From 5619ce81c5a2e1c88aed66bd54c902a3401528bc Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Mon, 29 Apr 2013 22:59:06 -0400
Subject: [PATCH] Add type inference for kernel arguments.

---
 examples/hello-loopy.py |   8 +--
 loopy/compiled.py       | 135 ++++++++++++++++++++++++++++++----------
 loopy/kernel.py         |  99 ++++++++++++++++++++++++++---
 test/test_loopy.py      |   8 +--
 4 files changed, 203 insertions(+), 47 deletions(-)

diff --git a/examples/hello-loopy.py b/examples/hello-loopy.py
index 629a6e6d0..a6b8e6d8b 100644
--- a/examples/hello-loopy.py
+++ b/examples/hello-loopy.py
@@ -19,9 +19,9 @@ knl = lp.make_kernel(ctx.devices[0],
     "{[i]: 0<=i<n}",   # "loop domain"-- what values does i take?
     "out[i] = 2*a[i]", # "instructions" to be executed across the domain
     [ # argument declarations
-        lp.GlobalArg("out", np.float32, shape="n"),
-        lp.GlobalArg("a", np.float32, shape="n"),
-        lp.ValueArg("n", np.int32),
+        lp.GlobalArg("out", shape="n"),
+        lp.GlobalArg("a", shape="n"),
+        lp.ValueArg("n"),
         ])
 
 # -----------------------------------------------------------------------------
@@ -35,4 +35,4 @@ knl = lp.split_iname(knl, "i", 128, outer_tag="g.0", inner_tag="l.0")
 cknl = lp.CompiledKernel(ctx, knl)
 evt, (out,) = cknl(queue, a=a, n=n)
 
-cknl.print_code()
+print cknl.get_highlighted_code({"a": np.float32})
diff --git a/loopy/compiled.py b/loopy/compiled.py
index 11a7b888e..84c525fd7 100644
--- a/loopy/compiled.py
+++ b/loopy/compiled.py
@@ -30,7 +30,7 @@ import pyopencl.array as cl_array
 
 import numpy as np
 
-from pytools import Record
+from pytools import Record, memoize_method
 
 AUTO_TEST_SKIP_RUN = False
 
@@ -109,18 +109,69 @@ class CompiledKernel:
 
         # }}}
 
+        self.context = context
         self.kernel = kernel
+        self.edit_code = edit_code
+        self.codegen_kwargs = codegen_kwargs
+        self.options = options
+
+        # {{{ precompile, store grid size functions
+
+        if size_args is None:
+            self.size_args = kernel.scalar_loop_args
+        else:
+            self.size_args = size_args
+
+        gsize_expr, lsize_expr = kernel.get_grid_sizes_as_exprs()
+
+        if not gsize_expr: gsize_expr = (1,)
+        if not lsize_expr: lsize_expr = (1,)
+
+        from pymbolic import compile
+        self.global_size_func = compile(
+                gsize_expr, self.size_args)
+        self.local_size_func = compile(
+                lsize_expr, self.size_args)
+
+        # }}}
+
+    @memoize_method
+    def get_kernel(self, dtype_mapping_set):
+        kernel = self.kernel
+
+        from loopy.kernel import (
+                add_argument_dtypes,
+                infer_argument_dtypes,
+                get_arguments_with_incomplete_dtype)
+
+        if get_arguments_with_incomplete_dtype(kernel):
+            if dtype_mapping_set is not None:
+                kernel = add_argument_dtypes(kernel, dict(dtype_mapping_set))
+
+            kernel = infer_argument_dtypes(kernel)
+
+            incomplete_args = get_arguments_with_incomplete_dtype(kernel)
+            if incomplete_args:
+                raise RuntimeError("not all argument dtypes are specified "
+                        "or could be inferred: " + ", ".join(incomplete_args))
+
+        return kernel
+
+    @memoize_method
+    def get_cl_kernel(self, dtype_mapping_set):
+        kernel = self.get_kernel(dtype_mapping_set)
+
         from loopy.codegen import generate_code
