ExponentialRK: remove per-step allocations from EPIRK/Exp4/EXPRB53s3 steppers - #3928
Merged
ChrisRackauckas merged 1 commit intoJul 23, 2026
Conversation
…steppers The in-place perform_step! implementations of Exp4, EPIRK4s3A/B, EPIRK5s3, EXPRB53s3, and EPIRK5P1/P2 allocated every step: - literal snapshot-time arrays ([dt/3, 2dt/3, dt], [g11, g21, g31], ...) constructed per step for phiv_timestep!; now a ts vector preallocated in the cache and filled in place - coefficient vectors for 3-column linear combinations built per step (mul!(rtmp, K, [-7/300, 97/150, -37/300])); now fused broadcasts over the columns of K - unfused broadcast right-hand sides that materialized N-sized temporaries (B[:, 4] .= (32 / dt^2) * rtmp and friends, and two broadcast-plus-broadcast sums); now fully fused with dots The N-sized temporaries in particular scale with the PDE size. Together with ExponentialUtilities#254/SciML#255 this takes per-step allocations on the Kuramoto-Sivashinsky pseudospectral benchmark (N=128, non-allocating nonlinearity) from 0.6-1.2 MB/step to 11-23 KB/step for the EPIRK family; solutions are identical. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
ChrisRackauckas
marked this pull request as ready for review
July 23, 2026 11:59
ChrisRackauckas
added a commit
that referenced
this pull request
Jul 23, 2026
…-step allocations (#3993) The EPIRK/Exp4/Exprb steppers still allocated per step through bare column-slice compound assignments of the form `X[:, i] .op= ...` (e.g. `K[:, i] ./= ts[i]`, `B[:, 4] .-= c .* rtmp`). Without `@views` the right-hand-side read of the slice materializes a fresh column copy each time (Base `_unsafe_getindex -> similar`), so each such line allocates one working vector per step. #3928 converted the large fused updates to `@views @.. broadcast=false` but left these smaller normalization/update slices unviewed. Wrapping them in `@views` makes the read a view instead of a copy; the operation is still an in-place elementwise write to the same memory, so the result is unchanged (each element reads only its own position -- no aliasing). Measured per-step allocations on a 1D semilinear reaction-diffusion problem (N=128, adaptive_krylov, m=30), via the `step!` interface after warmup, master vs this branch: Exp4 7647 -> 927 B/step (8.2x) EPIRK4s3B 5132 -> 652 B/step (7.9x) EPIRK4s3A 2879 -> 639 B/step (4.5x) EXPRB53s3 3180 -> 940 B/step (3.4x) EPIRK5P1 2047 -> 927 B/step (2.2x) Final solution is bit-identical before/after for every method (max|Δu| = 0 across Exp4/EPIRK4s3A/EPIRK4s3B/EXPRB53s3/EPIRK5P1/EPIRK5P2/ Exprb32/Exprb43). The remaining ~0.6-0.9 KB/step floor is in the ExponentialUtilities adaptive-Krylov path (a `v[1:end-1]` copy in `lanczos!` and a reshaped-subarray `similar`), tracked separately. Adds a runtime size-independence regression test (Exp4 on a reaction- diffusion problem): per-step allocations must not scale with the state size, which cleanly catches a reintroduced unviewed slice copy without a flaky absolute byte ceiling. Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The in-place
perform_step!implementations ofExp4,EPIRK4s3A/B,EPIRK5s3,EXPRB53s3, andEPIRK5P1/P2allocated on every step:[dt / 3, 2dt / 3, dt],[g11, g21, g31], …) constructed per step forphiv_timestep!— now atsvector preallocated in the cache and filled in place;mul!(rtmp, K, [-7 / 300, 97 / 150, -37 / 300])) — now fused broadcasts over the columns ofK;B[:, 4] .= (32 / dt^2) * rtmpand friends, plus two broadcast-plus-broadcast sums) — now fully fused. These scale with the PDE size, so they matter most exactly where the EPIRK methods are supposed to shine.Out-of-place (
ConstantCache) paths are untouched.Context / measurements
This is the OrdinaryDiffEq side of an allocation audit of the exponential integrators; the dominant per-step allocations were in ExponentialUtilities (SciML/ExponentialUtilities.jl#254 —
PhivCachereallocated on everyphiv!call — and SciML/ExponentialUtilities.jl#255 —exponential!workspace reuse). Combined per-step allocations on the Kuramoto–Sivashinsky pseudospectral benchmark problem (N=128, non-allocating nonlinearity, fixed dt, Julia 1.12.6):(*Exp4's remaining KB moved between categories as the two changes landed; the end state is equivalent.) Caching-mode ETDRK2/3/4, HochOst4, NorsettEuler, LawsonEuler, ETD2 were verified 0 bytes/step already. The remaining ~10–20 KB/step comes from small per-call allocations inside
phiv_timestep!(coeffs = ones(p), scalar-t[t]wrappers) and the finite-difference Jacobian.Solutions are identical before/after (endpoint norms bit-identical on the KS benchmark).
Tests
Ran
ODEDIFFEQ_TEST_GROUP=Core Pkg.test()locally on the sublibrary with the dev'd ExponentialUtilities fixes: Krylov tests + Linear-Nonlinear Convergence Tests pass (40 pass, 2 pre-existing@test_broken). Runic formatting applied.Version bumped 2.1.0 → 2.1.1 (internal cache structs only, no public API change).
Note: please ignore this PR until reviewed by @ChrisRackauckas.
🤖 Generated with Claude Code
https://claude.ai/code/session_01EDVGmmovzD5Aos3fFoPwMY