WIP

* About 8 tests currently failing, but they are non-essential for
  Firedrake.
* Substantially improve PetscMat implementation. Row and column axes
  must now be distinctly labelled. This avoids confusing failures.
* Add `Pack` class because sometimes `__getitem__` isn't enough (for
  example if you want a specific DoF layout).

Author: connorjward

pyop3/__init__.py

@@ -45,6 +45,7 @@
     Function,
     Loop,
     OpaqueKernelArgument,
+    Pack,
     ReplaceAssignment,
     do_loop,
     loop,

pyop3/array/petsc.py

@@ -105,6 +105,13 @@ def __init__(self, raxes, caxes, *, name=None):
         self.raxes = raxes
         self.caxes = caxes
 
+        axes = PartialAxisTree(raxes.parent_to_children)
+        for leaf_axis, leaf_cpt in raxes.leaves:
+            # do *not* uniquify, it makes indexing very complicated. Instead assert
+            # that external indices and axes must be sufficiently unique.
+            axes = axes.add_subtree(caxes, leaf_axis, leaf_cpt, uniquify_ids=True)
+        self.axes = AxisTree(axes.parent_to_children)
+
     def __getitem__(self, indices):
         # TODO also support context-free (see MultiArray.__getitem__)
         if len(indices) != 2:
@@ -151,6 +158,11 @@ def __getitem__(self, indices):
 
         arrays = {}
         for ctx, (rtree, ctree) in rcforest.items():
+            tree = rtree
+            for rleaf, clabel in rtree.leaves:
+                tree = tree.add_subtree(ctree, rleaf, clabel, uniquify_ids=True)
+            indexed_axes = _index_axes(tree, ctx, self.axes)
+
             indexed_raxes = _index_axes(rtree, ctx, self.raxes)
             indexed_caxes = _index_axes(ctree, ctx, self.caxes)
 
@@ -201,21 +213,41 @@ def __getitem__(self, indices):
             # breakpoint()
             packed = PackedPetscMat(self, rmap, cmap, shape)
 
-            indexed_axes = PartialAxisTree(indexed_raxes.parent_to_children)
-            for leaf_axis, leaf_cpt in indexed_raxes.leaves:
-                indexed_axes = indexed_axes.add_subtree(
-                    indexed_caxes, leaf_axis, leaf_cpt, uniquify=True
-                )
-            indexed_axes = indexed_axes.set_up()
+            # indexed_axes = PartialAxisTree(indexed_raxes.parent_to_children)
+            # for leaf_axis, leaf_cpt in indexed_raxes.leaves:
+            #     indexed_axes = indexed_axes.add_subtree(
+            #         indexed_caxes, leaf_axis, leaf_cpt, uniquify=True
+            #     )
+            # indexed_axes = indexed_axes.set_up()
+            # node_map = dict(indexed_raxes.parent_to_children)
+            # target_paths = dict(indexed_raxes.target_paths)
+            # index_exprs = dict(indexed_raxes.index_exprs)
+            # for leaf_axis, leaf_cpt in indexed_raxes.leaves:
+            #     for caxis in indexed_caxes.nodes:
+            #         if caxis.id not in indexed_raxes.parent_to_children:
+            #             cid = caxis.id
+            #         else:
+            #             cid = XXX
+            #
+            #         for ccpt in caxis.components:
+            #             node_map.update(...)
+            # indexed_axes = AxisTree(node_map, target_paths=???, index_exprs=???)
+            # can I make indexed_axes simply???
+            # breakpoint()
+
+            outer_loops = list(router_loops)
+            all_ids = [l.id for l in router_loops]
+            for ol in couter_loops:
+                if ol.id not in all_ids:
+                    outer_loops.append(ol)
 
             arrays[ctx] = HierarchicalArray(
                 indexed_axes,
                 data=packed,
                 target_paths=indexed_axes.target_paths,
                 index_exprs=indexed_axes.index_exprs,
                 # TODO ordered set?
-                outer_loops=router_loops
-                + tuple(filter(lambda l: l not in router_loops, couter_loops)),
+                outer_loops=outer_loops,
                 name=self.name,
             )
         return ContextSensitiveMultiArray(arrays)

pyop3/ir/lower.py

@@ -213,7 +213,7 @@ def add_argument(self, array):
             # Temporaries can have variable size, hence we allocate space for the
             # largest possible array
             # shape = (array.alloc_size,)
-            shape = self._temporary_shapes[array.name]
+            shape = self._temporary_shapes.get(array.name, (array.alloc_size,))
 
             # could rename array like the rest
             temp = lp.TemporaryVariable(array.name, dtype=array.dtype, shape=shape)
@@ -488,6 +488,11 @@ def _(expr: Assignment):
     return pmap()
 
 
