Trust JIT cache: skip DM rebuild on constant-only parameter changes

[lmoresi]

Summary

Fixes#123 — constant parameter updates (e.g. dt_elastic, scalar viscosity)
no longer trigger DM destruction and full solver rebuild.

Problem

The JIT cache handles constant-value changes correctly via _JITConstant
placeholders (same compiled code, different values in constants[]). But
_build() destroyed the DM before checking the cache, forcing expensive
DM recreation on every timestep.

Fix

_build() checks the JIT cache key before destroying the DM:

1. Run _setup_pointwise_functions to get the current cache key
2. Key matches last build → constants-only change → refresh PetscDS constants, skip DM rebuild
3. Key differs → structural change → full DM rebuild as before

Results (VE square-wave benchmark, 99 steps, 16x8 mesh)

	Before	After
Stokes _setup_pointwise calls	101	2
L2 error	1.89e-3	1.89e-3 (identical)

Combined with #124 (multicomponent projection): 515s → 56s (9.2x speedup).

Files changed

• src/underworld3/utilities/_jitextension.py — _GextResult includes cache_key
• src/underworld3/cython/petsc_generic_snes_solvers.pyx — _build() fast path

Underworld development team with AI support from Claude Code

[Copilot]

Pull request overview

This PR aims to avoid expensive PETSc DM destruction/rebuilds when only constant values change in JIT-compiled pointwise functions (no structural expression changes), by checking the JIT cache key first and refreshing PetscDS constants on a cache hit.

Changes:

• Extend getext()’s return value to include the computed JIT cache_key.
• Add a _build() fast path that checks the current vs previous JIT cache key before tearing down the DM, and on a match updates DS constants and skips DM rebuild.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
src/underworld3/utilities/_jitextension.py	Adds cache_key to the getext() result so callers can detect constants-only changes.
src/underworld3/cython/petsc_generic_snes_solvers.pyx	Implements a cache-key-based fast path in SolverBaseClass._build() and captures the key in SNES_Scalar.

Comments suppressed due to low confidence (1)

src/underworld3/utilities/_jitextension.py:505

• getext() now returns _GextResult with an additional cache_key field, but the Returns section of the getext() docstring below still documents only (ptrobj, fn_dicts, constants_manifest). Update the docstring (and any public-facing docs/types) to include cache_key so callers know it’s available and what it represents.

_GextResult = namedtuple("GextResult", ["ptrobj", "fn_dicts", "constants_manifest", "cache_key"])
@timing.routine_timer_decorator
def getext(
mesh,
callbacks: JITCallbackSet,
primary_field_list,
verbose=False,
debug=False,
debug_name=None,
cache=True,
):

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Fast-path cache-hit sets self.is_setup = True but does not restore self.constitutive_model._solver_is_setup. Since solve() forces self.is_setup = False whenever _solver_is_setup is false (see e.g. SNES_Scalar.solve and SNES_Stokes_SaddlePt.solve), leaving the flag unset will cause _build() to run every solve after a constant update. Consider setting self.constitutive_model._solver_is_setup = True (and any other setup flags that gate rebuilds) when taking this constants-only fast path.

Suggested change

	self.is_setup =True
	self.is_setup =True
	ifhasattr(self, "constitutive_model") and \
	self.constitutive_model isnotNoneand \
	hasattr(self.constitutive_model, "_solver_is_setup"):
	self.constitutive_model._solver_is_setup =True

[Copilot]

The fast path relies on _setup_pointwise_functions() setting _current_jit_cache_key, but in this file only SNES_Scalar._setup_pointwise_functions assigns it from getext(...).cache_key. For other solver subclasses, _current_jit_cache_key may never be set, preventing _last_jit_cache_key from being recorded and disabling the optimization for those solvers. Make sure every _setup_pointwise_functions() implementation that calls getext() stores the returned cache_key consistently (or centralize this in the base class).

[Copilot]

This change introduces new behavior (skipping DM destruction / rebuild on a cache-key hit) but there doesn’t appear to be a regression test asserting DM reuse when a constants-only update forces is_setup = False (the existing constants[] tests focus on avoiding JIT recompilation). Consider adding a test that (1) solves once, (2) mutates a constant-only parameter that invalidates setup, (3) solves again, and (4) asserts the solver DM object (or its underlying PETSc handle) is unchanged while the solution updates correctly.