Fix CalledMap.index() method

Allows for slope limiter code in thesis.

Author: connorjward

pyop3/axtree/tree.py

@@ -766,9 +766,12 @@ def datamap(self):
             for expr in index_exprs.values():
                 for array in MultiArrayCollector()(expr):
                     dmap.update(array.datamap)
-        for layout_expr in self.layouts.values():
-            for array in MultiArrayCollector()(layout_expr):
-                dmap.update(array.datamap)
+
+        # TODO: cleanup, indexed axis trees (from map.index()) do not have layouts
+        if not isinstance(self, IndexedAxisTree) or self.unindexed is not None:
+            for layout_expr in self.layouts.values():
+                for array in MultiArrayCollector()(layout_expr):
+                    dmap.update(array.datamap)
         return pmap(dmap)
 
     def as_tree(self):

pyop3/ir/lower.py

@@ -179,6 +179,7 @@ def add_assignment(self, assignee, expression, prefix="insn"):
             id=self._name_generator(prefix),
             within_inames=frozenset(self._within_inames),
             depends_on=self._depends_on,
+            depends_on_is_final=True,
         )
         self._add_instruction(insn)
 

pyop3/itree/tree.py

@@ -702,23 +702,23 @@ def connectivity(self):
 class ContextFreeCalledMap(Index):
     # FIXME this is clumsy
     # fields = Index.fields | {"map", "index", "leaf_target_paths"} - {"label", "component_labels"}
-    fields = {"map", "index", "leaf_target_paths", "label", "id"}
+    fields = {"map", "from_index", "leaf_target_paths", "label", "id"}
 
-    def __init__(self, map, index, leaf_target_paths, *, id=None, label=None):
+    def __init__(self, map, from_index, leaf_target_paths, *, id=None, label=None):
         super().__init__(id=id, label=label)
         self.map = map
         # better to call it "input_index"?
-        self.index = index
+        # self.index = index  # clash with method!
         self._leaf_target_paths = leaf_target_paths
 
         # alias for compat with ContextFreeCalledMap
-        self.from_index = index
+        self.from_index = from_index
 
     # TODO cleanup
     def with_context(self, context, axes=None):
         # maybe this line isn't needed?
         # cf_index = self.from_index.with_context(context, axes)
