diff --git a/src/KineticForces/Compute.jl b/src/KineticForces/Compute.jl index ada81bea0..42559149e 100644 --- a/src/KineticForces/Compute.jl +++ b/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 diff --git a/src/PerturbedEquilibrium/FieldReconstruction.jl b/src/PerturbedEquilibrium/FieldReconstruction.jl index ae0334e1f..6a0e71383 100644 --- a/src/PerturbedEquilibrium/FieldReconstruction.jl +++ b/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 diff --git a/src/PerturbedEquilibrium/SingularCoupling.jl b/src/PerturbedEquilibrium/SingularCoupling.jl index 2a2bec045..63772c95d 100644 --- a/src/PerturbedEquilibrium/SingularCoupling.jl +++ b/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