Symbolic-Numeric GPU Acceleration with ModelingToolkit

ModelingToolkit.jl is a symbolic-numeric computing system which allows for using symbolic transformations of equations before code generation. The goal is to improve numerical simulations by first turning them into the simplest set of equations to solve and exploiting things that normally cannot be done by hand. Those exact features are also potentially useful for GPU computing, and thus this tutorial showcases how to effectively use MTK with DiffEqGPU.jl.

The core aspect to doing this right is two things. First of all, MTK respects the types chosen by the user, and thus in order for GPU kernel generation to work the user needs to ensure that the problem that is built uses static structures. For example this means that the u0 and p specifications should use static arrays. This looks as follows:

using OrdinaryDiffEq, 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,
]

@named lorenz = System(eqs, t)
sys = mtkcompile(lorenz; split = false)

op = @SVector [
    D(x) => 2.0f0,
    x => 1.0f0,
    y => 0.0f0,
    z => 0.0f0,
    σ => 28.0f0,
    ρ => 10.0f0,
    β => 8.0f0 / 3.0f0,
]

tspan = (0.0f0, 100.0f0)
prob = ODEProblem{false}(sys, op, tspan)
sol = solve(prob, Tsit5())

retcode: Success
Interpolation: specialized 4th order "free" interpolation
t: 1498-element Vector{Float32}:
   0.0
   0.00014131522
   0.00070821616
   0.0022506418
   0.0054032984
   0.011556536
   0.022820365
   0.04294171
   0.07692823
   0.11737434
   ⋮
  99.505005
  99.576675
  99.642845
  99.72166
  99.79972
  99.86805
  99.933105
  99.99022
 100.0
u: 1498-element Vector{StaticArraysCore.SVector{4, Float32}}:
 [0.0, 0.0, 1.0, 2.0]
 [9.9860955f-8, 0.0014132521, 1.0002824, 1.9960454]
 [2.5092213f-6, 0.007084652, 1.0014095, 1.9802262]
 [2.537006f-5, 0.022531182, 1.0044309, 1.9375514]
 [0.0001465516, 0.05417118, 1.0104039, 1.852005]
 [0.000672939, 0.11615784, 1.0213017, 1.6916379]
 [0.0026377225, 0.23024164, 1.038805, 1.4212998]
 [0.00937996, 0.43503767, 1.0631571, 1.0161755]
 [0.029963775, 0.7798868, 1.0892379, 0.56920403]
 [0.068470694, 1.1825796, 1.1081138, 0.4366887]
 ⋮
 [6.380222, 0.22771135, -10.930152, 80.160675]
 [5.7150326, -1.9418658, -4.589691, 93.904305]
 [4.9683547, -2.2139828, 1.643521, 92.4817]
 [3.6729689, 0.26693404, 8.438807, 78.372765]
 [5.360412, 5.2036157, 13.830003, 59.637257]
 [10.916073, 6.850487, 17.324488, 42.0887]
 [14.709508, 2.5671852, 19.334965, 17.774876]
 [12.63868, -1.9967552, 19.496784, -13.401418]
 [11.898012, -2.4069638, 19.336784, -19.326712]

There are two things to notice here. The first is the split = false argument to mtkcompile. By default MTK builds an MTKParameters object, which stores the parameters in separate buffers grouped by how they are used. That object holds Vectors and is therefore not isbits, so it cannot be placed into a GPU kernel. split = false instead puts every parameter into a single flat buffer.

The second is that the operating point op is given as a single StaticArrays.jl vector of pairs, using Float32 values. MTK builds u0 and p in the same container type it was handed, so a static vector in gives static vectors out:

typeof(prob.u0), typeof(prob.p)

(StaticArraysCore.SVector{4, Float32}, StaticArraysCore.SVector{10, Float32})

Both are isbits and thus usable from a GPU kernel.

