Expunge ContextSensitiveMultiArray

Firedrake tests appear to be passing.

Author: connorjward

pyop3/array/base.py

@@ -1,10 +1,11 @@
 import abc
 
+from pyop3.axtree import ContextAware
 from pyop3.lang import FunctionArgument, ReplaceAssignment
 from pyop3.utils import UniqueNameGenerator
 
 
-class Array(FunctionArgument, abc.ABC):
+class Array(ContextAware, FunctionArgument, abc.ABC):
     _prefix = "array"
     _name_generator = UniqueNameGenerator()
 
@@ -16,3 +17,12 @@ def __init__(self, name=None, *, prefix=None) -> None:
     def assign(self, other, eager=True):
         expr = ReplaceAssignment(self, other)
         return expr() if eager else expr
+
+    @abc.abstractmethod
+    def with_context(self):
+        pass
+
+    @property
+    @abc.abstractmethod
+    def context_free(self):
+        pass

pyop3/array/harray.py

@@ -94,7 +94,7 @@ class FancyIndexWriteException(Exception):
     pass
 
 
-class HierarchicalArray(Array, ContextFree, KernelArgument):
+class HierarchicalArray(Array, KernelArgument):
     """Multi-dimensional, hierarchical array.
 
     Parameters
@@ -138,7 +138,7 @@ def __init__(
                 # always deal with flattened data
                 if len(data.shape) > 1:
                     data = data.flatten()
-                if data.size != axes.unindexed.global_size:
+                if data.size != axes.alloc_size:
                     raise ValueError("Data shape does not match axes")
 
             # IndexedAxisTrees do not currently have SFs, so create a dummy one here
@@ -147,10 +147,10 @@ def __init__(
             else:
                 assert isinstance(axes, (ContextSensitiveAxisTree, IndexedAxisTree))
                 # not sure this is the right thing to do
-                sf = serial_forest(axes.unindexed.global_size)
+                sf = serial_forest(axes.alloc_size)
 
             data = DistributedBuffer(
-                axes.unindexed.global_size,  # not a useful property anymore
+                axes.alloc_size,  # not a useful property anymore
                 sf,
                 dtype,
                 name=self.name,
@@ -210,6 +210,25 @@ def getitem(self, indices, *, strict=False):
     # to be iterable (which it's not). This avoids some confusing behaviour.
     __iter__ = None
 
+    def with_context(self, context):
+        return type(self)(
+            self.axes.with_context(context),
+            name=self.name,
+            data=self.buffer,
+            max_value=self.max_value,
+            constant=self.constant,
+        )
+
+    @property
+    def context_free(self, context):
+        return type(self)(
+            self.axes.context_free,
+            name=self.name,
+            data=self.buffer,
+            max_value=self.max_value,
+            constant=self.constant,
+        )
+
     @property
     def dtype(self):
         return self.buffer.dtype

pyop3/array/petsc.py

@@ -64,7 +64,8 @@ class PetscVecNest(PetscVec):
     ...
 
 
-class AbstractMat(Array, ContextFree):
+# BaseMat?
+class AbstractMat(Array):
     DEFAULT_MAT_TYPE = PETSc.Mat.Type.AIJ
 
     prefix = "mat"
@@ -216,6 +217,29 @@ def getitem(self, indices, *, strict=False):
         # self._cache[cache_key] = mat
         return mat
 
+    def with_context(self, context):
+        row_axes = self.raxes.with_context(context)
+        col_axes = self.caxes.with_context(context)
+        return type(self)(
+            row_axes,
+            col_axes,
+            name=self.name,
+            mat_type=self.mat_type,
+            mat=self.mat,
+        )
+
+    @property
+    def context_free(self):
+        row_axes = self.raxes.context_free
+        col_axes = self.caxes.context_free
+        return type(self)(
+            row_axes,
+            col_axes,
+            name=self.name,
+            mat_type=self.mat_type,
+            mat=self.mat,
+        )
+
     # like Dat, bad name? handle?
     @property
     def array(self):

pyop3/axtree/__init__.py

@@ -1,5 +1,6 @@
 from .tree import (
     Axis,
+    ContextMismatchException,
     AxisComponent,
     AxisTree,
     AxisVariable,

pyop3/axtree/tree.py

@@ -25,6 +25,7 @@
 from pyrsistent import freeze, pmap, thaw
 
 from pyop3.axtree.parallel import partition_ghost_points
+from pyop3.exceptions import Pyop3Exception
 from pyop3.dtypes import IntType
 from pyop3.sf import StarForest, serial_forest
 from pyop3.tree import (
@@ -51,10 +52,14 @@
 )
 
 
-class ExpectedLinearAxisTreeException(Exception):
+class ExpectedLinearAxisTreeException(Pyop3Exception):
     ...
 
 
+class ContextMismatchException(Pyop3Exception):
+    pass
+
+
 class ContextAware(abc.ABC):
     @abc.abstractmethod
     def with_context(self, context):
@@ -97,7 +102,10 @@ def keys(self):
             return frozenset(indices)
 
     def with_context(self, context):
-        return self.context_map[self.filter_context(context)]
+        try:
+            return self.context_map[self.filter_context(context)]
+        except KeyError:
+            raise ContextMismatchException
 
     def filter_context(self, context):
         key = {}
@@ -1144,18 +1152,22 @@ def __init__(
         layout_exprs,
         outer_loops,
     ):
+        if layout_exprs is None:
+            layout_exprs = pmap()
         if outer_loops is None:
             outer_loops = ()
-        else:
-            assert isinstance(outer_loops, tuple)
 
         super().__init__(node_map)
         self._unindexed = unindexed
-        self._target_paths = target_paths
-        self._index_exprs = index_exprs
-        self._layout_exprs = layout_exprs
+        self._target_paths = pmap(target_paths)
+        self._index_exprs = pmap(index_exprs)
+        self._layout_exprs = pmap(layout_exprs)
         self._outer_loops = tuple(outer_loops)
 
+    # @cached_property
+    # def _hash_key(self):
+    #     return super()._hash_key + (self.unindexed, self.target_paths, self.index_exprs, self.layout_exprs, self.outer_loops)
+
     @property
     def unindexed(self):
         return self._unindexed
@@ -1342,9 +1354,15 @@ def datamap(self):
     def sf(self):
         return single_valued([ax.sf for ax in self.context_map.values()])
 
+    # @cached_property
+    # def unindexed(self):
+        # this does not work because unindexed may have different IDs, so just return
+        # the first one.
+        # return single_valued([ax.unindexed for ax in self.context_map.values()])
+
     @cached_property
-    def unindexed(self):
-        return single_valued([ax.unindexed for ax in self.context_map.values()])
+    def context_free(self):
+        return just_one(self.context_map.values())
 
 
 @functools.singledispatch

pyop3/transform.py

@@ -9,7 +9,7 @@
 
 from pyop3.array import HierarchicalArray
 from pyop3.array.petsc import AbstractMat
-from pyop3.axtree import Axis, AxisTree, ContextFree, ContextSensitive
+from pyop3.axtree import Axis, AxisTree, ContextFree, ContextSensitive, ContextMismatchException, ContextAware
 from pyop3.buffer import DistributedBuffer, NullBuffer, PackedBuffer
 from pyop3.itree import Map, TabulatedMapComponent
 from pyop3.lang import (
@@ -162,7 +162,7 @@ def _(self, terminal: CalledFunction, *, context):
                 outer_context.update(
                     {
                         index: paths
-                        for ctx in arg.context_map.keys()
+                        for ctx in arg.axes.context_map.keys()
                         for index, paths in ctx.items()
                         if index not in context
                     }
@@ -179,38 +179,29 @@ def _(self, terminal: Assignment, *, context):
         # FIXME for now we assume an outer context of {}. In other words anything
         # context sensitive in the assignment is completely handled by the existing
         # outer loops.
+        # This is meaningful if the kernel accepts a loop index as an argument.
 
-        valid = True
         cf_args = []
         for arg in terminal.arguments:
-            try:
-                cf_arg = (
-                    arg.with_context(context)
-                    if isinstance(arg.axes, ContextSensitive)
-                    else arg
-                )
-            # FIXME We will hit issues here when we are missing outer context I think
-            except KeyError:
-                # assignment is not valid in this context, do nothing
-                valid = False
-                break
-            cf_args.append(cf_arg)
-
-        if valid:
-            return ((pmap(), terminal.with_arguments(cf_args)),)
-        else:
-            return ((pmap(), None),)
+            if isinstance(arg, ContextAware):
+                try:
+                    cf_args.append(arg.with_context(context))
+                except ContextMismatchException:
+                    # assignment is not valid in this context, do nothing
+                    return ((pmap(), None),)
+            else:
+                cf_args.append(arg)
+        return ((pmap(), terminal.with_arguments(cf_args)),)
 
     # TODO: this is just an assignment, fix inheritance
     @_apply.register
     def _(self, terminal: PetscMatInstruction, *, context):
-        if any(
-            isinstance(a.axes, ContextSensitive)
-            for a in {terminal.mat_arg, terminal.array_arg}
-        ):
-            raise NotImplementedError
-
-        return ((pmap(), terminal),)
+        try:
+            mat = terminal.mat_arg.with_context(context)
+            array = terminal.array_arg.with_context(context)
+            return ((pmap(), terminal.copy(mat_arg=mat, array_arg=array)),)
+        except ContextMismatchException:
+            return ((pmap(), None),)
 
 
 def expand_loop_contexts(expr: Instruction):
@@ -320,9 +311,9 @@ def _(self, terminal: CalledFunction):
         for (arg, intent), shape in checked_zip(
             terminal.function_arguments, terminal.argument_shapes
         ):
-            assert isinstance(
-                arg, ContextFree
-            ), "Loop contexts should already be expanded"
+            # assert isinstance(
+            #     arg, ContextFree
+            # ), "Loop contexts should already be expanded"
 
             if isinstance(arg, DummyKernelArgument):
                 arguments.append(arg)

pyop3/tree.py

@@ -251,6 +251,19 @@ def __init__(self, node_map=None):
         # post-init checks
         self._check_node_labels_unique_in_paths(self.node_map)
 
+    # This is arguably over-specific. Otherwise equivalent trees will currently
+    # fail this check if their IDs do not match. One way to resolve this would be
+    # to re-ID all of the nodes with a pre-order traversal. This is not a priority.
+    # def __eq__(self, other):
+    #     return type(other) is type(self) and other._hash_key == self._hash_key
+    #
+    # def __hash__(self):
+    #     return hash(self._hash_key)
+    #
+    # @cached_property
+    # def _hash_key(self):
+    #     return (self.node_map,)
+
     @classmethod
     def _check_node_labels_unique_in_paths(
         cls, node_map, node=None, seen_labels=frozenset()