Make every MPI rank compile the same generated C, rather than refuse

[lmoresi]

Closes nothing — #752 stays open. This stops the flaky abort; the non-determinism it is a symptom of is still there.

What was happening

The JIT derives the module name and the C symbol prefix from a hash of the generated source, so ranks that generate different text build disjoint artefacts and the rank-0-compiles/others-load protocol breaks. Agreement was required, and a mismatch raised.

It fires in practice. A Stokes solve with a power-law transversely isotropic viscosity under the consistent tangent trips it in roughly half of np=2 runs — measured 2 aborts in 4 on feature/discrete-adjoint. The abort lands in the forward solve, so it looks like a solver failure.

RuntimeError: JIT C-source hash differs across MPI ranks: {'471c0c17f9aba995', '5e640ad7875d8e9f'}.

Why adopting one rank's source is sound

Two measurements, not an argument:

• The sources differ only in the order of factors in commutative products. Diffing both ranks' output on a failing run: 12 of 225 lines differ, every one with an identical token multiset and identical length. Every rank's source is a correct kernel for the same equation.
• The solver's symbolic blocks agree. Hashing them at np=2 on runs that then abort: c_tensor, flux, _u_F1, _uu_G0–_uu_G3, _up_G0, _pu_G1 — all identical. What differs is produced inside generate_c_source, not handed to it.

So the disagreement is over which of several correct spellings to compile. Rank 0's is taken, every rank rehashes from it, and the invariant that matters is restored: one source, one hash, one module.

Measured on the reproducing branch: the repair fired in 4 of 6 runs and all 6 succeeded (baseline 2 aborts in 4).

This is a repair, not a fix

It warns, naming #752, rather than papering over the upstream bug silently. The common path is untouched — ranks that already agree keep their own objects and pay one allgather of a 16-character string.

Two dead ends, recorded so they are not retried

Both were implemented and measured before being discarded:

1. The printer is innocent. The C printers do not use StrPrinter._print_Mul — C89CodePrinter delegates to CodePrinter._print_Mul, which has no stored-order fast path and does sort (printer.order is None). Forcing that sort produces byte-identical C across all five kernel bundles, and a run still aborted.
2. It is not commutative ordering. Recursively sorting every Mul/Add argument by default_sort_key before hashing still leaves the ranks differing (2 of 3 runs).

The remaining lead for #752: the lowering phases between the solver's pointwise functions and printer.doprint — _reveal_constants, unwrap, the constants_subs_map xreplace, to_matrix, the free_syms walk, cse. Several consume .atoms(), which is a set; _stable_sorted already exists in that file.

Tests

tests/parallel/test_0022_jit_rank_source_agreement.py. The repair is extracted as _agree_source_across_ranks so it can be tested directly — reproducing a real disagreement means reproducing a non-deterministic bug at about one run in two, which is not a test: it would pass half the time with the repair removed. The tests construct the disagreement themselves, so they keep working once the lowering is made deterministic.

• ranks that disagree all adopt rank 0's source, and get the hash that text actually produces (not merely hashes that match each other);
• ranks that already agree are left exactly as they were — assert agreed is same — so the repair cannot quietly become an unconditional broadcast;
• the end-to-end case that trips it completes.

Writing that first test immediately caught a real bug in the extraction: hashlib is imported inside generate_c_source, not at module level, so the helper raised NameError on the path that matters.

Serial band tests/test_00[0-4]*py: 141 passed. Parallel at np=2 with test_1069: 10 passed, 1 skipped.

Note on scope

This is pre-existing on development, not something arriving with a feature branch. Controlled measurement — one worktree, one environment, one set of caches, branch as the only variable, counting repair firings rather than aborts:

arm	repair fired
development	2 / 10
feature/discrete-adjoint	2 / 10

An earlier uncontrolled comparison (0/12 vs 2/4, in two different checkouts) suggested a branch regression. It was wrong: the rate is environment-sensitive, so samples from different worktrees are not comparable. Any future comparison here needs the controlled form.

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.