Per-function JIT cache to avoid redundant recompilation by lmoresi · Pull Request #94 · underworldcode/underworld3

lmoresi · 2026-03-25T11:48:31Z

Summary

Adds per-function caching to the JIT compilation layer. Each compiled C function is cached individually by its structural hash, so when a solver recompiles (e.g. tensor projection changing only the residual RHS), only genuinely new functions are compiled — unchanged ones (Jacobians, smoothing terms) are reused via pointer copying.
Extends PtrContainer (Cython) with allocate() and copy_*_from() methods for cross-module function pointer assembly.
Preserves the existing whole-bundle fast path: if nothing changed, returns immediately with zero overhead.
Builds on PR Refactor JIT cache layer with JITCallbackSet #93 (JITCallbackSet structured cache keys).

Measured impact (VEP timing test, 8x8 mesh):

Tensor projection component loop: 2nd/3rd components go from full recompilation to 6/6 cache hits, 0 new compilations
RHS1 -> RHS2 -> RHS1 pattern: zero recompilation on return to RHS1
First VE solve: 22.0s -> 19.1s (-13%)
Total benchmark: 64.5s -> 57.6s (-11%)

Also includes timing/profiling test scripts (tests/minimal_vep_timing.py, tests/profile_jit_phases.py) and a bugfix for elastic_dt -> dt_elastic in VE_Stokes.solve().

Test plan

