Repeated rotated free-slip Stokes solves rebuild solver state and grow RSS until OOM

[gthyagi]

Summary

Repeated linear Stokes solves with rotated strong free-slip boundary conditions
grow resident memory substantially on every timestep. In an 8-rank Zhong et al.
(2008) A1 thermal-convection run on a 16 GiB Apple Silicon Mac, aggregate Python
RSS reached 54-59 GiB before macOS watchdog panics/restarts.

This is distinct from the constrained-multiplier scaling problem in #244. Single
rotated solves are accurate and fast; the failure appears during repeated
transient solves because the rotated path reconstructs PETSc matrices and the
fieldsplit/GAMG solver hierarchy for every call.

Reproducer

Benchmark script:

• https://github.com/gthyagi/uw3-mantle-convection-benchmarks/blob/main/benchmarks/bench_100_zhong2008_thermal_convection_A1.py

Detailed postmortem:

• https://github.com/gthyagi/uw3-mantle-convection-benchmarks/blob/main/docs/thermal_A1_memory_postmortem.md

The unsafe production case used cellsize=1/8, P2/P1, 8 MPI ranks, and rotated
free slip. The mesh had only about 48,202 Stokes unknowns, so solution-vector
storage does not explain the observed memory.

Crash evidence

Run	Aggregate Python RSS before failure	Per-rank RSS	System state
1	54.29 GiB	6.17-7.28 GiB	watchdog panic
2	59.14 GiB	7.05-7.60 GiB	watchdog panic
3	59.18 GiB	7.12-7.61 GiB	watchdog panic

At failure, the compressed-segment pool was effectively full and macOS had
created 41-42 swapfiles.

Bounded coarse probes

Serial cellsize=1/4 repeated-solve probes isolate the boundary-condition path:

BC / experiment	RSS increase	Approx. increase per subsequent solve
Rotated free slip	+102.03 MiB over initialization and 12 solves	8.54 MiB/solve
Rotated, explicitly destroy prior _rotated_freeslip_info	+88.65 MiB over 11 subsequent solves	8.06 MiB/solve
Native Nitsche	+25.57 MiB over 9 subsequent solves	2.84 MiB/solve and near plateau

Destroying the Python-visible prior result does not remove the slope. This points
below the retained result dictionary, most likely to repeated PETSc
matrix/preconditioner setup.

Current code path

For every linear solve, solve_rotated_freeslip() currently:

1. computes J and Jp;
2. copies the original operator;
3. builds Q and Q^T;
4. forms Ahat = Q J Q^T with PtAP;
5. creates the pressure mass matrix, nullspaces, fieldsplit KSP/PC, and GAMG
   hierarchy;
6. solves;
7. destroys only the immediate KSP context.

The nonlinear rotated path already demonstrates in-call reuse by updating
Ahat and the pressure block in place across Newton iterations. The linear
transient path has no persistent cross-call workspace.

Proposed fix

Add a persistent linear rotated-free-slip workspace owned by the Stokes solver:

• cache geometry-dependent rotation data (Q, Q^T, constrained rows);
• reuse/update Ahat and the pressure block in place;
• retain and reuse the fieldsplit KSP/PC and GAMG hierarchy when the operator is
  unchanged;
• refresh only the RHS when forcing changes;
• invalidate and destroy the cache when mesh topology, rotated BCs,
  discretisation, or operator coefficients change;
• retain the reaction vector required by rotated topography without retaining
  unnecessary copies of A and b.

Acceptance criteria

• A 25-50 solve coarse transient reproducer reaches a bounded post-warmup RSS
  plateau.
• The KSP/PC/GAMG objects are reused when only the RHS changes.
• Changing body force produces the correct updated velocity and pressure.
• Serial and MPI results remain numerically equivalent to the uncached
  implementation.
• Cache invalidation is covered for structural/operator changes.
• Rotated boundary traction/topography postprocessing remains correct.

Related issues

• Stokes_Constrained is slow in parallel despite low KSP iterations #244: constrained free-slip parallel setup/scaling; related solver-lifecycle
  concern but a different formulation.
• Compare free-slip methods on Kramer annulus and spherical benchmarks #320: free-slip method comparison; establishes rotated accuracy/performance
  for individual solves.
• Optimize solver settings for pure projection postprocessing #156: repeated GAMG setup in projection postprocessing; analogous setup-cost
  and memory behavior.

[gthyagi]

Implemented and pushed a focused fix on
gthyagi:feature/mantle-convection-benchmarks:

• UW3 commit: gthyagi@5c4b0716
• Detailed benchmark report:
  https://github.com/gthyagi/uw3-mantle-convection-benchmarks/blob/main/docs/rotated_freeslip_memory_fix.md

Root cause refined

There were two repeated allocations:

1. _residual_is_nonlinear() assembled/copied two trial Jacobians on every
   rotated solve.
2. The linear rotated path rebuilt Q, Ahat, pressure mass, nullspaces,
   fieldsplit KSP/PC, and GAMG on every timestep.

Stage isolation on a serial cellsize=1/4 spherical shell measured
+76.77 MiB over 12 nonlinearity probes and +117.23 MiB over 12 direct
rotated solves.

Fix

The Stokes solver now owns a persistent linear rotated workspace. RHS-only
field changes skip Jacobian assembly and reuse Q, Ahat, pressure mass,
nullspaces, KSP/PC, and GAMG. Constitutive coefficient
MeshVariable._state changes and explicit time= solves refresh and compare
the operator. Solver/DM/BC rebuilds explicitly destroy the workspace.

Two intermediate approaches were rejected:

• exact Mat.equal() was too strict for distributed assembly roundoff;
• tolerant CSR copying reintroduced +82.02 MiB over ten serial solves.

The final changed-operator check uses two persistent deterministic
matrix-vector probes and creates no per-step matrix/CSR copies.

Before / after

Run	Repeated solves	RSS result	Accuracy
Pre-fix serial spherical	12	+102.03 MiB	baseline
Post-fix serial spherical	30	peak +1.91 MiB; allocator later released memory	RHS scaling error <1e-9
Post-fix 8-rank coarse box	30	+11.84 MiB aggregate	<8e-14
Post-fix 8-rank spherical	10	+12.78 MiB aggregate	<3e-10
Doubled viscosity field	1 refresh	old KSP destroyed; no reuse	expected half-scale velocity, error 8.10e-11

Serial and eight-rank rotated velocity/leakage, reaction-topography,
boundary-traction, and dynamic-topography tests pass. The worktree rebuild
also passes.

MPI validation must use the Pixi environment mpirun; the shell's Homebrew
Open MPI 5.0.7 launcher mismatched the UW3 Open MPI 5.0.10 runtime and stalled
before model execution. No FI_PROVIDER=tcp workaround was used.

The full 8-rank cellsize=1/8 A1 continuation remains prohibited on the
16 GiB Mac. The next production validation should resume checkpoint step 500
on a larger-memory node with per-rank RSS logging and a hard abort threshold.