(draft PR for comparison only) Compare old loop nest constraint branch to new branch

doc/ref_kernel.rst

@@ -220,6 +220,8 @@ Tag                             Meaning
 Identifiers
 -----------
 
+.. _reserved-identifiers:
+
 Reserved Identifiers
 ^^^^^^^^^^^^^^^^^^^^
 

loopy/__init__.py

@@ -69,18 +69,20 @@
 from loopy.version import VERSION, MOST_RECENT_LANGUAGE_VERSION
 
 from loopy.transform.iname import (
-        set_loop_priority, prioritize_loops, untag_inames,
+        set_loop_priority, prioritize_loops, constrain_loop_nesting,
+        untag_inames,
         split_iname, chunk_iname, join_inames, tag_inames, duplicate_inames,
         rename_iname, remove_unused_inames,
         split_reduction_inward, split_reduction_outward,
         affine_map_inames, find_unused_axis_tag,
         make_reduction_inames_unique,
         has_schedulable_iname_nesting, get_iname_duplication_options,
-        add_inames_to_insn, add_inames_for_unused_hw_axes)
+        add_inames_to_insn, add_inames_for_unused_hw_axes,  map_domain)
 
 from loopy.transform.instruction import (
         find_instructions, map_instructions,
-        set_instruction_priority, add_dependency,
+        set_instruction_priority,
+        add_dependency, add_dependency_v2,
         remove_instructions,
         replace_instruction_ids,
         tag_instructions,
@@ -129,6 +131,8 @@
 from loopy.schedule import (
     generate_loop_schedules, get_one_scheduled_kernel,
     get_one_linearized_kernel, linearize)
+from loopy.schedule.checker import (
+    find_unsatisfied_dependencies)
 from loopy.statistics import (ToCountMap, ToCountPolynomialMap, CountGranularity,
         stringify_stats_mapping, Op, MemAccess, get_op_map, get_mem_access_map,
         get_synchronization_map, gather_access_footprints,
@@ -194,15 +198,16 @@
 
         # {{{ transforms
 
-        "set_loop_priority", "prioritize_loops", "untag_inames",
+        "set_loop_priority", "prioritize_loops", "constrain_loop_nesting",
+        "untag_inames",
         "split_iname", "chunk_iname", "join_inames", "tag_inames",
         "duplicate_inames",
         "rename_iname", "remove_unused_inames",
         "split_reduction_inward", "split_reduction_outward",
         "affine_map_inames", "find_unused_axis_tag",
         "make_reduction_inames_unique",
         "has_schedulable_iname_nesting", "get_iname_duplication_options",
-        "add_inames_to_insn", "add_inames_for_unused_hw_axes",
+        "add_inames_to_insn", "add_inames_for_unused_hw_axes", "map_domain",
 
         "add_prefetch", "change_arg_to_image",
         "tag_array_axes", "tag_data_axes",
@@ -212,7 +217,8 @@
         "rename_argument", "set_temporary_scope",
 
         "find_instructions", "map_instructions",
-        "set_instruction_priority", "add_dependency",
+        "set_instruction_priority",
+        "add_dependency", "add_dependency_v2",
         "remove_instructions",
         "replace_instruction_ids",
         "tag_instructions",
@@ -268,7 +274,7 @@
         "generate_loop_schedules",
         "get_one_scheduled_kernel", "get_one_linearized_kernel",
         "linearize",
-
+        "find_unsatisfied_dependencies",
         "GeneratedProgram", "CodeGenerationResult",
         "PreambleInfo",
         "generate_code", "generate_code_v2", "generate_body",

loopy/kernel/__init__.py

@@ -249,6 +249,7 @@ class LoopKernel(ImmutableRecordWithoutPickling, Taggable):
     .. automethod:: tagged
     .. automethod:: without_tags
     """
+    # TODO document loop_nest_constraints attribute
 
     # {{{ constructor
 
@@ -268,6 +269,7 @@ def __init__(self, domains, instructions, args=None,
 
             iname_slab_increments=None,
             loop_priority=frozenset(),
+            loop_nest_constraints=None,
             silenced_warnings=None,
 
             applied_iname_rewrites=None,
@@ -380,6 +382,7 @@ def __init__(self, domains, instructions, args=None,
                 assumptions=assumptions,
                 iname_slab_increments=iname_slab_increments,
                 loop_priority=loop_priority,
+                loop_nest_constraints=loop_nest_constraints,
                 silenced_warnings=silenced_warnings,
                 temporary_variables=temporary_variables,
                 local_sizes=local_sizes,
@@ -1543,6 +1546,7 @@ def __setstate__(self, state):
             "substitutions",
             "iname_slab_increments",
             "loop_priority",
+            "loop_nest_constraints",
             "silenced_warnings",
             "options",
             "state",

loopy/kernel/instruction.py

@@ -75,7 +75,7 @@ class UseStreamingStoreTag(Tag):
 # {{{ instructions: base class
 
 class InstructionBase(ImmutableRecord, Taggable):
-    """A base class for all types of instruction that can occur in
+    r"""A base class for all types of instruction that can occur in
     a kernel.
 
     .. attribute:: id
@@ -87,7 +87,7 @@ class InstructionBase(ImmutableRecord, Taggable):
 
     .. attribute:: depends_on
 
-        a :class:`frozenset` of :attr:`id` values of :class:`InstructionBase`
+        A :class:`frozenset` of :attr:`id` values of :class:`InstructionBase`
         instances that *must* be executed before this one. Note that
         :func:`loopy.preprocess_kernel` (usually invoked automatically)
         augments this by adding dependencies on any writes to temporaries read
@@ -106,6 +106,15 @@ class InstructionBase(ImmutableRecord, Taggable):
           :func:`loopy.make_kernel`. Note, that this is not meant as a user-facing
           interface.
 
+    .. attribute:: dependencies
+
+        A :class:`dict` mapping :attr:`id` values of :class:`InstructionBase`
+        instances (each referring to a statement with statement instances that
+        must be executed before instances of this statement) to lists (one list
+        per key) of class:`islpy.Map`\ s mapping each instance of the dependee
+        statement to all instances of this statement that must occur later. Note
+        that this dict will eventually replace the `depends_on` attribute.
+
     .. attribute:: depends_on_is_final
 
         A :class:`bool` determining whether :attr:`depends_on` constitutes
@@ -212,6 +221,8 @@ class InstructionBase(ImmutableRecord, Taggable):
     pymbolic_set_fields = {"predicates"}
 
     def __init__(self, id, depends_on, depends_on_is_final,
+            dependencies,
+            non_linearizing_deps,
             groups, conflicts_with_groups,
             no_sync_with,
             within_inames_is_final, within_inames,
@@ -241,6 +252,12 @@ def __init__(self, id, depends_on, depends_on_is_final,
         if depends_on is None:
             depends_on = frozenset()
 
+        if dependencies is None:
+            dependencies = {}
+        # TODO dependee ids for deps that don't affect cartoon dag
+        if non_linearizing_deps is None:
+            non_linearizing_deps = set()
+
         if groups is None:
             groups = frozenset()
 
@@ -297,6 +314,8 @@ def __init__(self, id, depends_on, depends_on_is_final,
                 id=id,
                 depends_on=depends_on,
                 depends_on_is_final=depends_on_is_final,
+                dependencies=dependencies,
+                non_linearizing_deps=non_linearizing_deps,  # TODO
                 no_sync_with=no_sync_with,
                 groups=groups, conflicts_with_groups=conflicts_with_groups,
                 within_inames_is_final=within_inames_is_final,
@@ -392,6 +411,11 @@ def get_str_options(self):
 
         if self.depends_on:
             result.append("dep="+":".join(self.depends_on))
+        if self.dependencies:
+            result.append("dependencies="+":".join(self.dependencies.keys()))
+        if self.non_linearizing_deps:
+            result.append(
+                "non_linearizing_deps="+":".join(self.non_linearizing_deps))
         if self.no_sync_with:
             result.append("nosync="+":".join(
                     "%s@%s" % entry for entry in self.no_sync_with))
@@ -461,6 +485,8 @@ def __setstate__(self, val):
         if self.id is not None:  # pylint:disable=access-member-before-definition
             self.id = intern(self.id)
         self.depends_on = intern_frozenset_of_ids(self.depends_on)
+        # TODO something with dependencies?
+        # TODO something with non_linearizing_deps?
         self.groups = intern_frozenset_of_ids(self.groups)
         self.conflicts_with_groups = (
                 intern_frozenset_of_ids(self.conflicts_with_groups))
@@ -887,6 +913,8 @@ def __init__(self,
             id=None,
             depends_on=None,
             depends_on_is_final=None,
+            dependencies=None,
+            non_linearizing_deps=None,  # TODO
             groups=None,
             conflicts_with_groups=None,
             no_sync_with=None,
@@ -903,6 +931,8 @@ def __init__(self,
                 id=id,
                 depends_on=depends_on,
                 depends_on_is_final=depends_on_is_final,
+                dependencies=dependencies,
+                non_linearizing_deps=non_linearizing_deps,  # TODO
                 groups=groups,
                 conflicts_with_groups=conflicts_with_groups,
                 no_sync_with=no_sync_with,
@@ -1038,6 +1068,8 @@ def __init__(self,
             id=None,
             depends_on=None,
             depends_on_is_final=None,
+            dependencies=None,
+            non_linearizing_deps=None,  # TODO
             groups=None,
             conflicts_with_groups=None,
             no_sync_with=None,
@@ -1051,6 +1083,8 @@ def __init__(self,
                 id=id,
                 depends_on=depends_on,
                 depends_on_is_final=depends_on_is_final,
+                dependencies=dependencies,
+                non_linearizing_deps=non_linearizing_deps,  # TODO
                 groups=groups,
                 conflicts_with_groups=conflicts_with_groups,
                 no_sync_with=no_sync_with,
@@ -1329,13 +1363,21 @@ class CInstruction(InstructionBase):
 
     def __init__(self,
             iname_exprs, code,
-            read_variables=frozenset(), assignees=tuple(),
-            id=None, depends_on=None, depends_on_is_final=None,
-            groups=None, conflicts_with_groups=None,
+            read_variables=frozenset(),
+            assignees=tuple(),
+            id=None,
+            depends_on=None,
+            depends_on_is_final=None,
+            dependencies=None,
+            non_linearizing_deps=None,  # TODO
+            groups=None,
+            conflicts_with_groups=None,
             no_sync_with=None,
-            within_inames_is_final=None, within_inames=None,
+            within_inames_is_final=None,
+            within_inames=None,
             priority=0,
-            predicates=frozenset(), tags=None):
+            predicates=frozenset(),
+            tags=None):
         """
         :arg iname_exprs: Like :attr:`iname_exprs`, but instead of tuples,
             simple strings pepresenting inames are also allowed. A single
@@ -1350,11 +1392,14 @@ def __init__(self,
                 id=id,
                 depends_on=depends_on,
                 depends_on_is_final=depends_on_is_final,
+                dependencies=dependencies,
+                non_linearizing_deps=non_linearizing_deps,  # TODO
                 groups=groups, conflicts_with_groups=conflicts_with_groups,
                 no_sync_with=no_sync_with,
                 within_inames_is_final=within_inames_is_final,
                 within_inames=within_inames,
-                priority=priority, predicates=predicates, tags=tags)
+                priority=priority, predicates=predicates,
+                tags=tags)
 
         # {{{ normalize iname_exprs
 
@@ -1495,16 +1540,27 @@ class NoOpInstruction(_DataObliviousInstruction):
         ... nop
     """
 
-    def __init__(self, id=None, depends_on=None, depends_on_is_final=None,
-            groups=None, conflicts_with_groups=None,
+    def __init__(
+            self,
+            id=None,
+            depends_on=None,
+            depends_on_is_final=None,
+            dependencies=None,  # TODO
+            non_linearizing_deps=None,
+            groups=None,
+            conflicts_with_groups=None,
             no_sync_with=None,
-            within_inames_is_final=None, within_inames=None,
+            within_inames_is_final=None,
+            within_inames=None,
             priority=None,
-            predicates=None, tags=None):
+            predicates=None,
+            tags=None):
         super().__init__(
                 id=id,
                 depends_on=depends_on,
                 depends_on_is_final=depends_on_is_final,
+                dependencies=dependencies,
+                non_linearizing_deps=non_linearizing_deps,  # TODO
                 groups=groups,
                 conflicts_with_groups=conflicts_with_groups,
                 no_sync_with=no_sync_with,
@@ -1554,12 +1610,22 @@ class BarrierInstruction(_DataObliviousInstruction):
     fields = _DataObliviousInstruction.fields | {"synchronization_kind",
                                                      "mem_kind"}
 
-    def __init__(self, id, depends_on=None, depends_on_is_final=None,
-            groups=None, conflicts_with_groups=None,
+    def __init__(
+            self,
+            id,
+            depends_on=None,
+            depends_on_is_final=None,
+            dependencies=None,  # TODO
+            non_linearizing_deps=None,
+            groups=None,
+            conflicts_with_groups=None,
             no_sync_with=None,
-            within_inames_is_final=None, within_inames=None,
+            within_inames_is_final=None,
+            within_inames=None,
             priority=None,
-            predicates=None, tags=None, synchronization_kind="global",
+            predicates=None,
+            tags=None,
+            synchronization_kind="global",
             mem_kind="local"):
 
         if predicates:
@@ -1569,6 +1635,8 @@ def __init__(self, id, depends_on=None, depends_on_is_final=None,
                 id=id,
                 depends_on=depends_on,
                 depends_on_is_final=depends_on_is_final,
+                dependencies=dependencies,
+                non_linearizing_deps=non_linearizing_deps,  # TODO
                 groups=groups,
                 conflicts_with_groups=conflicts_with_groups,
                 no_sync_with=no_sync_with,

loopy/options.py

@@ -242,6 +242,7 @@ def __init__(
                 disable_global_barriers=kwargs.get("disable_global_barriers",
                     False),
                 check_dep_resolution=kwargs.get("check_dep_resolution", True),
+                use_dependencies_v2=kwargs.get("use_dependencies_v2", False),
 
                 enforce_variable_access_ordered=kwargs.get(
                     "enforce_variable_access_ordered", True),