Tier A level 1 tests pass (33/33 verified locally)
CI test suite passes
VE shear box test still converges
No regression in Integral/BdIntegral cache behaviour (PR Fix Integral/BdIntegral JIT cache collisions #92 fix preserved)

Underworld development team with AI support from Claude Code

Refactor the JIT compilation pipeline to use a JITCallbackSet dataclass that groups the five PETSc callback lists (residual, bcs, jacobian, bd_residual, bd_jacobian) into a single structured container. This addresses the root cause of the cache-collision bug (PR #92) at an architectural level: the flat tuple hash that lost callback role information is replaced by a structured signature that preserves which slot each expression belongs to. Changes: - Add JITCallbackSet dataclass with flat(), signature(), map(), counts - Extract _structural_expand() as a module-level function (was inline) - Refactor getext() to accept JITCallbackSet (with backward compat) - Refactor _createext() to accept JITCallbackSet - Update all 6 call sites: 3 solvers (Scalar, Vector, Stokes) and 3 integrals (Integral, Integral._evaluate_integral, BdIntegral) - Include PR #92 regression tests (spherical shell cache collision) Incorporates the fix from PR #92 (gthyagi) which identified the bug and added the regression tests. Test results: 374 passed, 7 skipped, 1 xfailed (level_1 suite) Underworld development team with AI support from Claude Code

The existing JIT cache keys the entire bundle of functions (residuals + Jacobians + BCs) as one hash. If any single function changes, everything recompiles. This is wasteful for SNES_Tensor_Projection which loops over tensor components, changing only the residual RHS while the Jacobian (identity) never changes. New per-function cache: each compiled C function is cached individually by its structural hash + signature type. When getext() is called: - Fast path: whole-bundle hash hit → return immediately (unchanged) - Slow path: check per-function hashes, compile only new functions, assemble PtrContainer by copying cached function pointers PtrContainer gains allocate() and copy_*_from() methods to support cross-module pointer assembly. Also fixes elastic_dt → dt_elastic bug in VE_Stokes.solve(). Includes timing/profiling test scripts for JIT benchmarking. Underworld development team with AI support from Claude Code

Copilot

Pull request overview

This PR enhances Underworld3’s JIT extension layer by introducing per-function caching for compiled PETSc callback functions, aiming to avoid redundant recompilation when only some expressions change between solver setups.

Changes:

Add per-function JIT caching in getext() keyed by structural hashes, while keeping a whole-bundle cache fast path.
Extend PtrContainer (Cython) with allocation and pointer-copy helpers to assemble mixed-origin callback sets.
Update solvers/integrals/tests to use the structured JITCallbackSet interface; add two timing/profiling helper scripts; fix elastic_dt → dt_elastic in VE solver.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
`src/underworld3/utilities/_jitextension.py`	Implements per-function cache, partial compilation path, and pointer assembly via `PtrContainer`.
`src/underworld3/cython/petsc_types.pyx`	Adds `PtrContainer.allocate()` and `copy_*_from()` methods used for assembling callback pointer arrays.
`src/underworld3/cython/petsc_types.pxd`	Declares new `PtrContainer` methods for Cython consumers.
`src/underworld3/cython/petsc_maths.pyx`	Updates Integral/BdIntegral JIT calls to use `JITCallbackSet`.
`src/underworld3/cython/petsc_generic_snes_solvers.pyx`	Updates SNES solvers’ JIT calls to use `JITCallbackSet`.
`src/underworld3/systems/solvers.py`	Bugfix: use `Parameters.dt_elastic` instead of `elastic_dt`.
`tests/test_0004_pointwise_fns.py`	Updates unit tests to the new `getext(mesh, JITCallbackSet(...), ...)` signature.
`tests/minimal_vep_timing.py`	Adds a minimal runtime timing script for VEP solves / JIT caching.
`tests/profile_jit_phases.py`	Adds a profiling script to time derivative/JIT/solve phases.

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

Copilot · 2026-03-25T11:53:32Z

+                _fn_cache[cache_key] = entry
+                cached_hits[key] = entry
+
+    # Also register the full bundle in _ext_dict if ALL functions were new
+    # (common case: first compile of a solver)
+    if n_new > 0 and n_hits == 0:


When cache=False is passed, getext() still registers compiled functions into _fn_cache (and also adds new_jitname modules into _ext_dict). This means callers trying to force a clean rebuild can still pollute global caches and grow memory over time. Respect cache=False by skipping both reading and writing of _fn_cache/_ext_dict for per-function entries (or provide a separate flag controlling per-function caching).

Suggested change

_fn_cache[cache_key] = entry

cached_hits[key] = entry

# Also register the full bundle in _ext_dict if ALL functions were new

# (common case: first compile of a solver)

if n_new > 0 and n_hits == 0:

# Always make the entry available for this call via cached_hits,

# but only persist it in the global per-function cache when

# caching is enabled.

if cache:

_fn_cache[cache_key] = entry

cached_hits[key] = entry

# Also register the full bundle in _ext_dict if ALL functions were new

# (common case: first compile of a solver). Only do this when caching

# is enabled, so that callers using cache=False do not pollute the

# global extension registry.

if cache and n_new > 0 and n_hits == 0:

Copilot · 2026-03-25T11:53:32Z

+    # Also register the full bundle in _ext_dict if ALL functions were new
+    # (common case: first compile of a solver)
+    if n_new > 0 and n_hits == 0:
+        _ext_dict[jitname] = _ext_dict[new_jitname]
+
+    # ── Assemble PtrContainer from cached function pointers ─────────────
+    from underworld3.cython.petsc_types import PtrContainer
+
+    result_ptr = PtrContainer()
+    counts = callbacks.counts
+    result_ptr.allocate(*counts)
+
+    _copy_methods = {
+        "residual": result_ptr.copy_residual_from,
+        "bcs": result_ptr.copy_bcs_from,
+        "jacobian": result_ptr.copy_jacobian_from,
+        "bd_residual": result_ptr.copy_bd_residual_from,
+        "bd_jacobian": result_ptr.copy_bd_jacobian_from,
+    }

-    i_bd_res = {}
-    for index, fn in enumerate(fns_bd_residual):
-        i_bd_res[fn] = index
+    for sig_type, n_fns in zip(_SIG_TYPES, counts):
+        copy_fn = _copy_methods[sig_type]
+        for slot_idx in range(n_fns):
+            entry = cached_hits[(sig_type, slot_idx)]
+            copy_fn(slot_idx, entry.ptr_container, entry.index)

-    i_bd_jac = {}
-    for index, fn in enumerate(fns_bd_jacobian):
-        i_bd_jac[fn] = index
+    # ── Build fn_dicts (unchanged from original) ────────────────────────
+    i_res = {fn: i for i, fn in enumerate(callbacks.residual)}
+    i_ebc = {fn: i for i, fn in enumerate(callbacks.bcs)}
+    i_jac = {fn: i for i, fn in enumerate(callbacks.jacobian)}
+    i_bd_res = {fn: i for i, fn in enumerate(callbacks.bd_residual)}
+    i_bd_jac = {fn: i for i, fn in enumerate(callbacks.bd_jacobian)}

    extn_fn_dict = namedtuple(
-        "Functions",
-        ["res", "jac", "ebc", "bd_res", "bd_jac"],
+        "Functions", ["res", "jac", "ebc", "bd_res", "bd_jac"],
    )

-    extensions_functions_dicts = extn_fn_dict(i_res, i_jac, i_ebc, i_bd_res, i_bd_jac)
-
-    return _GextResult(ptrobj, extensions_functions_dicts, constants_manifest)
+    return _GextResult(
+        result_ptr,
+        extn_fn_dict(i_res, i_jac, i_ebc, i_bd_res, i_bd_jac),
+        constants_manifest,
+    )


For mixed cache-hit/miss cases, getext() assembles a new PtrContainer via allocate()+copy_*_from() every call, but the assembled container is not cached under jitname. That means repeated calls with the same final signature never reach the 'whole-bundle' fast path and repeatedly allocate/copy (and currently leak without PtrContainer freeing). Consider caching the assembled PtrContainer keyed by jitname (and constants manifest signature) once all required functions are available.

Copilot · 2026-03-25T11:53:32Z

+        # Compile the new functions
+        new_jitname = abs(hash((jitname, "partial", n_new, time.time())))
        _createext(
-            jitname,
+            new_jitname,
            mesh,
-            fns_residual,
-            fns_bcs,
-            fns_jacobian,
-            fns_bd_residual,
-            fns_bd_jacobian,
+            new_callbacks,
            primary_field_list,
            constants_subs_map=constants_subs_map,
            verbose=verbose,
            debug=debug,
            debug_name=debug_name,
        )
-    else:
-        if verbose and underworld3.mpi.rank == 0:
-            print(f"JIT compiled module cached ... {jitname} ", flush=True)
-
-    module = _ext_dict[jitname]
-    ptrobj = module.getptrobj()
-
-    i_res = {}
-    for index, fn in enumerate(fns_residual):
-        i_res[fn] = index
-
-    i_ebc = {}
-    for index, fn in enumerate(fns_bcs):
-        i_ebc[fn] = index



new_jitname is derived from abs(hash((jitname, "partial", n_new, time.time()))). Using Python's hash here risks collisions (rare but possible) and also makes module naming nondeterministic, which complicates debugging and reproducibility. Prefer a monotonic counter or uuid-based suffix to guarantee uniqueness without relying on hash randomness/time resolution.

Copilot · 2026-03-25T11:53:33Z

+# --- First solve (includes JIT compilation) ---
+# Set dt_elastic explicitly to work around elastic_dt alias bug in VE_Stokes.solve()
+t1 = time.time()
+stokes.solve(timestep=0.02, zero_init_guess=True)
+print(f"First solve (incl JIT): {time.time() - t1:.1f}s")


The comment says dt_elastic is set explicitly to work around the elastic_dt/dt_elastic bug, but the script does not actually set dt_elastic anywhere before calling stokes.solve(...). Either update the comment or set the parameter explicitly so the script matches its own documentation (especially now that VE_Stokes.solve() was fixed in this PR).

Copilot · 2026-03-25T11:53:33Z

+    cpdef allocate(self, int n_res, int n_bcs, int n_jac, int n_bd_res, int n_bd_jac):
+        """Allocate function pointer arrays of the given sizes."""
+        self.fns_residual    = <PetscDSResidualFn*>   malloc(n_res    * sizeof(PetscDSResidualFn))
+        self.fns_bcs         = <PetscDSResidualFn*>   malloc(n_bcs    * sizeof(PetscDSResidualFn))
+        self.fns_jacobian    = <PetscDSJacobianFn*>   malloc(n_jac    * sizeof(PetscDSJacobianFn))
+        self.fns_bd_residual = <PetscDSBdResidualFn*> malloc(n_bd_res * sizeof(PetscDSBdResidualFn))
+        self.fns_bd_jacobian = <PetscDSBdJacobianFn*> malloc(n_bd_jac * sizeof(PetscDSBdJacobianFn))


PtrContainer.allocate() uses malloc() for several pointer arrays but PtrContainer has no dealloc (or equivalent) to free them. With the new per-function assembly path creating a fresh PtrContainer frequently, this becomes an unbounded native-memory leak. Add free() calls in a dealloc method (and consider freeing/replacing existing pointers inside allocate() when called more than once).

Copilot · 2026-03-25T11:53:33Z

+    constants_manifest_key = tuple(
+        (str(expr), idx) for idx, expr in constants_manifest


The per-function cache key for constants uses str(expr), but UWexpression.str includes the current value (and units). This makes cache hits value-dependent and can force recompilation (or change constants[] ordering) when only constant values change—defeating the intended constants[] mechanism. Use a value-independent identifier instead (e.g., expr.name / a stable symbol id) for constants_manifest_key (and any related ordering).

Suggested change

constants_manifest_key = tuple(

(str(expr), idx) for idx, expr in constants_manifest

# Use a value-independent identifier for each constant expression so that

# changing the numerical value of a constant does not invalidate the cache.

def _constant_key(expr):

"""

Return a stable, value-independent identifier for a constant expression.

Prefer attributes that are expected to be stable across value changes,

falling back to the expression's type name if nothing better is

available.

"""

name = getattr(expr, "name", None)

if name is not None:

return name

symbol_id = getattr(expr, "symbol_id", None)

if symbol_id is not None:

return symbol_id

# Fallback: type name is at least independent of the current value.

return type(expr).__name__

constants_manifest_key = tuple(

(_constant_key(expr), idx) for idx, expr in constants_manifest

lmoresi added 2 commits March 25, 2026 16:53

Copilot AI review requested due to automatic review settings March 25, 2026 11:48

Copilot started reviewing on behalf of lmoresi March 25, 2026 11:49 View session

Copilot AI reviewed Mar 25, 2026

View reviewed changes

lmoresi merged commit 4f78272 into development Mar 26, 2026
5 checks passed

Navigation Menu

Uh oh!

Per-function JIT cache to avoid redundant recompilation - #94

Per-function JIT cache to avoid redundant recompilation#94
lmoresi merged 2 commits into
developmentfrom
feature/jit-fn-cache

lmoresi commented Mar 25, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Mar 25, 2026

Uh oh!

Copilot AI Mar 25, 2026

Uh oh!

Copilot AI Mar 25, 2026

Uh oh!

Copilot AI Mar 25, 2026

Uh oh!

Copilot AI Mar 25, 2026

Uh oh!

Copilot AI Mar 25, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		constants_manifest_key = tuple(
		(str(expr), idx) for idx, expr in constants_manifest

-    constants_manifest_key = tuple(
-        (str(expr), idx) for idx, expr in constants_manifest
+    # Use a value-independent identifier for each constant expression so that
+    # changing the numerical value of a constant does not invalidate the cache.
+    def _constant_key(expr):
+        """
+        Return a stable, value-independent identifier for a constant expression.
+        Prefer attributes that are expected to be stable across value changes,
+        falling back to the expression's type name if nothing better is
+        available.
+        """
+        name = getattr(expr, "name", None)
+        if name is not None:
+            return name
+        symbol_id = getattr(expr, "symbol_id", None)
+        if symbol_id is not None:
+            return symbol_id
+        # Fallback: type name is at least independent of the current value.
+        return type(expr).__name__
+    constants_manifest_key = tuple(
+        (_constant_key(expr), idx) for idx, expr in constants_manifest

Uh oh!

Conversation

lmoresi commented Mar 25, 2026

Summary

Test plan

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants