Parallelization transformation

loopy/__init__.py

@@ -199,6 +199,10 @@
     split_array_axis,
     split_array_dim,
 )
+from loopy.transform.parallelize import (
+    add_gbarrier_between_disjoint_loop_sets,
+    split_iteration_domain_across_work_items,
+)
 from loopy.transform.parameter import assume, fix_parameters
 from loopy.transform.precompute import precompute
 from loopy.transform.privatize import (
@@ -298,6 +302,7 @@
     "add_barrier",
     "add_dependency",
     "add_dtypes",
+    "add_gbarrier_between_disjoint_loop_sets",
     "add_inames_for_unused_hw_axes",
     "add_inames_to_insn",
     "add_nosync",
@@ -402,6 +407,7 @@
     "split_array_axis",
     "split_array_dim",
     "split_iname",
+    "split_iteration_domain_across_work_items",
     "split_reduction_inward",
     "split_reduction_outward",
     "t_unit_to_python",

loopy/transform/data.py

@@ -861,13 +861,15 @@ def set_temporary_scope(kernel, temp_var_names, address_space):
 
 @for_each_kernel
 def reduction_arg_to_subst_rule(
-        kernel, inames, insn_match=None, subst_rule_name=None):
+        kernel: LoopKernel, inames: str | Sequence[str],
+        insn_match: ToMatchConvertible = None,
+        subst_rule_name: str | None = None):
     if isinstance(inames, str):
         inames = [s.strip() for s in inames.split(",")]
 
     inames_set = frozenset(inames)
 
-    substs = kernel.substitutions.copy()
+    substs = dict(kernel.substitutions)
 
     var_name_gen = kernel.get_var_name_generator()
 
@@ -908,10 +910,13 @@ def map_reduction(expr, rec, nresults=1):
     cb_mapper = ReductionCallbackMapper(map_reduction)
 
     from loopy.kernel.data import MultiAssignmentBase
+    from loopy.match import parse_match
+    match = parse_match(insn_match)
 
     new_insns = []
     for insn in kernel.instructions:
-        if not isinstance(insn, MultiAssignmentBase):
+        if (not isinstance(insn, MultiAssignmentBase)
+                or not match(kernel, insn)):
             new_insns.append(insn)
         else:
             new_insns.append(insn.copy(expression=cb_mapper(insn.expression)))