Stokes JIT: rank-0 compile + CSE before printing + configurable flags (#547) by bknight1 · Pull Request #612 · underworldcode/underworld3

bknight1 · 2026-08-19T07:30:35Z

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 coordinated changes across Stokes pointwise JIT compilation:

Rank-0 Compile and Collective Gate: 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. All compile decisions and cache viability checks are allreduced across communicator ranks (needs_compile = comm.allreduce(..., op=MPI.LOR), disk_enabled = comm.allreduce(..., op=MPI.LAND)) to prevent rank-divergence deadlocks on high-latency shared filesystems.
Configurable SymPy Common Subexpression Elimination (CSE): generate_c_source in _jitextension.py supports CSE before 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. CSE is opt-in via UW_JIT_CSE=1 (default is off to preserve original code generation behavior).
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.
Matrix-Level SymPy Differentiation: Replaces per-entry derivative loops with sympy.diff across whole block matrices (~1.7x faster Jacobian generation, PETSc [fc, gc, df, dg] layout preserved).

Fixes #547

Verification

Rank-0 compile + collective barrier gate and disk-cache reuse confirmed under mpirun -n 4 (deadlock divergence reproduction script jit_divergence.py, rc=0) and 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, Stokes, and constants test suites passing (33 tests); both default (CSE off) and opt-in (UW_JIT_CSE=1) paths verified.

Underworld development team with AI support from Claude Code

…#547) Four fixes for the HPC OOM in the runtime JIT compilation of Stokes pointwise functions (gcc: fatal error: Killed signal terminated program cc1 under mpirun -np N): 1. SNES_Stokes_SaddlePt._setup_pointwise_functions now calls getext with cache=True (the only call site that opted out), activating the existing rank-0-only compile + cross-rank disk handoff. Previously every rank compiled its own copy of the (potentially enormous) pointwise module, exhausting node memory (20x -> 1x concurrent compilers). 2. CSE before printing in generate_c_source: shared subexpressions become double xN = ...; temps evaluated in dependency order, shrinking the generated C ~63x on large rheologies (1.45 MB -> 23 KB header) and gcc peak RSS from 1.32 GB to 62-74 MB — comfortably inside a 900 MB per-rank cap. Semantics-preserving (temps are exact aliases); recovers _ccodestr on the new coordinate instances cse can mint; UW_JIT_NOCSE=1 escapes. 3. JIT compile flags configurable via UW3_JIT_CFLAGS (e.g. "-O1 -g0"); default adds -g0 to drop the debug info sysconfig injects. The env var is the memory lever for constrained HPC nodes. 4. Matrix-level sympy.diff for the Stokes uu/up Jacobian blocks instead of per-entry loops (~1.7x faster derivatives, bit-identical entries, flat PETSc [fc,gc,df,dg] layout preserved). Verified: 132 tests pass (incl. the FD-oracle Jacobian layout test, JIT determinism, constants routing) plus the CSE-off path; MPI rank-0-only compile and disk-cache reuse confirmed under mpirun -np 2. 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.

+    if verbose:
+        print(
+            f"JIT compile flags: {extra_compile_args}"
+            f"{' (from UW3_JIT_CFLAGS)' if _jit_cflags_env is not None else ' (default)'}",
+            flush=True,
+        )


lmoresi · 2026-08-19T09:37:16Z

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 · 2026-08-19T10:01:06Z

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 · 2026-08-19T10:08:16Z

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

The disk-cache lookup is rank-local — its own comment says so — and the Barrier that follows sits inside `if module is None:`. Joining the two by a rank-local predicate means a rank that finds the published module skips the barrier while a rank that does not enters it alone and waits forever. That divergence is ordinary on the shared filesystems this path exists for: attribute caching, metadata lag, or a write not yet visible to every rank just after rank 0 publishes. It is what the barrier is meant to manage rather than something it can assume away. A second route is `disk_enabled`, which is `cache and get_cache_dir() is not None` — a rank whose cache directory cannot be resolved takes the else branch, which has no barrier, while its peers wait in one. Both predicates are now reduced. `needs_compile` is a global OR, so if any rank lacks the module every rank enters the branch and reaches the barrier; `disk_enabled` is a global AND, so a rank without a cache directory makes all ranks fall back together to local compiles, which is 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`. Demonstrated rather than argued. Monkeypatching the disk lookup so rank 0 sees the cache and the others do not — the shape NFS or Lustre lag produces — deadlocks this PR as it stands: at np=4 all four ranks enter the solve, none returns, mpirun times out at rc=241. With the reduction, rc=0 on four consecutive runs. This is not a defect introduced here: the gate predates the PR. What the PR changes is the exposure. Stokes was the one call site passing cache=False, so it always took the no-barrier branch; 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 filesystems #547 is about. Full ./uw test 1556 passed; the four JIT test files 14 passed. Underworld development team with AI support from Claude Code

Change the CSE preprocessing step in generate_c_source from opt-out (UW_JIT_NOCSE=1) to opt-in (UW_JIT_CSE=1), so that default code generation behavior and output are preserved verbatim. Users on memory-constrained HPC nodes running very large/complex rheologies can set UW_JIT_CSE=1 to activate CSE header shrinking. Underworld development team with AI support from Claude Code

bknight1 · 2026-08-21T00:29:08Z

Thanks @lmoresi — great catch on the rank-local predicate / collective barrier divergence hazard on shared filesystems.

Updates applied:

Collective Gate Applied (Option A): Cherry-picked aaba141 (9b8b3cd7 on this branch) to allreduce needs_compile and disk_enabled before the barrier branch. Verified locally with the jit_divergence.py reproduction script under mpirun -n 4 (rc=0) and ptest_jit_cache.py under mpirun -np 2.
CSE Made Opt-In (Default Off): To preserve original codegen behavior and exact output by default, CSE is now opt-in via UW_JIT_CSE=1 (commit 630ed8d2). Users on memory-constrained HPC nodes with large composite rheologies can activate CSE with export UW_JIT_CSE=1. Both default and CSE-enabled paths pass all test suites.
Findings 2 & 3 noted: We will track automatic harness collection of ptest_* parallel suites in a follow-up test harness cleanup PR.

Pushed and ready for CI.

Underworld development team with AI support from Claude Code

Copilot AI lite review requested due to automatic review settings August 19, 2026 07:30

bknight1 requested a review from lmoresi as a code owner August 19, 2026 07:30

Copilot started reviewing on behalf of bknight1 August 19, 2026 07:31 View session

Copilot AI reviewed Aug 19, 2026

View reviewed changes

Comment thread src/underworld3/utilities/_jitextension.py

Comment on lines +1216 to +1221

if verbose:

print(

f"JIT compile flags: {extra_compile_args}"

f"{' (from UW3_JIT_CFLAGS)' if _jit_cflags_env is not None else ' (default)'}",

flush=True,

)

lmoresi and others added 2 commits August 21, 2026 08:27

This was referenced Aug 21, 2026

adapt: snap_level_boundaries guards its allreduce on a rank-local mask (np>1 hang) #627

Open

Reach the boundary-snap reduction on every rank, and scan for the shape (#627) #628

Open

Navigation Menu

Uh oh!

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

Stokes JIT: rank-0 compile + CSE before printing + configurable flags (#547)#612
bknight1 wants to merge 3 commits into
developmentfrom
bugfix/stokes-jit-oom-cse-cache

bknight1 commented Aug 19, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

lmoresi commented Aug 19, 2026

Uh oh!

lmoresi commented Aug 19, 2026

Uh oh!

lmoresi commented Aug 19, 2026 •

edited

Loading

Uh oh!

bknight1 commented Aug 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

bknight1 commented Aug 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Verification

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

lmoresi commented Aug 19, 2026

Adversarial review

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

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

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

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

Uh oh!

lmoresi commented Aug 19, 2026

The collective, demonstrated — and a branch you can take

The deadlock is real

What the fix does

Scope

Uh oh!

lmoresi commented Aug 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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

Verifying it yourself

After applying

Uh oh!

bknight1 commented Aug 21, 2026

Updates applied:

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

bknight1 commented Aug 19, 2026 •

edited

Loading

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

lmoresi commented Aug 19, 2026 •

edited

Loading