Fix unsafe id()-keyed matrix cache and stale recompute

src/dolfinx_adjoint/blocks/interpolation.py

@@ -1,57 +1,89 @@
 from __future__ import annotations
 
 import typing
+import weakref
 from typing import Callable
 
 import dolfinx
+import dolfinx.fem.petsc
 from pyadjoint import Block
-from pyadjoint.overloaded_type import create_overloaded_object
 
 if typing.TYPE_CHECKING:
     from petsc4py import PETSc
 
-# Global cache to prevent redundant matrix assembly across multiple blocks/iterations
-_INTERPOLATION_MATRIX_CACHE: dict[tuple[int, int], dolfinx.la.MatrixCSR] = {}
+# Cache of assembled interpolation matrices, keyed by (id(space_from), id(space_to),
+# use_petsc), shared across all InterpolationBlocks to avoid redundant matrix assembly
+# (and, for PETSc matrices, MPI communicator exhaustion) when the same pair of spaces is
+# interpolated between many times (e.g. across optimization iterations).
+#
+# Keying on id() would normally risk a stale/wrong entry once a FunctionSpace is garbage
+# collected and its id is reused by an unrelated object (ids are only unique among
+# simultaneously-alive objects). We close that hole with weakref.finalize: each space
+# that contributes to a cache key gets a finalizer that purges every entry mentioning
+# its id. Finalizers run at the point the object is actually deallocated, which is
+# necessarily before CPython can hand that id out again, so a cache hit always
+# corresponds to spaces that are still alive.
+_INTERPOLATION_MATRIX_CACHE: dict[tuple[int, int, bool], "_MatrixCSRWorkspace | PETSc.Mat"] = {}
+_CACHE_KEYS_BY_SPACE_ID: dict[int, set[tuple[int, int, bool]]] = {}
 
 
-def attach_working_array(mat: dolfinx.la.MatrixCSR):
-    """Attach working arrays to a dolfinx.la.MatrixCSR for efficient matrix-vector multiplication."""
-    if not hasattr(mat, "_row_vec"):
-        mat._row_vec = dolfinx.la.vector(mat.index_map(0), mat.block_size[0], dtype=mat.data.dtype)
-    if not hasattr(mat, "_col_vec"):
-        mat._col_vec = dolfinx.la.vector(mat.index_map(1), mat.block_size[1], dtype=mat.data.dtype)
-    mat._row_vec.array[:] = 0.0
-    mat._col_vec.array[:] = 0.0
+def _purge_cache_entries_for(space_id: int) -> None:
+    for key in _CACHE_KEYS_BY_SPACE_ID.pop(space_id, ()):
+        _INTERPOLATION_MATRIX_CACHE.pop(key, None)
+
+
+def _register_cache_key(space: dolfinx.fem.FunctionSpace, key: tuple[int, int, bool]) -> None:
+    space_id = id(space)
+    if space_id not in _CACHE_KEYS_BY_SPACE_ID:
+        weakref.finalize(space, _purge_cache_entries_for, space_id)
+    _CACHE_KEYS_BY_SPACE_ID.setdefault(space_id, set()).add(key)
+
+
+class _MatrixCSRWorkspace:
+    """A dolfinx.la.MatrixCSR paired with pre-allocated working vectors.
+
+    dolfinx.la.MatrixCSR.mult requires vectors built from its own index maps
+    as scratch space. Wrapping them here (built once, alongside the matrix)
+    means they can be reused across repeated matrix-vector multiplications
+    without allocating on every adjoint/TLM/Hessian evaluation, and without
+    monkey-patching dynamic attributes onto the matrix object itself.
+    """
+
+    def __init__(self, mat: dolfinx.la.MatrixCSR):
+        self.mat = mat
+        self.row_vec = dolfinx.la.vector(mat.index_map(0), mat.block_size[0], dtype=mat.data.dtype)
+        self.col_vec = dolfinx.la.vector(mat.index_map(1), mat.block_size[1], dtype=mat.data.dtype)
 
 
 def get_mult(
-    mat: "PETSc.Mat" | dolfinx.la.MatrixCSR,
-    transpose: bool = False,  # type: ignore
+    mat: "PETSc.Mat" | _MatrixCSRWorkspace,
+    transpose: bool = False,
 ) -> Callable[[dolfinx.la.Vector, dolfinx.la.Vector], None]:
     """Return a function that performs matrix-vector multiplication with the given matrix."""
-    if isinstance(mat, dolfinx.la.MatrixCSR):
+    if isinstance(mat, _MatrixCSRWorkspace):
+        workspace = mat
 
         def mult(v_in: dolfinx.la.Vector, v_out: dolfinx.la.Vector):
-            # Need to use vectors from
             in_size_local = v_in.index_map.size_local * v_in.block_size
             out_size_local = v_out.index_map.size_local * v_out.block_size
-            attach_working_array(mat)  # Ensure working arrays are attached
+            # Zero the full working arrays (including ghosts) before each use
+            # to prevent double-counting in parallel.
+            workspace.row_vec.array[:] = 0.0
+            workspace.col_vec.array[:] = 0.0
             if transpose:
