Developer Interfaces

EnsembleArrayAlgorithm <: SciMLBase.EnsembleAlgorithm

Developer interface for ensemble algorithms that fuse a SciML ensemble into array-valued state and parameter problems before delegating each trajectory to an ordinary SciML differential equation solver.

Interface Rules

Subtypes are used as the third positional argument to solve(ensembleprob, alg, ensemblealg; trajectories, kwargs...) for SciMLBase.AbstractEnsembleProblems. A subtype must be accepted by DiffEqGPU's SciMLBase.__solve method and by the lower-level vectorized_map_solve path. It must define enough backend/device information for vectorized_map_solve_up to move the generated u0 and p arrays to the execution backend, and it must preserve the SciML ensemble semantics for prob_func, reduction, batch_size, trajectories, and solver keyword arguments.

Implementations

• EnsembleGPUArray: runs the fused array problem on a KernelAbstractions backend.
• EnsembleCPUArray: keeps the same fused-array code path on CPU for debugging.

Examples

solve(ensemble_prob, Tsit5(), EnsembleGPUArray(backend); trajectories = 10_000)

EnsembleKernelAlgorithm <: SciMLBase.EnsembleAlgorithm

Developer interface for ensemble algorithms that generate one GPU kernel for a complete fixed-size ODE or SDE solve.

Interface Rules

Subtypes are used as the ensemble algorithm in solve(ensembleprob, gpu_alg, ensemblealg; trajectories, kwargs...), where gpu_alg is a GPUODEAlgorithm or GPUSDEAlgorithm. The corresponding problem must be convertible to the kernel path with make_prob_compatible; in practice this means out-of-place dynamics over static state containers for EnsembleGPUKernel. Implementations must support the lower-level vectorized_solve and, for ODE algorithms, vectorized_asolve entry points used by batch_solve_up_kernel.

Implementations

• EnsembleGPUKernel: compiles a KernelAbstractions kernel for all trajectories in a batch.

Examples

solve(
    ensemble_prob, GPUTsit5(), EnsembleGPUKernel(backend);
    trajectories = 10_000, adaptive = false, dt = 0.1f0
)

maxthreads(backend)

Return the maximum work-group size used by DiffEqGPU kernels on backend.

This is a developer interface for backend extensions. A backend method must return a positive integer that is valid for the backend's kernel launch configuration.

Arguments

• backend: a KernelAbstractions backend supported by DiffEqGPU.

Returns

The backend-specific maximum number of threads in a work group.

Examples

maxthreads(CPU())

maybe_prefer_blocks(backend)

Return the backend configuration used for DiffEqGPU kernel launches.

This is a developer interface for backend extensions. A backend method may return a configuration with block-oriented execution enabled when that is required for efficient or correct kernel execution.

Arguments

• backend: a KernelAbstractions backend supported by DiffEqGPU.

Returns

The backend instance to pass to subsequent kernel allocation and launch operations.

Examples

maybe_prefer_blocks(CPU()) isa CPU

lufact!(backend, W)

Factorize each square matrix in a batched matrix array in place.

This is a developer interface implemented by backend extensions. The factorization is consumed by LinSolveGPUSplitFactorize; each slice W[:, :, i] must be a square matrix, and the backend implementation must provide the factorization operation callable from the selected execution environment.

Arguments

• backend: the execution backend.
• W: a three-dimensional array whose first two dimensions contain one matrix per batch index.

Returns

nothing; W is mutated in place.

Examples

W = reshape([2.0f0, 0.0f0, 0.0f0, 3.0f0], 2, 2, 1)
lufact!(CPU(), W)

LinSolveGPUSplitFactorize()
LinSolveGPUSplitFactorize(len, nfacts)

A parameter-parallel SciMLLinearSolveAlgorithm for applying pre-factorized per-trajectory linear systems on a KernelAbstractions backend.

Fields

• len::Int: the size of each factored linear system.
• nfacts::Int: the number of factorizations stored in the batched factorization array.

Arguments

• len::Int: the size of each factored linear system.
• nfacts::Int: the number of factorizations stored in the batched factorization array.

Most users do not need to construct this directly; EnsembleGPUArray installs it for compatible stiff ensemble solves.

Returns

