Tutorial utilising MTK to highlight the utility of DiffEqGPU.jl

[H-Sax]

Is your feature request related to a problem? Please describe.

It would be nice to see a tutorial where MTK examples show how to create a GPU compatible form to utilise the DiffEqGPU solvers

[ChrisRackauckas]

The simple thing:

using OrdinaryDiffEq, ModelingToolkit
using SymbolicIndexingInterface
using ModelingToolkit: t_nounits as t, D_nounits as D

@parameters σ ρ β
@variables x(t) y(t) z(t)

eqs = [D(D(x)) ~ σ * (y - x),
    D(y) ~ x * (ρ - z) - y,
    D(z) ~ x * y - β * z]

@mtkbuild sys = ODESystem(eqs, t)

u0 = [D(x) => 2f0,
    x => 1f0,
    y => 0f0,
    z => 0f0]

p = [σ => 28f0,
    ρ => 10f0,
    β => 8f0 / 3f0]

tspan = (0f0, 100f0)
prob = ODEProblem(sys, u0, tspan, p, jac = true)

prob_func = function (prob, i, repeat)
	prob = deepcopy(prob)
	prob.p.tunable .= rand(Float32, 3) .* prob.p.tunable
	prob
end

monteprob = EnsembleProblem(prob, prob_func = prob_func, safetycopy = false)
sol = solve(monteprob, Tsit5(), EnsembleGPUArray(CUDA.CUDABackend()), trajectories = 10_000,
    saveat = 1.0f0);

Should just work. But you probably don't want to remake using such low level stuff and instead use the SymbolicIndexingInterface. For that you probably want the following:

probchanger = SymbolicIndexingInterface.setp_oop(prob, [σ, ρ, β]);
prob_func = (prob, i, repeat) -> (probchanger(prob, rand(Float32, 3) .* [28f0, 10f0, 8f0/3f0]); prob)
monteprob = EnsembleProblem(prob, prob_func = prob_func, safetycopy = false)
sol = solve(monteprob, Tsit5(), EnsembleThreads(), trajectories = 10_000, saveat = 1.0f0);

using DiffEqGPU, CUDA
sol = solve(monteprob, Tsit5(), EnsembleGPUArray(CUDA.CUDABackend()), trajectories = 10_000,
    saveat = 1.0f0);

@AayushSabharwal is going to improve setsym_oop which allows for mixing initial conditions and parameters in what's set to make it even easier. I still need to test this works in a GPU kernel too, you might need to flat eval of the RGF so it's cache-free.

[utkarsh530]

https://github.com/utkarsh530/GPUODEBenchmarks/blob/1a90256220a8a557b238e37cb3101ba7fd5319fc/GPU_ODE_Julia/src/GPU_ODE_Julia.jl#L5

I had this written up to compile with just SArrays

[ChrisRackauckas]

using OrdinaryDiffEqTsit5, ModelingToolkit, StaticArrays
using ModelingToolkit: t_nounits as t, D_nounits as D

@parameters σ ρ β
@variables x(t) y(t) z(t)

eqs = [D(D(x)) ~ σ * (y - x),
    D(y) ~ x * (ρ - z) - y,
    D(z) ~ x * y - β * z]

@mtkbuild sys = ODESystem(eqs, t) split=false

u0 = SA[D(x) => 2f0,
    x => 1f0,
    y => 0f0,
    z => 0f0]

p = SA[σ => 28f0,
    ρ => 10f0,
    β => 8f0 / 3]

tspan = (0f0, 100f0)
_prob = ODEProblem{false}(sys, u0, tspan, p, jac = true, split=false)
prob = DiffEqGPU.ImmutableODEProblem{false, SciMLBase.FullSpecialize}(_prob.f.f, _prob.u0, _prob.tspan, Float32.(prob.p))

sol = solve(prob, Tsit5())

using SymbolicIndexingInterface
p_setter = setp(sys, [σ, ρ, β])