-                # Prevent double-counting in parallel by zeroing ghosts of input vector
                 # Calculate the exact number of local degrees of freedom
-                mat._row_vec.array[:in_size_local] = v_in.array[:in_size_local]
-                mat._row_vec.scatter_forward()  # Ensure ghost values are updated before multiplication
-                mat._col_vec.array[:out_size_local] = 0.0
-                mat.mult(mat._row_vec, mat._col_vec, transpose=True)
-                v_out.array[:out_size_local] = mat._col_vec.array[:out_size_local]
+                workspace.row_vec.array[:in_size_local] = v_in.array[:in_size_local]
+                workspace.row_vec.scatter_forward()  # Ensure ghost values are updated before multiplication
+                workspace.col_vec.array[:out_size_local] = 0.0
+                workspace.mat.mult(workspace.row_vec, workspace.col_vec, transpose=True)
+                v_out.array[:out_size_local] = workspace.col_vec.array[:out_size_local]
             else:
-                # Prevent double-counting in parallel by zeroing ghosts of input vector
                 # Calculate the exact number of local degrees of freedom
-                mat._row_vec.array[:out_size_local] = 0
-                mat._col_vec.array[:in_size_local] = v_in.array[:in_size_local]
-                mat._col_vec.scatter_forward()  # Ensure ghost values are updated before multiplication
-                mat.mult(mat._col_vec, mat._row_vec)
-                v_out.array[:out_size_local] = mat._row_vec.array[:out_size_local]
+                workspace.row_vec.array[:out_size_local] = 0
+                workspace.col_vec.array[:in_size_local] = v_in.array[:in_size_local]
+                workspace.col_vec.scatter_forward()  # Ensure ghost values are updated before multiplication
+                workspace.mat.mult(workspace.col_vec, workspace.row_vec)
+                v_out.array[:out_size_local] = workspace.row_vec.array[:out_size_local]
             v_out.scatter_forward()
 
         return mult
@@ -74,24 +106,29 @@ def mult(v_in: dolfinx.la.Vector, v_out: dolfinx.la.Vector):
         raise TypeError("Matrix type not supported. Expected dolfinx.la.MatrixCSR or PETSc.Mat, got {type(mat)=}.")
 
 
-def _get_interpolation_matrix(
+def _build_interpolation_matrix(
     space_from: dolfinx.fem.FunctionSpace, space_to: dolfinx.fem.FunctionSpace, use_petsc: bool = False
-) -> dolfinx.la.MatrixCSR | "PETSc.Mat":
-    """Retrieve or compute the interpolation matrix for a pair of spaces."""
-    key = (id(space_from), id(space_to))
+) -> _MatrixCSRWorkspace | "PETSc.Mat":
+    """Assemble the interpolation matrix for a pair of spaces."""
+    if use_petsc:
+        petsc_mat = dolfinx.fem.petsc.interpolation_matrix(space_from, space_to)
+        petsc_mat.assemble()
+        return petsc_mat
 
-    if key not in _INTERPOLATION_MATRIX_CACHE:
-        if use_petsc:
-            mat = dolfinx.fem.petsc.interpolation_matrix(space_from, space_to)
-            mat.assemble()
-        else:
-            mat = dolfinx.fem.interpolation_matrix(space_from, space_to)
-            mat.scatter_reverse()
-            # The built in interpolation matrix requires two working arrays
-            attach_working_array(mat)
+    mat = dolfinx.fem.interpolation_matrix(space_from, space_to)
+    mat.scatter_reverse()
+    return _MatrixCSRWorkspace(mat)
 
-        _INTERPOLATION_MATRIX_CACHE[key] = mat
 
+def _get_interpolation_matrix(
+    space_from: dolfinx.fem.FunctionSpace, space_to: dolfinx.fem.FunctionSpace, use_petsc: bool = False
+) -> _MatrixCSRWorkspace | "PETSc.Mat":
+    """Retrieve or assemble the interpolation matrix for a pair of spaces, cached for reuse."""
+    key = (id(space_from), id(space_to), use_petsc)
+    if key not in _INTERPOLATION_MATRIX_CACHE:
+        _INTERPOLATION_MATRIX_CACHE[key] = _build_interpolation_matrix(space_from, space_to, use_petsc=use_petsc)
+        _register_cache_key(space_from, key)
+        _register_cache_key(space_to, key)
     return _INTERPOLATION_MATRIX_CACHE[key]
 
 
@@ -116,15 +153,17 @@ def __init__(
         self._adj_output: dolfinx.fem.Function | None = None
         self._tlm_output: dolfinx.fem.Function | None = None
         self._hessian_output: dolfinx.fem.Function | None = None
-        self._recompute_output: dolfinx.fem.Function | None = None
 
     def __str__(self):
         return "interpolate_function"
 
+    def _get_interpolation_matrix(self) -> _MatrixCSRWorkspace | "PETSc.Mat":
+        return _get_interpolation_matrix(self.space_from, self.space_to, use_petsc=self._use_petsc)
+
     # --- Adjoint ---
 
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
-        return _get_interpolation_matrix(self.space_from, self.space_to, use_petsc=self._use_petsc)
+        return self._get_interpolation_matrix()
 
     def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None):
         adj_input = adj_inputs[0]
