JIT cache: key on generated C source (closes #121, #123)

docs/advanced/benchmarks/jit_cache_vs_recompile.py

@@ -0,0 +1,87 @@
+"""Benchmark for issue #123: parameter-update recompile.
+
+Reproduces the VE square-wave toggle pattern from
+https://github.com/underworldcode/underworld3/issues/123 — the original
+report measured ~459 s of JIT time vs ~3.7 s of actual SNES solve over
+99 timesteps that only changed dt_elastic and the top BC sign.
+
+With the C-source-hash JIT cache in place, the JIT time should collapse
+to roughly the cost of the first cold compile (~10 s on this hardware),
+because every later step generates the same C source and hits the cache.
+
+Run with:
+
+    pixi run -e amr-dev python -u docs/advanced/benchmarks/jit_cache_vs_recompile.py
+
+Optional: set ``UW_JIT_CACHE=0`` to disable the disk cache and time the
+in-memory cache only; or ``rm -rf ~/.cache/underworld3/jit`` first to
+force a full cold start.
+"""
+
+import time
+
+import sympy
+
+import underworld3 as uw
+import underworld3.timing as uw_timing
+from underworld3.function import expression
+from underworld3.utilities._jitextension import _ext_dict
+
+
+N_STEPS = 30
+ELEMENT_RES = (16, 8)
+
+
+def main():
+    mesh = uw.meshing.StructuredQuadBox(
+        elementRes=ELEMENT_RES,
+        minCoords=(0.0, 0.0),
+        maxCoords=(1.0, 0.5),
+    )
+
+    v = uw.discretisation.MeshVariable("V", mesh, mesh.dim, degree=2)
+    p = uw.discretisation.MeshVariable("P", mesh, 1, degree=1)
+
+    stokes = uw.systems.VE_Stokes(
+        mesh, velocityField=v, pressureField=p, order=2
+    )
+    stokes.constitutive_model = uw.constitutive_models.ViscoElasticPlasticFlowModel
+    stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0
+    stokes.constitutive_model.Parameters.shear_modulus = 1.0
+
+    V_top = expression("V_top", 0.5, "Top BC amplitude")
+    stokes.add_dirichlet_bc((V_top, 0.0), "Top")
+    stokes.add_dirichlet_bc((0.0, 0.0), "Bottom")
+    stokes.add_dirichlet_bc((sympy.oo, 0.0), "Left")
+    stokes.add_dirichlet_bc((sympy.oo, 0.0), "Right")
+
+    uw_timing.start()
+    t_total = time.time()
+
+    n_compiles_initial = len(_ext_dict)
+    step_times = []
+    for step in range(N_STEPS):
+        V_top.sym = sympy.Float(0.5 * (-1) ** step)
+        stokes.constitutive_model.Parameters.dt_elastic = 0.02
+        t_step = time.time()
+        stokes.solve(zero_init_guess=False, timestep=0.02)
+        step_times.append(time.time() - t_step)
+
+    elapsed = time.time() - t_total
+    n_compiles_after = len(_ext_dict)
+
+    uw.pprint(f"\n=== JIT cache benchmark — issue #123 reproducer ===")
+    uw.pprint(f"Steps                : {N_STEPS}")
+    uw.pprint(f"Mesh                 : {ELEMENT_RES}, BDF-2 VE Stokes")
+    uw.pprint(f"Wall time (total)    : {elapsed:.1f} s")
+    uw.pprint(f"First step (cold)    : {step_times[0]:.2f} s")
+    uw.pprint(f"Mean of remaining    : {sum(step_times[1:]) / (N_STEPS - 1):.2f} s/step")
+    uw.pprint(f"Compiled bundles     : {n_compiles_after - n_compiles_initial}")
+    uw.pprint(f"  (issue #123: would have shown ~3 new compiles per step "
+              f"= ~{N_STEPS * 3} total before this fix)")
+
+    uw_timing.print_summary()
+
+
+if __name__ == "__main__":
+    main()

docs/developer/index.md

@@ -165,6 +165,7 @@ subsystems/expressions-functions
 subsystems/containers
 subsystems/checkpointing-system
 subsystems/model-orchestration
+subsystems/jit-cache
 ```
 
 ```{toctree}

docs/developer/subsystems/jit-cache.md

@@ -0,0 +1,136 @@
+# JIT compilation cache
+
+Solvers in underworld3 compile their pointwise residual / Jacobian functions
+to native code via Cython + the system `cc`. For a non-trivial solver this
+takes seconds (simple Stokes) to minutes (VE Stokes BDF-2 with viscoplastic
+flow). The JIT cache eliminates that cost on every call after the first one
+within a session, and on every fresh Python process after the first one on a
+given machine.
+
+This document describes the on-disk layout, how cache entries are keyed, and
+how the cache is invalidated.
+
+## Layout
+
+```
+~/.cache/underworld3/jit/
+├── .env-fingerprint                    # written by `./uw build`
+├── 7e81af3b554b4126.so                 # the compiled extension
+├── 7e81af3b554b4126.manifest.json      # constants index → name mapping
+├── .7e81af3b554b4126.lock              # advisory lock used during writes
+├── 0202220e00e01df3.so
+├── 0202220e00e01df3.manifest.json
+└── ...
+```
+
+Locations honoured (first match wins):
+
+1. `$UW_JIT_CACHE_DIR` if set
+2. `$XDG_CACHE_HOME/underworld3/jit` if `XDG_CACHE_HOME` is set
+3. `~/.cache/underworld3/jit`
+
+Setting `UW_JIT_CACHE=0` (or `false` / `no`) disables the on-disk cache
+entirely — the in-memory dict still works, but nothing is persisted across
+processes.
+
+## Cache key
+
+The key is the SHA-256 of the **canonical** generated C source plus an
+**ABI salt**, truncated to 16 hex characters.
+
+```python
+canonical_source = "\n".join(setup_py, cy_ext_h, cy_ext_pyx)
+                   .replace(<modname>, "__UW_JIT_MOD__")
+                   .replace(<randstr>, "__UW_JIT_RS__")
+
+salt = f"petsc={PETSc.Sys.getVersion()}|uw={underworld3.__version__}"
+
+source_hash = sha256(canonical_source + "\n---\n" + salt).hexdigest()[:16]
+```
+
+Two design points worth understanding:
+
+**Canonicalisation**. The Cython module name and the symbol-prefix randomiser
+that the JIT generator picks vary between calls (random or counter-based).
+Hashing them directly would make every call produce a fresh key. We instead
+hash the source with both replaced by stable placeholders, and only after
+hashing do we substitute the real per-bundle names back in. Same C ⇒ same
+hash; different C ⇒ different hash; identifiers stay unique per bundle.
+
+**Why hash the C source rather than the sympy callback structure**. Two
+sympy expressions can be structurally distinct but emit identical C
+(constant placeholders, simplification, etc.); two structurally-identical
+expressions can emit different C if a path in the generator depends on
+something subtle. The compiled `.so` is what the solver actually executes —
+hashing exactly that makes the equivalence relation provable: same hash
+implies same `.so` implies same numerical behaviour.
+
+## Cache hit/miss flow
+
+`getext()` performs a three-tier lookup, cheapest first:
+
+1. **In-memory `_ext_dict`** — same Python process, just a dict access.
+2. **On-disk `{hash}.so` + `{hash}.manifest.json`** — same machine, fresh
+   process. `load_module` re-loads the `.so` via Python's import machinery
+   and verifies the saved constants list still matches the current call's
+   manifest (belt-and-braces: a name shift means an unnoticed ABI drift).
+3. **Cold compile** via `compile_and_load`. After a successful compile, the
+   `.so` and manifest are copied into the cache directory (rank 0 only),
+   guarded by an advisory `flock`.
+
+## Invalidation
+
+A cached entry stops being valid when *anything* it implicitly depended on
+changes. We handle the common categories like this:
+
+| Change                              | Handled by                                       |
+|-------------------------------------|--------------------------------------------------|
+| Solver expressions / topology       | `source_hash` differs ⇒ different entry          |
+| Constant **value**                  | `source_hash` is independent of value (uses `_JITConstant` placeholders); `PetscDSSetConstants` updates the array at solve time |
+| PETSc version, UW version           | Embedded in the ABI salt ⇒ different `source_hash` |
+| Compiler / CFLAGS / Python ABI      | `./uw build` writes `.env-fingerprint`; mismatch ⇒ wipe `*.so` and `*.manifest.json` |
+| Constant rename                     | `load_module` compares saved vs current names ⇒ miss |
+| Manifest schema change              | `MANIFEST_VERSION` bump ⇒ miss                   |
+
+The cache directory is safe to delete by hand at any time:
+
+```
+rm -rf ~/.cache/underworld3/jit
+```
+
+The next solve will repopulate.
+
+## MPI
+
+When `mpi.size > 1`:
+
+- Every rank computes the C-source hash independently. The hashes are
+  `comm.allgather`'d and compared — a mismatch raises immediately rather
+  than letting ranks diverge. Non-determinism in `generate_c_source` (e.g.
+  set/dict iteration order leaking into emitted C) would land here.
+- Rank 0 writes the cache entry; other ranks rely on the disk-cache hit
+  path on subsequent calls.
+- The `flock` on the per-hash lockfile serialises cross-shell concurrent
+  writes (e.g. a `mpirun -np 4` and a `mpirun -np 2` started seconds apart
+  on the same machine).
+
+A future refinement would have rank 0 compile while other ranks wait on a
+barrier and read the resulting `.so` directly; today every rank still
+performs the cold compile but only rank 0 publishes the result.
+
+## Environment variables
+
+| Variable             | Effect                                                  |
+|----------------------|---------------------------------------------------------|
+| `UW_JIT_CACHE`       | Set to `0`/`false`/`no` to disable disk cache           |
+| `UW_JIT_CACHE_DIR`   | Override the cache directory location                   |
+| `XDG_CACHE_HOME`     | Used when `UW_JIT_CACHE_DIR` is unset                   |
+
+## Code references
+
+- `src/underworld3/utilities/_jitextension.py` — `getext`, `generate_c_source`,
+  `compile_and_load`, `_abi_salt`, `_extract_constants`.
+- `src/underworld3/utilities/_jit_cache.py` — disk cache: `get_cache_dir`,
+  `load_module`, `store_module`, `_file_lock`.
+- `uw` (shell driver) — `.env-fingerprint` write + cache wipe inside
+  `run_build`.

src/underworld3/cython/petsc_generic_snes_solvers.pyx

@@ -519,6 +519,14 @@ class SolverBaseClass(uw_object):
             self._setup_solver(verbose, _rewire_only=True)
 
             self.is_setup = True
+
+            # _last_jit_cache_key must track "which bundle the solver is
+            # currently wired for", not "what the last full-build was".
+            # Otherwise a C → S → C transition can make fast path 1 trigger
+            # on a subsequent _build() call while the solver is still wired
+            # for S — returning S's result for a C solve.
+            if hasattr(self, '_current_jit_cache_key'):
+                self._last_jit_cache_key = self._current_jit_cache_key
             return
 
         # === Full rebuild path — teardown DM and reconstruct ===