using DiffEqGPU, CUDA
function prob_func2(prob, i, repeat)
    prob # = remake(prob, p = SVector{10, Float32}(@SVector(rand(Float32,3))..., prob.p[4:end]...))
end
monteprob = EnsembleProblem(prob, prob_func = prob_func2, safetycopy = false)
sol = solve(monteprob, GPUTsit5(), EnsembleGPUKernel(CUDA.CUDABackend()),
    trajectories = 10_000)

julia> sol = solve(monteprob, GPUTsit5(), EnsembleGPUKernel(CUDA.CUDABackend()),
           trajectories = 10_000)
ERROR: CUDA error: an illegal memory access was encountered (code 700, ERROR_ILLEGAL_ADDRESS)
Stacktrace:
  [1] throw_api_error(res::CUDA.cudaError_enum)
    @ CUDA C:\Users\accou\.julia\packages\CUDA\oymHm\lib\cudadrv\libcuda.jl:30
  [2] check(f::CUDA.var"#7#8"{Base.RefValue{UInt64}, Base.RefValue{UInt64}})
    @ CUDA C:\Users\accou\.julia\packages\CUDA\oymHm\lib\cudadrv\libcuda.jl:37
  [3] cuMemGetInfo_v2(free::Base.RefValue{UInt64}, total::Base.RefValue{UInt64})
    @ CUDA C:\Users\accou\.julia\packages\GPUToolbox\cZlg7\src\ccalls.jl:33
  [4] memory_info
    @ C:\Users\accou\.julia\packages\CUDA\oymHm\lib\cudadrv\memory.jl:869 [inlined]
  [5] free_memory
    @ C:\Users\accou\.julia\packages\CUDA\oymHm\lib\cudadrv\memory.jl:878 [inlined]
  [6] maybe_collect(will_block::Bool)
    @ CUDA C:\Users\accou\.julia\packages\CUDA\oymHm\src\memory.jl:80
  [7] maybe_collect
    @ C:\Users\accou\.julia\packages\CUDA\oymHm\src\memory.jl:65 [inlined]
  [8] pool_alloc
    @ C:\Users\accou\.julia\packages\CUDA\oymHm\src\memory.jl:621 [inlined]
  [9] (::CUDA.var"#1129#1131"{CUDA.DeviceMemory, Int64})()
    @ CUDA C:\Users\accou\.julia\packages\CUDA\oymHm\src\array.jl:76
 [10] cached_alloc(f::CUDA.var"#1129#1131"{CUDA.DeviceMemory, Int64}, key::Tuple{UnionAll, CuDevice, DataType, Int64})
    @ GPUArrays C:\Users\accou\.julia\packages\GPUArrays\uiVyU\src\host\alloc_cache.jl:36
 [11] CuArray{DiffEqGPU.ImmutableODEProblem{…}, 1, CUDA.DeviceMemory}(::UndefInitializer, dims::Tuple{Int64})
    @ CUDA C:\Users\accou\.julia\packages\CUDA\oymHm\src\array.jl:75
 [12] CuArray
    @ C:\Users\accou\.julia\packages\CUDA\oymHm\src\array.jl:421 [inlined]
 [13] CuArray
    @ C:\Users\accou\.julia\packages\CUDA\oymHm\src\array.jl:426 [inlined]
 [14] CuArray
    @ C:\Users\accou\.julia\packages\CUDA\oymHm\src\array.jl:435 [inlined]
 [15] convert
    @ C:\Users\accou\.julia\packages\GPUArrays\uiVyU\src\host\construction.jl:4 [inlined]
 [16] adapt_storage
    @ C:\Users\accou\.julia\packages\CUDA\oymHm\src\array.jl:716 [inlined]
 [17] adapt_structure
    @ C:\Users\accou\.julia\packages\Adapt\0LkKU\src\Adapt.jl:57 [inlined]
 [18] adapt
    @ C:\Users\accou\.julia\packages\Adapt\0LkKU\src\Adapt.jl:40 [inlined]
 [19] adapt_storage
    @ C:\Users\accou\.julia\packages\CUDA\oymHm\src\CUDAKernels.jl:34 [inlined]
 [20] adapt_structure
    @ C:\Users\accou\.julia\packages\Adapt\0LkKU\src\Adapt.jl:57 [inlined]
 [21] adapt
    @ C:\Users\accou\.julia\packages\Adapt\0LkKU\src\Adapt.jl:40 [inlined]
 [22] batch_solve_up_kernel(ensembleprob::EnsembleProblem{…}, probs::Vector{…}, alg::GPUTsit5, ensemblealg::EnsembleGPUKernel{…}, I::UnitRange{…}, adaptive::Bool; kwargs::@Kwargs{…})
    @ DiffEqGPU C:\Users\accou\.julia\dev\DiffEqGPU\src\solve.jl:276
 [23] batch_solve(ensembleprob::EnsembleProblem{…}, alg::GPUTsit5, ensemblealg::EnsembleGPUKernel{…}, I::UnitRange{…}, adaptive::Bool; kwargs::@Kwargs{…})
    @ DiffEqGPU C:\Users\accou\.julia\dev\DiffEqGPU\src\solve.jl:173
 [24] macro expansion
    @ .\timing.jl:421 [inlined]
 [25] __solve(ensembleprob::EnsembleProblem{…}, alg::GPUTsit5, ensemblealg::EnsembleGPUKernel{…}; trajectories::Int64, batch_size::Int64, unstable_check::Function, adaptive::Bool, kwargs::@Kwargs{})
    @ DiffEqGPU C:\Users\accou\.julia\dev\DiffEqGPU\src\solve.jl:55
 [26] solve(::EnsembleProblem{…}, ::GPUTsit5, ::Vararg{…}; kwargs::@Kwargs{…})
    @ DiffEqBase C:\Users\accou\.julia\dev\DiffEqBase\src\solve.jl:1215
 [27] top-level scope
    @ REPL[26]:1
