Skip to content

Pentrc statictyping preallocation#322

Merged
matt-pharr merged 7 commits into
developfrom
PENTRC_statictyping_preallocation
Jul 9, 2026
Merged

Pentrc statictyping preallocation#322
matt-pharr merged 7 commits into
developfrom
PENTRC_statictyping_preallocation

Conversation

@matt-pharr

@matt-pharr matt-pharr commented Jul 8, 2026

Copy link
Copy Markdown
Collaborator

There are a number of locations in pentrc that have poorly optimized dynamic typing and allocations happening close to operations. This means that when trying to multithread, we are needlessly multiplying our allocations by the number of threads we use, which is why the multithreading speedup for pentrc peaks at 2 threads. At time of making the PR the static typing and preallocations completed have taken the kinetic portion of the forcefreestates solve down from 619 s to 68.1 s in the d3d-like example case running with 8 threads on omega. Also made preallocations for the bounce averaging computations. This work is presently complete and ready for review.

Here is the final speedup table, running on omega:

threads develop branch speedup
1 1650.8 s 189.1 s 8.7×
2 1125.4 s 119.7 s 9.4×
4 745.0 s 72.1 s 10.3×
8 498.4 s 49.5 s 10.1×
16 ~520 s* 32.5 s ~16×
32 ~520 s* 24.3 s ~21×

*develop's 16- and 32-thread runs were run on a different cluster, but showed a plateau at ~520 s starting at ~8 threads.

