Automatically deduce domain parameters, if possible.

MEMO

@@ -107,9 +107,6 @@ Future ideas
 
 - Float4 joining on fetch/store?
 
-- How can one automatically generate something like microblocks?
-  -> Some sort of axis-adding transform?
-
 - Better for loop bound generation
   -> Try a triangular loop
 
@@ -134,6 +131,9 @@ Future ideas
 Dealt with
 ^^^^^^^^^^
 
+- How can one automatically generate something like microblocks?
+  -> Some sort of axis-adding transform?
+
 - ExpandingIdentityMapper
   extract_subst -> needs WalkMapper [actually fine as is]
   padding [DONE]

loopy/compiled.py

@@ -298,6 +298,9 @@ class CompiledKernel:
 
         import loopy as lp
 
+        kwargs.update(
+                kernel.domain_parameter_finder()(kwargs))
+
         domain_parameters = dict((name, int(kwargs[name]))
                 for name in kernel.scalar_loop_args)
 

loopy/kernel/__init__.py

@@ -861,6 +861,11 @@ class LoopKernel(Record):
 
     # }}}
 
+    @memoize_method
+    def domain_parameter_finder(self):
+        from loopy.kernel.tools import DomainParameterFinder
+        return DomainParameterFinder(self)
+
 # }}}
 
 # vim: foldmethod=marker

loopy/kernel/tools.py

@@ -291,4 +291,61 @@ def get_dot_dependency_graph(kernel, iname_cluster=False, iname_edge=True):
 
 # }}}
 
+# {{{ domain parameter finder
+
+class DomainParameterFinder:
+    """Finds parameters from shapes of passed arguments."""
+
+    def __init__(self, kernel):
+        # a mapping from parameter names to a list of tuples
+        # (arg_name, axis_nr, function), where function is a
+        # unary function of kernel.arg_dict[arg_name].shape[axis_nr]
+        # returning the desired parameter.
+        self.param_to_sources = param_to_sources = {}
+
+        param_names = kernel.all_params()
+
+        from loopy.kernel.data import GlobalArg
+        from loopy.symbolic import DependencyMapper
+        from pymbolic import compile
+        dep_map = DependencyMapper()
+
+        from pymbolic import var
+        for arg in kernel.args:
+            if isinstance(arg, GlobalArg):
+                for axis_nr, shape_i in enumerate(arg.shape):
+                    deps = dep_map(shape_i)
+                    if len(deps) == 1:
+                        dep, = deps
+
+                        if dep.name in param_names:
+                            from pymbolic.algorithm import solve_affine_equations_for
+                            try:
+                                # friggin' overkill :)
+                                param_expr = solve_affine_equations_for(
+                                        [dep.name], [(shape_i, var("shape_i"))]) \
+                                                [dep.name]
+                            except:
+                                # went wrong? oh well
+                                pass
+                            else:
+                                param_func = compile(param_expr, ["shape_i"])
+                                param_to_sources.setdefault(dep.name, []).append(
+                                        (arg.name, axis_nr, param_func))
+
+    def __call__(self, kwargs):
+        result = {}
+
+        for param_name, sources in self.param_to_sources.iteritems():
+            if param_name not in kwargs:
+                for arg_name, axis_nr, shape_func in sources:
+                    if arg_name in kwargs:
+                        result[param_name] = shape_func(
+                                kwargs[arg_name].shape[axis_nr])
+                        continue
+
+        return result
+
+# }}}
+
 # vim: foldmethod=marker

loopy/symbolic.py

@@ -47,6 +47,8 @@ from pymbolic.mapper.stringifier import \
         StringifyMapper as StringifyMapperBase
 from pymbolic.mapper.dependency import \
         DependencyMapper as DependencyMapperBase
+from pymbolic.mapper.coefficient import \
+        CoefficientCollector as CoefficientCollectorBase
 from pymbolic.mapper.unifier import UnidirectionalUnifier \
         as UnidirectionalUnifierBase
 
@@ -666,58 +668,8 @@ def parse(expr_str):
 
 # {{{ coefficient collector
 
-class CoefficientCollector(RecursiveMapper):
-    def map_sum(self, expr):
-        stride_dicts = [self.rec(ch) for ch in expr.children]
-
-        result = {}
-        for stride_dict in stride_dicts:
-            for var, stride in stride_dict.iteritems():
-                if var in result:
-                    result[var] += stride
-                else:
-                    result[var] = stride
-
-        return result
-
-    def map_product(self, expr):
-        result = {}
-
-        children_coeffs = [self.rec(child) for child in expr.children]
-
-        idx_of_child_with_vars = None
-        for i, child_coeffs in enumerate(children_coeffs):
-            for k in child_coeffs:
-                if isinstance(k, str):
-                    if (idx_of_child_with_vars is not None
-                            and idx_of_child_with_vars != i):
-                        raise RuntimeError(
-                                "nonlinear expression")
-                    idx_of_child_with_vars = i
-
-        other_coeffs = 1
-        for i, child_coeffs in enumerate(children_coeffs):
-            if i != idx_of_child_with_vars:
-                assert len(child_coeffs) == 1
-                other_coeffs *= child_coeffs[1]
-
-        if idx_of_child_with_vars is None:
-            return {1: other_coeffs}
-        else:
-            return dict(
-                    (var, other_coeffs*coeff)
-                    for var, coeff in
-                    children_coeffs[idx_of_child_with_vars].iteritems())
-
-        return result
-
-    def map_constant(self, expr):
-        return {1: expr}
-
-    def map_variable(self, expr):
-        return {expr.name: 1}
-
-    map_tagged_variable = map_variable
+class CoefficientCollector(CoefficientCollectorBase):
+    map_tagged_variable = CoefficientCollectorBase.map_variable
 
     def map_subscript(self, expr):
         raise RuntimeError("cannot gather coefficients--indirect addressing in use")

test/test_loopy.py

@@ -1201,7 +1201,7 @@ def test_array_with_offset(ctx_factory):
     a = a_full[3:10]
 
     print cknl.get_highlighted_code({"a": a.dtype}, {"a": True, "b": False})
-    evt, (b,) = cknl(queue, a=a, n=a.shape[0], m=a.shape[1])
+    evt, (b,) = cknl(queue, a=a)
 
     import numpy.linalg as la
     assert la.norm(b.get() - 2*a.get()) < 1e-13