Implements lp.rename_inames

loopy/__init__.py

@@ -71,7 +71,7 @@
 from loopy.transform.iname import (
         set_loop_priority, prioritize_loops, untag_inames,
         split_iname, chunk_iname, join_inames, tag_inames, duplicate_inames,
-        rename_iname, remove_unused_inames,
+        rename_iname, rename_inames, remove_unused_inames,
         split_reduction_inward, split_reduction_outward,
         affine_map_inames, find_unused_axis_tag,
         make_reduction_inames_unique,
@@ -198,7 +198,7 @@
         "set_loop_priority", "prioritize_loops", "untag_inames",
         "split_iname", "chunk_iname", "join_inames", "tag_inames",
         "duplicate_inames",
-        "rename_iname", "remove_unused_inames",
+        "rename_iname", "rename_inames", "remove_unused_inames",
         "split_reduction_inward", "split_reduction_outward",
         "affine_map_inames", "find_unused_axis_tag",
         "make_reduction_inames_unique",

loopy/transform/iname.py

@@ -58,6 +58,8 @@
 
 .. autofunction:: rename_iname
 
+.. autofunction:: rename_inames
+
 .. autofunction:: remove_unused_inames
 
 .. autofunction:: split_reduction_inward
@@ -1126,26 +1128,64 @@ def has_schedulable_iname_nesting(kernel):
 # {{{ rename_inames
 
 @for_each_kernel
-def rename_iname(kernel, old_iname, new_iname, existing_ok=False, within=None):
+def rename_inames(kernel, old_inames, new_iname, existing_ok=False, within=None):
     """
+    :arg old_inames: A collection of inames that must be renamed to **new_iname**.
     :arg within: a stack match as understood by
         :func:`loopy.match.parse_stack_match`.
     :arg existing_ok: execute even if *new_iname* already exists
     """
+    from collections.abc import Collection
+    if (isinstance(old_inames, str)
+            or not isinstance(old_inames, Collection)):
+        raise LoopyError("'old_inames' must be a collection of strings, "
+                         f"got '{type(old_inames)}'.")
+
+    if new_iname in old_inames:
+        raise LoopyError("new iname is part of inames being renamed")
+
+    if new_iname in (kernel.all_variable_names() - kernel.all_inames()):
+        raise LoopyError(f"New iname '{new_iname}' is already a variable in the"
+                         "kernel")
+
+    if any((len(insn.within_inames & frozenset(old_inames)) > 1)
+           for insn in kernel.instructions):
+        raise LoopyError("old_inames contains nested inames"
+                         " -- renaming is illegal.")
+
+    # sort to have deterministic implementation.
+    old_inames = sorted(old_inames)
 
     var_name_gen = kernel.get_var_name_generator()
 
     # FIXME: Distinguish existing iname vs. existing other variable
-    does_exist = var_name_gen.is_name_conflicting(new_iname)
+    does_exist = new_iname in kernel.all_inames()
 
-    if old_iname not in kernel.all_inames():
-        raise LoopyError("old iname '%s' does not exist" % old_iname)
+    if not (frozenset(old_inames) <= kernel.all_inames()):
+        raise LoopyError(f"old inames {frozenset(old_inames) - kernel.all_inames()}"
+                         " do not exist.")
 
     if does_exist and not existing_ok:
-        raise LoopyError("iname '%s' conflicts with an existing identifier"
-                "--cannot rename" % new_iname)
+        raise LoopyError(f"iname '{new_iname}' conflicts with an existing identifier"
+                         " --cannot rename")
 
-    if does_exist:
+    if not does_exist:
+        # {{{ rename old_inames[0] -> new_iname
+        # so that the code below can focus on "merging" inames that already exist
+
+        kernel = duplicate_inames(
+                kernel, old_inames[0], within=within, new_inames=[new_iname])
+        kernel = remove_unused_inames(kernel, old_inames[0])
+
+        # old_iname[0] is already renamed to new_iname => do not rename again.
+        old_inames = old_inames[1:]
+
+        # }}}
+
+    del does_exist
+    assert new_iname in kernel.all_inames()
+
+    for old_iname in old_inames:
         # {{{ check that the domains match up
 
         dom = kernel.get_inames_domain(frozenset((old_iname, new_iname)))
@@ -1177,42 +1217,45 @@ def rename_iname(kernel, old_iname, new_iname, existing_ok=False, within=None):
 
         # }}}
 
-        from pymbolic import var
-        subst_dict = {old_iname: var(new_iname)}
-
-        from loopy.match import parse_stack_match
-        within = parse_stack_match(within)
+    from pymbolic import var
+    subst_dict = {old_iname: var(new_iname) for old_iname in old_inames}
 
-        from pymbolic.mapper.substitutor import make_subst_func
-        rule_mapping_context = SubstitutionRuleMappingContext(
-                kernel.substitutions, var_name_gen)
-        smap = RuleAwareSubstitutionMapper(rule_mapping_context,
-                        make_subst_func(subst_dict), within)
+    from loopy.match import parse_stack_match
+    within = parse_stack_match(within)
 
-        kernel = rule_mapping_context.finish_kernel(
-                smap.map_kernel(kernel))
+    from pymbolic.mapper.substitutor import make_subst_func
+    rule_mapping_context = SubstitutionRuleMappingContext(
+            kernel.substitutions, var_name_gen)
+    smap = RuleAwareSubstitutionMapper(rule_mapping_context,
+                    make_subst_func(subst_dict), within)
 
-        new_instructions = []
-        for insn in kernel.instructions:
-            if (old_iname in insn.within_inames
-                    and within(kernel, insn, ())):
-                insn = insn.copy(
-                        within_inames=(
-                            (insn.within_inames - frozenset([old_iname]))
-                            | frozenset([new_iname])))
+    from loopy.kernel.instruction import MultiAssignmentBase
 
-            new_instructions.append(insn)
+    def does_insn_involve_iname(kernel, insn, *args):
+        return (not isinstance(insn, MultiAssignmentBase)
+                or frozenset(old_inames) & insn.dependency_names()
+                or frozenset(old_inames) & insn.reduction_inames())
 
-        kernel = kernel.copy(instructions=new_instructions)
+    kernel = rule_mapping_context.finish_kernel(
+            smap.map_kernel(kernel, within=does_insn_involve_iname))
 
-    else:
-        kernel = duplicate_inames(
-                kernel, [old_iname], within=within, new_inames=[new_iname])
+    new_instructions = [insn.copy(within_inames=((insn.within_inames
+                                                  - frozenset(old_inames))
+                                                 | frozenset([new_iname])))
+                        if ((len(frozenset(old_inames) & insn.within_inames) != 0)
+                            and within(kernel, insn, ()))
+                        else insn
+                        for insn in kernel.instructions]
 
-    kernel = remove_unused_inames(kernel, [old_iname])
+    kernel = kernel.copy(instructions=new_instructions)
+    kernel = remove_unused_inames(kernel, old_inames)
 
     return kernel
 
+
+def rename_iname(kernel, old_iname, new_iname, existing_ok=False, within=None):
+    return rename_inames(kernel, [old_iname], new_iname, existing_ok, within)
+
 # }}}
 
 

test/test_transform.py

@@ -1337,6 +1337,35 @@ def test_rename_inames_redn():
     assert "ifused" in t_unit.default_entrypoint.all_inames()
 
 
+def test_rename_inames(ctx_factory):
+    ctx = ctx_factory()
+
+    knl = lp.make_kernel(
+        "{[i1, i2]: 0<=i1, i2<10}",
+        """
+        y1[i1] = 2
+        y2[i2] = 3
+        """)
+    ref_knl = knl
+    knl = lp.rename_inames(knl, ["i1", "i2"], "ifused")
+    lp.auto_test_vs_ref(knl, ctx, ref_knl)
+
+
+def test_rename_inames_existing_ok(ctx_factory):
+    ctx = ctx_factory()
+
+    knl = lp.make_kernel(
+        "{[i1, i2, i3]: 0<=i1, i2, i3<10}",
+        """
+        y1[i1] = 2
+        y2[i2] = 3
+        y3[i3] = 4
+        """)
+    ref_knl = knl
+    knl = lp.rename_inames(knl, ["i1", "i2"], "i3", existing_ok=True)
+    lp.auto_test_vs_ref(knl, ctx, ref_knl)
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])