Solving ModelingToolkit models with EnsembleGPUArray

Specific Domains GPU
,
1

Hi everyone,

I’m running an ensemble of small (2–6 state) stiff/DAE thermal RC-network models (ModelingToolkit-built, many independent “buildings” of the same model structure with their own physical parameters) and trying to parallelize the ensemble solve on GPU via DiffEqGPU.jl using EnsembleGPUArray.

Setup: MTK-built RC network → structural_simplify(...; split=false) → ODEProblem → EnsembleProblem → solve(..., Rodas5P(autodiff=AutoFiniteDiff()), EnsembleGPUArray(CUDA.CUDABackend()); trajectories=n, saveat=...).

Symptom: a single trajectory (n=1) solves within a ~3s. The moment I batch ≥2 trajectories, cost jumps dramatically — e.g. ~300 - 2000s real execution time for 6 trajectories.

MWE (2-state analog: one differential + one algebraic state, small per-trajectory forcing table, partly vibe-coded to be honest):

using OrdinaryDiffEq, DiffEqGPU, CUDA, StaticArrays

const N_TIMEPOINTS = 24
const DT  = 3600.0f0
const ncases = 2

function make_table(n)
    SMatrix{N_TIMEPOINTS, n, Float32}(Float32.(5 .+ 10 .* rand(N_TIMEPOINTS, n)))
end

const TABLE = make_table(ncases)

function lookup(t, case_id)
    normalized_t = t / DT
    lower_idx = clamp(floor(Int, normalized_t) + 1, 1, N_TIMEPOINTS - 1)
    frac = normalized_t - floor(normalized_t)
    col = Int(case_id)
    tbl = TABLE
    return tbl[lower_idx, col] + frac * (tbl[lower_idx + 1, col] - tbl[lower_idx, col])
end

function thermal_rhs!(dT, T, p, t)
    R1, C1, Cond, case_id = p[1], p[2], p[3], p[4]
    Tout = lookup(t, case_id)
    dT[1] = (Tout - T[1]) / (R1 * C1)   # differential state
    dT[2] = Cond * (T[1] - T[2])        # algebraic constraint
end

Mass_mat = [1.0 0.0; 0.0 0.0]

u0 = Float32[20.0, 20.0]
tspan = (0.0, 24 * 3600.0 * 30)
p0 = Float32[1.0, 1000.0, 5.0, 1.0]

func = ODEFunction(thermal_rhs!, mass_matrix = Mass_mat)
prob = ODEProblem(func, u0, tspan, p0)

function assign_case(prob, ctx)
    case_id = Float32(getfield(ctx, :sim_id))
    newp = copy(prob.p)
    newp[4] = Float32(case_id)
    remake(prob; p = newp)
end

ensprob = EnsembleProblem(prob; prob_func = assign_case, safetycopy = false)

println("Starting single solve.")
@time sol1 = solve(ensprob, Rodas5P(autodiff = AutoFiniteDiff()),
                    EnsembleGPUArray(CUDA.CUDABackend());
                    trajectories = 1, saveat = 3600.0f0)
println("Single done, converged: $(sol1.converged), steps: ", sol1.stats.naccept + sol1.stats.nreject)

@time sol6 = solve(ensprob, Rodas5P(autodiff = AutoFiniteDiff()),
                    EnsembleGPUArray(CUDA.CUDABackend());
                    trajectories = ncases, saveat = 3600.0f0)
println("$(ncases) cases done, converged: $(sol6.converged)")

On my machine n=1 finishes in ~3 s (~11 M allocations); n=2 ranges between ~22 and 90s (140 / 550 M allocations), n=10 gives 50 - 100s (270 - 520 M allocations) and n=100 gives ~280 s (1.3 G allocations). I am aware that compared to CPU parallelized versions, this is likely to only produce comparable results once I scale the problem size up to a few thousand, however these high execution times were making me sceptical about the feasibility of scaling this up further. Where is this extreme jump between a single case and 2 cases coming from? Is this expected/known behaviour? Is there any way to improve these runtimes?
I am still rather new to Julia and GPU computing in general, so I’d be very grateful for any pointers you might have! I am working with Julia 1.12.6, DiffEqGPU 3.15.4 and CUDA.jl 6.2.1.

2

that 1-to-2 cliff is the giveaway, its not a scaling problem, its your rhs doing something the gpu path hates. my money is on allocations: MTK-simplified functions often allocate under the hood, and on EnsembleGPUArray that either errors or silently tanks performance. fastest check is @allocated on a single thermal_rhs! call, plus one trajectory on cpu vs gpu for a baseline. if the rhs allocates, thats your cliff: make it fully non-allocating (static arrays all the way down, no heap captures in the lookup table) and generate the jacobian symbolically instead of AutoFiniteDiff, finite differencing on device is rough. for 2-6 state systems id honestly just use EnsembleGPUKernel, the Array path only wins once per-trajectory work is big enough to hide launch overhead.

3

Is there any reason to not use EnsembleGPUKernel here?

4

I’ve already had a bunch of discussions about this and there are two main reasons that lead to focusing on EnsembleGPUArray:

  1. In the long run, we might want to introduce some coupling between buildings which would potentially lead to less uniform model structures/bigger systems and we hoped to keep a bit more flexibility here using the Array version.
  2. Coming from a much more complex model than this example here, I was simply a bit scared of the restructuring for my initial conditions and parameters I would probably need to do to use the Kernel version.

But if the answer to my questions simply is that the kernel launches create this type of overhead, I will definitely reconsider this choice.

5

The kernel generation is actually a bit easier for an MTK model given how the code generation has control over it, and most MTK models would be well-sized for it, so I made sure to get the kernel generation approach working with MTK-generated models and that’s how it’s documented. Since that’s generally a lot faster, I’d recommend going that route.

I’m not sure EnsembleGPUArray will buy you much in this case. I think that just needs some MTK GPU codegen targeting, which is something that is moving pretty fast.