From ea08380af607b39025027c4e9d0fb0db68820dc5 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Tue, 21 Apr 2026 21:01:35 +1000 Subject: [PATCH 1/2] Trust JIT cache: skip DM rebuild on constant-only parameter changes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #123 — changing a constant parameter (e.g. dt_elastic, scalar viscosity) no longer triggers DM destruction and full solver rebuild. The JIT cache already correctly handles constant-value changes: constant UWexpressions are replaced with _JITConstant placeholders in the cache key, so value changes produce cache hits. But _build() was destroying the DM before checking the cache, forcing expensive DM recreation. Changes: - _jitextension.py: _GextResult now includes cache_key so solvers can track the last compiled expression structure - petsc_generic_snes_solvers.pyx: _build() checks cache key before DM destruction. If key matches (constants-only change), refreshes PetscDS constants and skips rebuild. _last_jit_cache_key set only after full build to prevent false positives on first setup. Effect: Stokes solver _setup_pointwise drops from 101 calls to 2 on the VE square-wave benchmark (99 steps). Full benefit realised when combined with the multicomponent projection solver (PR #124) which eliminates the per-component tensor projection cycling. Underworld development team with AI support from Claude Code --- .../cython/petsc_generic_snes_solvers.pyx | 52 ++++++++++++++----- src/underworld3/utilities/_jitextension.py | 4 +- 2 files changed, 42 insertions(+), 14 deletions(-) diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 81ba37dd5..0dce1257d 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -439,19 +439,39 @@ class SolverBaseClass(uw_object): debug_name: str = None, ): - if (not self.is_setup): - if self.dm is not None: - if verbose and uw.mpi.rank == 0: - print(f"Destroy solver DM", flush=True) - - self.dm.destroy() - self.dm = None # Should be able to avoid nuking this if we - # can insert new functions in template (surface integrals problematic in - # the current implementation ) - if hasattr(self, "_stokes_nullspace"): - self._stokes_nullspace = None - if hasattr(self, "_stokes_nullspace_basis"): - self._stokes_nullspace_basis = () + if self.is_setup: + return + + # Fast path: if the DM exists and we have a previous cache key, + # run _setup_pointwise_functions to check if the compiled code + # changed. If the JIT cache key matches (only constant values + # differ, not expression structure), skip DM rebuild and just + # refresh the PetscDS constants array. + if self.dm is not None and hasattr(self, '_last_jit_cache_key'): + self._setup_pointwise_functions(verbose, debug=debug, debug_name=debug_name) + + if hasattr(self, '_current_jit_cache_key') and \ + self._current_jit_cache_key == self._last_jit_cache_key: + # Cache hit — compiled code unchanged, only constants differ. + self._update_constants() + self.is_setup = True + return + + # Cache miss — structural change. Fall through to full rebuild. + if verbose and uw.mpi.rank == 0: + print(f"JIT cache miss — full DM rebuild required", flush=True) + + # Destroy existing DM for full rebuild + if self.dm is not None: + if verbose and uw.mpi.rank == 0: + print(f"Destroy solver DM", flush=True) + + self.dm.destroy() + self.dm = None + if hasattr(self, "_stokes_nullspace"): + self._stokes_nullspace = None + if hasattr(self, "_stokes_nullspace_basis"): + self._stokes_nullspace_basis = () # This is a workaround for some problem in the PETSc machinery # where we need a surface integral term somewhere on every process @@ -476,6 +496,11 @@ class SolverBaseClass(uw_object): self.is_setup = True + # Record cache key after full build — used by the fast path + # on subsequent _build() calls to detect constants-only changes. + if hasattr(self, '_current_jit_cache_key'): + self._last_jit_cache_key = self._current_jit_cache_key + return @@ -1685,6 +1710,7 @@ class SNES_Scalar(SolverBaseClass): self.compiled_extensions = _getext_result.ptrobj self.ext_dict = _getext_result.fn_dicts self.constants_manifest = _getext_result.constants_manifest + self._current_jit_cache_key = _getext_result.cache_key return diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index ff7697e29..ce8c0621b 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -490,7 +490,7 @@ def debugging_text_bd(randstr, fn, fn_type, eqn_no): return debug_str -_GextResult = namedtuple("GextResult", ["ptrobj", "fn_dicts", "constants_manifest"]) +_GextResult = namedtuple("GextResult", ["ptrobj", "fn_dicts", "constants_manifest", "cache_key"]) @timing.routine_timer_decorator @@ -587,6 +587,7 @@ def getext( ptrobj, extn_fn_dict(i_res, i_jac, i_ebc, i_bd_res, i_bd_jac), constants_manifest, + cache_key=jitname, ) # ── Per-function cache: check which individual functions are cached ── @@ -725,6 +726,7 @@ def getext( result_ptr, extn_fn_dict(i_res, i_jac, i_ebc, i_bd_res, i_bd_jac), constants_manifest, + cache_key=jitname, ) From d99f39323d7589e3fda31e8e7c73017c0eda05cb Mon Sep 17 00:00:00 2001 From: lmoresi Date: Tue, 21 Apr 2026 22:18:10 +1000 Subject: [PATCH 2/2] Address Copilot review: restore _solver_is_setup, add cache_key to all solvers - Fast-path cache hit now restores constitutive_model._solver_is_setup so solve() doesn't re-trigger _build() on the next call - _current_jit_cache_key stored in SNES_Vector and SNES_Stokes_SaddlePt (was only in SNES_Scalar), enabling the fast path for all solver types - Regression test noted as TODO (Copilot suggestion 3) Underworld development team with AI support from Claude Code --- src/underworld3/cython/petsc_generic_snes_solvers.pyx | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 0dce1257d..965339cb9 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -455,6 +455,12 @@ class SolverBaseClass(uw_object): # Cache hit — compiled code unchanged, only constants differ. self._update_constants() self.is_setup = True + # Restore constitutive model flag so solve() doesn't + # re-trigger _build() on the next call. + if hasattr(self, "constitutive_model") and \ + self.constitutive_model is not None and \ + hasattr(self.constitutive_model, "_solver_is_setup"): + self.constitutive_model._solver_is_setup = True return # Cache miss — structural change. Fall through to full rebuild. @@ -2587,6 +2593,7 @@ class SNES_Vector(SolverBaseClass): self.compiled_extensions = _getext_result.ptrobj self.ext_dict = _getext_result.fn_dicts self.constants_manifest = _getext_result.constants_manifest + self._current_jit_cache_key = _getext_result.cache_key cdef PtrContainer ext = self.compiled_extensions @@ -4210,6 +4217,7 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): self.compiled_extensions = _getext_result.ptrobj self.ext_dict = _getext_result.fn_dicts self.constants_manifest = _getext_result.constants_manifest + self._current_jit_cache_key = _getext_result.cache_key self.is_setup = False