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

src/KineticForces/Compute.jl

@@ -47,14 +47,21 @@ function integrate_psi_quadgk(
         return (total=ComplexF64(0.0), torque_profile=nothing, matrix_integrated=nothing, psi_nsteps=0)
     end
 
-    # Buffers for the batch callback. The outer ψ-integral is intentionally
-    # serial: QuadGK.BatchIntegrand's refine loop invokes the callback many
-    # times with small batches (~15 nodes per Kronrod rule), and Threads.@threads
-    # fork-join overhead at this granularity dominates the ~40 ms per-ψ work,
-    # producing catastrophic slowdowns at ≥2 threads. Serial `for` inside the
-    # batch matches 1-thread wall time and is the fastest correct option found.
-    tpsi_val = Ref{ComplexF64}(0.0im)
-    wtw_l = is_matrix_method ? zeros(ComplexF64, mpert, mpert, 6) : nothing
+    # The outer ψ-integral (the QuadGK batch / ψ-node loop) stays serial: QuadGK's refine
+    # loop invokes the callback with small batches (~15 nodes), so threading it is
+    # fork-join-bound. Instead thread the inner bounce-harmonic loop (2·nl+1 harmonics),
+    # which carries the ~40 ms of per-ψ work. `tpsi!` writes per-`intr` scratch buffers and
+    # spline hints, so each thread gets its own `deepcopy(intr)` (the same per-thread-scratch
+    # pattern the self-consistent kinetic matrix build uses). Per-harmonic results are stored
+    # by index and reduced in fixed ℓ order, so the sum is identical to the serial reduction
+    # and independent of thread count.
+    nharm = 1 + 2 * nl
+    nth = Threads.maxthreadid()
+    thread_intrs = [deepcopy(intr) for _ in 1:nth]
+    thread_tpsi = [Ref{ComplexF64}(0.0im) for _ in 1:nth]
+    thread_wtw = is_matrix_method ? [zeros(ComplexF64, mpert, mpert, 6) for _ in 1:nth] : nothing
+    harm_vals = Vector{ComplexF64}(undef, nharm)
+    harm_elems = is_matrix_method ? [zeros(ComplexF64, mpert, mpert, 6) for _ in 1:nharm] : nothing
 
     logged_psi = Float64[]
     logged_dtdpsi = ComplexF64[]
@@ -64,35 +71,36 @@ function integrate_psi_quadgk(
         for k in eachindex(x)
             psi = Float64(x[k])
 
-            if is_matrix_method && !isnothing(wtw_l)
-                wtw_l .= 0
-            end
-            elems_accum = is_matrix_method ? zeros(ComplexF64, mpert, mpert, 6) : nothing
-
-            total = ComplexF64(0.0)
-            for ell_idx in 1:(1 + 2 * nl)
+            # Each bounce harmonic is independent; thread over them, each thread using its own
+            # `intr` copy and writing its own harm_vals[ell_idx] / harm_elems[ell_idx] slot.
+            Threads.@threads :static for ell_idx in 1:nharm
+                tid = Threads.threadid()
                 l = ell_idx - 1 - nl
-                if is_matrix_method && !isnothing(wtw_l)
-                    wtw_l .= 0
-                end
-
-                tpsi!(tpsi_val, psi, n, l, zi, mi, wdfac, divxfac,
-                      electron, method, equil, intr, kinetic_profiles;
-                      op_wmats=wtw_l,
+                w = is_matrix_method ? thread_wtw[tid] : nothing
+                is_matrix_method && fill!(w, 0)
+                tpsi!(thread_tpsi[tid], psi, n, l, zi, mi, wdfac, divxfac,
+                      electron, method, equil, thread_intrs[tid], kinetic_profiles;
+                      op_wmats=w,
                       atol_xlmda=ctrl.atol_xlmda, rtol_xlmda=ctrl.rtol_xlmda)
-                total += tpsi_val[]
-
-                if is_matrix_method && !isnothing(wtw_l) && !isnothing(elems_accum)
-                    elems_accum .+= wtw_l
-                end
+                harm_vals[ell_idx] = thread_tpsi[tid][]
+                is_matrix_method && (harm_elems[ell_idx] .= w)
             end
 
+            # Reduce in fixed ℓ order (bit-identical to the serial sum).
+            total = ComplexF64(0.0)
+            for ell_idx in 1:nharm
+                total += harm_vals[ell_idx]
+            end
             y[k] = total
 
             push!(logged_psi, psi)
             push!(logged_dtdpsi, total)
-            if is_matrix_method && !isnothing(elems_accum)
-                push!(logged_elems, copy(elems_accum))
+            if is_matrix_method
+                elems_accum = zeros(ComplexF64, mpert, mpert, 6)
+                for ell_idx in 1:nharm
+                    elems_accum .+= harm_elems[ell_idx]
+                end
+                push!(logged_elems, elems_accum)
             end
         end
     end

src/PerturbedEquilibrium/FieldReconstruction.jl

@@ -80,6 +80,13 @@ function reconstruct_physical_fields(
     npsi = size(ForceFreeStates_results.u_store, 4)
     psi_grid = ForceFreeStates_results.psi_store[1:npsi]
 
+    # Pin BLAS to a single thread for the per-surface reconstruction below. Each threaded
+    # loop over ψ calls only small BLAS kernels (per-surface mpert×mpert solves and mode
+    # DFTs), so leaving BLAS multi-threaded oversubscribes the cores against the Julia
+    # `@threads` over ψ and destroys thread scaling. Restored before returning.
+    _blas_nthreads = BLAS.get_num_threads()
+    BLAS.set_num_threads(1)
+
     # Sum weighted eigenmode contributions to get ξ_ψ, dξ_ψ/dψ, and ξ_s in mode space
     xi_psi_modes, xi_psi1_modes, xi_s_modes = sum_eigenmode_contributions(
         response_vector,
@@ -192,6 +199,7 @@ function reconstruct_physical_fields(
         R=b_R, Z=b_Z, phi=b_phi
     )
 
+    BLAS.set_num_threads(_blas_nthreads)
     return xi_modes, b_modes
 end
 
@@ -235,17 +243,18 @@ function sum_eigenmode_contributions(
     xi_s_modes = zeros(ComplexF64, npsi, mpert)
     # Surfaces are independent; threaded over ψ (run with `julia -t N` or JULIA_NUM_THREADS).
     Threads.@threads :static for ipsi in 1:npsi
-        # u_store[:,:,1] = Ξ_ψ (radial displacement)
+        # u_store[:,:,1] = Ξ_ψ (radial displacement). @view avoids copying the mpert×mpert
+        # eigenmode-matrix slice on every surface (mul! takes the view directly).
         mul!(view(xi_psi_modes, ipsi, :),
-            ForceFreeStates_results.u_store[:, :, 1, ipsi],
+            @view(ForceFreeStates_results.u_store[:, :, 1, ipsi]),
             alpha)
         # ud_store[:,:,1] = dΞ_ψ/dψ (radial derivative)
         mul!(view(xi_psi1_modes, ipsi, :),
-            ForceFreeStates_results.ud_store[:, :, 1, ipsi],
+            @view(ForceFreeStates_results.ud_store[:, :, 1, ipsi]),
             alpha)
         # ud_store[:,:,2] = Ξ_s = -A⁻¹(B·Ξ'_ψ + C·Ξ_ψ) (toroidal displacement, Glasser 2016 eq. 18)
         mul!(view(xi_s_modes, ipsi, :),
-            ForceFreeStates_results.ud_store[:, :, 2, ipsi],
+            @view(ForceFreeStates_results.ud_store[:, :, 2, ipsi]),
             alpha)
     end
 
@@ -393,8 +402,10 @@ function compute_clebsch_displacements(
         xsp_vec = view(xi_psi_modes, ipsi, :)
         mul!(xms_vec, bmat, xmp1_vec)                     # xms = B*xmp1
         mul!(xms_vec, cmat_buf, xsp_vec, 1.0+0.0im, 1.0+0.0im)  # xms += C*xsp
-        # amat is positive-definite by construction (Newcomb kinetic-energy form), so cholesky is safe.
-        amat_fact = cholesky(Hermitian(amat, :L))
+        # amat is positive-definite by construction (Newcomb kinetic-energy form), so cholesky is
+        # safe. cholesky! factorizes in place (amat is a per-thread scratch buffer, refilled by
+        # ffit.amats each surface), avoiding a fresh factorization allocation per surface.
+        amat_fact = cholesky!(Hermitian(amat, :L))
         ldiv!(amat_fact, xms_vec)                          # xms = A\(B*xmp1 + C*xsp)
         xms_vec .*= -1                                     # xms = -A\(B*xmp1 + C*xsp)
 
@@ -963,19 +974,27 @@ function _apply_rzphi_transform(
     Z_modes = zeros(ComplexF64, npsi, mpert)
     phi_modes = zeros(ComplexF64, npsi, mpert)
 
-    # Per-thread theta-space buffers; the immutable `ft` functor and `geom` are shared read-only.
-    fun_bufs = [(R=zeros(ComplexF64, mtheta), Z=zeros(ComplexF64, mtheta), P=zeros(ComplexF64, mtheta)) for _ in 1:Threads.maxthreadid()]
+    # Per-thread scratch (the immutable `ft` functor and `geom` are shared read-only): θ-space
+    # transform inputs/outputs (length mtheta) and mode-space forward-DFT outputs (length mpert),
+    # so the DFTs run in place with no per-surface allocation.
+    bufs = [(R=zeros(ComplexF64, mtheta), Z=zeros(ComplexF64, mtheta), P=zeros(ComplexF64, mtheta),
+             psi=zeros(ComplexF64, mtheta), th=zeros(ComplexF64, mtheta), ze=zeros(ComplexF64, mtheta),
+             Ro=zeros(ComplexF64, mpert), Zo=zeros(ComplexF64, mpert), Po=zeros(ComplexF64, mpert))
+            for _ in 1:Threads.maxthreadid()]
 
     Threads.@threads :static for ipsi in 1:npsi
-        buf = fun_bufs[Threads.threadid()]
+        buf = bufs[Threads.threadid()]
         R_fun = buf.R
         Z_fun = buf.Z
         phi_fun = buf.P
 
-        # Inverse DFT: modes → theta-space
-        psi_fun = Utilities.FourierTransforms.inverse(ft, view(psi_input, ipsi, :))
-        theta_fn = Utilities.FourierTransforms.inverse(ft, view(theta_input, ipsi, :))
-        zeta_fn = Utilities.FourierTransforms.inverse(ft, view(cova_zeta_input, ipsi, :))
+        # Inverse DFT: modes → theta-space (in place)
+        psi_fun = buf.psi
+        theta_fn = buf.th
+        zeta_fn = buf.ze
+        Utilities.FourierTransforms.inverse_transform!(psi_fun, ft, view(psi_input, ipsi, :))
+        Utilities.FourierTransforms.inverse_transform!(theta_fn, ft, view(theta_input, ipsi, :))
+        Utilities.FourierTransforms.inverse_transform!(zeta_fn, ft, view(cova_zeta_input, ipsi, :))
 
         # Pointwise transformation (Fortran gpeq_rzphi, gpeq.f:484-489)
         for itheta in 1:mtheta
@@ -994,10 +1013,13 @@ function _apply_rzphi_transform(
             phi_fun[itheta] = geom.t33[itheta, ipsi] * xvz
         end
 
-        # Forward DFT: theta-space → modes
-        R_modes[ipsi, :] .= ft(R_fun)
-        Z_modes[ipsi, :] .= ft(Z_fun)
-        phi_modes[ipsi, :] .= ft(phi_fun)
+        # Forward DFT: theta-space → modes (in place)
+        Utilities.FourierTransforms.transform!(buf.Ro, ft, R_fun)
+        Utilities.FourierTransforms.transform!(buf.Zo, ft, Z_fun)
+        Utilities.FourierTransforms.transform!(buf.Po, ft, phi_fun)
+        R_modes[ipsi, :] .= buf.Ro
+        Z_modes[ipsi, :] .= buf.Zo
+        phi_modes[ipsi, :] .= buf.Po
     end
 
     return R_modes, Z_modes, phi_modes

src/PerturbedEquilibrium/SingularCoupling.jl

@@ -152,10 +152,17 @@ function compute_singular_coupling_metrics!(
     edge_mn = intr.plasma_response ./ (chi1 * 2π * im .* reshape(singfac_lim, :, 1))
     C_coeffs = u_bnd \ edge_mn  # mpert × numpert_total
 
-    # Phase 3: Compute full coupling matrix rows
+    # Phase 3: Compute full coupling matrix rows. Each rational surface is independent -- it
+    # builds its own Green's functions (a fresh, allocation-local vacuum solve) and writes
+    # disjoint coupling-matrix rows / sing[s] -- so thread the loop over surfaces. BLAS is
+    # pinned to one thread across it (each surface's solve is small; multi-threaded BLAS would
+    # oversubscribe against the Julia threads). Restored after the loop.
     psi_store_all = ForceFreeStates_results.psi_store
     nstep = ForceFreeStates_results.step
-    for (row, (s, nn)) in enumerate(resonant_pairs)
+    _blas_nthreads = BLAS.get_num_threads()
+    BLAS.set_num_threads(1)
+    Threads.@threads :static for row in 1:length(resonant_pairs)
+        (s, nn) = resonant_pairs[row]
         sing_surf = ffs_intr.sing[s]
         m_res = round(Int, sing_surf.q * nn)
 
@@ -280,6 +287,7 @@ function compute_singular_coupling_metrics!(
             @info "Row $row: q=$(@sprintf("%.3f", sing_surf.q)), ψ=$(@sprintf("%.3f", sing_surf.psifac)), m=$m_res, n=$nn, Δ'(diag)=$(@sprintf("%.3e", dp_diag))"
         end
     end
+    BLAS.set_num_threads(_blas_nthreads)
 
     # Phase 4: Apply forcing amplitudes → R = C · Φ_x. The applied resonant scalars are
     # physical, coordinate-invariant quantities, so evaluate them from the flux-space C and