EnsembleGPUKernel

EnsembleGPUKernel(backend, cpu_offload = 0.0)

A massively parallel ensemble algorithm that generates one GPU kernel for the complete fixed-step ODE or SDE solve. It minimizes kernel-launch overhead, at the cost of stricter requirements on the problem and solver.

Fields

• dev: the KernelAbstractions backend used to launch the fused kernel.
• cpu_offload: the fraction of trajectories solved on the CPU. The default is 0.0.

Arguments

• backend: a KernelAbstractions backend, such as CUDA.CUDABackend() or KernelAbstractions.CPU().
• cpu_offload: the fraction of trajectories to offload to CPU execution. The two-argument constructor stores this value as a Float64; the one-argument constructor defaults to 0.0.

Returns

An EnsembleGPUKernel algorithm selector.

Throws

The constructor itself does not validate kernel compatibility. Unsupported state containers, right-hand sides, callbacks, or solver combinations throw when the ensemble is solved.

Limitations

• Not all standard Julia f functions are allowed. Only Julia f functions which are capable of being compiled into a GPU kernel are allowed. This notably means that certain features of Julia can cause issues inside a kernel, like:

  • Allocating memory (building arrays)
  • Linear algebra (anything that calls BLAS)
  • Broadcast

• Only out-of-place f definitions are allowed. Coupled with the requirement of not allowing for memory allocations, this means that the ODE must be defined with StaticArray initial conditions.

• Only specific ODE solvers are allowed. This includes:

  • GPUTsit5
  • GPUVern7
  • GPUVern9

• To use multiple GPUs over clusters, one must manually set up one process per GPU. See the multi-GPU tutorial for more details.

Examples

using DiffEqGPU, CUDA, OrdinaryDiffEq, StaticArrays

function lorenz(u, p, t)
    σ = p[1]
    ρ = p[2]
    β = p[3]
    du1 = σ * (u[2] - u[1])
    du2 = u[1] * (ρ - u[3]) - u[2]
    du3 = u[1] * u[2] - β * u[3]
    return SVector{3}(du1, du2, du3)
end

u0 = @SVector [1.0f0; 0.0f0; 0.0f0]
tspan = (0.0f0, 10.0f0)
p = @SVector [10.0f0, 28.0f0, 8 / 3.0f0]
prob = ODEProblem{false}(lorenz, u0, tspan, p)
prob_func = (prob, ctx) -> remake(prob, p = (@SVector rand(Float32, 3)) .* p)
monteprob = EnsembleProblem(prob; prob_func, safetycopy = false)

@time sol = solve(
    monteprob, GPUTsit5(), EnsembleGPUKernel(CUDA.CUDABackend()), trajectories = 10_000,
    adaptive = false, dt = 0.1f0
)

GPUTsit5()

Fifth-order Tsitouras Runge-Kutta method specialized for EnsembleGPUKernel ODE solves.

Use GPUTsit5 as the ODE algorithm when solving an EnsembleProblem with EnsembleGPUKernel:

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

GPUTsit5 supports the EnsembleGPUKernel restrictions, including out-of-place ODE functions over GPU-compatible static state containers. For a similar CPU implementation, see SimpleATsit5 from SimpleDiffEq.jl.

Returns

A GPUTsit5 algorithm selector.

Throws

The constructor has no runtime checks. solve throws when the problem is not compatible with the EnsembleGPUKernel restrictions.

GPUTsit5IController()

Fifth-order Tsitouras method with a current-error-only integral (I) step-size controller for EnsembleGPUKernel. The Runge-Kutta tableau, error estimate, and interpolation are the same as GPUTsit5; fixed-step solves are identical.

For adaptive solves, the next step is multiplied by clamp(0.9 * E^(-1/5), 0.2, 5), where E is the scaled error norm. Zero error uses the maximum growth factor. The same rule applies after rejected steps. Use EnsembleGPUKernel(backend, 0.0); CPU offloading is not supported.

Generally prefer the PI-controlled GPUTsit5, which has more stable step-size control. On simple nonstiff benchmarks such as Lorenz, an I controller can appear faster at identical tolerances partly because it achieves lower accuracy. Compare performance at matched achieved error.