A LinSolveGPUSplitFactorize selector configured for the supplied factorization layout.

Examples

linsolve = LinSolveGPUSplitFactorize(3, 256)

make_prob_compatible(prob)

Prepare a problem for the lower-level EnsembleGPUKernel interface.

For an ODEProblem, this updates any initialization problem on the host, converts the initialization problem and its maps to immutable static representations for device-side solving, and adapts parameter and mass-matrix storage. Other problem-like values are returned unchanged. The resulting problem must still satisfy the selected backend's GPU-compatibility requirements.

Arguments

• prob: an ODEProblem or another problem value accepted by the lower-level ensemble interface.

Returns

An immutable, backend-compatible representation for ODE problems, or prob unchanged for other values.

Examples

using DiffEqGPU, SciMLBase, StaticArrays
f(u, p, t) = u
prob = ODEProblem{false}(f, SVector(1.0f0), (0.0f0, 1.0f0), SVector(1.0f0))
gpu_prob = DiffEqGPU.make_prob_compatible(prob)

DiffEqGPU.make_static_storage(x)

Convert x into storage a GPU kernel can hold: arrays become SArrays, and Tuples and NamedTuples are rebuilt element by element. Everything else is returned unchanged, so a value that is already isbits passes straight through.

This is the extension point for types that a kernel cannot take as-is. A package owning such a type adds a method converting it to an isbits stand-in, and EnsembleGPUKernel picks it up when converting u0 and the parameters:

# In a package that owns `MyInterpolation`
function DiffEqGPU.make_static_storage(itp::MyInterpolation)
    return MyStaticInterpolation(
        DiffEqGPU.make_static_storage(itp.t), DiffEqGPU.make_static_storage(itp.u)
    )
end

Without such a method the value is left alone, and the problem is rejected with an error naming it rather than failing inside the kernel.

GPUODEAlgorithm <: SciMLBase.AbstractODEAlgorithm

Developer interface for ODE algorithms supported by EnsembleGPUKernel.

Interface Rules

Subtypes must be immutable algorithm selectors with all per-solve state allocated by the kernel integrator constructors. They are passed as the second positional argument to solve(ensembleprob, alg, EnsembleGPUKernel(backend); kwargs...) and to the lower-level vectorized_solve/vectorized_asolve functions. A subtype must have kernel integrator support in batch_solve_up_kernel, an order from alg_order, and any required tableau, interpolation, callback, nonlinear-solve, or mass-matrix support implemented in GPU-safe code.

The ODE function must be GPU compilable. The kernel path is tested through the generic lower-level API by constructing compatible ODEProblems and calling vectorized_solve and vectorized_asolve without reaching into solver internals.

Examples

DiffEqGPU.vectorized_solve(gpu_probs, prob, GPUTsit5(); dt = 0.1f0)

GPUSDEAlgorithm <: SciMLBase.AbstractSDEAlgorithm

Developer interface for SDE algorithms supported by EnsembleGPUKernel.

Interface Rules

Subtypes are passed as the SDE algorithm in solve(ensembleprob, alg, EnsembleGPUKernel(backend); kwargs...) and through the lower-level vectorized_solve entry point. Implementations must provide GPU-safe stepping code, static state support, and noise compatibility checks before launching kernels. Current kernel SDE algorithms are fixed-step methods and must reject unsupported noise structures rather than falling back to host-side behavior.

Examples

DiffEqGPU.vectorized_solve(
    gpu_probs, sde_prob, GPUEM();
    dt = 0.01f0, save_everystep = false
)

GPUODEImplicitAlgorithm{AD} <: GPUODEAlgorithm

Developer interface for stiff ODE algorithms in the EnsembleGPUKernel path.

Type Parameters

• AD: Boolean-like type parameter indicating whether the algorithm may derive missing Jacobian and time-gradient information with automatic differentiation. Constructors accept autodiff = Val(true) or Val(false).

Interface Rules

Subtypes must implement the GPUODEAlgorithm rules and additionally provide GPU-safe linear/nonlinear solve support. They must either receive analytical Jacobian/time-gradient functions through the ODEFunction or use the AD parameter to select automatic or finite difference derivative construction. Their nonlinear solve path is expected to use AbstractNLSolver state built by build_nlsolver and to solve static linear systems with DiffEqGPU's GPU-compatible linear algebra utilities.