Symbolic problem transformations are inherently dynamic, so changes to u0 and p should be made on the CPU before the per-trajectory problems are sent to the GPU. This can be done using the SymbolicIndexingInterface.jl. For example, let's define a problem which randomizes the choice of (σ, ρ, β). We do this by first constructing the function that will change a prob.p object into the updated form by changing those 3 values by using the setsym_oop as follows:

using SymbolicIndexingInterface
sym_setter = setsym_oop(sys, [σ, ρ, β])

The return sym_setter is our optimized function, let's see it in action:

u0, p = sym_setter(prob, SVector{3}(rand(Float32, 3)))

(Float32[0.0, 0.0, 1.0, 2.0], Float32[0.45456356, 0.8650038, 0.41637892, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 2.0])

Notice it takes in the vector of values for [σ, ρ, β] and spits out the new u0, p. So we can build and solve an MTK generated ODE on the GPU using the following:

using DiffEqGPU, CUDA
function prob_func2(prob, ctx)
    u0, p = sym_setter(prob, SVector{3}(rand(Float32, 3)))
    return remake(prob; u0, p)
end

monteprob = EnsembleProblem(prob, prob_func = prob_func2, safetycopy = false)
sol = solve(
    monteprob, GPUTsit5(), EnsembleGPUKernel(CUDA.CUDABackend()),
    trajectories = 10_000
)

We can then using symbolic indexing on the result to inspect it:

[sol.u[i][y] for i in 1:length(sol.u)]

DAE initialization

ModelingToolkit can generate a nonlinear initialization problem for a mass-matrix DAE. EnsembleGPUKernel converts that problem to static storage on the host, then solves one copy per trajectory inside the GPU kernel. Square systems use SimpleTrustRegion, while rectangular nonlinear least-squares systems use SimpleGaussNewton, both from SimpleNonlinearSolve.jl. The resulting consistent state and parameters are used to start the ODE solve. ModelingToolkit's default SCCNonlinearProblem representation is lowered to an immutable nonlinear problem plus statically typed SCC block metadata. The blocks are solved sequentially in the kernel: nonlinear blocks use SimpleTrustRegion, and linear blocks use DiffEqGPU's device-compatible static square solve (closed-form for blocks of size three or smaller and pivoted LU otherwise). The mutable SCC caches and linear-problem update wrappers are not placed in the kernel.

For example, the Cartesian pendulum can be initialized and solved as follows:

using CUDA, DiffEqGPU, ModelingToolkit, OrdinaryDiffEq, SciMLBase, StaticArrays
using ModelingToolkit: t_nounits as t, D_nounits as D

@parameters g = 9.81 L = 1.0
@variables px(t) py(t) [state_priority = 10] pλ(t)

eqs = [
    D(D(px)) ~ pλ * px / L
    D(D(py)) ~ pλ * py / L - g
    px^2 + py^2 ~ L^2
]

@mtkcompile pendulum = ODESystem(eqs, t, [px, py, pλ], [g, L])

static_constructor(values) = SVector{length(values)}(values)

prob = ODEProblem{false, SciMLBase.FullSpecialize}(
    pendulum,
    [py => 0.99, D(px) => 0.0],
    (0.0, 1.0);
    guesses = [pλ => 0.0, px => 0.1, D(py) => 0.0],
    u0_constructor = static_constructor,
    p_constructor = static_constructor,
)

ensemble_prob = EnsembleProblem(prob; safetycopy = false)
sol = solve(
    ensemble_prob,
    GPURodas5P(),
    EnsembleGPUKernel(CUDA.CUDABackend());
    trajectories = 10_000,
    dt = 0.01,
    adaptive = false,
)

How the problem must be built

Kernel initialization evaluates ModelingToolkit's own state and parameter initialization maps on the device, so the problem has to be built so that those maps are usable there. Three things are needed together, as in the example above:

