PerturbedEquilibrium & KineticForces: parallelize serial workflow loops (coupling, NTV torque) + reduce allocations

[logan-nc]

PerturbedEquilibrium & KineticForces: parallelize serial workflow loops (+ allocation cleanup)

Three performance improvements to the perturbed-equilibrium / kinetic-diagnostic workflow on
the DIII-D-like reference case. Each is verified to leave results unchanged (bit-identical or
within regression tolerance) and independent of thread count.

1. Parallelize the NTV-torque bounce-harmonic loop (KineticForces)

The NTV-torque outer ψ-integral (QuadGK BatchIntegrand) stays serial — its small per-batch
node counts make threading the ψ loop fork-join-bound — but the inner bounce-harmonic loop
(2·nl+1 harmonics per ψ node) carries the ~40 ms of per-node work and is embarrassingly
parallel. Threaded it: each harmonic's tpsi! writes per-intr scratch buffers/hints, so
each thread gets its own deepcopy(intr) (the pattern the self-consistent matrix build uses),
and per-harmonic results are reduced in fixed harmonic order.

• fgar NTV torque: 35.3 s → 8.2 s at 8 threads (4.3×).
• Bit-identical: 82.1808181112855 − 14.6098675986818i to every digit across the serial
  baseline, threaded t=1, and threaded t=8.

2. Parallelize the singular-coupling loop over rational surfaces (PerturbedEquilibrium)

Each rational surface builds its own vacuum Green's functions (allocation-local
compute_vacuum_response) and writes disjoint coupling rows / sing[s], so the loop is now
@threads :static over resonant pairs (BLAS pinned to 1 across it).

• Singular coupling: ~0.54 s → ~0.18 s at 8 threads (~3×); results identical across thread
  counts and equal to serial to ~13–15 significant figures (BLAS-pin reordering).

3. Reduce allocations in field reconstruction (PerturbedEquilibrium)

reconstruct_physical_fields: in-place mode DFTs, @view eigenmode slices, cholesky!, BLAS
pin. ~30% fewer allocations (594 → 420 MB), bit-identical. (A memory cleanup — reconstruction
is bandwidth-bound, not GC-bound.)

Verification

Regression fingerprints of every affected output (b_n/xi_n, C_resonant_*, resonant_*,
and the NTV torque) match the serial baseline and are thread-count-independent; the repo's
coordinate-invariance, coils, full-run, and HDF5-replay suites pass.

[logan-nc]

@claude review this pr

[claude]

Review: PR #325 (perf/perturbed-equilibrium-threading)

Three independent optimizations (NTV bounce-harmonic threading, singular-coupling threading, field-reconstruction allocation cleanup). Good use of per-thread scratch buffers and fixed-order reduction to preserve determinism where it has been checked -- but one of the three has a real correctness gap the PR's own testing would not have caught.

Bug: data race in the singular-coupling loop (SingularCoupling.jl)

compute_singular_coupling_metrics! now runs Threads.@threads :static for row in 1:length(resonant_pairs), and inside the loop body (unchanged from before, now executed concurrently):

ffs_intr.sing[s].grri = grri
ffs_intr.sing[s].grre = grre

resonant_pairs is built as (s, nn) for nn in nlow:nhigh, s in 1:msing (SingularCoupling.jl Phase 1) -- it is unique per (s, nn) pair, not per s. Any surface with an (near-)integer q resonates with every nn in range (m_res = q*nn is an integer for all nn when q is exactly integer), so the same s legitimately appears in multiple rows for multi-n runs. Those rows can land on different threads and now race on the same ffs_intr.sing[s].grri/.grre write -- a genuine data race, and whichever thread finishes last wins nondeterministically. That directly contradicts the PR description's claim that 'each [surface] writes disjoint coupling-matrix rows / sing[s]' -- it is disjoint in n=1 cases only.

This is why it was not caught: examples/DIIID-like_ideal_example (the diiid_n1 regression case referenced in the PR) is single-n (nn = 1), so resonant_pairs never has a repeated s, and the race cannot manifest there. It is also invisible to CI regardless, since test.yaml does not set JULIA_NUM_THREADS/-t, so these @threads loops always execute on 1 thread in CI -- none of the 'bit-identical across thread counts' claims in this PR are actually exercised by automated testing.