-        self.code = generate_code(kernel, **codegen_kwargs)
+        code = generate_code(kernel, **self.codegen_kwargs)
 
-        if edit_code:
+        if self.edit_code:
             from pytools import invoke_editor
-            self.code = invoke_editor(self.code, "code.cl")
+            code = invoke_editor(code, "code.cl")
 
         try:
-            self.cl_program = cl.Program(context, self.code)
-            self.cl_kernel = getattr(
-                    self.cl_program.build(options=options),
+            cl_program = cl.Program(self.context, code)
+            cl_kernel = getattr(
+                    cl_program.build(options=self.options),
                     kernel.name)
         except KeyboardInterrupt:
             raise
@@ -128,7 +179,7 @@ class CompiledKernel:
             print "[Loopy] ----------------------------------------------------"
             print "[Loopy] build failed, here's the source code:"
             print "[Loopy] ----------------------------------------------------"
-            print self.code
+            print code
             print "[Loopy] ----------------------------------------------------"
             print "[Loopy] end source code"
             print "[Loopy] ----------------------------------------------------"
@@ -143,23 +194,33 @@ class CompiledKernel:
             else:
                 arg_types.append(None)
 
-        self.cl_kernel.set_scalar_arg_dtypes(arg_types)
+        cl_kernel.set_scalar_arg_dtypes(arg_types)
 
-        if size_args is None:
-            self.size_args = kernel.scalar_loop_args
-        else:
-            self.size_args = size_args
+        return kernel, cl_kernel
 
-        gsize_expr, lsize_expr = kernel.get_grid_sizes_as_exprs()
+    # {{{ debugging aids
 
-        if not gsize_expr: gsize_expr = (1,)
-        if not lsize_expr: lsize_expr = (1,)
+    def get_code(self, dtype_dict=None):
+        if dtype_dict is not None:
+            dtype_dict = frozenset(dtype_dict.items())
 
-        from pymbolic import compile
-        self.global_size_func = compile(
-                gsize_expr, self.size_args)
-        self.local_size_func = compile(
-                lsize_expr, self.size_args)
+        kernel = self.get_kernel(dtype_dict)
+
+        from loopy.codegen import generate_code
+        return generate_code(kernel, **self.codegen_kwargs)
+
+    def get_highlighted_code(self, dtype_dict=None):
+        return get_highlighted_code(self.get_code(dtype_dict))
+
+    @property
+    def code(self):
+        from warnings import warn
+        warn("CompiledKernel.code is deprecated. Use .get_code() instead.",
+                DeprecationWarning, stacklevel=2)
+
+        return self.get_code()
+
+    # }}}
 
     def __call__(self, queue, **kwargs):
         """If all array arguments are :mod:`numpy` arrays, defaults to returning
@@ -172,15 +233,28 @@ class CompiledKernel:
         no_run = kwargs.pop("no_run", None)
         warn_numpy = kwargs.pop("warn_numpy", None)
 
-        import loopy as lp
+        # {{{ process arg types, get cl kernel
+
+        dtype_dict = {}
+        for arg in self.kernel.args:
+            val = kwargs.get(arg.name)
+            if val is not None:
+                try:
+                    dtype = val.dtype
+                except AttributeError:
+                    pass
+                else:
+                    dtype_dict[arg.name] = dtype
+
+        kernel, cl_kernel = self.get_cl_kernel(frozenset(dtype_dict.iteritems()))
+        del dtype_dict
 
-        if self.needs_check:
-            assert len(list(lp.check_kernels([self.kernel], kwargs))) == 1
+        # }}}
 
-            self.needs_check = False
+        import loopy as lp
 
         domain_parameters = dict((name, int(kwargs[name]))
-                for name in self.kernel.scalar_loop_args)
+                for name in kernel.scalar_loop_args)
 
         args = []
         outputs = []
@@ -188,8 +262,8 @@ class CompiledKernel:
 
         kwargs_copy = kwargs.copy()
 
-        for arg in self.kernel.args:
-            is_written = arg.name in self.kernel.get_written_variables()
+        for arg in kernel.args:
+            is_written = arg.name in kernel.get_written_variables()
 
             val = kwargs_copy.pop(arg.name, None)
 
@@ -247,7 +321,7 @@ class CompiledKernel:
         if no_run:
             evt = cl.enqueue_marker(queue)
         else:
-            evt = self.cl_kernel(queue,
+            evt = cl_kernel(queue,
                     self.global_size_func(**domain_parameters),
                     self.local_size_func(**domain_parameters),
                     *args,
@@ -260,9 +334,6 @@ class CompiledKernel:
 
         return evt, outputs
 
-    def print_code(self):
-        print get_highlighted_code(self.code)
-
 # }}}
 
 
diff --git a/loopy/kernel.py b/loopy/kernel.py
index 80aa93842..f0c799e5d 100644
--- a/loopy/kernel.py
+++ b/loopy/kernel.py
@@ -147,7 +147,7 @@ def parse_tag(tag):
 # {{{ arguments
 
 class _ShapedArg(Record):
-    def __init__(self, name, dtype, shape=None, strides=None, order="C",
+    def __init__(self, name, dtype=None, shape=None, strides=None, order="C",
             offset=0):
         """
         All of the following are optional. Specify either strides or shape.
@@ -159,7 +159,8 @@ class _ShapedArg(Record):
         :arg offset: Offset from the beginning of the vector from which
             the strides are counted.
         """
-        dtype = np.dtype(dtype)
+        if dtype is not None:
+            dtype = np.dtype(dtype)
 
         def parse_if_necessary(x):
             if isinstance(x, str):
@@ -227,7 +228,7 @@ class ConstantArg(_ShapedArg):
                 self.name, self.dtype, ",".join(str(i) for i in self.shape))
 
 class ImageArg(Record):
-    def __init__(self, name, dtype, dimensions=None, shape=None):
+    def __init__(self, name, dtype=None, dimensions=None, shape=None):
         dtype = np.dtype(dtype)
         if shape is not None:
             if dimensions is not None and dimensions != len(shape):
@@ -248,17 +249,19 @@ class ImageArg(Record):
     def __repr__(self):
         return "<ImageArg '%s' of type %s>" % (self.name, self.dtype)
 
-
 class ValueArg(Record):
-    def __init__(self, name, dtype, approximately=None):
-        Record.__init__(self, name=name, dtype=np.dtype(dtype),
+    def __init__(self, name, dtype=None, approximately=None):
+        if dtype is not None:
+            dtype = np.dtype(dtype)
+
+        Record.__init__(self, name=name, dtype=dtype,
                 approximately=approximately)
 
     def __repr__(self):
         return "<ValueArg '%s' of type %s>" % (self.name, self.dtype)
 
 class ScalarArg(ValueArg):
-    def __init__(self, name, dtype, approximately=None):
+    def __init__(self, name, dtype=None, approximately=None):
         from warnings import warn
         warn("ScalarArg is a deprecated name of ValueArg",
                 DeprecationWarning, stacklevel=2)
@@ -1351,6 +1354,16 @@ class LoopKernel(Record):
             result.update(dom.get_var_names(dim_type.set))
         return frozenset(result)
 
+    @memoize_method
+    def all_params(self):
+        all_inames = self.all_inames()
+
+        result = set()
+        for dom in self.domains:
+            result.update(set(dom.get_var_names(dim_type.param)) - all_inames)
+
+        return frozenset(result)
+
     @memoize_method
     def all_insn_inames(self):
         """Return a mapping from instruction ids to inames inside which
@@ -1678,6 +1691,78 @@ class LoopKernel(Record):
 
 # }}}
 
+# {{{ add and infer argument dtypes
+
+def add_argument_dtypes(knl, dtype_dict):
+    dtype_dict = dtype_dict.copy()
+    new_args = []
+
+    for arg in knl.args:
+        new_dtype = dtype_dict.pop(arg.name, None)
+        if new_dtype is not None:
+            new_dtype = np.dtype(new_dtype)
+            if arg.dtype is not None and arg.dtype != new_dtype:
+                raise RuntimeError(
+                        "argument '%s' already has a different dtype "
+                        "(existing: %s, new: %s)"
+                        % (arg.name, arg.dtype, new_dtype))
+            arg = arg.copy(dtype=new_dtype)
+
+        new_args.append(arg)
+
+    knl = knl.copy(args=new_args)
+
+    if dtype_dict:
+        raise RuntimeError("unused argument dtypes: %s"
+                % ", ".join(dtype_dict))
+
+    return knl.copy(args=new_args)
+
+def infer_argument_dtypes(knl):
+    new_args = []
+
+    writer_map = knl.writer_map()
+
+    from loopy.codegen.expression import (
+            TypeInferenceMapper, TypeInferenceFailure)
+    tim = TypeInferenceMapper(knl)
+
+    for arg in knl.args:
+        if arg.dtype is None:
+            new_dtype = None
+
+            if arg.name in knl.all_params():
+                new_dtype = knl.index_dtype
+            else:
+                try:
+                    for write_insn_id in writer_map.get(arg.name, ()):
+                        write_insn = knl.id_to_insn[write_insn_id]
+                        new_tim_dtype = tim(write_insn.expression)
+                        if new_dtype is None:
+                            new_dtype = new_tim_dtype
+                        elif new_dtype != new_tim_dtype:
+                            # Now we know *nothing*.
+                            new_dtype = None
+                            break
+
+                except TypeInferenceFailure:
+                    # Even one type inference failure is enough to
+                    # make this dtype not safe to guess. Don't.
+                    pass
+
+            if new_dtype is not None:
+                arg = arg.copy(dtype=new_dtype)
+
+        new_args.append(arg)
+
+    return knl.copy(args=new_args)
+
+def get_arguments_with_incomplete_dtype(knl):
+    return [arg.name for arg in knl.args
+            if arg.dtype is None]
+
+# }}}
+
 # {{{ find_all_insn_inames fixed point iteration
 
 def find_all_insn_inames(kernel):
diff --git a/test/test_loopy.py b/test/test_loopy.py
index c9648a46e..16a096c19 100644
--- a/test/test_loopy.py
+++ b/test/test_loopy.py
@@ -204,7 +204,7 @@ def test_wg_too_small(ctx_factory):
 
     for gen_knl in kernel_gen:
         try:
-            lp.CompiledKernel(ctx, gen_knl)
+            lp.CompiledKernel(ctx, gen_knl).get_code()
         except RuntimeError, e:
             assert "implemented and desired" in str(e)
             pass # expected!
@@ -655,7 +655,7 @@ def test_dependent_loop_bounds(ctx_factory):
 
     cknl = lp.CompiledKernel(ctx, knl)
     print "---------------------------------------------------"
-    cknl.print_code()
+    print cknl.get_highlighted_code()
     print "---------------------------------------------------"
 
 
@@ -689,7 +689,7 @@ def test_dependent_loop_bounds_2(ctx_factory):
             inner_tag="l.0")
     cknl = lp.CompiledKernel(ctx, knl)
     print "---------------------------------------------------"
-    cknl.print_code()
+    print cknl.get_highlighted_code()
     print "---------------------------------------------------"
 
 
@@ -727,7 +727,7 @@ def test_dependent_loop_bounds_3(ctx_factory):
 
     cknl = lp.CompiledKernel(ctx, knl)
     print "---------------------------------------------------"
-    cknl.print_code()
+    print cknl.get_highlighted_code()
     print "---------------------------------------------------"
 
     knl_bad = lp.split_iname(knl, "jj", 128, outer_tag="g.1",
-- 
GitLab