1. SciMLBase.FullSpecialize, passed as the problem type parameter (ODEProblem{iip, SciMLBase.FullSpecialize}(sys, ...); ModelingToolkit ignores a specialize keyword argument). This is the level at which ModelingToolkit emits the initialization maps as isbits RuntimeGeneratedFunctions (ModelingToolkit.jl#5043); at the default SciMLBase.AutoDespecialize they are host closures that cannot be uploaded. DiffEqGPU uses the generated maps exactly as ModelingToolkit produces them.
2. Static storage, through u0_constructor and p_constructor. The maps rebuild u0 and p inside the kernel, which cannot allocate, so those buffers have to be StaticArrays. ModelingToolkit fixes the container when the problem is built, so it cannot be corrected afterwards.
3. Out-of-place (ODEProblem{false, ...}). This follows from the other two rather than being an independent choice: an MVector is a mutable struct and so is not isbits, which rules it out of a device array entirely, while an in-place problem cannot write into an immutable SVector. Out-of-place with SVector storage is the only combination that is both isbits and self-consistent.

A problem that misses any of these is rejected with an error saying which, rather than failing inside the kernel.

Parameters that are not plain numbers

Everything reaching a kernel has to be isbits. Converting buffer storage to SArray cannot rescue a parameter whose contents are not — an interpolation object, a callable closing over an array, a type. Such a problem is rejected by name:

These `MTKParameters` cannot be used by EnsembleGPUKernel: the nonnumeric portion holds
values that are not isbits, so the problem cannot be uploaded to the device.

A package owning such a type can make it work by giving it an isbits stand-in, via a DiffEqGPU.make_static_storage method. DiffEqGPU calls that hook while converting u0 and the parameters, so no change to DiffEqGPU is needed:

# In the package that owns `MyInterpolation`
function DiffEqGPU.make_static_storage(itp::MyInterpolation)
    return MyStaticInterpolation(
        DiffEqGPU.make_static_storage(itp.t), DiffEqGPU.make_static_storage(itp.u)
    )
end

The stand-in has to be isbits and has to implement whatever the model calls on it.

The current initialization path has the following restrictions:

• Square, underdetermined, and overdetermined nonlinear least-squares initialization problems are supported. Square systems use SimpleTrustRegion; rectangular systems use the equivalent static normal-equation Gauss-Newton step because rectangular StaticArray factorizations are not device-compatible. Their Jacobian must have the rank required by that step.
• Lower and upper bounds are supported through a smooth transformation to unconstrained variables. A solution exactly on a finite bound is represented by a limiting unconstrained value and can therefore converge less robustly than an interior solution.
• ModelingToolkit's state and parameter initialization maps are used as generated, so entries computed from the solved initialization system (for example states that are observed variables of the torn initialization system) are supported alongside entries copied straight from the ODE and initialization problems.
• Ordinary nonlinear and linear SCC initialization blocks are supported, including all-linear SCC problems that carry no initial state: missing linear-block states are seeded with zeros, which the exact one-step linear solve does not depend on. SCC initialization containing Modelica homotopy blocks is not supported; compile the system with homotopy = false to remove them (see Systems using homotopy).

Structured MTKParameters storage is converted recursively to static storage, so this path does not require split = false.

Systems using homotopy

A system whose initialization carries Modelica homotopy(actual, simplified) operators builds a HomotopyProblem, which is solved by continuation. Continuation is a host-side sweep over a solver, so it has no device-compatible lowering and EnsembleGPUKernel rejects it:

ArgumentError: SCC nonlinear initialization problems containing homotopy blocks are not
supported by EnsembleGPUKernel.

Pass homotopy = false to mtkcompile to compile the system without them:

pendulum = mtkcompile(pendulum; homotopy = false)

Every homotopy(actual, simplified) node is replaced by actual before compilation, so the compiled system contains no homotopy nodes, initialization builds a plain nonlinear problem, and the result lowers to the kernel like any other system. This is exact rather than an approximation: actual is what the operator evaluates to numerically anyway, per the Modelica specification. What is given up is only the simplified starting heuristic — the continuation's easier warm start — so a system that relied on it to converge from a cold start may need better guesses instead.