Fix MG 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", )
@@ -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):
+        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[:2]))
+            return caches.setdefault(key, TransferManager.Cache(*key[:2]))
 
     def cache_key(self, V):
-        return (V.dim(),)
+        _, level = get_level(V.mesh())
+        return (level,)
 
     def V_dof_weights(self, V):
         """Dof weights for averaging projection.
@@ -143,7 +146,7 @@ def DG_inv_mass(self, DG):
         :returns: A PETSc Mat.
         """
         cache = self.cache(DG)
-        key = DG.dim()
+        key = self.cache_key(DG)
         try:
             return cache._DG_inv_mass[key]
         except KeyError:

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)
@@ -1209,14 +1208,16 @@ def __init__(self, Vc, Vf, Vc_bcs, Vf_bcs):
             self.Vc_bcs = [bc.reconstruct(V=self.Vc, g=0) for bc in self.Vc_bcs]
             self.Vf_bcs = [bc.reconstruct(V=self.Vf, g=0) for bc in self.Vf_bcs]
 
-    def work_function(self, V):
+    @staticmethod
+    def work_function(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):

tests/firedrake/multigrid/test_p_multigrid.py

@@ -73,6 +73,17 @@ def test_reconstruct_degree(tp_mesh, mixed_family):
         assert e == PMGPC.reconstruct_degree(elist[0], degree)
 
 
+def test_work_function_cache(tp_mesh):
+    from firedrake.preconditioners.pmg import StandaloneInterpolationMatrix
+
+    V1 = FunctionSpace(tp_mesh, "Lagrange", 1)
+    V2 = FunctionSpace(tp_mesh, "Lagrange", 1)
+    assert V1 is not V2
+    w1 = StandaloneInterpolationMatrix.work_function(V1)
+    w2 = StandaloneInterpolationMatrix.work_function(V2)
+    assert w1 is w2
+
+
 @pytest.mark.parametrize("family", ["Q", "NCE", "NCF", "DQ"])
 def test_prolong_basic(tp_mesh, family):
     """ Interpolate a constant function between low-order and high-order spaces

tests/firedrake/multigrid/test_transfer_manager.py

@@ -131,3 +131,73 @@ def test_transfer_manager_dat_version_cache(action, transfer_op, spaces):
 
     else:
         raise ValueError(f"Unrecognized action {action}")
+
+
+@pytest.fixture
+def DG_spaces(hierarchy):
+    return tuple(VectorFunctionSpace(mesh, "DG", 1) for mesh in hierarchy)
+
+
+@pytest.fixture
+def RT_spaces(hierarchy):
+    return tuple(RestrictedFunctionSpace(FunctionSpace(mesh, "RT", 1), ["on_boundary"]) for mesh in hierarchy)
+
+
+@pytest.mark.parametrize("action", [
+    "DG_work_function",
+    "DG_work_cofunction",
+    "V_dof_weights",
+    "work_vec",
+    "V_DG_mass",
+    "DG_inv_mass",
+    "V_approx_inv_mass",
+    "V_inv_mass_ksp",
+])
+def test_transfer_manager_cache(action, DG_spaces, RT_spaces):
+    V1 = RT_spaces[0]
+    DG1 = DG_spaces[0]
+
+    V2 = V1.reconstruct(name="V2")
+    DG2 = DG1.reconstruct(name="DG2")
+    assert V2 is not V1
+    assert DG2 is not DG1
+
+    if action == "DG_work_function":
+        w1 = transfer.DG_work(DG1)
+        w2 = transfer.DG_work(DG2)
+        assert w1 is w2
+
+    elif action == "DG_work_cofunction":
+        w1 = transfer.DG_work(DG1.dual())
+        w2 = transfer.DG_work(DG2.dual())
+        assert w1 is w2
+
+    elif action == "V_dof_weights":
+        w1 = transfer.V_dof_weights(V1)
+        w2 = transfer.V_dof_weights(V2)
+        assert w1 is w2
+
+    elif action == "work_vec":
+        w1 = transfer.work_vec(V1)
+        w2 = transfer.work_vec(V2)
+        assert w1 is w2
+
+    elif action == "V_DG_mass":
+        M1 = transfer.V_DG_mass(V1, DG1)
+        M2 = transfer.V_DG_mass(V2, DG2)
+        assert M1 is M2
+
+    elif action == "DG_inv_mass":
+        M1 = transfer.DG_inv_mass(DG1)
+        M2 = transfer.DG_inv_mass(DG2)
+        assert M1 is M2
+
+    elif action == "V_approx_inv_mass":
+        M1 = transfer.V_approx_inv_mass(V1, DG1)
+        M2 = transfer.V_approx_inv_mass(V2, DG2)
+        assert M1 is M2
+
+    elif action == "V_inv_mass_ksp":
+        ksp1 = transfer.V_inv_mass_ksp(V1)
+        ksp2 = transfer.V_inv_mass_ksp(V2)
+        assert ksp1 is ksp2