This library is a component package of the DifferentialEquations.jl ecosystem. It includes functionality for making use of GPUs in the differential equation solvers.
There are two very different ways that one can
accelerate an ODE solution with GPUs. There is one case where u
is very big and f
is very expensive but very structured, and you use GPUs to accelerate the computation
of said f
. The other use case is where u
is very small but you want to solve the ODE
f
over many different initial conditions (u0
) or parameters p
. In that case, you can
use GPUs to parallelize over different parameters and initial conditions. In other words:
Type of Problem | SciML Solution |
---|---|
Accelerate a big ODE | Use CUDA.jl's CuArray as u0 |
Solve the same ODE with many u0 and p |
Use DiffEqGPU.jl's EnsembleGPUArray and EnsembleGPUKernel |
using OrdinaryDiffEq, CUDA, LinearAlgebra
u0 = cu(rand(1000))
A = cu(randn(1000, 1000))
f(du, u, p, t) = mul!(du, A, u)
prob = ODEProblem(f, u0, (0.0f0, 1.0f0)) # Float32 is better on GPUs!
sol = solve(prob, Tsit5())
using DiffEqGPU, 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, i, repeat) -> remake(prob, p = (@SVector rand(Float32, 3)) .* p)
monteprob = EnsembleProblem(prob, prob_func = prob_func, safetycopy = false)
@time sol = solve(monteprob, GPUTsit5(), EnsembleGPUKernel(), trajectories = 10_000,
adaptive = false, dt = 0.1f0)