Some type information was truncated. Use `show(err)` to see complete types.

Some things may need to strip better. @AayushSabharwal this is related to the static array not carrying forward.

[ChrisRackauckas]

using DiffEqGPU
using OrdinaryDiffEqTsit5, ModelingToolkit, StaticArrays
using ModelingToolkit: t_nounits as t, D_nounits as D

@parameters σ ρ β
@variables x(t) y(t) z(t)

eqs = [D(D(x)) ~ σ * (y - x),
    D(y) ~ x * (ρ - z) - y,
    D(z) ~ x * y - β * z]

@mtkbuild sys = ODESystem(eqs, t) split=false

u0 = SA[D(x) => 2f0,
    x => 1f0,
    y => 0f0,
    z => 0f0]

p = SA[σ => 28f0,
    ρ => 10f0,
    β => 8f0 / 3f0]

tspan = (0f0, 100f0)
prob = ODEProblem{false}(sys, u0, tspan, p, split=true)
prob = remake(prob, p = p = SVector{10, Float32}(prob.p...))
sol = solve(prob, Tsit5())

using SymbolicIndexingInterface
p_setter = setp_oop(sys, [σ, ρ, β])

using DiffEqGPU, CUDA
function prob_func2(prob, i, repeat)
    remake(prob, p = p_setter(prob,@SVector(rand(Float32,3))))
end
monteprob = EnsembleProblem(prob, prob_func = prob_func2, safetycopy = false)
sol = solve(monteprob, GPUTsit5(), EnsembleGPUKernel(CUDA.CUDABackend()),
    trajectories = 10_000)

It's basically working now, but:

1. split=true goes to the non-split form?
2. It's converting the parameters to Float64 in the backend, hence the weird prob = remake(prob, p = p = SVector{10, Float32}(prob.p...))