Remove reference cycle in VecAccessMixin (#4033)

* Remove reference cycle in VecAccessMixin

With an associated PETSc Vec, VecAccessMixin deferred its version
property to a lambda to avoid allocating the storage until necessary.
Unfortunately, this lambda creates a reference cycle to self for all
users of the VecAccessMixin. Given that counter accesses should be
relatively infrequent, it seems fine to look up the counter type within
the method itself.

* Don't cache topological on CoordinatelessFunction

Doesn't make sense to cache a reference to self, just return self.

* Convert increment_dat_version to abstract on DataCarrier

The inheritance chain for Dat and Global puts VecAccessMixin
(rightly) behind DataCarrier. This means that by MRO, the
increment_dat_version method provided on DataCarrier will be used,
which is a null operation. I think this makes sense, given that not all
classes use this. However, because we're providing
increment_dat_version as an override through VecAccessMixin, we need to
explicitly refer to it in the inheriting classes.

* Update pyop2/types/dat.py

Otherwise updating a dat through a view doesn't increment the version.

Co-authored-by: Connor Ward <c.ward20@imperial.ac.uk>

---------

Co-authored-by: David A. Ham <david.ham@imperial.ac.uk>
Co-authored-by: Connor Ward <c.ward20@imperial.ac.uk>

Author: 3 people

firedrake/function.py

@@ -84,7 +84,7 @@ def __init__(self, function_space, val=None, name=None, dtype=ScalarType):
         else:
             self.dat = function_space.make_dat(val, dtype, self.name())
 
-    @utils.cached_property
+    @property
     def topological(self):
         r"""The underlying coordinateless function."""
         return self

firedrake/preconditioners/patch.py

@@ -127,6 +127,9 @@ def _wrapper_cache_key_(self):
     def __call__(self, access, map_=None):
         return LocalDatLegacyArg(self, map_, access)
 
+    def increment_dat_version(self):
+        pass
+
 
 register_petsc_function("MatSetValues")
 

pyop2/types/dat.py

@@ -725,6 +725,9 @@ def __init__(self, dat, index):
                                       dtype=dat.dtype,
                                       name="view[%s](%s)" % (index, dat.name))
 
+    def increment_dat_version(self):
+        self._parent.increment_dat_version()
+
     @utils.cached_property
     def _kernel_args_(self):
         return self._parent._kernel_args_
@@ -841,6 +844,9 @@ def vec_context(self, access):
         if access is not Access.READ:
             self.halo_valid = False
 
+    def increment_dat_version(self):
+        VecAccessMixin.increment_dat_version(self)
+
 
 class MixedDat(AbstractDat, VecAccessMixin):
     r"""A container for a bag of :class:`Dat`\s.

pyop2/types/data_carrier.py

@@ -44,6 +44,7 @@ def cdim(self):
         the product of the dim tuple."""
         return self._cdim
 
+    @abc.abstractmethod
     def increment_dat_version(self):
         pass
 
@@ -80,24 +81,21 @@ def _is_allocated(self):
 class VecAccessMixin(abc.ABC):
 
     def __init__(self, petsc_counter=None):
-        if petsc_counter:
-            # Use lambda since `_vec` allocates the data buffer
-            # -> Dat/Global should not allocate storage until accessed
-            self._dat_version = lambda: self._vec.stateGet()
-            self.increment_dat_version = lambda: self._vec.stateIncrease()
-        else:
-            # No associated PETSc Vec if incompatible type:
-            # -> Equip Dat/Global with their own counter.
-            self._version = 0
-            self._dat_version = lambda: self._version
-
-            def _inc():
-                self._version += 1
-            self.increment_dat_version = _inc
+        self._petsc_counter = petsc_counter
+        self._version = 0
 
     @property
     def dat_version(self):
-        return self._dat_version()
+        if self._petsc_counter:
+            return self._vec.stateGet()
+
+        return self._version
+
+    def increment_dat_version(self):
+        if self._petsc_counter:
+            self._vec.stateIncrease()
+        else:
+            self._version += 1
 
     @abc.abstractmethod
     def vec_context(self, access):

pyop2/types/glob.py

@@ -399,6 +399,9 @@ def vec_context(self, access):
             data = self._data
             self.comm.Bcast(data, 0)
 
+    def increment_dat_version(self):
+        VecAccessMixin.increment_dat_version(self)
+
 
 # has no comm, can only be READ
 class Constant(SetFreeDataCarrier):
@@ -464,3 +467,6 @@ def duplicate(self):
             dtype=self.dtype,
             name=self.name
         )
+
+    def increment_dat_version(self):
+        pass

pyop2/types/mat.py

@@ -550,6 +550,9 @@ def __repr__(self):
         return "Mat(%r, %r, %r)" \
                % (self._sparsity, self._datatype, self._name)
 
+    def increment_dat_version(self):
+        pass
+
 
 class Mat(AbstractMat):
     """OP2 matrix data. A Mat is defined on a sparsity pattern and holds a value