Options for solving ODE systems on GPUs?

Question

I would like to farm out solving systems of ODEs onto GPUs, in a 'trivially parallelisable' setting. For example, doing a sensitivity analysis with 512 different parameter sets.

Ideally I want to do ODE solving with a smart adaptive timestep solver like CVODE, rather than a fixed timestep like Forward Euler, but running it on an NVIDIA GPU instead of CPU.

Has anyone done this? Are there libraries for it?

Answer 1

You may want to look into Boost's odeint library and Thrust. They can be combined as discussed here.

Answer 2

DifferentialEquations.jl library is a library for a high level language (Julia) which has tools for automatically transforming the ODE system to an optimized version for parallel solution on GPUs. There are two forms of parallelism that can be employed: array-based parallelism for large ODE systems and parameter parallelism for parameter studies on relatively small (<100) ODE systems. It supports high order implicit and explicit methods and routinely outperforms or matches other systems in benchmarks (at the very least, it wraps the others so it's easy to check and use them!)

For this specific functionality, you may want to take a look at DiffEqGPU.jl which is the module for automated parameter parallelism. The DifferentialEquations.jl library has functionality for parallel parameter studies, and this module augments the existing configurations to make the study happen automatically in parallel. What one does is transforms their existing ODEProblem (or other DEProblem like SDEProblem) into an EnsembleProblem and specify with a prob_func how the other problems are generated from the prototype. The following solves 10,000 trajectories of the Lorenz equation on the GPU with a high order explicit adaptive method:

using OrdinaryDiffEq, DiffEqGPU
function lorenz(du,u,p,t)
 @inbounds begin
     du[1] = p[1]*(u[2]-u[1])
     du[2] = u[1]*(p[2]-u[3]) - u[2]
     du[3] = u[1]*u[2] - p[3]*u[3]
 end
 nothing
end

u0 = Float32[1.0;0.0;0.0]
tspan = (0.0f0,100.0f0)
p = (10.0f0,28.0f0,8/3f0)
prob = ODEProblem(lorenz,u0,tspan,p)
prob_func = (prob,i,repeat) -> remake(prob,p=rand(Float32,3).*p)
monteprob = EnsembleProblem(prob, prob_func = prob_func)
@time sol = solve(monteprob,Tsit5(),EnsembleGPUArray(),trajectories=10_000,saveat=1.0f0)

Notice the user needs to write no GPU code, and with a single RTX 2080 this benchmarks as a 5x improvement over using a 16 core Xeon machine with multithreaded parallelism. One can then check out the README for how to do things like utilize multiple GPUs and doing multiprocessing+GPUs for utilizing a full cluster of GPUs simultaneously. Note that switching to multithreading instead of GPUs is one line change: EnsembleThreads() instead of EnsembleGPUArray().

Then for implicit solvers, the same interface holds. For example, the following uses high order Rosenbrock and implicit Runge-Kutta methods:

function lorenz_jac(J,u,p,t)
 @inbounds begin
     σ = p[1]
     ρ = p[2]
     β = p[3]
     x = u[1]
     y = u[2]
     z = u[3]
     J[1,1] = -σ
     J[2,1] = ρ - z
     J[3,1] = y
     J[1,2] = σ
     J[2,2] = -1
     J[3,2] = x
     J[1,3] = 0
     J[2,3] = -x
     J[3,3] = -β
 end
 nothing
end

function lorenz_tgrad(J,u,p,t)
 nothing
end

func = ODEFunction(lorenz,jac=lorenz_jac,tgrad=lorenz_tgrad)
prob_jac = ODEProblem(func,u0,tspan,p)
monteprob_jac = EnsembleProblem(prob_jac, prob_func = prob_func)

@time solve(monteprob_jac,Rodas5(linsolve=LinSolveGPUSplitFactorize()),EnsembleGPUArray(),dt=0.1,trajectories=10_000,saveat=1.0f0)
@time solve(monteprob_jac,TRBDF2(linsolve=LinSolveGPUSplitFactorize()),EnsembleGPUArray(),dt=0.1,trajectories=10_000,saveat=1.0f0)

While this form requires that you give a Jacobian in order to be used on the GPU (currently, will be fixed soon), the DifferentialEquations.jl documentation demonstrates how to do automatic symbolic Jacobian calculations on numerically defined functions, so there is still no manual labor here. I would highly recommend these algorithms because the branching logic of a method like CVODE generally causes thread desync and doesn't seem to perform as well as a Rosenbrock method in these types of scenarios anyways.

By using DifferentialEquations.jl, you also get access to the full library, which includes functionality like global sensitivity analysis which can make use of this GPU-acceleration. It is also compatible with dual numbers for fast local sensitivity analysis. The GPU-based code gets all of the features of DifferentialEquations.jl, like the event handling and the large set of ODE solvers which are optimized for different types of problems, meaning it's not just a simple one-off GPU ODE solver but instead a part of a fully featured system which also happens to have efficient GPU support.