SciML / DiffEqGPU.jl

GPU-acceleration routines for DifferentialEquations.jl and the broader SciML scientific machine learning ecosystem
https://docs.sciml.ai/DiffEqGPU/stable/
MIT License
274 stars 28 forks source link

Attempt to isolate cause of dynamic allocation in continuous callbacks #221

Closed utkarsh530 closed 9 months ago

utkarsh530 commented 1 year ago

The PR attempts to remove closure to better isolate the issue.

using DiffEqGPU, SimpleDiffEq, StaticArrays, CUDA, BenchmarkTools, NonlinearSolve
#using OrdinaryDiffEq

CUDA.allowscalar(false)

function f(u, p, t)
    du1 = u[2]
    du2 = -p[1]
    return SVector{2}(du1, du2)
end

u0 = @SVector[45.0f0, 0.0f0]
tspan = (0.0f0, 1.0f0)
p = @SVector [10.0f0]
prob = ODEProblem{false}(f, u0, tspan, p)
prob_func = (prob, i, repeat) -> remake(prob, p = prob.p)
monteprob = EnsembleProblem(prob, safetycopy = false)

function affect!(integrator)
    integrator.u += @SVector[0.0f0, -2.0f0] .* integrator.u
end

function condition(u, t, integrator)
   u[1]
end

cb = ContinuousCallback(condition, affect!; save_positions = (false, false))
sol = solve(monteprob, GPUTsit5(), EnsembleGPUKernel(),
            trajectories = 2,
            adaptive = false, dt = 0.1f0, callback = cb, merge_callbacks = true)

@device_code_llvm dump_module=true sol = solve(monteprob, GPUTsit5(), EnsembleGPUKernel(),
                                               trajectories = 2,
                                               adaptive = false, dt = 0.1f0, callback = cb, merge_callbacks = true)