Caching mapper changes

doc/ref_transform.rst

@@ -141,4 +141,10 @@ TODO: Matching instruction tags
 
 .. automodule:: loopy.match
 
+
+Fusing Loops
+------------
+
+.. automodule:: loopy.transform.loop_fusion
+
 .. vim: tw=75:spell

doc/tutorial.rst

@@ -610,7 +610,7 @@ commonly called 'loop tiling':
     ...     assumptions="n mod 16 = 0 and n >= 1")
     >>> knl = lp.split_iname(knl, "i", 16)
     >>> knl = lp.split_iname(knl, "j", 16)
-    >>> knl = lp.prioritize_loops(knl, "i_outer,j_outer,i_inner")
+    >>> knl = lp.prioritize_loops(knl, "i_outer,j_outer,i_inner,j_inner")
     >>> knl = lp.set_options(knl, "write_code")
     >>> evt, (out,) = knl(queue, a=a_mat_dev)
     #define lid(N) ((int) get_local_id(N))
@@ -1029,8 +1029,8 @@ transformation exists in :func:`loopy.add_prefetch`:
     >>> evt, (out,) = knl_pf(queue, a=x_vec_dev)
     #define lid(N) ((int) get_local_id(N))
     ...
-        acc_k = 0.0f;
         a_fetch = a[16 * gid(0) + lid(0)];
+        acc_k = 0.0f;
         for (int k = 0; k <= 15; ++k)
           acc_k = acc_k + a_fetch;
         out[16 * gid(0) + lid(0)] = acc_k;
@@ -1053,10 +1053,10 @@ earlier:
     >>> evt, (out,) = knl_pf(queue, a=x_vec_dev)
     #define lid(N) ((int) get_local_id(N))
     ...
-      if (-1 + -16 * gid(0) + -1 * lid(0) + n >= 0)
-        acc_k = 0.0f;
       if (-1 + -16 * gid(0) + -1 * lid(0) + n >= 0)
         a_fetch[lid(0)] = a[16 * gid(0) + lid(0)];
+      if (-1 + -16 * gid(0) + -1 * lid(0) + n >= 0)
+        acc_k = 0.0f;
       barrier(CLK_LOCAL_MEM_FENCE) /* for a_fetch (insn_k_update depends on a_fetch_rule) */;
       if (-1 + -16 * gid(0) + -1 * lid(0) + n >= 0)
       {
@@ -1908,18 +1908,16 @@ Now to make things more interesting, we'll create a kernel with barriers:
     {
       __local int c[50 * 10 * 99];
     <BLANKLINE>
-      {
-        int const k_outer = 0;
-    <BLANKLINE>
+      for (int i = 0; i <= 49; ++i)
         for (int j = 0; j <= 9; ++j)
-          for (int i = 0; i <= 49; ++i)
-          {
-            barrier(CLK_LOCAL_MEM_FENCE) /* for c (insn rev-depends on insn_0) */;
-            c[990 * i + 99 * j + lid(0) + 1] = 2 * a[980 * i + 98 * j + lid(0) + 1];
-            barrier(CLK_LOCAL_MEM_FENCE) /* for c (insn_0 depends on insn) */;
-            e[980 * i + 98 * j + lid(0) + 1] = c[990 * i + 99 * j + 1 + lid(0) + 1] + c[990 * i + 99 * j + -1 + lid(0) + 1];
-          }
-      }
+        {
+          int const k_outer = 0;
+    <BLANKLINE>
+          barrier(CLK_LOCAL_MEM_FENCE) /* for c (insn rev-depends on insn_0) */;
+          c[990 * i + 99 * j + lid(0) + 1] = 2 * a[980 * i + 98 * j + lid(0) + 1];
+          barrier(CLK_LOCAL_MEM_FENCE) /* for c (insn_0 depends on insn) */;
+          e[980 * i + 98 * j + lid(0) + 1] = c[990 * i + 99 * j + 1 + lid(0) + 1] + c[990 * i + 99 * j + -1 + lid(0) + 1];
+        }
     }
 
 In this kernel, when a work-item performs the second instruction it uses data

loopy/__init__.py

@@ -161,6 +161,9 @@
 
 from loopy.tools import Optional, t_unit_to_python, memoize_on_disk
 
+from loopy.transform.loop_fusion import (get_kennedy_unweighted_fusion_candidates,
+                                         rename_inames_in_batch)
+
 
 __all__ = [
         "TaggedVariable", "Reduction", "LinearSubscript", "TypeCast",
@@ -253,6 +256,9 @@
 
         "pack_and_unpack_args_for_call",
 
+        "rename_inames_in_batch",
+        "get_kennedy_unweighted_fusion_candidates",
+
         # }}}
 
         "get_dot_dependency_graph",
@@ -320,6 +326,15 @@
         # }}}
         ]
 
+
+try:
+    import loopy.relations as relations
+except ImportError:
+    # catching ImportErrors to avoid making minikanren a hard-dep
+    pass
+else:
+    __all__ += ["relations"]
+
 # }}}
 
 

loopy/codegen/result.py

@@ -195,6 +195,28 @@ def with_new_ast(self, codegen_state, new_ast):
                 self.current_program(codegen_state).copy(
                     ast=new_ast))
 
