Reuse the rotated free-slip solver workspace across solves — reworked with a constants-aware key (#417, supersedes #418)

[lmoresi]

Reuse the rotated free-slip solver workspace across solves (rework of #418, fixes #417)

This is the landing rework of #418 (gthyagi's "Reuse linear rotated free-slip
solver workspace", bec76bb), which went CONFLICTING after the rotated solve
loop was rewritten (#437, #458, #465/#471, #469, #493, #500, #502, #530/#534),
and which the 2026-07-27 adversarial review held on one structural finding.
gthyagi's commit is cherry-picked with his authorship preserved; the caching is
re-derived at the current seam rather than force-fitting the old text.

What was ported, and where it lives now

The original PR split a linear path from a nonlinear one and cached the linear
path's workspace whole — operator, KSP, and (implicitly) the solution state —
keyed on MeshVariable._state counters, with a matrix probe as safety net.
Since the fork, rotated_bc.solve_rotated_freeslip became ONE manual
Newton/Picard loop for linear and nonlinear models alike, which already reuses
its operator and KSP context between its own iterations. The port extends
exactly that in-loop pattern across solves, split by what each piece
actually depends on:

• Geometry tier (reused on a matching geometry signature: boundary specs,
  fault registration, DM identity): the rotation Q/Qt, constrained normal
  rows, fault pair blocks, custom-FMG prolongation.
• Structure tier (always correct — objects reused, values reassembled in
  place): Ahat via ptap-with-result, the Schur pmat via
  createSubMatrix-with-submat, the fieldsplit KSP/PC via a setOperators
  poke. This is the identical operation sequence the Newton loop performs
  between iterations, so it carries the production-validated risk profile.
• Iteration-0 fast path (skips Jacobian assembly, ptap and PCSetUp): only
  when the operator key proves nothing operator-relevant changed AND the last
  solve on this workspace behaved linearly (converged in ≤ 1 increment — a
  self-measured hint, no up-front nonlinearity probe).

Nothing about the rotated BC's discrete equations changed. All current result
keys (rnorm/rnorm0 for the solve report, the #534 tolerance/report
plumbing, the #502 fault keys) are preserved; rotation_reused and
workspace_reused are added.

Teardown hooks: _reset_rotated_solver_cache() runs from _reset() and from
the _build() full-rebuild branch before the SNES/DM are destroyed
(mesh.deform funnels there). The in-place rewire fast path keeps the cache;
its new kernels are caught by the JIT-key half of the operator signature.
The cache is forfeited outright for direct-LU, prescribed-datum and fault
interface-law solves. Stokes.solve gains the time= pass-through from the
original PR; an explicit time vetoes the fast path (petsc_t bypasses every
counter and constant).

The blind-spot fix (the finding that held the PR)

The review probed that a 2× viscosity ramp via a rampable UWexpression
constant (the #416 contract: value changes bump NO state counter) returned a
bit-identical stale solution flagged workspace_reused=True, and that the
same "unchanged" verdict short-circuited the matrix-probe safety net that was
supposed to catch it. Both halves are fixed, structurally:

1. The key covers the kernels' constants. The operator signature now
   includes the packed constants[] values the compiled kernels will actually
   assemble with, plus the JIT bundle key. If the manifest or the coefficient
   enumeration cannot be read, the fast path is forfeited — correctness first.
   (This deliberately over-invalidates on RHS-only constant changes;
   RHS-only field changes — the production temperature pattern — still ride
   the fast path.)
2. The safety net is independent of the key it guards. In the reworked
   loop there is no separate one-shot linear path to poison: the loop measures
   the TRUE residual (fresh kernels, current constants) at every iterate,
   declares convergence only on that, and always reassembles from iteration 1
   on. A wrong fast-path verdict therefore costs one extra increment — it
   cannot return a stale solution. Measured on the naive-key port (negative
   control, house rule): the ramp probe showed reused=True with the WRONG
   flag but the CORRECT halved velocity (4.1e-11 vs a fresh control) — the
   old code returned the stale field bit-identically. The key fix makes the
   flag honest and removes the wasted work.
3. The verdict is collective. State counters follow rank-local writes, so
   the match verdict is allgathered and must be unanimous before it gates any
   collective PETSc call (a rank-divergent verdict is a hang, not a wrong
   answer).

The matrix probe from the original PR is dropped: its role is subsumed by the
loop's exact-residual verification, which is stronger (it validates the
solution, not two matvecs) and cannot be disabled by its own trigger.

Fault composition (#502)

solve_with_fault drives the same rotated loop, so the decision is explicit
and tested:

• Frictionless pair blocks cache — the coincident-node pairing and fault
  normals live in Q and are geometry, keyed by the fault registration in
  the geometry signature. Warm repeat reuses the rotation; a cold re-solve
  rides the iteration-0 fast path; both match a fresh-solver control.
• Interface-law solvers opt out entirely (cache_allowed excludes them):
  the interface tangent is reassembled per iterate at the current slip rates
  and the reaction-fed normal stress is Picard-lagged solver state — neither
  is keyable registration state.

Verification

All serial runs sequential on the pr418-rework worktree (amr-dev env).

Probe	Result
Ramp probe (2× UWexpression constant, warm solve, fresh control)	fast path correctly defeated (reused=False); error vs fresh control 4.8e-08–4.1e-11 across configs; fail-before validated: naive key reports reused=True
Reuse actually happens	34/35 solves reused on the RHS-change pattern; warm solve 1.000 s vs 2.10 s fresh-solver control = 2.10× speedup (gthyagi measured 1.92× in production); same Q/Ahat/KSP handles asserted
Deform between solves	workspace invalidated (rotation_reused=False), post-deform solve matches fresh control < 1e-6 (test)
Memory (#417)	35 solves one process: RSS 347.7 → 351.3 MiB over solves 5–34, +0.12 MiB/solve (was: unbounded growth to OOM); bound used: < 1 MiB/solve with margin
np2	4-solve sequence (fresh, RHS×2, ramp): reuse verdicts rank-identical at every solve [(F,F),(T,T),(T,T),(F,T)]×2 ranks; RHS scaling 9.7e-08, ramp scaling 4.8e-08
test_1018_rotated_freeslip	22 passed (19 existing + reuse + ramp + deform)
Fault set 0845–0848	30 passed (incl. new test_fault_repeat_solve_composes_with_workspace_reuse)
Full gate -m "level_1 and tier_a" (–test_0050)	0 failed

Production evidence (gthyagi)

The original mechanism carries a substantial production validation trail that
this rework inherits and that deserves explicit credit: 310 guarded Zhong A1
steps (cellsize=1/8, 8 ranks, checkpoints 500→810 in guarded 50-step
segments, UW_MEMPROBE=1, 0.25 s RSS sampling, hard memory/time stops that
never fired), flat warm RSS with periodic collective releases, a continuous
physical trajectory (Vrms 57.83→57.97, Nu_surf 2.694→2.706 over 810 steps),
and a ~1.92× per-step speedup over the pre-fix rate. Since the loop internals
moved under the cache in this rework, gthyagi may wish to re-run an A1
segment on the reworked branch (a 50-step guarded restart from an existing
checkpoint would do) before merge.

Commits

1. 3eba6aa4 — Reuse the rotated free-slip solver workspace across solves
   (Repeated rotated free-slip Stokes solves rebuild solver state and grow RSS until OOM #417) — author: Tyagi, cherry-pick of bec76bb re-derived at the
   current seam.
2. 78338211 — Close the rotated-workspace blind spot: rampable constants
   join the key, the verdict goes collective.
3. f8c2b726 — Fault contact composes with the rotated workspace cache:
   pair blocks cache, interface laws opt out.

Underworld development team with AI support from Claude Code

[Copilot]

Pull request overview

This PR introduces a cross-solve cache for the rotated strong free-slip Stokes solve, reusing the PETSc rotation/operator/KSP workspace across repeated solves to reduce rebuild cost and bound memory growth, while adding a constants-aware invalidation key (to catch rampable UWexpression constants) and a collective reuse verdict for MPI safety.

Changes:

• Add a persistent rotated-workspace cache (geometry + structure tiers) and an iteration-0 “fast path” that can skip Jacobian/PtAP/PCSetUp when the operator is provably unchanged.
• Extend invalidation to include packed constants[] values + JIT bundle key (and make reuse verdict collective across ranks); add workspace_reused / rotation_reused reporting.
• Add time= pass-through on Stokes.solve() and veto the fast path for explicit time evaluation; add focused regression tests and a design note.

Reviewed changes

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

Show a summary per file

File	Description
tests/test_1018_rotated_freeslip.py	Adds regression coverage for workspace reuse, constants-ramp invalidation, and deform invalidation.
tests/test_0846_fault_contact.py	Adds coverage that fault contact (frictionless vs interface-law) composes correctly with the rotated workspace cache.
src/underworld3/utilities/rotated_bc.py	Implements cross-solve workspace caching, constants-aware operator signature, and collective reuse gating.
src/underworld3/systems/solvers.py	Exposes time= on solve() and forwards it to the underlying solver call sites.
src/underworld3/cython/petsc_generic_snes_solvers.pyx	Adds cache teardown hooks and vetoes the rotated fast path when time= is provided.
docs/developer/design/ROTATED_FREESLIP_LINEAR_REUSE.md	Documents the reuse tiers, safety net, invalidation rules, and validation approach.

---

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

[lmoresi]

Adversarial review — PR #543: rotated free-slip workspace reuse (rework of #418, fixes #417)

Reviewed independently of the implementer at f8c2b72 (3 commits on 689523f),
isolated worktree r543-review, amr-dev env, uw.__file__ verified into the
worktree's own site-packages, probes serial and sequential, np2 via the
worktree's own mpirun. Probe scripts: session scratchpad
probe1_staleness.py … probe6d_true_window.py.

The structural claim we held #418 on is now real: the safety net
(true-residual convergence + always-reassemble-from-iteration-1) is
independent of the key it guards, and we could not make the branch return a
wrong answer — not with a deliberately corrupted key, not with BC/rheology/
geometry mutation, not with a fault-normal change, not at np2. Every finding
below is an honesty/latency wart or a hardening suggestion, none blocks.

MERGE-BLOCKERS

None. One condition we would attach: gthyagi's A1 re-validation segment on
this branch (the PR itself proposes it) — the loop internals moved under
the cache since the 310-step production trail was measured; a 50-step
guarded restart from an existing checkpoint should land in the PR thread.

Findings (severity-ranked)

m1 (minor) — solver configuration latches into the cached KSP

_solve_rotated_iterative reads solver.tolerance (and the FMG/GAMG
choice, smoother variant) only at ctx creation; pre-PR the ctx died each
solve, now it persists until a geometry/BC teardown. Tightening
stokes.tolerance between solves leaves the inner Krylov at the OLD rtol:
the outer loop re-reads the new tolerance and still converges on the true
residual — measured (probe 1-D) 1e-4→1e-10 honoured, final rnorm 4.5e-13,
err vs a tight fresh control 1.8e-11 — but at extra increments, with the
Citcom inner-margin discipline (inner = 0.1×tol) silently violated. Same
class: a set_custom_fmg registration after the first solve is ignored
until the cache dies. Correctness unaffected. Suggest a config tuple
(tolerance, custom-mg identity, smoother variant) in the signature, or one
sentence in the design doc.

m2 (minor) — persist can store a key describing values Ahat does not hold

If the operator key mismatches (poison correctly fired) but the loop exits
at iteration 0 without ever assembling (a warm start whose TRUE residual
already passes under the changed constants — reachable when operator and
RHS co-scale), the persist block stores the NEW operator_sig against the
UNREFRESHED Ahat, violating its own comment's invariant ("the stored key
describes the operator values Ahat now holds"). Bounded by the safety net
(+1 increment, lying flag, never a wrong answer). Hardening: persist
operator_sig only if assemble ran at least once this solve, else keep
it poisoned (None).

m3 (hardening) — coefficient-variable enumeration reuse is safe by coincidence, not construction

The enumeration (cache["coefficient_variables"]) is reused whenever the
GEOMETRY signature matches (~line 866), yet it encodes an OPERATOR
property (which variables the kernels read). We hunted the stale window
hard (probes 6/6b/6c/6d): every dependency-changing route we could build —
Parameters constant→field, field→field with the new field already in the
aux layout — forced a full _build teardown before the solve (the cache
was already None; probe 6c shows _build() alone kills it), so the
enumeration was always re-derived and the ramp of the new field was always
detected (final err 3.6e-11). But nothing in the cache logic ENFORCES
that: the safety is inherited from _build's aux-layout invalidation
behaviour, which this module does not control. One-line hardening:
re-derive the enumeration whenever the JIT half of operator_sig
mismatches, not only on geometry change. (Worst case if the coincidence
ever breaks: permanent lying workspace_reused=True + one wasted
increment per solve — the safety net still guarantees the answer.)

m4 (nit) — destroy-then-null ordering in teardown

_destroy_rotated_linear_cache(cache) runs before
solver._rotated_linear_cache = None (rotated_bc.py lines 895–896 and the
pyx hook). An exception mid-destroy leaves the solver holding a
half-destroyed dict. Null-first is the safer order. (Inside the
destructors the #492 release discipline is right: destroyed entries are
None'd, double-calls are no-ops — verified by repeated
_reset_rotated_solver_cache()/_reset() in probe 3.)

Observation (positive) — rank-divergent state counters are unreachable via the sanctioned API

Our adversarial np2 variant (rank-0-only T.data write, then solve)
deadlocked INSIDE the data-access layer, not in the PR: the stack sample
shows the .data write hook itself calling
DM.createSubDM → DMClone → MPI_Allreduce while rank 1 sat at the probe's
own barrier. A sanctioned field write is already collective, so the
counter divergence the PR's allgather defends against cannot be produced
through .data at all. The allgather is belt-and-braces — cheap and
correct; keep it (.array/swarm/exception paths remain conceivable).

Observation — in practice the cache does not survive constitutive rewires

The PR text says the in-place rewire keeps the cache and the JIT key
catches the new kernels. In every constitutive-change flow we probed,
_build ended up on the full-rebuild branch (aux-layout hash) and the
cache was torn down before the solve — over-invalidation, the safe
direction, and irrelevant to the target time-stepping pattern (no rewires
there: 34/35 solves reused per the PR, 25/30 in our probe 3). The JIT-key
half of the signature is defensive depth today, not the working
invalidation path. No action; recorded so nobody later "optimises" the
teardown away believing the JIT key alone carries rewires.

Attack log

1. Staleness beyond the ramp (probe 1, serial, Q2/P1disc):

Change between solves	invalidates / forfeits?	err vs fresh control
second solver, same mesh	caches independent (distinct handles); interleaved drift 0.0	0.0; ¼-viscosity ratio 0.0
add a Dirichlet wall	full rebuild, rotation_reused=False	0.0
add a second rotated boundary	full rebuild + boundaries in geometry sig	0.0
tolerance tightened 1e-4→1e-10	honoured on the true residual (m1)	1.8e-11
nonlinear rheology switched on, η(v)	fast path defeated; linear_hint withdrawn after the 2-increment solve	5.7e-16
bodyforce expression REPLACED	rebuild/rewire; no stale answer	3.4e-10
rampable-constant 2× ramp (THE #418 killer)	their test: fast path defeated, halved velocity exact	(in test, passes)
direct mesh.X.coords writes	not a live path: the property returns scaled/wrapped arrays and writes never reach the PETSc coordinate vec; deform is the sanctioned route and tears the cache down (their test + probe 3)	n/a

2. Iteration-0 fast path (probe 2-G): after a fast-path solve, poisoned
the key to force a slow-path in-place reassembly of the SAME Mat: every
constrained row identical, worst abs diff 0.0, |Ahat| unchanged to all
digits. The skipped zeroRowsColumns/PCSetUp is genuinely redundant.
Poisoning order verified in code AND by probe 2-I (hand-poison, ramp, no
fast path, err 3.5e-11): an interrupt in the poison→refresh window can only
under-reuse.

3. Collective discipline (probe 4, np2): code audit — every verdict
input is registration state, packed-constants (global), or goes through
the allgather; exception paths inside both signature helpers swallow to
None symmetrically (→ forfeit). Probe: fresh / RHS×2 / rank-VARYING RHS
values / constant ramp — flags allgathered at every solve, rank-identical:
[(F,F),(T,T),(T,T),(T,F)], ramp defeated on both ranks, no hang, final
err vs fresh control 1.2e-10. (The rank-asymmetric-write variant is
unreachable — see the positive observation above.)

4. The structural safety net (probe 2-H): monkeypatched
_operator_constants_signature to a frozen value (equivalent to reverting
the blind-spot fix), ramped 2×: workspace_reused=True (the lie fires),
halved-velocity error 1.7e-11, exactly +1 increment (2 vs 1). The net
that replaced the dropped matrix probe does what the PR claims — a
corrupted key costs work, never an answer.

5. Fault composition (their test + probes 5/5b): fault set passes (30).
Changing the fault's normal= override between solves is caught by the
geometry signature (rotation_reused=False); a COLD post-change solve
matches a fresh override control bit-identically (0.0, floor between
two fresh controls also 0.0). For the record: a WARM start straight across
the normal change landed 6.3e-2 off while reporting converged — with the
cache already destroyed (pre-PR behavior identical), in a deliberately
flux-incompatible override configuration, and with the loop's
mass-conservation gauge warning printed loudly. Pre-existing rotated-loop
warm-start semantics, not a cache defect.

6. Teardown (probe 3): 30 solves with 5 interleaved mesh.deform
teardown/rebuild cycles: 25/30 fast-pathed, ru_maxrss slope
0.79 MiB/solve including the five full rebuilds (PR's pure-reuse figure
0.12; high-water RSS overstates); repeated _reset() /
_reset_rotated_solver_cache() idempotent, no double-destroy; σ_nn
traction recovery off a fast-path solve doubles with the load (1.3e-10) —
verified in code that the recoveries read only reaction+boundaries, so
the "result dict survives cache teardown" contract holds. Custom-FMG
prolongation deref-not-destroy is correct (shared coarse Mats; the rotated
Pfine/Qv fall to petsc4py GC).

7. Tests & negative control: test_1018 22 passed; fault set
0845/0846/0847/0848 30 passed. House-rule negative control: swapped
commit 1's rotated_bc.py (naive key) into the INSTALLED copy → the ramp
test fails exactly at the lying-flag assertion, restored and re-verified.
The regression test has teeth, and its own internal negative control (fast
path proven armed before the ramp assertion) fires. Full gate
-m "level_1 and tier_a" (–test_0050): 599 passed, 0 failed, 17 skipped
(MPI-gated), 1 xfailed, in 8:13.

Notes for the record

• PR base is 2 commits behind origin/development; the divergence touches
  the pyx by 6 lines (Boundary flux on degree-3 traces: per-slot DOF identity, true edge-node coordinates, and the consistent line mass (#459) #537), disjoint from this change — expect a clean
  merge.
• Commit 1 carries gthyagi's authorship, as it should.
• time= semantics: petsc_t persists on the DM after a time= solve; a
  later no-time solve fast-paths against the operator assembled at that
  time, which is CONSISTENT (operator and residual agree). The veto covers
  the transition solve — the only place it must.
• _pack_constants packs 0.0 for an un-floatable constant, but signature
  and kernels share the same packing, so the key cannot desync from the
  assembly.
• Suggest (non-blocking): an np2 ptest_ for the reuse verdicts — the
  parallel evidence currently lives only in probe scripts, and
  tests/parallel/test_1064/1066 predate the cache.

Verdict

APPROVE (with the A1 re-run segment requested of gthyagi, and m2/m3 as
cheap hardenings here or in a follow-up). The #418 failure class — a silent
stale answer — is structurally closed: under deliberate key corruption the
branch still produced the correct field at +1 increment, and every
mutation we threw at the key (BCs, rheology, bodyforce, fault normals,
tolerances, np2 rank-varying state) either invalidated, forfeited, or was
rescued by the always-reassemble net. Memory (#417) stays bounded through
deform/rebuild cycles.

[lmoresi]

Response commit 2f17743 takes m2 (operator_sig persisted only when Ahat provably holds the keyed values) and m4 (null-before-destroy at both teardown sites); test_1018 22/22 after rebuild. m1 (config latched into the cached KSP) and m3 (coefficient-enumeration reuse safe by teardown coincidence, not construction) are tracked as a follow-up issue. The attached condition stands: @gthyagi — when you have a slot, a 50-step guarded A1 restart segment on this branch would re-anchor your 310-step production trail against the reworked loop; the memory numbers here (+0.12 MiB/solve over 35 solves, 2.10x warm-solve speedup) match your original measurements.

Underworld development team with AI support from Claude Code