Within-Method GPU Parallelism of Ordinary Differential Equation Solves

Within-Method GPU Parallelism for ODE solvers is a method for accelerating large ODE solves with regularity, i.e., only using array-based “vectorized” operations like linear algebra, maps, and broadcast statements. In these cases, the solve can be GPU accelerated simply by placing the initial condition array on the GPU. As a quick example:

using OrdinaryDiffEq, CUDA, LinearAlgebra
function f(du, u, p, t)
    return mul!(du, A, u)
end

A = cu(-rand(3, 3))
u0 = cu([1.0; 0.0; 0.0])
tspan = (0.0f0, 100.0f0)

prob = ODEProblem(f, u0, tspan)
sol = solve(prob, Tsit5())
sol = solve(prob, Rosenbrock23())

retcode: Success
Interpolation: specialized 2nd order "free" stiffness-aware interpolation
t: 145-element Vector{Float32}:
   0.0
   0.00012846051
   0.0021160643
   0.008777656
   0.019103168
   0.039619245
   0.06823447
   0.113024265
   0.17060909
   0.248942
   ⋮
  94.23676
  95.0316
  95.82644
  96.62128
  97.416115
  98.21095
  99.00579
  99.80063
 100.0
u: 145-element Vector{CUDACore.CuArray{Float32, 1, CUDACore.DeviceMemory}}:
 Float32[1.0, 0.0, 0.0]
 Float32[0.9999266, -0.00011696152, -6.706105f-5]
 Float32[0.9987933, -0.0019246978, -0.001103668]
 Float32[0.99502146, -0.007956853, -0.0045643635]
 Float32[0.9892553, -0.017226433, -0.009887526]
 Float32[0.9780834, -0.035359725, -0.020319209]
 Float32[0.9631196, -0.06003515, -0.03455526]
 Float32[0.9410907, -0.097276665, -0.056136005]
 Float32[0.9151422, -0.14281873, -0.082696095]
 Float32[0.88383615, -0.20087713, -0.11685908]
 ⋮
 Float32[4.363863f11, -4.573449f11, -3.8751814f11]
 Float32[5.507078f11, -5.77157f11, -4.890375f11]
 Float32[6.949785f11, -7.283567f11, -6.171522f11]
 Float32[8.7704424f11, -9.191666f11, -7.788295f11]
 Float32[1.1068063f12, -1.1599637f12, -9.82862f11]
 Float32[1.3967599f12, -1.4638432f12, -1.2403456f12]
 Float32[1.7626736f12, -1.8473307f12, -1.5652829f12]
 Float32[2.2244468f12, -2.331282f12, -1.9753451f12]
 Float32[2.3578597f12, -2.4711025f12, -2.093818f12]

Notice that both stiff and non-stiff ODE solvers were used here.

Note

Time span was changed to Float32 types, as GPUs generally have very slow Float64 operations, usually around 1/32 of the speed of Float32. cu(x) on an array automatically changes an Array{Float64} to a CuArray{Float32}. If this is not intended, use the CuArray constructor directly. For more information on GPU Float64 performance issues, search around Google for discussions like this.

Warn

Float32 precision is sometimes not enough precision to accurately solve a stiff ODE. Make sure that the precision is necessary by investigating the condition number of the Jacobian. If this value is well-above 1e8, use Float32 with caution!

Restrictions of CuArrays

Note that all the rules of CUDA.jl apply when CuArrays are being used in the solver. While for most of the AbstractArray interface they act similarly to Arrays, such as having valid broadcasting operations (x .* y) defined, they will work on GPUs. For more information on the rules and restrictions of CuArrays, see this page from the CUDA.jl documentation.