+    def get_idis_for_subkernel(self, kernel, name):
+        """
+        Returns a :class:`list` of :class:`~loopy.codegen.ImplementedDataInfo` for
+        the subkernel named *name*.
+
+        :arg kernel: An instance of :class:`loopy.LoopKernel`.
+        """
+        from loopy.schedule.tools import get_callkernel_dependencies
+        from loopy.kernel.data import InameArg
+        name2idi = {idi.name: idi for idi in self.implemented_data_info}
+        subknl_deps = get_callkernel_dependencies(kernel, name)
+        return [idi
+                for idi in self.implemented_data_info
+                if (idi.name in subknl_deps
+                    or idi.arg_class is InameArg
+                    or idi.base_name in subknl_deps
+                    or (idi.offset_for_name is not None
+                        and name2idi[idi.offset_for_name].base_name in subknl_deps)
+                    or (idi.stride_for_name_and_axis is not None
+                        and idi.stride_for_name_and_axis[0] in subknl_deps))]
+
+
 # }}}
 
 

loopy/kernel/tools.py

@@ -2070,4 +2070,67 @@ def get_outer_params(domains):
 # }}}
 
 
+# {{{ get access map from an instruction
+
+class _IndexCollector(CombineMapper):
+    def __init__(self, var):
+        self.var = var
+        super().__init__()
+
+    def combine(self, values):
+        import operator
+        return reduce(operator.or_, values, frozenset())
+
+    def map_subscript(self, expr):
+        if expr.aggregate.name == self.var:
+            return (super().map_subscript(expr) | frozenset([expr.index_tuple]))
+        else:
+            return super().map_subscript(expr)
+
+    def map_algebraic_leaf(self, expr):
+        return frozenset()
+
+    map_constant = map_algebraic_leaf
+
+
+def _project_out_inames_from_maps(amaps, inames_to_project_out):
+    new_amaps = []
+    for amap in amaps:
+        for iname in inames_to_project_out:
+            dt, pos = amap.get_var_dict()[iname]
+            amap = amap.project_out(dt, pos, 1)
+
+        new_amaps.append(amap)
+
+    return new_amaps
+
+
+def _union_amaps(amaps):
+    import islpy as isl
+    return reduce(isl.Map.union, amaps[1:], amaps[0])
+
+
+def get_insn_access_map(kernel, insn_id, var, inner_inames):
+    from loopy.transform.subst import expand_subst
+    from loopy.match import Id
+    from loopy.symbolic import get_access_map
+
+    insn = kernel.id_to_insn[insn_id]
+
+    kernel = expand_subst(kernel, within=Id(insn_id))
+    indices = list(_IndexCollector(var)((insn.expression,
+                                         insn.assignees,
+                                         tuple(insn.predicates))))
+
+    amaps = _project_out_inames_from_maps(
+        [get_access_map(kernel.get_inames_domain(insn.within_inames),
+                        idx, kernel.assumptions)
+
+         for idx in indices],
+        inner_inames)
+
+    return _union_amaps(amaps)
+
+# }}}
+
 # vim: foldmethod=marker

loopy/relations.py

@@ -0,0 +1,122 @@
+from kanren import Relation, facts
+
+
+def get_inameo(kernel):
+    inameo = Relation()
+    for iname in kernel.all_inames():
+        facts(inameo, (iname,))
+    return inameo
+
+
+def get_argo(kernel):
+    argo = Relation()
+    for arg in kernel.args:
+        facts(argo, (arg.name,))
+
+    return argo
+
+
+def get_tempo(kernel):
+    tempo = Relation()
+    for tv in kernel.temporary_variables:
+        facts(tempo, (tv,))
+
+    return tempo
+
+
+def get_insno(kernel):
+    insno = Relation()
+    for insn in kernel.instructions:
+        facts(insno, (insn.id,))
+
+    return insno
+
+
+def get_taggedo(kernel):
+    taggedo = Relation()
+
+    for arg_name, arg in kernel.arg_dict.items():
+        for tag in arg.tags:
+            facts(taggedo, (arg_name, tag))
+
+    for iname_name, iname in kernel.inames.items():
+        for tag in iname.tags:
+            facts(taggedo, (iname_name, tag))
+
+    for insn in kernel.instructions:
+        for tag in insn.tags:
+            facts(taggedo, (insn.id, tag))
+
+    return taggedo
+
+
+def get_taggedo_of_type(kernel, tag_type):
+    taggedo = Relation()
+
+    for arg_name, arg in kernel.arg_dict.items():
+        for tag in arg.tags_of_type(tag_type):
+            facts(taggedo, (arg_name, tag))
+
+    for iname_name, iname in kernel.inames.items():
+        for tag in iname.tags_of_type(tag_type):
+            facts(taggedo, (iname_name, tag))
+
+    for insn in kernel.instructions:
+        for tag in insn.tags_of_type(tag_type):
+            facts(taggedo, (insn.id, tag))
+
+    return taggedo
+
+
+def get_producero(kernel):
+    producero = Relation()
+
+    for insn in kernel.instructions:
+        for var in insn.assignee_var_names():
+            facts(producero, (insn.id, var))
+
+    return producero
+
+
+def get_consumero(kernel):
+    consumero = Relation()
+
+    for insn in kernel.instructions:
+        for var in insn.read_dependency_names():
+            facts(consumero, (insn.id, var))
+
+    return consumero
+
+
+def get_withino(kernel):
+    withino = Relation()
+
+    for insn in kernel.instructions:
+        facts(withino, (insn.id, insn.within_inames))
+
+    return withino
+
+
+def get_reduce_insno(kernel):
+    reduce_insno = Relation()
+
+    for insn in kernel.instructions:
+        if insn.reduction_inames():
+            facts(reduce_insno, (insn.id,))
+
+    return reduce_insno
+
+
+def get_reduce_inameo(kernel):
+    from functools import reduce
+    reduce_inameo = Relation()
+
+    for iname in reduce(frozenset.union,
+                        (insn.reduction_inames()
+                         for insn in kernel.instructions),
+                        frozenset()):
+        facts(reduce_inameo, (iname,))
+
+    return reduce_inameo
+
+# vim: fdm=marker