+@_collect_temporary_shapes.register
+def _(expr: PetscMatInstruction):
+    return pmap()
+
+
 @_collect_temporary_shapes.register
 def _(call: CalledFunction):
     return freeze(
@@ -496,7 +501,6 @@ def _(call: CalledFunction):
             for lp_arg, arg in checked_zip(
                 call.function.code.default_entrypoint.args, call.arguments
             )
-            if lp_arg.shape is not None
         }
     )
 

pyop3/itree/tree.py

@@ -324,11 +324,14 @@ def datamap(self):
 # FIXME class hierarchy is very confusing
 class ContextFreeLoopIndex(ContextFreeIndex):
     def __init__(self, iterset: AxisTree, source_path, path, *, id=None):
-        super().__init__(id=id)
+        super().__init__(id=id, label=id, component_labels=("XXX",))
         self.iterset = iterset
         self.source_path = freeze(source_path)
         self.path = freeze(path)
 
+        # if self.label == "_label_ContextFreeLoopIndex_15":
+        #     breakpoint()
+
     def with_context(self, context, *args):
         return self
 
@@ -437,7 +440,8 @@ class Slice(ContextFreeIndex):
 
     """
 
-    fields = Index.fields | {"axis", "slices", "numbering"} - {"label"}
+    # fields = Index.fields | {"axis", "slices", "numbering"} - {"label", "component_labels"}
+    fields = {"axis", "slices", "numbering"}
 
     def __init__(self, axis, slices, *, numbering=None, id=None):
         super().__init__(label=axis, id=id)
@@ -496,9 +500,10 @@ def datamap(self):
         return pmap(data)
 
 
-class CalledMap(Identified, LoopIterable):
-    def __init__(self, map, from_index, *, id=None):
+class CalledMap(Identified, Labelled, LoopIterable):
+    def __init__(self, map, from_index, *, id=None, label=None):
         Identified.__init__(self, id=id)
+        Labelled.__init__(self, label=label)
         self.map = map
         self.from_index = from_index
 
@@ -596,7 +601,9 @@ def with_context(self, context, axes=None):
         )
         if len(leaf_target_paths) == 0:
             raise RuntimeError
-        return ContextFreeCalledMap(self.map, cf_index, leaf_target_paths, id=self.id)
+        return ContextFreeCalledMap(
+            self.map, cf_index, leaf_target_paths, id=self.id, label=self.label
+        )
 
     @property
     def name(self):
@@ -609,8 +616,12 @@ def connectivity(self):
 
 # class ContextFreeCalledMap(Index, ContextFree):
 class ContextFreeCalledMap(Index):
-    def __init__(self, map, index, leaf_target_paths, *, id=None):
-        super().__init__(id=id)
+    # FIXME this is clumsy
+    # fields = Index.fields | {"map", "index", "leaf_target_paths"} - {"label", "component_labels"}
+    fields = {"map", "index", "leaf_target_paths", "label", "id"}
+
+    def __init__(self, map, 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
@@ -1274,7 +1285,7 @@ def _make_leaf_axis_from_called_map(
             {map_cpt.target_axis: map_cpt.target_component}
         )
 
-        axisvar = AxisVariable(called_map.id)
+        axisvar = AxisVariable(called_map.label)
 
         if not isinstance(map_cpt, TabulatedMapComponent):
             raise NotImplementedError("Currently we assume only arrays here")
@@ -1297,7 +1308,7 @@ def _make_leaf_axis_from_called_map(
         # a replacement
         map_leaf_axis, map_leaf_component = map_axes.leaf
         old_inner_index_expr = map_array.index_exprs[
-            map_leaf_axis.id, map_leaf_component.label
+            map_leaf_axis.id, map_leaf_component
         ]
 
         my_index_exprs = {}
@@ -1322,7 +1333,10 @@ def _make_leaf_axis_from_called_map(
         raise RuntimeError("map does not target any relevant axes")
 
     axis = Axis(
-        components, label=called_map.id, id=axis_id, numbering=called_map.map.numbering
+        components,
+        label=called_map.label,
+        id=axis_id,
+        numbering=called_map.map.numbering,
     )
 
     return (
@@ -1358,9 +1372,18 @@ def _index_axes(
         debug=debug,
     )
 
-    # index trees should track outer loops, I think?
     outer_loops += indices.outer_loops
 
+    # drop duplicates
+    outer_loops_ = []
+    allids = set()
+    for ol in outer_loops:
+        if ol.id in allids:
+            continue
+        outer_loops_.append(ol)
+        allids.add(ol.id)
+    outer_loops = tuple(outer_loops_)
+
     # check that slices etc have not been missed
     assert not include_loop_index_shape, "old option"
     if axes is not None:

pyop3/lang.py

@@ -31,6 +31,7 @@
     checked_zip,
     just_one,
     merge_dicts,
+    single_valued,
     unique,
 )
 
@@ -80,6 +81,21 @@ def kernel_dtype(self):
         pass
 
 
+# this is an expression, like passing an array through to a kernel
+# but it is transformed first.
+class Pack(KernelArgument, ContextFree):
+    def __init__(self, big, small):
+        self.big = big
+        self.small = small
+
+    @property
+    def kernel_dtype(self):
+        try:
+            return single_valued([self.big.dtype, self.small.dtype])
+        except ValueError:
+            raise ValueError("dtypes must match")
+
+
 class Instruction(UniqueRecord, abc.ABC):
     pass
 

pyop3/transform.py

@@ -24,6 +24,7 @@
     DummyKernelArgument,
     Instruction,
     Loop,
+    Pack,
     PetscMatAdd,
     PetscMatLoad,
     PetscMatStore,
@@ -276,7 +277,6 @@ def _(self, assignment: Assignment):
                 axes = AxisTree(arg.axes.parent_to_children)
                 new_arg = HierarchicalArray(
                     axes,
-                    layouts=arg.layouts,
                     data=NullBuffer(arg.dtype),  # does this need a size?
                     name=self._name_generator("t"),
                 )
@@ -319,42 +319,56 @@ def _(self, terminal: CalledFunction):
             # this is a separate stage to the assignment operations because one
             # can index a packed mat. E.g. mat[p, q][::2] would decompose into
             # two calls, one to pack t0 <- mat[p, q] and another to pack t1 <- t0[::2]
-            if isinstance(arg.buffer, PackedBuffer):
+            if (
+                isinstance(arg, Pack)
+                and isinstance(arg.big.buffer, PackedBuffer)
+                or not isinstance(arg, Pack)
+                and isinstance(arg.buffer, PackedBuffer)
+            ):
+                if isinstance(arg, Pack):
+                    myarg = arg.big
+                else:
+                    myarg = arg
+
                 # TODO add PackedPetscMat as a subclass of buffer?
-                if not isinstance(arg.buffer.array, PetscMat):
+                if not isinstance(myarg.buffer.array, PetscMat):
                     raise NotImplementedError("Only handle Mat at the moment")
 
-                axes = AxisTree(arg.axes.parent_to_children)
+                axes = AxisTree(myarg.axes.parent_to_children)
                 new_arg = HierarchicalArray(
                     axes,
-                    data=NullBuffer(arg.dtype),  # does this need a size?
+                    data=NullBuffer(myarg.dtype),  # does this need a size?
                     name=self._name_generator("t"),
                 )
 
                 if intent == READ:
-                    gathers.append(PetscMatLoad(arg, new_arg))
+                    gathers.append(PetscMatLoad(myarg, new_arg))
                 elif intent == WRITE:
-                    scatters.insert(0, PetscMatStore(arg, new_arg))
+                    scatters.insert(0, PetscMatStore(myarg, new_arg))
                 elif intent == RW:
-                    gathers.append(PetscMatLoad(arg, new_arg))
-                    scatters.insert(0, PetscMatStore(arg, new_arg))
+                    gathers.append(PetscMatLoad(myarg, new_arg))
+                    scatters.insert(0, PetscMatStore(myarg, new_arg))
                 else:
                     assert intent == INC
                     gathers.append(ReplaceAssignment(new_arg, 0))
-                    scatters.insert(0, PetscMatAdd(arg, new_arg))
+                    scatters.insert(0, PetscMatAdd(myarg, new_arg))
 
                 # the rest of the packing code is now dealing with the result of this
                 # function call
                 arg = new_arg
 
             # unpick pack/unpack instructions
             if intent != NA and _requires_pack_unpack(arg):
-                axes = AxisTree(arg.axes.parent_to_children)
-                temporary = HierarchicalArray(
-                    axes,
-                    data=NullBuffer(arg.dtype),  # does this need a size?
-                    name=self._name_generator("t"),
-                )
+                if isinstance(arg, Pack):
+                    temporary = arg.small
+                    arg = arg.big
+                else:
+                    axes = AxisTree(arg.axes.parent_to_children)
+                    temporary = HierarchicalArray(
+                        axes,
+                        data=NullBuffer(arg.dtype),  # does this need a size?
+                        name=self._name_generator("t"),
+                    )
 
                 if intent == READ:
                     gathers.append(ReplaceAssignment(temporary, arg))
@@ -426,7 +440,7 @@ def _requires_pack_unpack(arg):
     # however, it is overly restrictive since we could pass something like dat[i0, :] directly
     # to a local kernel
     # return isinstance(arg, HierarchicalArray) and arg.subst_layouts != arg.layouts
-    return isinstance(arg, HierarchicalArray)
+    return isinstance(arg, HierarchicalArray) or isinstance(arg, Pack)
 
 
 # *below is old untested code*