Kernels for SimpleTauLeaping

[sivasathyaseeelan]

using JumpProcesses, DiffEqBase, Plots
using Test, LinearAlgebra
using KernelAbstractions, Adapt, CUDA
using StableRNGs
rng = StableRNG(12345)

# SEIR model with exposed compartment
β = 0.3 / 1000.0
σ = 0.2
ν = 0.01
p = (β, σ, ν)

regular_rate = (out, u, p, t) -> begin
    out[1] = p[1] * u[1] * u[3]  # β*S*I (infection)
    out[2] = p[2] * u[2]         # σ*E (progression)
    out[3] = p[3] * u[3]         # ν*I (recovery)
end

regular_c = (dc, u, p, t, counts, mark) -> begin
    dc .= 0.0
    dc[1] = -counts[1]           # S: -infection
    dc[2] = counts[1] - counts[2] # E: +infection - progression
    dc[3] = counts[2] - counts[3] # I: +progression - recovery
    dc[4] = counts[3]            # R: +recovery
end

 # Initial state
u0 = [999.0, 0.0, 10.0, 0.0]  # S, E, I, R
tspan = (0.0, 250.0)

# Create JumpProblem
prob_disc = DiscreteProblem(u0, tspan, p)
rj = RegularJump(regular_rate, regular_c, 3)
jump_prob = JumpProblem(prob_disc, Direct(), rj; rng=StableRNG(12345))
sol = solve(EnsembleProblem(jump_prob), SimpleTauLeaping(), EnsembleGPUKernel(CUDABackend()); trajectories = 10, dt = 1.0)
plot(sol)

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

Add any other context about the problem here.

[isaacsas]

Can Github CI test GPU code for free?

[ChrisRackauckas]

No we have to setup the Buildkite

[ChrisRackauckas]

https://buildkite.com/julialang/jumpprocesses-dot-jl

[ChrisRackauckas]

Rebase onto latest master and accept the new runtests.yml and buildkite changes. With that, then add gpu tests into the gpu folder

[isaacsas]

https://docs.nvidia.com/cuda/curand/device-api-overview.html

CUDA has a builtin Poisson RNG, wouldn't we expect that to be more appropriate for GPU use?

[sivasathyaseeelan]

but in that case we can't use kernels for other GPUs

[isaacsas]

AMD offers it too in ROCRAND: differentiable programming for odes review. Can we not dispatch between NVDIA vs. AMD, and otherwise then fall back to pois_rand if in neither of those? I'd imagine there might be substantial performance benefits to using the vendor provided versions.

[isaacsas]

This needs some actual correctness tests for statistics of the obtained solutions (compare to CPU tau-leaping or Gillespie). All this is testing right now is that the code doesn't crash.

[sivasathyaseeelan]

https://github.com/SciML/JumpProcesses.jl/blob/master/test/regular_jumps.jl SimpleTauLeaping needed to be tested too

[isaacsas]

Yup, that needs updating too.

[ChrisRackauckas]

This part should probably just be on CPU

[sivasathyaseeelan]

This can't be done in CPU as we are allocating GPU memory (for current_u_buf) which can only be accessed inside kernels

[isaacsas]

Sorry, I don’t really have much gpu experience so I’ll have to leave this for @ChrisRackauckas to review.

[ChrisRackauckas]

Did you make the upstream issue on PoissonRandom.jl?

[sivasathyaseeelan]

No, I thought @oscardssmith will take that. I will create an issue now

[sivasathyaseeelan]

SciML/PoissonRandom.jl#51

[ChrisRackauckas]

Buildkite doesn't run on forks, so what we can do here is merge as the code seems finished now, but I'll open a master test to double check the test passes.