diff --git a/src/dolfinx_adjoint/blocks/interpolation.py b/src/dolfinx_adjoint/blocks/interpolation.py index 8cf8b58..dcd0c23 100644 --- a/src/dolfinx_adjoint/blocks/interpolation.py +++ b/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 diff --git a/tests/test_interpolate.py b/tests/test_interpolate.py index 1328412..1106f8c 100644 --- a/tests/test_interpolate.py +++ b/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