I did not find any current read of sing[s].grri/.grre elsewhere in src/ (grep for .grri/.grre only turns up Vacuum.jl's unrelated vac_data.grri and this file), so today it may be write-only/dead. But that is fragile -- please either guard the write so it is skipped/idempotent-safe for duplicate s (or key it by (s, nn) instead of overwriting the shared sing[s] field), or confirm it is genuinely dead and remove the mutation from the hot loop rather than leave a landmine for the next reader who wires it into output.

Given the regression suite only has a single-n perturbed-equilibrium case, please add (or point to) a multi-n PerturbedEquilibrium regression case with an integer-q rational surface in range -- that is the actual reproducer for this bug and is exactly the gap the regression harness policy in CLAUDE.md exists to catch.

Robustness: BLAS thread count not restored on exception

Both reconstruct_physical_fields (FieldReconstruction.jl) and compute_singular_coupling_metrics! (SingularCoupling.jl) save BLAS.get_num_threads(), call BLAS.set_num_threads(1), run a threaded loop that can throw (e.g. vacuum solve, cholesky!, indexing), then restore the thread count -- all without a try/finally. If anything in the loop throws, BLAS stays pinned at 1 thread for the rest of the process, silently degrading every subsequent linear-algebra call in that Julia session. Worth wrapping in try/finally.

Minor

• KineticForces/Compute.jl: the harmonic loop is threaded inside the per-batch for k in eachindex(x) loop, so a Threads.@threads team is spawned/joined once per psi-node (~15x per QuadGK batch) rather than once per batch. Given the measured 4.3x speedup this clearly is not a problem in practice, but flattening to a single @threads over (k, ell_idx) pairs per batch would cut fork-join overhead further if it is ever worth it.
• thread_intrs = [deepcopy(intr) for _ in 1:nth] runs once per call to integrate_psi_quadgk, i.e. once per n in multi-n NTV runs. Fine for the current DIII-D case, but worth a passing mention of memory scaling (nthreads x sizeof(intr) x npert) if KineticForcesInternal carries large per-instance buffers, since multi-n kinetic runs with many threads could add up.
• The cholesky!/@view/in-place DFT changes in FieldReconstruction.jl look correct: amat is thread-local scratch fully overwritten by ffit.amats(...) before each cholesky!, and the @view slices into mul! avoid the copy without changing semantics.

Test coverage

No new tests, consistent with a 'results should be identical' performance PR, but per CLAUDE.md's regression harness policy this needs a regress --cases diiid_n1,solovev_multi_n --refs develop,local run (or equivalent) attached before merge, and -- per the bug above -- ideally a multi-n perturbed-equilibrium case added to the harness so this class of threading bug is regression-tested going forward, not just the ideal-stability multi-n path.

[logan-nc]

Addressed the review findings in 424463c (threading-correctness pass):

• Removed the sing[s].grri/.grre data race — these were write-only (never read anywhere in src/ or test/), and a surface s legitimately repeats across rows in multi-n runs (integer q resonates at several n), so the shared-field write was a genuine race. Dropped the dead store; behavior-neutral.
• Guarded both BLAS pins with try/finally (SingularCoupling, FieldReconstruction) so an exception in a threaded loop can't leak BLAS.set_num_threads(1) into the rest of the session.
• Documented the top-level-@threads requirement of the threadid()-indexed buffers and the nthreads × sizeof(intr) memory scaling of the per-thread deepcopy.

Verified thread-safety of the remaining shared reads: the vacuum solve uses a task-local pool (@with_pool → get_task_local_pool), and FastInterpolations' hint is an external caller kwarg (never internal mutable state), so concurrent hint-less spline calls are race-free.

Regression: regress --cases diiid_n1,solovev_kinetic_calculated --refs develop,local is bit-identical to develop (all quantities 0.0e+00), covering the singular-coupling, field-reconstruction, and NTV-torque paths.

Follow-up (separate PR): a multi-n perturbed-equilibrium regression case to exercise the repeated-s path and a threaded CI job — tracked in a new issue.