Performance of ODE solver on GPU in the presence of parameters

[ghost]

using CUDA, CUDA.CUSPARSE
using DifferentialEquations
using LinearAlgebra, MKLSparse, SparseArrays
using BenchmarkTools

N = 2^8

a = sprand(Float64, N, N, 0.2)
x = rand(N)

cu_a = CUSPARSE.CuSparseMatrixCSC(a)
cu_x = CuArray(x)

function func(du, u, p, t)
    CUSPARSE.mv!('T', 1.0, p, u, 0.0, du, 'O')
end

tspan = (0.0, 1.0)
prob = ODEProblem(func, cu_x, tspan, cu_a)

@btime sol = solve(prob, Tsit5(), save_everystep=false, save_start=false)

global cu_b = cu_a

function func_global(du, u, p, t)
    CUSPARSE.mv!('T', 1.0, cu_b, u, 0.0, du, 'O')
end

prob_global = ODEProblem(func_global, cu_x, tspan)

@btime sol_global = solve(prob_global, Tsit5(), save_everystep=false, save_start=false)

I am benchmarking the performance of ODE solver on a GPU with Julia 1.4.2. The above piece of code reproduces the behavior. The benchmarks indicate that the performance is significantly better when I declare the parameters needed for ODE as global variables. On GTX 1060, to solve 'prob' it takes 6.363 s (4197131 allocations: 154.08 MiB) and 'prob_global' takes 14.548 ms (30185 allocations: 1.21 MiB).

[ChrisRackauckas]

Thanks for the report! Fixed in SciML/DiffEqBase.jl#535 . Quite an edge case!