Skip to content

Add GPURodas3 and GPURodas42 support to EnsembleGPUKernel #429

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

Open
AJ0070 wants to merge 4 commits into
base: master
Choose a base branch
from
Open

Add GPURodas3 and GPURodas42 support to EnsembleGPUKernel #429

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
998cd6f
Add GPURodas3 and GPURodas42 support to EnsembleGPUKernel
AJ0070 Apr 6, 2026
8fd9c51
doc update
AJ0070 Apr 6, 2026
6b9dd34
fix format check
AJ0070 Apr 7, 2026
b35e14b
ran runic
AJ0070 Apr 7, 2026
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: 2 additions & 0 deletions docs/src/manual/ensemblegpukernel.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,9 @@ GPUVern9
GPUEM
GPUSIEA
GPURosenbrock23
GPURodas3
GPURodas4
GPURodas42
GPURodas5P
GPUKvaerno3
GPUKvaerno5
Expand Down
4 changes: 3 additions & 1 deletion src/DiffEqGPU.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ include("ensemblegpukernel/perform_step/gpu_vern9_perform_step.jl")
include("ensemblegpukernel/perform_step/gpu_em_perform_step.jl")
include("ensemblegpukernel/perform_step/gpu_siea_perform_step.jl")
include("ensemblegpukernel/perform_step/gpu_rosenbrock23_perform_step.jl")
include("ensemblegpukernel/perform_step/gpu_rodas3_perform_step.jl")
include("ensemblegpukernel/perform_step/gpu_rodas4_perform_step.jl")
include("ensemblegpukernel/perform_step/gpu_rodas5P_perform_step.jl")
include("ensemblegpukernel/perform_step/gpu_kvaerno3_perform_step.jl")
Expand All @@ -78,7 +79,8 @@ export EnsembleCPUArray, EnsembleGPUArray, EnsembleGPUKernel, LinSolveGPUSplitFa

export GPUTsit5, GPUVern7, GPUVern9, GPUEM, GPUSIEA
## Stiff ODE solvers
export GPURosenbrock23, GPURodas4, GPURodas5P, GPUKvaerno3, GPUKvaerno5
export GPURosenbrock23, GPURodas3, GPURodas4, GPURodas42, GPURodas5P, GPUKvaerno3,
GPUKvaerno5

export terminate!

Expand Down
2 changes: 2 additions & 0 deletions src/ensemblegpukernel/alg_utils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,9 @@ alg_order(alg::GPUTsit5) = 5
alg_order(alg::GPUVern7) = 7
alg_order(alg::GPUVern9) = 9
alg_order(alg::GPURosenbrock23) = 2
alg_order(alg::GPURodas3) = 3
alg_order(alg::GPURodas4) = 4
alg_order(alg::GPURodas42) = 4
alg_order(alg::GPURodas5P) = 5
alg_order(alg::GPUKvaerno3) = 3
alg_order(alg::GPUKvaerno5) = 5
Expand Down
18 changes: 17 additions & 1 deletion src/ensemblegpukernel/gpukernel_algorithms.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,14 @@ generation with EnsembleGPUKernel.
"""
struct GPURosenbrock23{AD} <: GPUODEImplicitAlgorithm{AD} end

"""
GPURodas3()

A specialized implementation of the `Rodas3` method specifically for kernel
generation with EnsembleGPUKernel.
"""
struct GPURodas3{AD} <: GPUODEImplicitAlgorithm{AD} end

"""
GPURodas4()

Expand All @@ -39,6 +47,14 @@ generation with EnsembleGPUKernel.
"""
struct GPURodas4{AD} <: GPUODEImplicitAlgorithm{AD} end

"""
GPURodas42()

A specialized implementation of the `Rodas42` method specifically for kernel
generation with EnsembleGPUKernel.
"""
struct GPURodas42{AD} <: GPUODEImplicitAlgorithm{AD} end

"""
GPURodas5P()