Parallel efficiency (speedup vs. that version's own 1-thread):

2t 4t 8t
develop 73% 55% 41%
branch 79% 66% 48%

…rams to reduce multithreading overhead and add a pre-allocated quadgk buffer
…memory allocation in QuadGK inner loop and update Energy Integration to support QuadGK buffer pre-allocation
@matt-pharr matt-pharr requested a review from logan-nc July 8, 2026 16:30
…-surface scratch buffers to reduce memory allocation during bounce-averaging computations
@matt-pharr matt-pharr marked this pull request as ready for review July 8, 2026 19:57
@logan-nc

logan-nc commented Jul 8, 2026

Copy link
Copy Markdown
Collaborator

@claude review this PR

@claude

claude Bot commented Jul 8, 2026

Copy link
Copy Markdown

Review: Pentrc static typing / preallocation

Solid, well-documented allocation-reduction pass on the KineticForces hot loops (bounce averaging, energy integration, pitch integration). This will be useful for scaling the multithreaded pentrc solve past 2 threads.

Potential bugs / correctness

  • BounceScratch.vpar_fine sizing is only coincidentally correct. BounceScratch(ntheta, mpert; nfine::Int=256) sizes vpar_fine as nfine+1 (BounceAveraging.jl), while _find_bounce_points_and_grid's trapped branch independently hardcodes its own nfine = 256 and indexes scr.vpar_fine[i] for i in 1:(nfine+1). These two 256 literals aren't derived from one another — before this PR, vpar_fine was built inline from the local nfine, so it was self-consistent by construction. Now, if either literal changes in the future without the other, the @inbounds loop will silently read/write past the buffer end instead of erroring. Worth threading one nfine value through instead of duplicating it (e.g. size from length(scr.vpar_fine) - 1 at the call site, or pass nfine explicitly).

Style / minor

  • PitchGARParams.fvals::Vector{ComplexF64} buffers what is always real data — fbnce's underlying Series is built from zeros(Float64, ...) in Torque.jl and is never complex. The in-place call still works (Float64 promotes fine into a ComplexF64 slot), so this isn't a functional bug, but it doubles the buffer's footprint for no benefit; Vector{Float64} would be tighter and match what's actually stored.
  • The new BounceScratch struct documents its fields via trailing inline comments only, with no ## Fields section in the docstring (per CLAUDE.md's struct-doc convention). Minor — it matches this file's existing PitchGARParams precedent, so not a new departure, but worth fixing going forward.

Performance

  • The core change (hoisting the per-θ Vector{Float64}(undef,5) / Vector{ComplexF64}(undef,mpert) allocations out to a per-surface BounceScratch, plus reusing a QuadGK segbuf across the λ sweep in PitchGARParams) is a reasonable way to cut per-thread allocation pressure and matches the reported speedup numbers.
  • Pulling _resonant_integrand out of a closure into a named @inline function is a good change independent of the buffer work — avoids boxed-capture overhead in the QuadGK hot loop.

Test coverage

  • No unit tests were added for BounceScratch, _bounce_integrate, or the new PitchGARParams fields. runtests_kinetic.jl covers integrate_energy/find_resonance_energies directly but not compute_bounce_data or the pitch-integration path; those are only exercised indirectly via the full-run smoke tests in runtests_fullruns.jl, which check finiteness/sign, not exact values. Given the stated goal of bit-for-bit parity with the previous per-call-allocation behavior, a small targeted regression (or at least a regression-harness run) pinning a kinetic quantity before/after would give more confidence.
  • Per CLAUDE.md, the regression harness should be run against develop before merging any substantive change and its report included in the PR — that isn't present here yet. Since this touches the dJdJ/kinetic-matrix hot path directly, it's worth running (e.g. against the Solovev kinetic case) before merge.

@logan-nc logan-nc left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good and great speedup results 🥇
Only minor suggestions from me

Comment thread src/KineticForces/BounceAveraging.jl
Comment thread src/KineticForces/EnergyIntegration.jl
@matt-pharr matt-pharr requested a review from logan-nc July 8, 2026 21:29
@matt-pharr matt-pharr enabled auto-merge July 8, 2026 21:30
@matt-pharr

Copy link
Copy Markdown
Collaborator Author

@logan-nc added the NaN fill. Ready for merge pending your thumbs up.

…buffer, BounceScratch field docs

fvals stays ComplexF64: the kinetic-matrix (wt) path builds fbnce from complex
op_wmats data, so a real buffer throws InexactError in the pitch kernel.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TTu6sczod1Qtjy9e9pB6jP
@logan-nc

logan-nc commented Jul 8, 2026

Copy link
Copy Markdown
Collaborator

Review feedback addressed (38b5f42)

All review items are in:

  • NaN fill of wmats_lmda — done in 0007e62 (kept as-is).
  • ν=0 pole annotation — restored in 17a441a.
  • nfine duplication_find_bounce_points_and_grid now derives nfine = length(scr.vpar_fine) - 1 from the BounceScratch buffer instead of a second hardcoded 256, so the @inbounds loop can't silently go out of sync with the constructor.
  • BounceScratch field docs — moved into a ## Fields docstring section per convention.
  • fvals::Vector{Float64} suggestion — rejected, with evidence. The suggestion only holds for the scalar-torque path (Torque.jl builds that fbnce from real data). The kinetic-matrix path builds fbnce from complex op_wmats data ("fbnce is COMPLEX on this path", Torque.jl), and with a Float64 buffer the wt pitch kernel throws InexactError on the Solovev kinetic cases — a path the unit tests don't reach (all 255 pass either way). The buffer stays ComplexF64 with a comment explaining why. Caught by the regression harness, which is a good advertisement for running it.

Regression harness vs develop

  • diiid_n1: all 47 tracked quantities unchanged (0.0 diff), and the tracked runtime drops 677 s → 426 s (−37 %) from this PR's allocation work.
  • solovev_kinetic_calculated / solovev_kinetic_nuzero: both run successfully at the branch tip and are now cached as baselines. They cannot run on develop — the develop copy of the example still lacks the [SOL_INPUT] section in gpec.toml (pre-existing; fixed here by 6f49409), so no cross-ref diff is possible. Note the new baseline values (et[1] = 1.831−1.485i calculated, 2.189−2.251i nuzero) differ by a few percent from the ≈-values quoted in the case descriptions, which likely reflects the sol.toml→gpec.toml parameter consolidation — worth a quick confirm that the consolidated parameters match the originals.
Full diiid_n1 regression table
Regression Report: diiid_n1
===================================================================================================================
Ref 1: develop  @ 94e2ea31 (2026-07-06)
Ref 2: PENTRC_statictyping_preallocation  @ 38b5f425 (2026-07-08)
-------------------------------------------------------------------------------------------------------------------
Quantity                                      develop          PENTRC_statictyping_preallocation  Diff       Status
-------------------------------------------------------------------------------------------------------------------
root-area-weighted total energy Re(et[1])     9.107359e-14     9.107359e-14                       0.0e+00    OK    
root-area-weighted total energy Im(et[1])     -2.460303e-20    -2.460303e-20                      0.0e+00    OK    
root-area-weighted plasma energy Re(ep[1])    4.574355e-14     4.574355e-14                       0.0e+00    OK    
root-area-weighted vacuum energy Re(ev[1])    4.532993e-14     4.532993e-14                       0.0e+00    OK    
vacuum matrix min eigenvalue (XiNorm)         1.873975e-01     1.873975e-01                       0.0e+00    OK    
root-area-weighted plasma energy (all)        [35 elem]        [35 elem]                          0.0e+00    OK    
root-area-weighted vacuum energy (all)        [35 elem]        [35 elem]                          0.0e+00    OK    
root-area-weighted total energy (all)         [35 elem]        [35 elem]                          0.0e+00    OK    
ODE steps (saved)                             1463             1463                               0.0e+00    OK    
ODE steps (total)                             2486             2486                               0.0e+00    OK    
q0                                            1.204206e+00     1.204206e+00                       0.0e+00    OK    
q95                                           4.781722e+00     4.781722e+00                       0.0e+00    OK    
beta_t                                        1.327021e-02     1.327021e-02                       0.0e+00    OK    
beta_n                                        1.372507e+00     1.372507e+00                       0.0e+00    OK    
internal inductance li1                       8.842447e-01     8.842447e-01                       0.0e+00    OK    
internal inductance li2                       7.080909e-01     7.080909e-01                       0.0e+00    OK    
internal inductance li3                       7.304497e-01     7.304497e-01                       0.0e+00    OK    
poloidal beta betap1                          6.680728e-01     6.680728e-01                       0.0e+00    OK    
poloidal beta betap2                          5.349834e-01     5.349834e-01                       0.0e+00    OK    
poloidal beta betap3                          5.518761e-01     5.518761e-01                       0.0e+00    OK    
# singular surfaces                           5                5                                  0.0e+00    OK    
singular psi locations                        [5 elem]         [5 elem]                           0.0e+00    OK    
singular q values                             [5 elem]         [5 elem]                           0.0e+00    OK    
current beta betaj                            4.236473e-01     4.236473e-01                       0.0e+00    OK    
plasma volume                                 1.829476e+01     1.829476e+01                       0.0e+00    OK    
plasma current                                1.152130e+00     1.152130e+00                       0.0e+00    OK    
mpert                                         35               35                                 0.0e+00    OK    
npert                                         1                1                                  0.0e+00    OK    
toroidal field bt0                            2.006573e+00     2.006573e+00                       0.0e+00    OK    
wall field bwall                              3.880145e-01     3.880145e-01                       0.0e+00    OK    
aspect ratio                                  2.845746e+00     2.845746e+00                       0.0e+00    OK    
elongation kappa                              1.708350e+00     1.708350e+00                       0.0e+00    OK    
q profile (checksum)                          a1b1c2215ac6...  a1b1c2215ac6...                    identical  OK    
pressure profile (checksum)                   f1484167f5bd...  f1484167f5bd...                    identical  OK    
Mercier D_I profile (checksum)                472a5914dcb6...  472a5914dcb6...                    identical  OK    
resistive interchange D_R profile (checksum)  7918d44452d8...  7918d44452d8...                    identical  OK    
ballooning Delta' profile (checksum)          7f6b347c5fc8...  7f6b347c5fc8...                    identical  OK    
delta prime (BVP diagonal)                    [5 elem]         [5 elem]                           0.0e+00    OK    
island half-widths                            [5 elem]         [5 elem]                           0.0e+00    OK    
Chirikov parameter                            [5 elem]         [5 elem]                           0.0e+00    OK    
||resonant area-weighted field||              3.800199e-04     3.800199e-04                       0.0e+00    OK    
PE plasma energy                              3.423709e+00     3.423709e+00                       0.0e+00    OK    
PE vacuum energy                              3.174509e+00     3.174509e+00                       0.0e+00    OK    
PE surface energy                             5.853353e+00     5.853353e+00                       0.0e+00    OK    
PE toroidal torque                            -5.060600e-02    -5.060600e-02                      0.0e+00    OK    
NTV torque FGAR (Re,Im) [N·m]                 [2 elem]         [2 elem]                           1.4e-14    OK    
Runtime (s)                                   677.0s           426.4s                                        --    
resonant area-weighted field b^r              [5 elem]         [5 elem]                           0.0e+00    OK    
===================================================================================================================
Summary: 47 unchanged

🤖 Generated with Claude Code

https://claude.ai/code/session_01TTu6sczod1Qtjy9e9pB6jP

@matt-pharr matt-pharr merged commit 5044073 into develop Jul 9, 2026
6 checks passed
@matt-pharr matt-pharr deleted the PENTRC_statictyping_preallocation branch July 9, 2026 02:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants