diff --git a/doc/reference.rst b/doc/reference.rst index 839851d1796068a0c48e052c49dea654ac250389..412fad92b59320950b4e7c913797f061c211276d 100644 --- a/doc/reference.rst +++ b/doc/reference.rst @@ -442,6 +442,8 @@ Wrangling inames .. autofunction:: realize_ilp +.. autofunction:: find_unused_axis_tag + Dealing with Parameters ^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/loopy/__init__.py b/loopy/__init__.py index 07d536eb3b680492a5a4827b61f44b7389a4c31e..3b946f8f89c5440894ac726bbe45db321233939e 100644 --- a/loopy/__init__.py +++ b/loopy/__init__.py @@ -57,7 +57,7 @@ from loopy.transform.iname import ( assume, set_loop_priority, split_iname, chunk_iname, join_inames, tag_inames, duplicate_inames, rename_iname, link_inames, remove_unused_inames, - affine_map_inames) + affine_map_inames, find_unused_axis_tag) from loopy.transform.instruction import ( find_instructions, map_instructions, @@ -131,7 +131,7 @@ __all__ = [ "split_iname", "chunk_iname", "join_inames", "tag_inames", "duplicate_inames", "rename_iname", "link_inames", "remove_unused_inames", - "affine_map_inames", + "affine_map_inames", "find_unused_axis_tag", "add_prefetch", "change_arg_to_image", "tag_data_axes", "set_array_dim_names", "remove_unused_arguments", diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 8927098cf7f1111b591040440b7839732fce602c..bea37d56135618b5bbce84f7de416ab9fa382ae1 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1173,4 +1173,52 @@ def affine_map_inames(kernel, old_inames, new_inames, equations): # }}} +# {{{ find unused axes + +def find_unused_axis_tag(kernel, kind, insn_match=None): + """For one of the hardware-parallel execution tags, find an unused + axis. + + :arg insn_match: An instruction match as understood by + :func:`loopy.context_matching.parse_match`. + :arg kind: may be "l" or "g", or the corresponding tag class name + + :returns: an :class:`GroupIndexTag` or :class:`LocalIndexTag` + that is not being used within the instructions matched by + *insn_match*. + """ + used_axes = set() + + from looopy.kernel.data import GroupIndexTag, LocalIndexTag + + if isinstance(kind, str): + found = False + for cls in [GroupIndexTag, LocalIndexTag]: + if kind == cls.print_name: + kind = cls + found = True + break + + if not found: + raise LoopyError("invlaid tag kind: %s" % kind) + + from loopy.context_matching import parse_match + match = parse_match(insn_match) + insns = [insn for insn in kernel.instructions if match(kernel, insn)] + + for insn in insns: + for iname in kernel.insn_inames(insn): + dim_tag = kernel.iname_to_tag.get(iname) + + if isinstance(dim_tag, kind): + used_axes.add(kind.axis) + + i = 0 + while i in used_axes: + i += 1 + + return kind(i) + +# }}} + # vim: foldmethod=marker