Examples

solve(
    ensemble_prob, GPURodas4(autodiff = Val(false)),
    EnsembleGPUKernel(backend); trajectories = 10_000
)

AbstractNLSolver

Developer interface for nonlinear solver state used by stiff EnsembleGPUKernel integrators.

Interface Rules

Subtypes are mutable-by-replacement state containers used inside GPU kernel integrators. They must store the current stage correction, temporary stage state, derivative scaling, Jacobian/W-operator constructors, time-step metadata, parameters, and iteration counters needed by nlsolve. A subtype must be usable from GPU-compiled code: all fields and all functions called from nlsolve must be concrete and GPU compatible, and the state update must return the updated solver value rather than relying on host mutation.

Implicit kernel algorithms build this state with build_nlsolver and then call nlsolve(nlsolver, integrator) from the generic step implementation. The required public behavior is tested through generic stiff EnsembleGPUKernel solves and the lower-level vectorized_solve/vectorized_asolve paths rather than by inspecting solver fields.

Fields

Concrete subtypes are expected to provide the fields read by nlsolve, including z, tmp, γ, c, J, W, dt, t, p, iter, and maxiters.

AbstractNLSolverCache

Developer interface for nonlinear solver cache objects used by future EnsembleGPUKernel nonlinear solver implementations.

Interface Rules

Subtypes are reserved for GPU-compatible nonlinear solver caches. Cache fields must be static or otherwise acceptable to KernelAbstractions kernels, and methods using the cache must be callable from device code. A cache must not depend on host-only allocation, reflection, dynamic dispatch, BLAS/LAPACK calls, or non-bitstype closures in the kernel step path.

This is developer-facing API for DiffEqGPU solver implementations. User code should select documented algorithms such as GPURodas4 or GPUKvaerno5 instead of constructing nonlinear solver caches directly.

NLSolver{uType, gamType, tmpType, tType, JType, WType, pType} <: AbstractNLSolver

Concrete Newton-style nonlinear solver state used by stiff EnsembleGPUKernel integrators.

Fields

• z: current nonlinear correction.
• tmp: stage state used by DIRK and multistep methods.
• tmp2: additional temporary state for methods that need a second work vector.
• ztmp: temporary correction storage.
• γ: method coefficient multiplying the nonlinear correction.
• c: stage abscissa.
• α: method-specific stage coefficient.
• κ: nonlinear convergence damping parameter.
• J: callable Jacobian builder J(u, p, t).
• W: callable W-operator builder W(u, p, t).
• dt: current step size.
• t: current step start time.
• p: parameters for the current trajectory.
• iter: current nonlinear iteration count.
• maxiters: maximum nonlinear iterations.

Interface Rules

NLSolver is constructed by build_nlsolver, consumed by nlsolve, and updated by returning a new value with modified fields. It is not intended as a direct user constructor; users select a stiff GPU algorithm and provide GPU-compatible derivative functions or enable algorithm-controlled differentiation.

Device-compatible representation of an SCC-split initialization problem.

problem is the out-of-place residual for the complete initialization system. blocks stores the statically known row/state range of each SCC and whether the block is linear. Keeping only immutable, fully specialized data avoids the mutable cache writers used by the host SCCNonlinearProblem representation.

vectorized_solve(
    probs, prob::Union{ODEProblem, SDEProblem}, alg; dt,
    saveat = nothing, save_everystep = true, debug = false,
    callback = CallbackSet(nothing), tstops = nothing
)

Run a fixed-step EnsembleGPUKernel solve on a batch of compatible problems. This is a developer-facing entry point for packages that need the batched time and state arrays instead of a collection of EnsembleSolutions.

Arguments

• probs: a batch of problems adapted to the backend returned by get_backend(probs). For ODE problems, each element must be compatible with the static, GPU-compilable kernel representation.
• prob: a representative ODEProblem or SDEProblem whose time span and state type determine the output layout. For an ensemble batch this is normally the original problem or probs[1].
• alg: a supported GPUODEAlgorithm or GPUSDEAlgorithm, such as GPUTsit5() or GPUEM().

