Skip to content

Commit 479de67

Browse files
Add a tutorial which showcases ModelingToolkit and SII
Fixes #330. Currently the MWE works: ```julia using DiffEqGPU using OrdinaryDiffEqTsit5, 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] @mtkbuild sys = ODESystem(eqs, t) split=false u0 = SA[D(x) => 2f0, x => 1f0, y => 0f0, z => 0f0] p = SA[σ => 28f0, ρ => 10f0, β => 8f0 / 3f0] tspan = (0f0, 100f0) prob = ODEProblem{false}(sys, u0, tspan, p, split=true) prob = remake(prob, p = p = SVector{10, Float32}(prob.p...)) sol = solve(prob, Tsit5()) using SymbolicIndexingInterface p_setter = setp_oop(sys, [σ, ρ, β]) using DiffEqGPU, CUDA function prob_func2(prob, i, repeat) remake(prob, p = p_setter(prob,@svector(rand(Float32,3)))) end monteprob = EnsembleProblem(prob, prob_func = prob_func2, safetycopy = false) sol = solve(monteprob, GPUTsit5(), EnsembleGPUKernel(CUDA.CUDABackend()), trajectories = 10_000) ``` But you need to `#prob = get_updated_symbolic_problem(_get_root_indp(prob), prob; kwargs...)` in DiffEqBase. What's in here drops the `split=false` part. We need to fix `get_updated_symbolic_problem` to not promote to `Float64` and fix static array outputs in `split=true`, i.e. SciML/ModelingToolkit.jl#3585, in order to finish this tutorial.
1 parent 3d94277 commit 479de67

File tree

2 files changed

+85
-0
lines changed

2 files changed

+85
-0
lines changed

docs/pages.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ pages = ["index.md",
55
"Tutorials" => Any[
66
"GPU Ensembles" => Any["tutorials/gpu_ensemble_basic.md",
77
"tutorials/parallel_callbacks.md",
8+
"tutorials/modelingtoolkit.md",
89
"tutorials/multigpu.md",
910
"tutorials/lower_level_api.md",
1011
"tutorials/weak_order_conv_sde.md"],

docs/src/tutorials/modelingtoolkit.md

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
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+
The core aspect to doing this right is two things. First of all, MTK respects the types
11+
chosen by the user, and thus in order for GPU kernel generation to work the user needs
12+
to ensure that the problem that is built uses static structures. For example this means
13+
that the `u0` and `p` specifications should use static arrays. This looks as follows:
14+
15+
```@example mtk
16+
using OrdinaryDiffEqTsit5, ModelingToolkit, StaticArrays
17+
using ModelingToolkit: t_nounits as t, D_nounits as D
18+
19+
@parameters σ ρ β
20+
@variables x(t) y(t) z(t)
21+
22+
eqs = [D(D(x)) ~ σ * (y - x),
23+
D(y) ~ x * (ρ - z) - y,
24+
D(z) ~ x * y - β * z]
25+
26+
@mtkbuild sys = ODESystem(eqs, t)
27+
28+
u0 = SA[D(x) => 2f0,
29+
x => 1f0,
30+
y => 0f0,
31+
z => 0f0]
32+
33+
p = SA[σ => 28f0,
34+
ρ => 10f0,
35+
β => 8f0 / 3f0]
36+
37+
tspan = (0f0, 100f0)
38+
prob = ODEProblem{false}(sys, u0, tspan, p)
39+
sol = solve(prob, Tsit5())
40+
```
41+
42+
with the core aspect to notice are the `SA`
43+
[StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl) annotations on the parts
44+
for the problem construction, along with the use of `Float32`.
45+
46+
Now one of the difficulties for building an ensemble problem is that we must make a kernel
47+
for how to construct the problems, but the use of symbolics is inherently dynamic. As such,
48+
we need to make sure that any changes to `u0` and `p` are done on the CPU and that we
49+
compile an optimized function to run on the GPU. This can be done using the
50+
[SymbolicIndexingInterface.jl](https://docs.sciml.ai/SymbolicIndexingInterface/stable/).
51+
For example, let's define a problem which randomizes the choice of `(σ, ρ, β)`. We do this
52+
by first constructing the function that will change a `prob.p` object into the updated
53+
form by changing those 3 values by using the `setsym_oop` as follows:
54+
55+
```@example mtk
56+
using SymbolicIndexingInterface
57+
sym_setter = setsym_oop(sys, [σ, ρ, β])
58+
```
59+
60+
The return `sym_setter` is our optimized function, let's see it in action:
61+
62+
```@example mtk
63+
u0, p = sym_setter(prob,@SVector(rand(Float32,3)))
64+
```
65+
66+
Notice it takes in the vector of values for `[σ, ρ, β]` and spits out the new `u0, p`. So
67+
we can build and solve an MTK generated ODE on the GPU using the following:
68+
69+
```@example mtk
70+
using DiffEqGPU, CUDA
71+
function prob_func2(prob, i, repeat)
72+
u0, p = sym_setter(prob,@SVector(rand(Float32,3)))
73+
remake(prob, u0 = u0, p = p)
74+
end
75+
monteprob = EnsembleProblem(prob, prob_func = prob_func2, safetycopy = false)
76+
sol = solve(monteprob, GPUTsit5(), EnsembleGPUKernel(CUDA.CUDABackend()),
77+
trajectories = 10_000)
78+
```
79+
80+
We can then using symbolic indexing on the result to inspect it:
81+
82+
```@example mtk
83+
[sol.u[i][y] for i in 1:length(sol.u)]
84+
```

0 commit comments

Comments
 (0)