Skip to content

ExponentialRK: remove per-step allocations from EPIRK/Exp4/EXPRB53s3 steppers #3928

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ChrisRackauckas merged 1 commit into from
Jul 23, 2026
Merged

ExponentialRK: remove per-step allocations from EPIRK/Exp4/EXPRB53s3 steppers #3928

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion lib/OrdinaryDiffEqExponentialRK/Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "OrdinaryDiffEqExponentialRK"
uuid = "e0540318-69ee-4070-8777-9e2de6de23de"
authors = ["ParamThakkar123 <paramthakkar864@gmail.com>"]
version = "2.1.0"
version = "2.1.1"

[deps]
FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
Expand Down
36 changes: 24 additions & 12 deletions lib/OrdinaryDiffEqExponentialRK/src/exponential_rk_caches.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -443,7 +443,7 @@ for (Alg, Cache) in [
end
end

@cache struct Exp4Cache{uType, rateType, JCType, FType, matType, JType, KsType} <:
@cache struct Exp4Cache{uType, rateType, JCType, FType, matType, JType, KsType, tsType} <:
ExpRKCache
u::uType
uprev::uType
Expand All @@ -458,6 +458,7 @@ end
J::JType
B::matType
KsCache::KsType
ts::tsType
end
function alg_cache(
alg::Exp4, u, rate_prototype, ::Type{uEltypeNoUnits},
Expand Down Expand Up @@ -489,10 +490,11 @@ function alg_cache(
# Allocate caches for phiv_timestep
maxiter = min(alg.m, n)
KsCache = _phiv_timestep_caches(u, maxiter, 1)
return Exp4Cache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache)
ts = Vector{typeof(dt)}(undef, 3)
return Exp4Cache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache, ts)
end

@cache struct EPIRK4s3ACache{uType, rateType, JCType, FType, matType, JType, KsType} <:
@cache struct EPIRK4s3ACache{uType, rateType, JCType, FType, matType, JType, KsType, tsType} <:
ExpRKCache
u::uType
uprev::uType
Expand All @@ -507,6 +509,7 @@ end
J::JType
B::matType
KsCache::KsType
ts::tsType
end
function alg_cache(
alg::EPIRK4s3A, u, rate_prototype, ::Type{uEltypeNoUnits},
Expand Down Expand Up @@ -537,10 +540,11 @@ function alg_cache(
# Allocate caches for phiv_timestep
maxiter = min(alg.m, n)
KsCache = _phiv_timestep_caches(u, maxiter, 4)
return EPIRK4s3ACache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache)
ts = Vector{typeof(dt)}(undef, 2)
return EPIRK4s3ACache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache, ts)
end

@cache struct EPIRK4s3BCache{uType, rateType, JCType, FType, matType, JType, KsType} <:
@cache struct EPIRK4s3BCache{uType, rateType, JCType, FType, matType, JType, KsType, tsType} <:
ExpRKCache
u::uType
uprev::uType
Expand All @@ -555,6 +559,7 @@ end
J::JType
B::matType
KsCache::KsType
ts::tsType
end
function alg_cache(
alg::EPIRK4s3B, u, rate_prototype, ::Type{uEltypeNoUnits},
Expand Down Expand Up @@ -585,7 +590,8 @@ function alg_cache(
# Allocate caches for phiv_timestep
maxiter = min(alg.m, n)
KsCache = _phiv_timestep_caches(u, maxiter, 4)
return EPIRK4s3BCache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache)
ts = Vector{typeof(dt)}(undef, 2)
return EPIRK4s3BCache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache, ts)
end

@cache struct EPIRK5s3Cache{uType, rateType, JCType, FType, matType, JType, KsType} <:
Expand Down Expand Up @@ -635,7 +641,7 @@ function alg_cache(
return EPIRK5s3Cache(u, uprev, tmp, dz, k, rtmp, rtmp2, du1, jac_config, uf, J, B, KsCache)
end

@cache struct EXPRB53s3Cache{uType, rateType, JCType, FType, matType, JType, KsType} <:
@cache struct EXPRB53s3Cache{uType, rateType, JCType, FType, matType, JType, KsType, tsType} <:
ExpRKCache
u::uType
uprev::uType
Expand All @@ -650,6 +656,7 @@ end
J::JType
B::matType
KsCache::KsType
ts::tsType
end
function alg_cache(
alg::EXPRB53s3, u, rate_prototype, ::Type{uEltypeNoUnits},
Expand Down Expand Up @@ -680,10 +687,11 @@ function alg_cache(
# Allocate caches for phiv_timestep
maxiter = min(alg.m, n)
KsCache = _phiv_timestep_caches(u, maxiter, 4)
return EXPRB53s3Cache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache)
ts = Vector{typeof(dt)}(undef, 2)
return EXPRB53s3Cache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache, ts)
end

@cache struct EPIRK5P1Cache{uType, rateType, JCType, FType, matType, JType, KsType} <:
@cache struct EPIRK5P1Cache{uType, rateType, JCType, FType, matType, JType, KsType, tsType} <:
ExpRKCache
u::uType
uprev::uType
Expand All @@ -698,6 +706,7 @@ end
J::JType
B::matType
KsCache::KsType
ts::tsType
end
function alg_cache(
alg::EPIRK5P1, u, rate_prototype, ::Type{uEltypeNoUnits},
Expand Down Expand Up @@ -728,10 +737,11 @@ function alg_cache(
# Allocate caches for phiv_timestep
maxiter = min(alg.m, n)
KsCache = _phiv_timestep_caches(u, maxiter, 3)
return EPIRK5P1Cache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache)
ts = Vector{typeof(dt)}(undef, 3)
return EPIRK5P1Cache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache, ts)
end

@cache struct EPIRK5P2Cache{uType, rateType, JCType, FType, matType, JType, KsType} <:
@cache struct EPIRK5P2Cache{uType, rateType, JCType, FType, matType, JType, KsType, tsType} <:
ExpRKCache
u::uType
uprev::uType
Expand All @@ -747,6 +757,7 @@ end
J::JType
B::matType
KsCache::KsType
ts::tsType
end
function alg_cache(
alg::EPIRK5P2, u, rate_prototype, ::Type{uEltypeNoUnits},
Expand Down Expand Up @@ -777,7 +788,8 @@ function alg_cache(
# Allocate caches for phiv_timestep
maxiter = min(alg.m, n)
KsCache = _phiv_timestep_caches(u, maxiter, 3)
return EPIRK5P2Cache(u, uprev, tmp, dz, rtmp, rtmp2, dR, du1, jac_config, uf, K, J, B, KsCache)
ts = Vector{typeof(dt)}(undef, 3)
return EPIRK5P2Cache(u, uprev, tmp, dz, rtmp, rtmp2, dR, du1, jac_config, uf, K, J, B, KsCache, ts)
end

####################################
Expand Down
74 changes: 45 additions & 29 deletions lib/OrdinaryDiffEqExponentialRK/src/exponential_rk_perform_step.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -801,11 +801,13 @@ end

function perform_step!(integrator, cache::Exp4Cache, repeat_step = false)
(; t, dt, uprev, u, f, p) = integrator
(; tmp, rtmp, rtmp2, K, J, B, KsCache) = cache
(; tmp, rtmp, rtmp2, K, J, B, KsCache, ts) = cache
calc_J!(J, integrator, cache)
alg = unwrap_alg(integrator, true)
f0 = integrator.fsalfirst # f(u0) is fsaled
ts = [dt / 3, 2dt / 3, dt]
ts[1] = dt / 3
ts[2] = 2dt / 3
ts[3] = dt
kwargs = (
tol = integrator.opts.reltol, iop = alg.iop,
opnorm = integrator.opts.internalopnorm,
Expand All @@ -819,29 +821,31 @@ function perform_step!(integrator, cache::Exp4Cache, repeat_step = false)
@inbounds for i in 1:3
K[:, i] ./= ts[i]
end
mul!(rtmp, K, [-7 / 300, 97 / 150, -37 / 300]) # rtmp is now w4
@views @.. broadcast = false rtmp = (-7 / 300) * K[:, 1] + (97 / 150) * K[:, 2] +
(-37 / 300) * K[:, 3] # rtmp is now w4
@muladd @.. broadcast = false tmp = uprev + dt * rtmp # tmp is now u4
mul!(rtmp2, J, rtmp)
f(rtmp, tmp, p, t + dt) # TODO: what should be the time?
OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
@muladd @.. broadcast = false @view(B[:, 2]) = rtmp - f0 - dt * rtmp2 # B[:,2] is now d4
# Partially update entities that use k1, k2, k3
mul!(rtmp, K, [59 / 300, -7 / 75, 269 / 300]) # rtmp is now w7
@views @.. broadcast = false rtmp = (59 / 300) * K[:, 1] + (-7 / 75) * K[:, 2] +
(269 / 300) * K[:, 3] # rtmp is now w7
@muladd @.. broadcast = false u = uprev + dt * @view(K[:, 3])
# Krylov for the first remainder d4
phiv_timestep!(K, ts, J, B; kwargs...)
@inbounds for i in 1:3
K[:, i] ./= ts[i]
end
mul!(rtmp2, K, [2 / 3, 2 / 3, 2 / 3])
@views @.. broadcast = false rtmp2 = (2 / 3) * (K[:, 1] + K[:, 2] + K[:, 3])
rtmp .+= rtmp2 # w7 fully updated
@muladd @.. broadcast = false tmp = uprev + dt * rtmp # tmp is now u7
mul!(rtmp2, J, rtmp)
f(rtmp, tmp, p, t + dt) # TODO: what should be the time?
OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
@muladd @.. broadcast = false @view(B[:, 2]) = rtmp - f0 - dt * rtmp2 # B[:,2] is now d7
# Partially update entities that use k4, k5, k6
mul!(rtmp, K, [1.0, -4 / 3, 1.0])
@views @.. broadcast = false rtmp = K[:, 1] + (-4 / 3) * K[:, 2] + K[:, 3]
axpy!(dt, rtmp, u)
# Krylov for the second remainder d7
k7 = @view(K[:, 1])
Expand Down Expand Up @@ -892,7 +896,7 @@ end

function perform_step!(integrator, cache::EPIRK4s3ACache, repeat_step = false)
(; t, dt, uprev, u, f, p) = integrator
(; tmp, rtmp, rtmp2, K, J, B, KsCache) = cache
(; tmp, rtmp, rtmp2, K, J, B, KsCache, ts) = cache
calc_J!(J, integrator, cache)
alg = unwrap_alg(integrator, true)
f0 = integrator.fsalfirst # f(u0) is fsaled
Expand All @@ -905,21 +909,23 @@ function perform_step!(integrator, cache::EPIRK4s3ACache, repeat_step = false)

# Compute U2 and U3 vertically
B[:, 2] .= f0
phiv_timestep!(K, [dt / 2, 2dt / 3], J, @view(B[:, 1:2]); kwargs...)
ts[1] = dt / 2
ts[2] = 2dt / 3
phiv_timestep!(K, ts, J, @view(B[:, 1:2]); kwargs...)
## U2 and R2
@.. broadcast = false tmp = uprev + @view(K[:, 1]) # tmp is now U2
f(rtmp, tmp, p, t + dt / 2)
mul!(rtmp2, J, @view(K[:, 1]))
@.. broadcast = false rtmp = rtmp - f0 - rtmp2 # rtmp is now R2
B[:, 4] .= (32 / dt^2) * rtmp
B[:, 5] .= (-144 / dt^3) * rtmp
B[:, 4] .= (32 / dt^2) .* rtmp
B[:, 5] .= (-144 / dt^3) .* rtmp
## U3 and R3
@.. broadcast = false tmp = uprev + @view(K[:, 2]) # tmp is now U3
f(rtmp, tmp, p, t + 2dt / 3)
mul!(rtmp2, J, @view(K[:, 2]))
@.. broadcast = false rtmp = rtmp - f0 - rtmp2 # rtmp is now R3
B[:, 4] .-= (13.5 / dt^2) * rtmp
B[:, 5] .+= (81 / dt^3) * rtmp
B[:, 4] .-= (13.5 / dt^2) .* rtmp
B[:, 5] .+= (81 / dt^3) .* rtmp
OrdinaryDiffEqCore.increment_nf!(integrator.stats, 2)

# Update u
Expand Down Expand Up @@ -972,7 +978,7 @@ end

function perform_step!(integrator, cache::EPIRK4s3BCache, repeat_step = false)
(; t, dt, uprev, u, f, p) = integrator
(; tmp, rtmp, rtmp2, K, J, B, KsCache) = cache
(; tmp, rtmp, rtmp2, K, J, B, KsCache, ts) = cache
calc_J!(J, integrator, cache)
alg = unwrap_alg(integrator, true)
f0 = integrator.fsalfirst # f(u0) is fsaled
Expand All @@ -986,23 +992,25 @@ function perform_step!(integrator, cache::EPIRK4s3BCache, repeat_step = false)
# Compute U2 and U3 vertically
fill!(@view(B[:, 2]), zero(eltype(B)))
B[:, 3] .= f0
phiv_timestep!(K, [dt / 2, 3dt / 4], J, @view(B[:, 1:3]); kwargs...)
ts[1] = dt / 2
ts[2] = 3dt / 4
phiv_timestep!(K, ts, J, @view(B[:, 1:3]); kwargs...)
K[:, 1] .*= 8 / (3 * dt)
K[:, 2] .*= 16 / (9 * dt)
## U2 and R2
@.. broadcast = false tmp = uprev + @view(K[:, 1]) # tmp is now U2
f(rtmp, tmp, p, t + dt / 2)
mul!(rtmp2, J, @view(K[:, 1]))
@.. broadcast = false rtmp = rtmp - f0 - rtmp2 # rtmp is now R2
B[:, 4] .= (54 / dt^2) * rtmp
B[:, 5] .= (-324 / dt^3) * rtmp
B[:, 4] .= (54 / dt^2) .* rtmp
B[:, 5] .= (-324 / dt^3) .* rtmp
## U3 and R3
@.. broadcast = false tmp = uprev + @view(K[:, 2]) # tmp is now U3
f(rtmp, tmp, p, t + 3dt / 4)
mul!(rtmp2, J, @view(K[:, 2]))
@.. broadcast = false rtmp = rtmp - f0 - rtmp2 # rtmp is now R3
B[:, 4] .-= (16 / dt^2) * rtmp
B[:, 5] .+= (144 / dt^3) * rtmp
B[:, 4] .-= (16 / dt^2) .* rtmp
B[:, 5] .+= (144 / dt^3) .* rtmp
OrdinaryDiffEqCore.increment_nf!(integrator.stats, 2)

# Update u
Expand Down Expand Up @@ -1091,7 +1099,7 @@ function perform_step!(integrator, cache::EPIRK5s3Cache, repeat_step = false)
# Compute U3 horizontally
B[:, 2] .= (53 / 5) .* f0
B[:, 3] .= (-648 / (5 * dt)) .* f0
B[:, 4] .= (2916 / (5 * dt^2)) .* f0 + (32065 / (1152 * dt^2)) .* rtmp
B[:, 4] .= (2916 / (5 * dt^2)) .* f0 .+ (32065 / (1152 * dt^2)) .* rtmp
phiv_timestep!(k, 4dt / 9, J, @view(B[:, 1:4]); kwargs...)
## Update B matrix using R2
B[:, 2] .= f0
Expand Down Expand Up @@ -1161,7 +1169,7 @@ end

function perform_step!(integrator, cache::EXPRB53s3Cache, repeat_step = false)
(; t, dt, uprev, u, f, p) = integrator
(; tmp, rtmp, rtmp2, K, J, B, KsCache) = cache
(; tmp, rtmp, rtmp2, K, J, B, KsCache, ts) = cache
calc_J!(J, integrator, cache)
alg = unwrap_alg(integrator, true)
f0 = integrator.fsalfirst # f(u0) is fsaled
Expand All @@ -1174,7 +1182,9 @@ function perform_step!(integrator, cache::EXPRB53s3Cache, repeat_step = false)

# Compute the first group for U2 and U3
B[:, 2] .= f0
phiv_timestep!(K, [dt / 2, 9dt / 10], J, @view(B[:, 1:2]); kwargs...)
ts[1] = dt / 2
ts[2] = 9dt / 10
phiv_timestep!(K, ts, J, @view(B[:, 1:2]); kwargs...)
## U2 and R2
@.. broadcast = false tmp = uprev + @view(K[:, 1]) # tmp is now U2
f(rtmp, tmp, p, t + dt / 2)
Expand All @@ -1186,21 +1196,21 @@ function perform_step!(integrator, cache::EXPRB53s3Cache, repeat_step = false)
# Compute the second group for U3
fill!(@view(B[:, 2]), zero(eltype(B)))
B[:, 4] .= rtmp
phiv_timestep!(K, [dt / 2, 9dt / 10], J, @view(B[:, 1:4]); kwargs...)
phiv_timestep!(K, ts, J, @view(B[:, 1:4]); kwargs...)
## Update B using R2
B[:, 2] .= f0
B[:, 4] .= (18 / dt^2) .* rtmp
B[:, 5] .= (-60 / dt^3) .* rtmp
## U3 and R3
@views tmp .+= 216 / (25 * dt^2) .* K[:, 1] + 8 / dt^2 .* K[:, 2] # tmp is now U3
@views tmp .+= 216 / (25 * dt^2) .* K[:, 1] .+ 8 / dt^2 .* K[:, 2] # tmp is now U3
f(rtmp, tmp, p, t + 9dt / 10)
OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
tmp .-= uprev
mul!(rtmp2, J, tmp)
@.. broadcast = false rtmp = rtmp - f0 - rtmp2 # rtmp is now R3
## Update B using R3
B[:, 4] .-= (250 / (81 * dt^2)) * rtmp
B[:, 5] .+= (500 / (27 * dt^3)) * rtmp
B[:, 4] .-= (250 / (81 * dt^2)) .* rtmp
B[:, 5] .+= (500 / (27 * dt^3)) .* rtmp

# Update u
du = @view(K[:, 1])
Expand Down Expand Up @@ -1268,7 +1278,7 @@ end

function perform_step!(integrator, cache::EPIRK5P1Cache, repeat_step = false)
(; t, dt, uprev, u, f, p) = integrator
(; tmp, rtmp, rtmp2, K, J, B, KsCache) = cache
(; tmp, rtmp, rtmp2, K, J, B, KsCache, ts) = cache
calc_J!(J, integrator, cache)
alg = unwrap_alg(integrator, true)
f0 = integrator.fsalfirst # f(u0) is fsaled
Expand All @@ -1291,7 +1301,10 @@ function perform_step!(integrator, cache::EPIRK5P1Cache, repeat_step = false)

# Compute the first column (f0)
B[:, 2] .= f0
phiv_timestep!(K, [g11, g21, g31], J, @view(B[:, 1:2]); kwargs...)
ts[1] = g11
ts[2] = g21
ts[3] = g31
phiv_timestep!(K, ts, J, @view(B[:, 1:2]); kwargs...)
## U1 and R1
@.. broadcast = false tmp = uprev + @view(K[:, 1]) # tmp is now U1
f(rtmp, tmp, p, t + g11)
Expand Down Expand Up @@ -1384,7 +1397,7 @@ end

function perform_step!(integrator, cache::EPIRK5P2Cache, repeat_step = false)
(; t, dt, uprev, u, f, p) = integrator
(; tmp, rtmp, rtmp2, dR, K, J, B, KsCache) = cache
(; tmp, rtmp, rtmp2, dR, K, J, B, KsCache, ts) = cache
calc_J!(J, integrator, cache)
alg = unwrap_alg(integrator, true)
f0 = integrator.fsalfirst # f(u0) is fsaled
Expand All @@ -1409,7 +1422,10 @@ function perform_step!(integrator, cache::EPIRK5P2Cache, repeat_step = false)

# Compute the first column (f0)
B[:, 2] .= f0
phiv_timestep!(K, [g11, g21, g31], J, @view(B[:, 1:2]); kwargs...)
ts[1] = g11
ts[2] = g21
ts[3] = g31
phiv_timestep!(K, ts, J, @view(B[:, 1:2]); kwargs...)
## U1 and R1
@.. broadcast = false tmp = uprev + @view(K[:, 1]) # tmp is now U1
f(rtmp, tmp, p, t + g11)
Expand Down
Loading