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

src/underworld3/cython/petsc_generic_snes_solvers.pyx

@@ -439,19 +439,45 @@ 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
+                # 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.
+            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 +502,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 +1716,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
 
@@ -2561,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
 
@@ -4184,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
 

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,
     )