JIT: generate kernels from the shared graph of named quantities, old route kept switchable (#823, tier 2)

[lmoresi]

Tier 2 of #823. Closes #752.

What changes

The JIT no longer expands every named sub-expression (UWexpression) into one tree before it differentiates and prints. src/underworld3/utilities/_jit_graph.py:

• Lowers each non-constant atom to a node: an applied function of the leaves its value depends on (field values and gradients, coordinates, constant atoms).
• Differentiates through nodes by SymPy's own chain rule. A node's fdiff is the partial derivative of its body, taken against real placeholders and itself a node.
• Emits one C temporary per distinct computation, ordered and merged by a hash of the C each computes, with every leaf written as the C the kernel reads.

Two routes, switchable. The graph is the default. The JIT before tier 2 is kept beside it as the "expanded" route: a fallback, and a reference for telling a JIT regression from a fault in the model. If a model misbehaves on both routes, the model is likely wrong.

• Process default: uw.use_jit_route("expanded"), or UW_JIT_ROUTE=expanded in the environment.
• One solver: solver.jit_route = "expanded". The solver rebuilds its kernels at the next solve and keeps its state; None returns to the default.
• Exactness: the expanded route is development's code, moved and not edited (_expanded_equations, _extract_constants, _reveal_constants, the class patching, the UW_JIT_CSE option). Its C is byte-identical to development's on all six fixtures (linear, power law, box, VEP, TI, notch).
• Cache: the two routes never share a cache entry for a law with a named quantity, because their C differs. For a law without one they generate the same C, so they share a module, which is the same function.

On the graph route:

• An explicit map from leaf to C, built for each compile (_leaf_spellings), replaces the C names patched onto the field classes.
• The two unwrappers expand nodes, for code that evaluates a lowered Jacobian block.
• The unused prepare_for_cache_key and _createext are removed.

Design note: docs/developer/design/jit-shared-graph-codegen.md (measurements, risks and the test that closes each, decisions). Also updated:

• jit-cache.md, whose MPI section was already stale;
• expressions-functions.md and UW3_Developers_MathematicalObjects.md;
• jacobian-consistent-tangent.md;
• the plasticity capability guide.

Measurements

Against development with tier 1 and -fno-math-errno (#834, PR #835). Mac, single runs on a loaded machine.

	notch	VEP	box	power law	TI	linear
pointwise setup, s	16.6 / 1.95	3.3 / 2.1	2.5 / 1.7	1.5 / 1.5	1.6 / 1.6	—
generated C	4.0 MB / 22 KB	370 / 35 KB	104 / 16 KB	28 / 14 KB	37 / 25 KB	byte-identical
Jacobian assembly, ms	184 / 122	2.37 / 2.32	1.52 / 1.47	1.48 / 1.40	2.42 / 2.46	equal
solver path	chaotic (see below)	identical	identical	identical	identical	identical

On the notch, the Jacobian callbacks cost 2.67 → 0.49 µs per quadrature point. The rest of each assembly is PETSc's finite-element machinery, which is the same on both routes.

Linux (gcc 14.3): the same ratios hold. gcc's default -fmath-errno accounts for part of the tree's callback cost on some laws (hence #834), but none of it on the VEP.

Agreement:

• The assembled residual and Jacobian agree to ≤5e-16 of each row's largest entry, at rest and at solved states, for the Newton, Picard and continuation tangents.
• The notch's Newton iteration count is a property of the problem, not the route: over 17 round-off-sized perturbations of the body force, tier 1 took 60–111 iterations and once failed to converge in 300; tier 2 took 43–86.

Behaviour that changes

• Round-off on the default path: the default (Picard) kernels and the residuals of any law with a named non-constant quantity change at round-off, because the evaluation order is new (decided 2026-10-07). A law without one emits byte-identical C.
• Cache: every cached JIT module recompiles once.
• Fields from another mesh: a field the compile does not own is now an unconvertible-symbol error. Before, it silently printed whatever slot the last compile had patched onto its class.
• Constants by structure: constancy is decided by structure rather than by current value. (1 + T**2)**(-m) + 1 at m = 0 used to bank as one constants[] slot and raise when m ramped. m now has its own slot and ramps without a recompile, which fixes the "rampable constant in exponent position does not ramp" report; test_0104 is updated to pin it.
• Guard at rest: a power law on a named invariant (η = ε̇^(1/n−1), ε̇ = √g) had a NaN Newton flux at rest on tier 1: SymPy merged the powers past the sqrt guard. On tier 2 it is finite.
• Printing Jacobian blocks: stokes._uu_G3 and friends hold node applications, which print as their quantity's name. unwrap_expression and evaluate expand them; lambdify needs unwrap_expression first.
• Verbose output: the "Processing JIT" line prints the lowered kernel.
• -fp_trap: temporaries are computed even on an untaken Piecewise branch. Values are unaffected; under PETSc -fp_trap such a branch could trap, and nothing in the repository traps.

Tests

• test_0026_jit_route_switch.py covers the switch. Both routes give the same nonlinear and linear iteration counts and the same solution on a viscoplastic box, Newton and Picard.
• The expanded route keeps its own tests: test_0022's tree guard test, test_0104's collapse-and-raise behaviour, and test_0105's hash-seed sweep, which now runs on both routes.
• test_0024_jit_graph_lowering.py (15 tests, pinned to the graph route) covers the design note's risks. Every test was checked against a deliberate break of the lowering, and each of the second review's fixes has a test shown to fail first.
• test_0105 adds a hash-seed sweep over a Newton viscoplastic law that emits temporaries.
• test_0022 keeps its guard identity test, now pointed at the graph's guard function.
• test_0023, test_0103 and test_0104 are updated for the removed internals and the structural constancy rule.
• The full serial suite (scripts/test.sh batches, not tier C) passed before the routes were made switchable: 2,910 passed, 0 failed. Runs on both routes are in progress and will be posted. It also passed after step 4 alone (3,082) and with the graph behind a switch (3,023, once two fault-network defects were fixed).

Review

Two adversarial reviews ran before this PR. Neither found a blocker; their findings and what was done are in comments below.

Open questions for the maintainer

1. Should each C temporary carry its quantity's name as a comment? It makes kernels readable; a rename would then recompile.
2. Which tier for new tests: test_0024 is tier B. TESTING-RELIABILITY-SYSTEM.md says tier C, and recent practice is tier A.
3. Should the adjoint branches' _peel_except adopt the nodes when they rebase?

Merge PR #835 first, so that this is measured against the improved JIT.

Underworld development team with AI support from Claude Code

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[lmoresi]

Adversarial review 1 of 2: mechanism and correctness (before opening), and what we did

No blocker.

Should-fix, all fixed, each with a test in test_0024 shown failing first:

• S1. Coordinate derivative. A mesh.X coordinate beside a field in one named quantity lost its explicit partial derivative. The per-body cse rebuilt the UWCoordinate with a cloned coordinate system, so it was no longer equal to mesh.N.x. Solver Jacobians differentiate by fields only, so no tangent was wrong; coordinate derivatives of lowered blocks were. Fix: leaves and child nodes are hidden behind placeholders while cse runs, and lowering finds UWCoordinates by type.
• S2. Matrix and vector atoms. A matrix- or vector-valued atom became a scalar node (isinstance(Matrix, Expr) is true) and failed to print. Fix: such bodies are expanded in place (_is_scalar).
• S3. Constancy cost. Constancy was decided by a complete unwrap per atom (_is_truly_constant), exponential in the nesting depth while the C stays small: 4.4 s to lower a 12-layer chain. Fix: decided bottom-up on the graph. The rule is structural; the consequence for test_0104 is in the PR body.

Nits:

• N1. Recursion. About 23 Python frames per named layer. Newton differentiation of a 48-layer chain hits the recursion limit; tier 1 failed between 64 and 96 layers. Noted, not changed.
• N2. Number symbols. A number symbol in a temporary (EulerGamma) wrote its declaration into the initialiser. Fixed: declarations are hoisted, once each.
• N3. Unreachable refusal. The // Not supported in C: check cannot fire on SymPy 1.14, whose printer raises instead. Unchanged from tier 1; documented.
• N4. Untaken branches. Temporaries are computed even on an untaken Piecewise branch, which matters only under -fp_trap. Recorded in the design note's risk table.

Attacks that found nothing:

• Derivatives through Min, Max, Piecewise, Abs and Heaviside, by value and gradient, through diff_wrt_field and sympy.diff: worst relative error 0. Second derivatives 5.6e-15; sensitivities to constants exact.
• Lost cancellations across names: ViscoPlastic (sqrt, harmonic, min, powermean, floors), VEP (all yield modes, anchor) and TI, at rest and at random states. Residual, Picard and Newton agree to 2.6e-16 or better, and the graph is never NaN where the tree is finite.
• Determinism: the header hash is identical under four hash seeds, after a 37-object preamble, and across ranks at np = 2.
• Constants: same-named slots stay distinct, ramping keeps the cache key, units-active values are non-dimensional, and no constant is baked.
• The leaf map: auxiliary offsets for scalar, vector and symmetric-tensor variables; the integration-point gradient is refused; a field on another mesh is refused; petsc_n, petsc_t and mesh.t are spelled correctly.
• DiracDelta dropping, analytic _ccode functions with their libraries, pickle and deepcopy of node-bearing blocks, N-d arrays.

[lmoresi]

Adversarial review 2 of 2: tests, CI, docs, Charter, behaviour (before opening), and what we did

No blocker. Every should-fix is fixed.

1. Seed sweep. No suite test swept hash seeds over a law that emits temporaries; test_0105's fixture emitted none. Added a Newton viscoplastic sweep that asserts temporaries were emitted.
2. Design note and comments. They promised tests that did not exist and were inconsistent on JIT C-source generation is non-deterministic across MPI ranks (flaky hard abort at np>1) #752.
   • Added: a power law in the guarded-lowering test, a twelve-layer law, a coordinate-spelling unit test, the guard identity test.
   • Risk rows now name the test that closes each.
   • JIT C-source generation is non-deterministic across MPI ranks (flaky hard abort at np>1) #752 is stated one way, in the note, the getext comment and jit-cache.md.
   • What was not measured is listed: 3-D, Integral/BdIntegral, set_jacobian_F1_source, SolCx with Nitsche.
3. Guard test. The deleted guard test's protection is restored against _jit_graph.guard_half_integer_powers, by srepr on eight laws.
4. Stale docs. expressions-functions.md and UW3_Developers_MathematicalObjects.md described the two-phase unwrap; rewritten.
5. Solver docstrings. The consistent_jacobian docstring promised bit-identical Picard kernels; corrected, together with the comments and the pointers to a design doc that does not exist.
6. Opaque names. Lowered blocks printed as hashed names (N06e56...). Nodes now print as their quantity's name; srepr and the C are unchanged.
7. _JITConstant. Its docstring and test_0103 claimed rank-safety it no longer provides; both now say what it does.
8. Dead code. Removed: unreachable fallbacks in generate_c_source, nodes_exist, an unused import, and the _ccodestr constancy branch. One manifest helper replaces three copies.
9. Allowlist. A stale entry in deprecated_pattern_allowlist.txt is removed.

Nits fixed: the design-note references (uwt_0, i_jac, "memoised once per lowering", line counts), session-script docstrings, uw.Params in the compare scripts, a docstring on test_second_derivatives, and a constant-law test with named constants.

Behaviour changes the review asked to be stated are in the PR body.

Checks that found nothing:

• Mutations: every one tried was caught by the intended test.
• No tautological test.
• Reach: check_test_coverage.py passes, and test_0024 gates as tier B.
• No leftover references to the removed names.
• describe(), view(), transcripts, pickle and evaluate() all handle node-bearing blocks.
• No memory accumulation over 40 rebuilds.
• The session scripts run in both builds.