rename_inames: replace old inames that appear as params in other domains

Author: isuruf

loopy/transform/iname.py

@@ -2327,6 +2327,7 @@ def rename_inames(kernel, old_inames, new_iname, existing_ok=False,
         raise LoopyError(f"iname '{new_iname}' conflicts with an existing identifier"
                          " --cannot rename")
 
+    orig_old_inames = old_inames
     if not does_exist:
         # {{{ rename old_inames[0] -> new_iname
         # so that the code below can focus on "merging" inames that already exist
@@ -2404,6 +2405,16 @@ def does_insn_involve_iname(kernel, insn, *args):
             smap.map_kernel(kernel, within=does_insn_involve_iname,
                             map_tvs=False, map_args=False))
 
+    # replace instances where the old inames appear as a param
+    new_domains = []
+    for dom in kernel.domains:
+        for old_iname in orig_old_inames:
+            d = dom.get_var_dict()
+            if old_iname in d and new_iname not in d:
+                var_type, var_num = d[old_iname]
+                dom = dom.set_dim_name(var_type, var_num, new_iname)
+        new_domains.append(dom)
+
     new_instructions = [insn.copy(within_inames=((insn.within_inames
                                                   - frozenset(old_inames))
                                                  | frozenset([new_iname])))
@@ -2412,7 +2423,7 @@ def does_insn_involve_iname(kernel, insn, *args):
                         else insn
                         for insn in kernel.instructions]
 
-    kernel = kernel.copy(instructions=new_instructions)
+    kernel = kernel.copy(instructions=new_instructions, domains=new_domains)
 
     return kernel
 

test/test_transform.py

@@ -1468,6 +1468,22 @@ def test_rename_inames(ctx_factory):
     lp.auto_test_vs_ref(knl, ctx, ref_knl)
 
 
+def test_rename_inames_with_params(ctx_factory):
+    # https://github.com/inducer/loopy/issues/726
+    ctx = ctx_factory()
+
+    knl = lp.make_kernel(
+        [
+            "{ [i]: 0<=i<10 }",
+            "{ [k]: 0<=k<i }",
+        ],
+        "out[i, k] = 2"
+    )
+    ref_knl = knl
+    knl = lp.rename_inames(knl, ["i"], "j")
+    lp.auto_test_vs_ref(knl, ctx, ref_knl)
+
+
 def test_buffer_array_preserves_rev_deps(ctx_factory):
     # See https://github.com/inducer/loopy/issues/546
     ctx = ctx_factory()