GPUVern7()

Seventh-order Verner Runge-Kutta method specialized for EnsembleGPUKernel ODE solves.

Use GPUVern7 for non-stiff ODE ensemble solves that satisfy the EnsembleGPUKernel kernel-generation restrictions:

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

Returns

A GPUVern7 algorithm selector.

Throws

solve throws when the problem is not compatible with the EnsembleGPUKernel restrictions.

GPUVern9()

Ninth-order Verner Runge-Kutta method specialized for EnsembleGPUKernel ODE solves.

Use GPUVern9 for high-accuracy non-stiff ODE ensemble solves that satisfy the EnsembleGPUKernel kernel-generation restrictions:

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

Returns

A GPUVern9 algorithm selector.

Throws

solve throws when the problem is not compatible with the EnsembleGPUKernel restrictions.

GPUEM()

Euler-Maruyama method with weak order 1.0 specialized for EnsembleGPUKernel SDE solves.

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

Returns

A GPUEM algorithm selector.

Throws

solve throws when the SDE noise structure is not supported by the kernel implementation.

GPUSIEA()

Weak order 2.0 SIEA method for Ito SDEs specialized for EnsembleGPUKernel SDE solves.

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

Returns

A GPUSIEA algorithm selector.

Throws

solve throws when the SDE noise structure is not diagonal or otherwise unsupported.

GPURosenbrock23(; autodiff = Val{true}())

Second/third-order Rosenbrock-W method specialized for stiff EnsembleGPUKernel ODE solves.

Keyword Arguments

• autodiff: whether automatic differentiation is used for derivative generation. This accepts either a Bool or Val{Bool}; pass false or Val{false}() when providing the required derivatives manually.

solve(
    ensemble_prob, GPURosenbrock23(), EnsembleGPUKernel(backend);
    trajectories = 10_000
)

Returns

A GPURosenbrock23 algorithm selector.

Throws

solve throws when the problem does not provide GPU-compatible derivative information.

GPURodas4(; autodiff = Val{true}())

Fourth-order Rosenbrock method specialized for stiff EnsembleGPUKernel ODE solves.

Keyword Arguments

• autodiff: whether automatic differentiation is used for derivative generation. This accepts either a Bool or Val{Bool}; pass false or Val{false}() when providing the required derivatives manually.

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

Returns

A GPURodas4 algorithm selector.

Throws

solve throws when the problem does not provide GPU-compatible derivative information.

GPURodas5P(; autodiff = Val{true}())

Fifth-order Rosenbrock method specialized for stiff EnsembleGPUKernel ODE solves.

Keyword Arguments

• autodiff: whether automatic differentiation is used for derivative generation. This accepts either a Bool or Val{Bool}; pass false or Val{false}() when providing the required derivatives manually.

solve(ensemble_prob, GPURodas5P(), EnsembleGPUKernel(backend); trajectories = 10_000)

Returns

A GPURodas5P algorithm selector.

Throws

solve throws when the problem does not provide GPU-compatible derivative information.

GPUKvaerno3(; autodiff = Val{true}())

Third-order ESDIRK method specialized for stiff EnsembleGPUKernel ODE solves.

Keyword Arguments

• autodiff: whether automatic differentiation is used for derivative generation. This accepts either a Bool or Val{Bool}; pass false or Val{false}() when providing the required derivatives manually.

solve(ensemble_prob, GPUKvaerno3(), EnsembleGPUKernel(backend); trajectories = 10_000)

Returns

A GPUKvaerno3 algorithm selector.

Throws

solve throws when the problem does not provide GPU-compatible derivative information.

GPUKvaerno5(; autodiff = Val{true}())

Fifth-order ESDIRK method specialized for stiff EnsembleGPUKernel ODE solves.

Keyword Arguments

• autodiff: whether automatic differentiation is used for derivative generation. This accepts either a Bool or Val{Bool}; pass false or Val{false}() when providing the required derivatives manually.

solve(ensemble_prob, GPUKvaerno5(), EnsembleGPUKernel(backend); trajectories = 10_000)

Returns

A GPUKvaerno5 algorithm selector.

Throws

solve throws when the problem does not provide GPU-compatible derivative information.