From e855c21f8c734377d0b1fcf12a68d86f5bec75e3 Mon Sep 17 00:00:00 2001
From: Nicholas Christensen <11543181+nchristensen@users.noreply.github.com>
Date: Wed, 21 Jun 2023 14:59:38 +0000
Subject: [PATCH] Preserve tags in rename_iname (#786)

* Preserve tags in rename_iname

* Make preserving the tags optional

* Preserve tags by default

* Remove redundant line

* Doc fix, maintain order of pre-existing arguments

* add for_each_kernel decorator

* flake8 fix
---
 loopy/transform/iname.py | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py
index 870f6c843..673920364 100644
--- a/loopy/transform/iname.py
+++ b/loopy/transform/iname.py
@@ -2502,8 +2502,23 @@ def does_insn_involve_iname(kernel, insn, *args):
     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)
+@for_each_kernel
+def rename_iname(kernel, old_iname, new_iname, existing_ok=False,
+                 within=None, preserve_tags=True):
+    """
+    Single iname version of :func:`loopy.rename_inames`.
+    :arg existing_ok: execute even if *new_iname* already exists
+    :arg within: a stack match understood by :func:`loopy.match.parse_stack_match`.
+    :arg preserve_tags: copy the tags on the old iname to the new iname
+    """
+    from itertools import product
+    from loopy import tag_inames
+
+    tags = kernel.inames[old_iname].tags
+    kernel = rename_inames(kernel, [old_iname], new_iname, existing_ok, within)
+    if preserve_tags:
+        kernel = tag_inames(kernel, product([new_iname], tags))
+    return kernel
 
 # }}}