Add find_unused_axis_tag

doc/reference.rst

@@ -442,6 +442,8 @@ Wrangling inames
 
 .. autofunction:: realize_ilp
 
+.. autofunction:: find_unused_axis_tag
+
 Dealing with Parameters
 ^^^^^^^^^^^^^^^^^^^^^^^
 

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",

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