Stokes JIT: rank-0 compile + CSE before printing + configurable flags (#547)

[bknight1]

Summary

Fixes the HPC compiler OOM and wall-clock timeout during runtime JIT compilation of Stokes pointwise functions (e.g. gcc: fatal error: Killed signal terminated program cc1 under mpirun -np N) described in #547.

This PR implements four coordinated changes:

1. Rank-0 Compile and Disk-Cache Activation: SNES_Stokes_SaddlePt._setup_pointwise_functions now calls getext with cache=True (the only call site that previously opted out). Under MPI, this activates the rank-0-only compile gate and cross-rank disk handoff, preventing N concurrent compiler invocations per node.

2. Common Subexpression Elimination (CSE): generate_c_source in _jitextension.py performs CSE prior to C printing. Shared subexpressions become intermediate double xN = ...; statements evaluated in topological dependency order. This shrinks generated C headers by up to ~63x on complex rheologies (1.45 MB → 23 KB) and lowers peak gcc RSS from ~1.32 GB to ~62–74 MB. Escapable via UW_JIT_NOCSE=1.

3. Configurable JIT CFLAGS: Defaults to -O3 -g0 to eliminate debug symbol bloat from Python sysconfig. Customizable via UW3_JIT_CFLAGS (e.g. -O1 -g0) for memory-constrained HPC nodes.

4. Matrix-Level SymPy Differentiation: Replaces per-entry derivative loops with sympy.diff across whole block matrices (~1.7x faster Jacobian generation, bit-identical entries, PETSc [fc, gc, df, dg] layout preserved).

Fixes #547

Verification

• Rank-0 compile + disk-cache reuse confirmed under mpirun -np 2 (ptest_jit_cache.py).
• Jacobian layout validated against the finite-difference oracle (test_1066_stokes_jacobian_layout.py).
• JIT determinism and constants manifest verified (test_jit_cache.py, test_jit_deterministic_ordering.py, test_0103_jit_rampable_constants.py).
• Core and Stokes test suites passing (79 tests); CSE-off path exercised.

Underworld development team with AI support from Claude Code

[Copilot]

Pull request overview

This PR addresses MPI/HPC failures during runtime JIT compilation of Stokes pointwise functions by reducing redundant compilation work, shrinking generated C via SymPy common-subexpression elimination, and improving JIT compile configurability/performance.

Changes:

• Enable Stokes JIT disk caching so MPI runs use rank-0 compilation with cross-rank reuse.
• Add optional SymPy CSE prior to C printing to dramatically reduce generated code size for large expressions.
• Make JIT compile flags configurable via UW3_JIT_CFLAGS (defaulting to -O3 -g0) and speed up Jacobian generation via matrix-level differentiation.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
src/underworld3/utilities/_jitextension.py	Adds lowering gate, CSE-based code shrinking, and configurable JIT compile flags for generated extensions.
src/underworld3/cython/petsc_generic_snes_solvers.pyx	Switches Stokes pointwise-function JIT to cached/rank-0 compile and optimizes Jacobian construction with matrix-level sympy.diff.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[lmoresi]

Adversarial review

Reviewed against the PR head. CI is green and the change does what it says; one
finding is worth acting on before it lands, because this PR makes a latent
deadlock live on the main solver path.

1. The Barrier() is reached conditionally, on a rank-local predicate

In _jitextension.getext the cold-compile block is entered per rank:

module = None
if cache:
    module = _jc.load_module(...)     # "Disk lookup is cheap and rank-local
                                      #  — every rank checks independently"
if module is None:                    # <-- rank-local decision
    ...
    if multi_rank and disk_enabled:
        if rank == 0: compile + store
        underworld3.mpi.comm.Barrier()   # <-- COLLECTIVE, inside the branch
        if rank != 0: load

The comment above the disk lookup states it is rank-local. The Barrier two
lines below is not. If any rank finds the published module on disk and any other
does not, the second enters the barrier alone and the job hangs.

Two ways they diverge, both realistic on the systems #547 is about:

• Shared-filesystem visibility. NFS attribute caching, Lustre metadata lag,
  or a partially-visible write means rank 3 sees the .so and rank 5 does not.
  This is the same class of hazard the barrier exists to manage, and the barrier
  sits inside the branch it is meant to protect.
• disk_enabled differing. disk_enabled = cache and _jc.get_cache_dir() is not None. A rank whose cache directory cannot be created or resolved takes the
  else branch — local compile, no barrier — while its peers take the
  barrier branch.

Why this PR changes the exposure. Before it, Stokes was the one call site
passing cache=False, so disk_enabled was False and Stokes always took the
no-barrier branch. Item 1 switches it to cache=True, which is the point of the
change — and puts the main solver path onto the branch containing a conditional
collective, on precisely the shared high-latency filesystems the PR targets.

The fix is small: decide collectively. Allreduce the "does anyone need a
compile?" predicate before branching so every rank takes the same branch, or lift
the barrier out of the conditional. Either makes the discipline true by
construction rather than by every rank happening to see the same directory
contents at the same instant.

2. The gate's verification is in a file nothing runs

tests/parallel/ptest_jit_cache.py is named in the verification list.
tests/pytest.ini sets python_files = test_*.py, so ptest_* is never
collected, and neither scripts/test.sh nor scripts/test_levels.sh mentions
ptest at all. The rank-0 compile gate is therefore covered by a manual
invocation and by nothing automatic. Given finding 1, that is the coverage worth
having. (#615 fixes the directory globs; it does not make ptest_ files
collectible, which is a separate decision.)

3. "Bit-identical" is claimed, and the evidence is a finite-difference oracle

Item 4 says the matrix-level sympy.diff produces "bit-identical entries", and
the verification for it is test_1066_stokes_jacobian_layout.py against a
finite-difference oracle. An FD comparison validates layout and approximate
values; it cannot establish bit-identity. If the claim matters — and it is the
reason to believe a 1.7x codegen speedup is free — the check is a comparison of
the generated C, or of evaluated entries, against the per-entry path.

4. Checked and clean: the CSE is bit-safe

Worth stating because it is the part that looks riskiest. Naming a repeated
subexpression and reusing it yields the same bits as recomputing it, and
sympy.cse with the default optimizations=None does no algebraic rewriting, so
this is structural rather than a simplification. The coordinate patch is also the
right call: cse minting BaseScalar instances without _ccodestr is a known
trap, and _patch_coords is applied to both the replacement expressions and the
reduced one.

Underworld development team with AI support from Claude Code

[lmoresi]

The collective, demonstrated — and a branch you can take

Following finding 1 above, we reproduced it and fixed it rather than leaving it
as an argument. The fix is one commit on top of this branch, at
bugfix/jit-collective-gate (aaba141).

The deadlock is real

Monkeypatching the disk lookup so rank 0 sees the cache and the other ranks do
not — the shape NFS attribute caching or Lustre metadata lag produces just after
rank 0 publishes — then running an ordinary Stokes solve at np=4:

	ranks entering the solve	ranks returning	result
this branch as it stands	4	0	mpirun timeout, rc=241
with the reduction	4	4	rc=0, four consecutive runs

Nothing exotic is needed to provoke it: the divergence is in which branch each
rank takes, and the barrier is inside the branch.

What the fix does

Both predicates are reduced before they are used to choose a branch:

• needs_compile becomes a global OR of module is None, so if any rank
  lacks the module every rank enters the branch and reaches the Barrier.
• disk_enabled becomes a global AND, so a rank whose get_cache_dir()
  returns None makes every rank fall back together to local compiles rather
  than some waiting in a barrier the others never enter. Wasteful, and the safe
  direction.

Ranks that already hold the module keep it — the rank-0 compile, the
post-barrier load and the no-disk fallback are each guarded on module is None,
so nothing recompiles or reloads what it already has.

Scope

The gate predates this PR; what this PR changes is the exposure. Stokes was the
one call site passing cache=False, so it always took the no-barrier branch.
Item 1 switching it to cache=True is the point of the change, and it puts the
main solver path onto the branch with the conditional collective — on exactly
the shared filesystems #547 is about.

Verified: full ./uw test 1556 passed, 32 skipped, 2 xfailed; the four JIT test
files 14 passed.

@bknight1 — the branch is there to cherry-pick or ignore as you prefer; say the
word and we can push it onto this one instead. Findings 2 and 3 above (the
ptest_ file nothing collects, and "bit-identical" evidenced by a
finite-difference oracle) are unaddressed and are not blockers.

Underworld development team with AI support from Claude Code

[lmoresi]

How to take the fix — pick one, all three are fine

The commit is aaba141 on bugfix/jit-collective-gate, sitting directly on top
of this PR's head (8978ad8), so it applies cleanly with no rebase.

Option A — cherry-pick onto this branch (recommended).

git fetch origin bugfix/jit-collective-gate
git checkout bugfix/stokes-jit-oom-cse-cache
git cherry-pick aaba1417b7d9ecd359d5407ceedd12987a406cf5
./uw build
git push origin bugfix/stokes-jit-oom-cse-cache

Option B — merge the branch in. Same result, keeps the fix as its own
commit with the merge recorded:

git fetch origin bugfix/jit-collective-gate
git checkout bugfix/stokes-jit-oom-cse-cache
git merge origin/bugfix/jit-collective-gate
./uw build
git push origin bugfix/stokes-jit-oom-cse-cache

Option C — say so here and we will push it onto this branch for you. It adds
one commit and changes nothing else; we have not done it unasked because it is
your PR.

Verifying it yourself

The reproduction is a monkeypatch, not a real filesystem race, so it is
deterministic and takes a few seconds. Save as jit_divergence.py:

import sympy
import underworld3 as uw
from underworld3.utilities import _jit_cache as _jc

_real = _jc.load_module
# Rank 0 sees the published module, the others do not — what NFS attribute
# caching or Lustre metadata lag produces just after rank 0 publishes.
_jc.load_module = lambda *a, **k: _real(*a, **k) if uw.mpi.rank == 0 else None

mesh = uw.meshing.StructuredQuadBox(elementRes=(4, 4))
v = uw.discretisation.MeshVariable("Vj", mesh, mesh.dim, degree=2)
p = uw.discretisation.MeshVariable("Pj", mesh, 1, degree=1)
s = uw.systems.Stokes(mesh, velocityField=v, pressureField=p)
s.constitutive_model = uw.constitutive_models.ViscousFlowModel
s.constitutive_model.Parameters.shear_viscosity_0 = 1
s.bodyforce = sympy.Matrix([0, -1])
s.add_dirichlet_bc((0.0, 0.0), "Bottom")
s.add_dirichlet_bc((0.0, 0.0), "Top")

print(f"[{uw.mpi.rank}] entering solve", flush=True)
s.solve()
print(f"[{uw.mpi.rank}] SOLVE RETURNED", flush=True)

mpirun --timeout 120 -n 4 python -u jit_divergence.py ; echo "rc=$?"

rc=241 is the deadlock (mpirun's timeout). rc=0 is the fix. Read the exit
code, not the printed lines — mpirun drops a rank's final stdout at teardown,
so the "SOLVE RETURNED" count reads 3 or 4 on a healthy run at np=4 and is not a
reliable signal.

Clear the JIT cache between runs if you want a genuinely cold start
(_jit_cache.get_cache_dir() tells you where it is).

After applying

Nothing else in the PR needs to change. ./uw test was 1556 passed with the fix
applied, and the four JIT test files 14 passed.

Findings 2 and 3 in the review above — ptest_jit_cache.py being in a file
nothing collects, and "bit-identical" resting on a finite-difference oracle — are
not blockers and can be follow-ups.

Underworld development team with AI support from Claude Code