Implements a Loop Fusion Transformation

doc/ref_transform.rst

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

loopy/__init__.py

@@ -180,6 +180,10 @@
     simplify_indices,
     tag_instructions,
 )
+from loopy.transform.loop_fusion import (
+    get_kennedy_unweighted_fusion_candidates,
+    rename_inames_in_batch,
+)
 from loopy.transform.pack_and_unpack_args import pack_and_unpack_args_for_call
 from loopy.transform.padding import (
     add_padding,
@@ -325,6 +329,7 @@
     "get_dot_dependency_graph",
     "get_global_barrier_order",
     "get_iname_duplication_options",
+    "get_kennedy_unweighted_fusion_candidates",
     "get_mem_access_map",
     "get_one_linearized_kernel",
     "get_one_scheduled_kernel",
@@ -371,6 +376,7 @@
     "rename_callable",
     "rename_iname",
     "rename_inames",
+    "rename_inames_in_batch",
     "replace_instruction_ids",
     "save_and_reload_temporaries",
     "set_argument_order",

loopy/kernel/tools.py

@@ -2151,4 +2151,67 @@ def get_hw_axis_base_for_codegen(kernel: LoopKernel, iname: str) -> isl.Aff:
                                        constants_only=False)
     return lower_bound
 
+
+# {{{ get access map from an instruction
+
+class _IndexCollector(CombineMapper):
+    def __init__(self, var):
+        self.var = var
+
+    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.match import Id
+    from loopy.symbolic import get_access_map
+    from loopy.transform.subst import expand_subst
+
+    insn = kernel.id_to_insn[insn_id]
+
+    kernel = expand_subst(kernel, within=Id(insn_id))
+    indices = list(_IndexCollector(var)((insn.expression,
+                                         insn.assignees,
+                                         list(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