Expand All @@ -63,7 +79,7 @@ generation with EnsembleGPUKernel.
"""
struct GPUKvaerno5{AD} <: GPUODEImplicitAlgorithm{AD} end

for Alg in [:GPURosenbrock23, :GPURodas4, :GPURodas5P, :GPUKvaerno3, :GPUKvaerno5]
for Alg in [:GPURosenbrock23, :GPURodas3, :GPURodas4, :GPURodas42, :GPURodas5P, :GPUKvaerno3, :GPUKvaerno5]
@eval begin
function $Alg(; autodiff = Val{true}())
return $Alg{SciMLBase._unwrap_val(autodiff)}()
Expand Down
309 changes: 309 additions & 0 deletions src/ensemblegpukernel/integrators/stiff/types.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -365,6 +365,295 @@ const GPUARodas4I = GPUARodas4Integrator
)
end

@inline function init(
alg::GPURodas42, f::F, IIP::Bool, u0::S, t0::T, dt::T,
p::P, tstops::TS,
callback::CB,
save_everystep::Bool,
saveat::ST
) where {
F, P, T,
S,
TS, CB, ST,
}
!IIP && @assert S <: SArray
event_last_time = 1
vector_event_last_time = 0
last_event_error = zero(T)

tab = Rodas42Tableau(T, T)

return integ = GPURodas4I{IIP, S, T, ST, P, F, TS, CB, typeof(tab), typeof(alg)}(
alg, f,
copy(u0),
copy(u0),
copy(u0), t0,
t0,
t0,
dt,
sign(dt), p,
true,
tstops, 1,
callback,
save_everystep,
saveat, 1, 1,
event_last_time,
vector_event_last_time,
last_event_error,
copy(u0),
copy(u0), tab,
DiffEqBase.ReturnCode.Default
)
end

@inline function init(
alg::GPURodas42, f::F, IIP::Bool, u0::S, t0::T, tf::T,
dt::T, p::P,
abstol::TOL, reltol::TOL,
internalnorm::N, tstops::TS,
callback::CB,
saveat::ST
) where {
F, P, S, T, N, TOL, TS,
CB, ST,
}
!IIP && @assert S <: SArray
qoldinit = T(1.0e-4)
event_last_time = 1
vector_event_last_time = 0
last_event_error = zero(T)

tab = Rodas42Tableau(T, T)

return integ = GPUARodas4I{
IIP, S, T, ST, P, F, N, TOL, typeof(qoldinit), TS, CB, typeof(tab),
typeof(alg),
}(
alg,
f,
copy(u0),
copy(u0),
copy(u0),
t0,
t0,
t0,
tf,
dt,
dt,
sign(
tf -
t0
),
p,
true,
tstops,
1,
callback,
false,
saveat,
1, 1,
event_last_time,
vector_event_last_time,
last_event_error,
copy(u0),
copy(u0),
tab,
qoldinit,
abstol,
reltol,
internalnorm,
DiffEqBase.ReturnCode.Default
)
end

##########################
# Rodas 3
##########################

mutable struct GPURodas3Integrator{IIP, S, T, ST, P, F, TS, CB, TabType, AlgType} <:
DiffEqBase.AbstractODEIntegrator{AlgType, IIP, S, T}
alg::AlgType
f::F
uprev::S
u::S
tmp::S
tprev::T
t::T
t0::T
dt::T
tdir::T
p::P
u_modified::Bool
tstops::TS
tstops_idx::Int
callback::CB
save_everystep::Bool
saveat::ST
cur_t::Int
step_idx::Int
event_last_time::Int
vector_event_last_time::Int
last_event_error::T
k1::S
k2::S
tab::TabType
retcode::DiffEqBase.ReturnCode.T
end
const GPURodas3I = GPURodas3Integrator

@inline function init(
alg::GPURodas3, f::F, IIP::Bool, u0::S, t0::T, dt::T,
p::P, tstops::TS,
callback::CB,
save_everystep::Bool,
saveat::ST
) where {
F, P, T,
S,
TS, CB, ST,
}
!IIP && @assert S <: SArray
event_last_time = 1
vector_event_last_time = 0
last_event_error = zero(T)

tab = Rodas3Tableau(T, T)

return integ = GPURodas3I{IIP, S, T, ST, P, F, TS, CB, typeof(tab), typeof(alg)}(
alg, f,
copy(u0),
copy(u0),
copy(u0), t0,
t0,
t0,
dt,
sign(dt), p,
true,
tstops, 1,
callback,
save_everystep,
saveat, 1, 1,
event_last_time,
vector_event_last_time,
last_event_error,
copy(u0),
copy(u0), tab,
DiffEqBase.ReturnCode.Default
)
end

mutable struct GPUARodas3Integrator{
IIP,
S,
T,
ST,
P,
F,
N,
TOL,
Q,
TS,
CB,
TabType,
AlgType,
} <:
DiffEqBase.AbstractODEIntegrator{AlgType, IIP, S, T}
alg::AlgType
f::F
uprev::S
u::S
tmp::S
tprev::T
t::T
t0::T
tf::T
dt::T
dtnew::T
tdir::T
p::P
u_modified::Bool
tstops::TS
tstops_idx::Int
callback::CB
save_everystep::Bool
saveat::ST
cur_t::Int
step_idx::Int
event_last_time::Int
vector_event_last_time::Int
last_event_error::T
k1::S
k2::S
tab::TabType
qold::Q
abstol::TOL
reltol::TOL
internalnorm::N
retcode::DiffEqBase.ReturnCode.T
end

const GPUARodas3I = GPUARodas3Integrator

@inline function init(
alg::GPURodas3, f::F, IIP::Bool, u0::S, t0::T, tf::T,
dt::T, p::P,
abstol::TOL, reltol::TOL,
internalnorm::N, tstops::TS,
callback::CB,
saveat::ST
) where {
F, P, S, T, N, TOL, TS,
CB, ST,
}
!IIP && @assert S <: SArray
qoldinit = T(1.0e-4)
event_last_time = 1
vector_event_last_time = 0
last_event_error = zero(T)

tab = Rodas3Tableau(T, T)

return integ = GPUARodas3I{
IIP, S, T, ST, P, F, N, TOL, typeof(qoldinit), TS, CB, typeof(tab),
typeof(alg),
}(
alg,
f,
copy(u0),
copy(u0),
copy(u0),
t0,
t0,
t0,
tf,
dt,
dt,
sign(
tf -
t0
),
p,
true,
tstops,
1,
callback,
false,
saveat,
1, 1,
event_last_time,
vector_event_last_time,
last_event_error,
copy(u0),
copy(u0),
tab,
qoldinit,
abstol,
reltol,
internalnorm,
DiffEqBase.ReturnCode.Default
)
end

##########################
# Rodas 5P
##########################
Expand Down Expand Up @@ -937,6 +1226,26 @@ end
return integrator.u_modified = bool
end

# GPURodas3Integrator
@inline function (integrator::GPURodas3I)(t)
Θ = (t - integrator.tprev) / integrator.dt
return _ode_interpolant(Θ, integrator.dt, integrator.uprev, integrator)
end

@inline function DiffEqBase.u_modified!(integrator::GPURodas3I, bool::Bool)
return integrator.u_modified = bool
end

# GPUARodas3Integrator
@inline function (integrator::GPUARodas3I)(t)
Θ = (t - integrator.tprev) / integrator.dt
return _ode_interpolant(Θ, integrator.dt, integrator.uprev, integrator)
end

@inline function DiffEqBase.u_modified!(integrator::GPUARodas3I, bool::Bool)
return integrator.u_modified = bool
end

# GPURodas4Integrator
@inline function (integrator::GPURodas4I)(t)
Θ = (t - integrator.tprev) / integrator.dt
Expand Down
Loading
Loading