Keyword Arguments

• dt: required fixed time step.
• saveat: optional scalar, vector, or range of output times. nothing uses the regular time grid implied by dt.
• save_everystep: whether to retain every fixed-step state when saveat === nothing.
• debug: reserved debugging option; it is accepted for compatibility with the solver interface.
• callback: a GPU-compatible callback set. The default is CallbackSet(nothing).
• tstops: optional additional stopping times.

Returns

A pair (ts, us). ts contains the saved times and us contains the corresponding states, with one batch trajectory per column. The arrays remain on the selected backend.

Throws

An ArgumentError or MethodError can be raised when the batch, problem, algorithm, or callback is not compatible with GPU kernel execution. SDE algorithms also throw when the noise structure is unsupported.

Examples

ts, us = DiffEqGPU.vectorized_solve(
    gpu_probs, prob, GPUTsit5(); dt = 0.1f0, save_everystep = false
)

vectorized_asolve(
    probs, prob::ODEProblem, alg; dt = 0.1f0,
    saveat = nothing, save_everystep = false, abstol = 1.0f-6,
    reltol = 1.0f-3, debug = false, callback = CallbackSet(nothing),
    tstops = nothing
)

Run an adaptive EnsembleGPUKernel ODE solve on a batch of compatible problems. This is a developer-facing entry point for packages that need the batched time and state arrays instead of a collection of EnsembleSolutions.

Arguments

• probs: a batch of ODE problems adapted to the backend returned by get_backend(probs). Each element must be compatible with the static, GPU-compilable kernel representation.
• prob: a representative ODEProblem whose time span and state type determine the output layout. For an ensemble batch this is normally the original problem or probs[1].
• alg: a supported GPUODEAlgorithm, such as GPUTsit5() or GPURodas4().

Keyword Arguments

• dt: initial time step; defaults to 0.1f0.
• saveat: optional scalar, vector, or range of output times. nothing uses adaptive internal steps and the value of save_everystep.
• save_everystep: whether to retain the adaptive internal steps. Defaults to false.
• abstol: absolute error tolerance. Defaults to 1.0f-6.
• reltol: relative error tolerance. Defaults to 1.0f-3.
• debug: accepted for compatibility with the vectorized solver interface.
• callback: a GPU-compatible callback set. The default is CallbackSet(nothing).
• tstops: optional additional stopping times.

Returns

A pair (ts, us) containing the saved times and states, with one batch trajectory per column. The arrays remain on the selected backend.

Throws

An ArgumentError or MethodError can be raised when the batch, problem, algorithm, or callback is not compatible with adaptive GPU kernel execution.

Examples

ts, us = DiffEqGPU.vectorized_asolve(
    gpu_probs, prob, GPUTsit5(); dt = 0.1f0, abstol = 1.0f-6, reltol = 1.0f-3
)

vectorized_map_solve(
    probs, alg, ensemblealg::EnsembleArrayAlgorithm, I, adaptive;
    kwargs...
)

Run the EnsembleArrayAlgorithm path directly for a selected set of trajectories. The function is intended for solver and ensemble package developers that need the lower-level solution collection without the EnsembleSolution construction performed by solve.

Arguments

• probs: the collection of problems generated by the ensemble's prob_func.
• alg: a differential-equation algorithm supported by the selected array ensemble implementation. OrdinaryDiffEq algorithms are the usual choice.
• ensemblealg: an EnsembleArrayAlgorithm, normally EnsembleGPUArray(backend) or EnsembleCPUArray().
• I: an iterator of trajectory indices, such as 1:10_000.
• adaptive: whether to use adaptive time stepping.

Keyword Arguments

Additional solver keywords are forwarded to the underlying differential-equation solve. Only keywords supported by the selected algorithm and ensemble implementation are valid.

Returns

The solution returned by the underlying batched solve. Its state arrays contain the selected trajectories in their second dimension.

Throws

An ArgumentError or MethodError can be raised when the problems, algorithm, backend, or solver keywords are incompatible with the selected array implementation.

Examples

sols = DiffEqGPU.vectorized_map_solve(
    probs, Tsit5(), EnsembleCPUArray(), 1:10, false;
    dt = 0.1f0, save_everystep = false
)