@@ -143,7 +182,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
     # --- Tangent Linear Model (TLM) ---
 
     def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs):
-        return _get_interpolation_matrix(self.space_from, self.space_to, use_petsc=self._use_petsc)
+        return self._get_interpolation_matrix()
 
     def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None):
         tlm_input = tlm_inputs[0]
@@ -165,7 +204,7 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
     # --- Hessian ---
 
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_dependencies):
-        return _get_interpolation_matrix(self.space_from, self.space_to, use_petsc=self._use_petsc)
+        return self._get_interpolation_matrix()
 
     def evaluate_hessian_component(
         self, inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None
@@ -191,12 +230,17 @@ def prepare_recompute_component(self, inputs, relevant_outputs):
     def recompute_component(self, inputs, block_variable, idx, prepared):
         func_from = inputs[0]
 
-        if self._recompute_output is None:
-            self._recompute_output = dolfinx.fem.Function(self.space_to)
-
-        self._recompute_output.interpolate(func_from)
-        self._recompute_output.x.scatter_forward()
-
-        # Overload the object to ensure PyAdjoint tracks it properly
-        output = create_overloaded_object(self._recompute_output)
+        # Update the tape's actual output object in-place (rather than a
+        # separately cached Function) so that any Python reference the user
+        # is holding to the interpolated Function stays in sync after the
+        # tape is replayed, matching FunctionAssignBlock's convention. Note
+        # this only holds until the output is first used as a dependency of
+        # another block: pyadjoint then freezes a private checkpoint copy
+        # (see OverloadedType._ad_will_add_as_dependency), so the live
+        # Python object can still go stale after that point. This is an
+        # existing pyadjoint/dolfinx_adjoint-wide characteristic (identical
+        # behavior in FunctionAssignBlock), not specific to interpolation.
+        output = block_variable.saved_output
+        output.interpolate(func_from)
+        output.x.scatter_forward()
         return output

tests/test_interpolate.py

@@ -1,3 +1,5 @@
+import gc
+
 from mpi4py import MPI
 
 import dolfinx
@@ -7,7 +9,12 @@
 import ufl
 
 from dolfinx_adjoint import Function, assemble_scalar, interpolate
-from dolfinx_adjoint.blocks.interpolation import InterpolationBlock
+from dolfinx_adjoint.blocks.interpolation import (
+    _CACHE_KEYS_BY_SPACE_ID,
+    _INTERPOLATION_MATRIX_CACHE,
+    InterpolationBlock,
+    _get_interpolation_matrix,
+)
 
 # Dynamically determine available matrix backends
 petsc_options = [False]
@@ -148,3 +155,61 @@ def u_ex(mod, x_coords):
     assert np.isclose(min_rate, 3.0, rtol=1e-3, atol=1e-3)
 
     pyadjoint.get_working_tape().clear_tape()
+
+
+# ==============================================================================
+# Test 3: Interpolation Matrix Cache
+# ==============================================================================
+
+
+def test_interpolation_matrix_cache_reused_for_same_spaces(mesh_1D):
+    """The same (space_from, space_to) pair should reuse one assembled matrix."""
+    V_from = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 1))
+    V_to = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 2))
+
+    mat1 = _get_interpolation_matrix(V_from, V_to)
+    mat2 = _get_interpolation_matrix(V_from, V_to)
+    assert mat1 is mat2
+
+    # A different space pair must not share the cached matrix.
+    V_other = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 3))
+    mat3 = _get_interpolation_matrix(V_from, V_other)
+    assert mat3 is not mat1
+
+
+def test_interpolation_matrix_cache_purged_when_space_is_garbage_collected(mesh_1D):
+    """Regression test for the id()-collision bug: a cache entry keyed on
+    id(space) must be dropped once that space is garbage collected, otherwise
+    a later, unrelated FunctionSpace object that happens to reuse the same id
+    could silently be served a stale matrix built for a completely different
+    pair of spaces.
+    """
+    V_to = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 2))
+
+    def build_and_return_id():
+        # V_from only lives inside this function, so it is eligible for
+        # collection as soon as it returns.
+        v_from = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 1))
+        _get_interpolation_matrix(v_from, V_to)
+        return id(v_from)
+
+    space_id = build_and_return_id()
+    gc.collect()
+
+    assert space_id not in _CACHE_KEYS_BY_SPACE_ID
+    assert all(key[0] != space_id for key in _INTERPOLATION_MATRIX_CACHE)
+
+
+def test_interpolation_matrix_cache_does_not_leak_across_transient_spaces(mesh_1D):
+    """The cache must not grow without bound as short-lived FunctionSpaces
+    (e.g. created once per optimization iteration) are used and discarded.
+    """
+    baseline = len(_INTERPOLATION_MATRIX_CACHE)
+
+    for _ in range(20):
+        v_from = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 1))
+        v_to = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 2))
+        _get_interpolation_matrix(v_from, v_to)
+
+    gc.collect()
+    assert len(_INTERPOLATION_MATRIX_CACHE) <= baseline + 1