-        cf_index = self.index
+        cf_index = self.from_index
         leaf_target_paths = tuple(
             freeze({mcpt.target_axis: mcpt.target_component})
             for path in cf_index.leaf_target_paths
@@ -737,6 +737,23 @@ def with_context(self, context, axes=None):
     def name(self) -> str:
         return self.map.name
 
+    def index(self) -> LoopIndex | ContextFreeLoopIndex:
+        index_forest = as_index_forest(self)
+        assert index_forest.keys() == {pmap()}
+        index_tree = index_forest[pmap()]
+        iterset = index_axes(index_tree, pmap())
+
+        # The loop index from a context-free map can be context-sensitive if it
+        # has multiple components.
+        if len(iterset.leaves) == 1:
+            path = iterset.path(*iterset.leaf)
+            target_path = {}
+            for ax, cpt in iterset.path_with_nodes(*iterset.leaf).items():
+                target_path.update(iterset.target_paths.get((ax.id, cpt), {}))
+            return ContextFreeLoopIndex(iterset, path, target_path)
+        else:
+            return LoopIndex(iterset)
+
     # is this ever used?
     # @property
     # def components(self):
@@ -1450,7 +1467,7 @@ def _(
         outer_loops,
         prior_extra_index_exprs,
     ) = collect_shape_index_callback(
-        called_map.index,
+        called_map.from_index,
         indices,
         prev_axes=prev_axes,
         **kwargs,
@@ -1553,7 +1570,7 @@ def _make_leaf_axis_from_called_map(
 
         all_skipped = False
         if isinstance(map_cpt.arity, HierarchicalArray):
-            arity = map_cpt.arity[called_map.index]
+            arity = map_cpt.arity[called_map.from_index]
         else:
             arity = map_cpt.arity
         cpt = AxisComponent(arity, label=map_cpt.label)
@@ -1658,7 +1675,8 @@ def index_axes(
         prev_axes=axes,
     )
 
-    outer_loops += axes.outer_loops
+    if axes is not None:
+        outer_loops += axes.outer_loops
 
     # drop duplicates
     outer_loops_ = []
@@ -1686,7 +1704,7 @@ def index_axes(
 
     return IndexedAxisTree(
         indexed_axes.node_map,
-        axes.unindexed,
+        axes.unindexed if axes else None,
         target_paths=mytpaths,
         index_exprs=myindex_expr_per_target,
         # layout_exprs=mylayout_expr_per_target,
@@ -2260,47 +2278,52 @@ def partition_iterset(index: LoopIndex, arrays):
     # for p in index.iterset.iter():
     #     # hack because I wrote bad code and mix up loop indices and itersets
     #     p = dataclasses.replace(p, index=index)
-    for p in index.iter():
-        parindex = p.source_exprs[paraxis.label]
-        assert isinstance(parindex, numbers.Integral)
-
-        for array in arrays:
-            # same nasty hack
-            if isinstance(array, (Mat, Sparsity)) or not hasattr(array, "buffer"):
-                continue
-            # skip purely local arrays
-            if not array.buffer.is_distributed:
-                continue
-            if labels[parindex] == IterationPointType.LEAF:
-                continue
-
-            # loop over stencil
-            array = array.with_context({index.id: (p.source_path, p.target_path)})
-
-            for q in array.axes.iter({p}):
-                # offset = array.axes.offset(q.target_exprs, q.target_path)
-                offset = array.axes.offset(q.source_exprs, q.source_path, loop_exprs=p.replace_map)
-
-                point_label = is_root_or_leaf_per_array[array.name][offset]
-                if point_label == ArrayPointLabel.LEAF:
-                    labels[parindex] = IterationPointType.LEAF
-                    break  # no point doing more analysis
-                elif point_label == ArrayPointLabel.ROOT:
-                    assert labels[parindex] != IterationPointType.LEAF
-                    labels[parindex] = IterationPointType.ROOT
-                else:
-                    assert point_label == ArrayPointLabel.CORE
-                    pass
+    # for p in index.iter():
+    #     parindex = p.source_exprs[paraxis.label]
+    #     assert isinstance(parindex, numbers.Integral)
+    #
+    #     for array in arrays:
+    #         # same nasty hack
+    #         if isinstance(array, (Mat, Sparsity)) or not hasattr(array, "buffer"):
+    #             continue
+    #         # skip purely local arrays
+    #         if not array.buffer.is_distributed:
+    #             continue
+    #         if labels[parindex] == IterationPointType.LEAF:
+    #             continue
+    #
+    #         # loop over stencil
+    #         array = array.with_context({index.id: (p.source_path, p.target_path)})
+    #
+    #         for q in array.axes.iter({p}):
+    #             # offset = array.axes.offset(q.target_exprs, q.target_path)
+    #             offset = array.axes.offset(q.source_exprs, q.source_path, loop_exprs=p.replace_map)
+    #
+    #             point_label = is_root_or_leaf_per_array[array.name][offset]
+    #             if point_label == ArrayPointLabel.LEAF:
+    #                 labels[parindex] = IterationPointType.LEAF
+    #                 break  # no point doing more analysis
+    #             elif point_label == ArrayPointLabel.ROOT:
+    #                 assert labels[parindex] != IterationPointType.LEAF
+    #                 labels[parindex] = IterationPointType.ROOT
+    #             else:
+    #                 assert point_label == ArrayPointLabel.CORE
+    #                 pass
 
     parcpt = just_one(paraxis.components)  # for now
 
     # I don't think this is working - instead everything touches a leaf
     # core = just_one(np.nonzero(labels == IterationPointType.CORE))
     # root = just_one(np.nonzero(labels == IterationPointType.ROOT))
     # leaf = just_one(np.nonzero(labels == IterationPointType.LEAF))
-    core = np.asarray([], dtype=IntType)
-    root = np.asarray([], dtype=IntType)
-    leaf = np.arange(paraxis.size, dtype=IntType)
+    # core = np.asarray([], dtype=IntType)
+    # root = np.asarray([], dtype=IntType)
+    # leaf = np.arange(paraxis.size, dtype=IntType)
+
+    # hack to check things
+    core = np.asarray([0], dtype=IntType)
+    root = np.asarray([1], dtype=IntType)
+    leaf = np.arange(2, paraxis.size, dtype=IntType)
 
     subsets = []
     for data in [core, root, leaf]:
@@ -2313,6 +2336,7 @@ def partition_iterset(index: LoopIndex, arrays):
         )
         subsets.append(subset)
     subsets = tuple(subsets)
+    return "not used", subsets
 
     # make a new iteration set over just these indices
     # index with just core (arbitrary)