|
| 1 | +# Symbolic-Numeric GPU Acceleration with ModelingToolkit |
| 2 | + |
| 3 | +[ModelingToolkit.jl](https://docs.sciml.ai/ModelingToolkit/stable/) is a symbolic-numeric |
| 4 | +computing system which allows for using symbolic transformations of equations before |
| 5 | +code generation. The goal is to improve numerical simulations by first turning them into |
| 6 | +the simplest set of equations to solve and exploiting things that normally cannot be done |
| 7 | +by hand. Those exact features are also potentially useful for GPU computing, and thus this |
| 8 | +tutorial showcases how to effectively use MTK with DiffEqGPU.jl. |
| 9 | + |
| 10 | +!!! warn |
| 11 | + This tutorial currently only works for ODEs defined by ModelingToolkit. More work |
| 12 | + will be required to support DAEs in full. This is work that is ongoing and expected |
| 13 | + to be completed by the summer of 2025. |
| 14 | + |
| 15 | +The core aspect to doing this right is two things. First of all, MTK respects the types |
| 16 | +chosen by the user, and thus in order for GPU kernel generation to work the user needs |
| 17 | +to ensure that the problem that is built uses static structures. For example this means |
| 18 | +that the `u0` and `p` specifications should use static arrays. This looks as follows: |
| 19 | + |
| 20 | +```@example mtk |
| 21 | +using OrdinaryDiffEqTsit5, ModelingToolkit, StaticArrays |
| 22 | +using ModelingToolkit: t_nounits as t, D_nounits as D |
| 23 | +
|
| 24 | +@parameters σ ρ β |
| 25 | +@variables x(t) y(t) z(t) |
| 26 | +
|
| 27 | +eqs = [D(D(x)) ~ σ * (y - x), |
| 28 | + D(y) ~ x * (ρ - z) - y, |
| 29 | + D(z) ~ x * y - β * z] |
| 30 | +
|
| 31 | +@mtkbuild sys = ODESystem(eqs, t) |
| 32 | +u0 = SA[D(x) => 2f0, |
| 33 | + x => 1f0, |
| 34 | + y => 0f0, |
| 35 | + z => 0f0] |
| 36 | +
|
| 37 | +p = SA[σ => 28f0, |
| 38 | + ρ => 10f0, |
| 39 | + β => 8f0 / 3f0] |
| 40 | +
|
| 41 | +tspan = (0f0, 100f0) |
| 42 | +prob = ODEProblem{false}(sys, u0, tspan, p) |
| 43 | +sol = solve(prob, Tsit5()) |
| 44 | +``` |
| 45 | + |
| 46 | +with the core aspect to notice are the `SA` |
| 47 | +[StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl) annotations on the parts |
| 48 | +for the problem construction, along with the use of `Float32`. |
| 49 | + |
| 50 | +Now one of the difficulties for building an ensemble problem is that we must make a kernel |
| 51 | +for how to construct the problems, but the use of symbolics is inherently dynamic. As such, |
| 52 | +we need to make sure that any changes to `u0` and `p` are done on the CPU and that we |
| 53 | +compile an optimized function to run on the GPU. This can be done using the |
| 54 | +[SymbolicIndexingInterface.jl](https://docs.sciml.ai/SymbolicIndexingInterface/stable/). |
| 55 | +For example, let's define a problem which randomizes the choice of `(σ, ρ, β)`. We do this |
| 56 | +by first constructing the function that will change a `prob.p` object into the updated |
| 57 | +form by changing those 3 values by using the `setsym_oop` as follows: |
| 58 | + |
| 59 | +```@example mtk |
| 60 | +using SymbolicIndexingInterface |
| 61 | +sym_setter = setsym_oop(sys, [σ, ρ, β]) |
| 62 | +``` |
| 63 | + |
| 64 | +The return `sym_setter` is our optimized function, let's see it in action: |
| 65 | + |
| 66 | +```@example mtk |
| 67 | +u0, p = sym_setter(prob,@SVector(rand(Float32,3))) |
| 68 | +``` |
| 69 | + |
| 70 | +Notice it takes in the vector of values for `[σ, ρ, β]` and spits out the new `u0, p`. So |
| 71 | +we can build and solve an MTK generated ODE on the GPU using the following: |
| 72 | + |
| 73 | +```@example mtk |
| 74 | +using DiffEqGPU, CUDA |
| 75 | +function prob_func2(prob, i, repeat) |
| 76 | + u0, p = sym_setter(prob,@SVector(rand(Float32,3))) |
| 77 | + remake(prob, u0 = u0, p = p) |
| 78 | +end |
| 79 | +
|
| 80 | +monteprob = EnsembleProblem(prob, prob_func = prob_func2, safetycopy = false) |
| 81 | +sol = solve(monteprob, GPUTsit5(), EnsembleGPUKernel(CUDA.CUDABackend()), |
| 82 | + trajectories = 10_000) |
| 83 | +``` |
| 84 | + |
| 85 | +We can then using symbolic indexing on the result to inspect it: |
| 86 | + |
| 87 | +```@example mtk |
| 88 | +[sol.u[i][y] for i in 1:length(sol.u)] |
| 89 | +``` |
0 commit comments