Fix multigrid caches

Author: pbrubeck

firedrake/mg/embedded.py

@@ -5,6 +5,7 @@
 from enum import IntEnum
 from firedrake.petsc import PETSc
 from firedrake.embedding import get_embedding_dg_element
+from .utils import get_level
 
 
 __all__ = ("TransferManager", )
@@ -40,7 +41,7 @@ class Cache(object):
         """A caching object for work vectors and matrices.
 
         :arg element: The element to use for the caching."""
-        def __init__(self, ufl_element, value_shape):
+        def __init__(self, ufl_element, value_shape, boundary_set):
             self.embedding_element = get_embedding_dg_element(ufl_element, value_shape)
             self._dat_versions = {}
             self._V_DG_mass = {}
@@ -65,7 +66,6 @@ def __init__(self, *, native_transfers=None, use_averaging=True):
         """
         self.native_transfers = native_transfers or {}
         self.use_averaging = use_averaging
-        self.caches = {}
 
     def is_native(self, element, op):
         if element in self.native_transfers.keys():
@@ -87,14 +87,17 @@ def _native_transfer(self, element, op):
         return None
 
     def cache(self, V):
-        key = (V.ufl_element(), V.value_shape)
+        mh, _ = get_level(V.mesh())
+        caches = mh._shared_data_cache["transfer_manager_cache"]
+        key = (V.ufl_element(), V.value_shape, V.boundary_set)
         try:
-            return self.caches[key]
+            return caches[key]
         except KeyError:
-            return self.caches.setdefault(key, TransferManager.Cache(*key))
+            return caches.setdefault(key, TransferManager.Cache(*key))
 
     def cache_key(self, V):
-        return (V.dim(), V.boundary_set)
+        _, level = get_level(V.mesh())
+        return level
 
     def V_dof_weights(self, V):
         """Dof weights for averaging projection.

firedrake/preconditioners/pmg.py

@@ -1188,7 +1188,6 @@ class StandaloneInterpolationMatrix(object):
     """
 
     _cache_kernels = {}
-    _cache_work = {}
 
     def __init__(self, Vc, Vf, Vc_bcs, Vf_bcs):
         self.uc = self.work_function(Vc)
@@ -1212,11 +1211,12 @@ def __init__(self, Vc, Vf, Vc_bcs, Vf_bcs):
     def work_function(self, V):
         if isinstance(V, firedrake.Function):
             return V
-        key = (V.ufl_element(), V.mesh(), V.boundary_set)
+        cache = V.mesh()._shared_data_cache["pmg_work_function"]
+        key = (V.ufl_element(), V.value_shape, V.boundary_set)
         try:
-            return self._cache_work[key]
+            return cache[key]
         except KeyError:
-            return self._cache_work.setdefault(key, firedrake.Function(V))
+            return cache.setdefault(key, firedrake.Function(V))
 
     @cached_property
     def _weight(self):