From 965ed5edec4eab104f2fabc2f3cd8054e6b640a7 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Sat, 13 Nov 2021 11:32:25 -0600 Subject: [PATCH 1/2] implements rename_inames --- loopy/__init__.py | 4 +- loopy/transform/iname.py | 109 +++++++++++++++++++++++++++------------ 2 files changed, 78 insertions(+), 35 deletions(-) diff --git a/loopy/__init__.py b/loopy/__init__.py index 9bd01534b..7e6ee5234 100644 --- a/loopy/__init__.py +++ b/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", diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index e55bad50c..3712d678b 100644 --- a/loopy/transform/iname.py +++ b/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) + # }}} From 66b8f3c76ee73c66ca00980da35d485759cb7530 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Sat, 13 Nov 2021 11:32:35 -0600 Subject: [PATCH 2/2] tests rename inames --- test/test_transform.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/test/test_transform.py b/test/test_transform.py index 3915ce161..e42eeb498 100644 --- a/test/test_transform.py +++ b/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])