Persistent JIT cache: avoid recompilation across sessions

[lmoresi]

Summary

The JIT compiler (_jitextension.py) recompiles every pointwise function from scratch on each Python session start. For complex models (VE Stokes with BDF-2, TI rheology on geographic meshes), this takes 3-8 minutes of wall time before the first solve step can execute.

A persistent disk cache keyed by the C source hash would eliminate repeat compilation across sessions.

Observed compilation times

Model	Pointwise functions	First-solve wall time	Solve-only time
VE Stokes BDF-2 (square wave)	~15	~5 min	~1s/step
TI Stokes (geographic mesh)	~10	~2 min	~160s total
AdvDiffusion (geographic, DDt)	~8	~90s	~7s/step
Stokes (Cartesian, simple)	~5	~10s	fast

The JIT cost dominates short runs and makes interactive exploration painful. Users often think the code is hung.

Current behaviour

• _ext_dict in _jitextension.py caches compiled extensions in-memory
• Cache is lost when the Python process exits
• Every new session recompiles the same expressions
• Each compilation: sympy → C code → write .c → call cc → load .so
• Compiler subprocess shows as ~12% CPU on the Python process (waiting on cc)

Proposed solution

Hash the generated C source (or the sympy expression tree) and check a persistent directory (e.g. ~/.cache/underworld3/jit/) before compiling:

import hashlib

def get_or_compile(c_source, ...):
    h = hashlib.sha256(c_source.encode()).hexdigest()[:16]
    cached_so = CACHE_DIR / f"{h}.so"
    if cached_so.exists():
        return load_dynamic(cached_so)
    # compile as usual, then copy .so to cache
    ...

Invalidation: the hash includes the C source (which encodes the full expression), so any change in the expression produces a different hash. Library ABI changes (PETSc/underworld version) could be handled by including a version salt in the hash.

Impact

• First session: same as today (compile + cache)
• Subsequent sessions: instant load from cache
• Interactive workflows: parameter studies restart in seconds instead of minutes
• Parallel runs: all ranks can share the cache directory

[lmoresi]

Fixed by PR #129 (JIT cache keyed on generated C source).