diff --git a/docs/pages.jl b/docs/pages.jl index c11f4c24..9217da7f 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -5,6 +5,7 @@ pages = ["index.md", "Tutorials" => Any[ "GPU Ensembles" => Any["tutorials/gpu_ensemble_basic.md", "tutorials/parallel_callbacks.md", + "tutorials/modelingtoolkit.md", "tutorials/multigpu.md", "tutorials/lower_level_api.md", "tutorials/weak_order_conv_sde.md"], diff --git a/docs/src/tutorials/modelingtoolkit.md b/docs/src/tutorials/modelingtoolkit.md new file mode 100644 index 00000000..af2b3957 --- /dev/null +++ b/docs/src/tutorials/modelingtoolkit.md @@ -0,0 +1,89 @@ +# Symbolic-Numeric GPU Acceleration with ModelingToolkit + +[ModelingToolkit.jl](https://docs.sciml.ai/ModelingToolkit/stable/) 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. + +!!! warn + This tutorial currently only works for ODEs defined by ModelingToolkit. More work + will be required to support DAEs in full. This is work that is ongoing and expected + to be completed by the summer of 2025. + +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: + +```@example mtk +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) +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) +sol = solve(prob, Tsit5()) +``` + +with the core aspect to notice are the `SA` +[StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl) annotations on the parts +for the problem construction, along with the use of `Float32`. + +Now one of the difficulties for building an ensemble problem is that we must make a kernel +for how to construct the problems, but the use of symbolics is inherently dynamic. As such, +we need to make sure that any changes to `u0` and `p` are done on the CPU and that we +compile an optimized function to run on the GPU. This can be done using the +[SymbolicIndexingInterface.jl](https://docs.sciml.ai/SymbolicIndexingInterface/stable/). +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: + +```@example mtk +using SymbolicIndexingInterface +sym_setter = setsym_oop(sys, [σ, ρ, β]) +``` + +The return `sym_setter` is our optimized function, let's see it in action: + +```@example mtk +u0, p = sym_setter(prob,@SVector(rand(Float32,3))) +``` + +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: + +```@example mtk +using DiffEqGPU, CUDA +function prob_func2(prob, i, repeat) + u0, p = sym_setter(prob,@SVector(rand(Float32,3))) + remake(prob, u0 = u0, p = 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: + +```@example mtk +[sol.u[i][y] for i in 1:length(sol.u)] +``` diff --git a/src/DiffEqGPU.jl b/src/DiffEqGPU.jl index 79613416..71c492fd 100644 --- a/src/DiffEqGPU.jl +++ b/src/DiffEqGPU.jl @@ -37,6 +37,7 @@ include("ensemblegpuarray/kernels.jl") include("ensemblegpuarray/problem_generation.jl") include("ensemblegpuarray/lowerlevel_solve.jl") +include("ensemblegpukernel/problems/ode_problems.jl") include("ensemblegpukernel/callbacks.jl") include("ensemblegpukernel/lowerlevel_solve.jl") include("ensemblegpukernel/gpukernel_algorithms.jl") @@ -67,7 +68,7 @@ include("ensemblegpukernel/tableaus/verner_tableaus.jl") include("ensemblegpukernel/tableaus/rodas_tableaus.jl") include("ensemblegpukernel/tableaus/kvaerno_tableaus.jl") -include("ensemblegpukernel/problems/ode_problems.jl") + include("utils.jl") include("algorithms.jl") diff --git a/src/ensemblegpukernel/integrators/nonstiff/types.jl b/src/ensemblegpukernel/integrators/nonstiff/types.jl index 22f7ca69..0faa46d0 100644 --- a/src/ensemblegpukernel/integrators/nonstiff/types.jl +++ b/src/ensemblegpukernel/integrators/nonstiff/types.jl @@ -1,13 +1,24 @@ ## Fixed TimeStep Integrator -function Adapt.adapt_structure(to, prob::ODEProblem{<:Any, <:Any, iip}) where {iip} - ODEProblem{iip, true}(adapt(to, prob.f), +function Adapt.adapt_structure(to, prob::Union{ODEProblem{<:Any, <:Any, iip}, ImmutableODEProblem{<:Any, <:Any, iip}}) where {iip} + ImmutableODEProblem{iip, true}(adapt(to, prob.f), adapt(to, prob.u0), adapt(to, prob.tspan), adapt(to, prob.p); adapt(to, prob.kwargs)...) end +function Adapt.adapt_structure(to, f::ODEFunction{iip}) where {iip} + if f.mass_matrix !== I && f.initialization_data !== nothing + error("Adaptation to GPU failed: DAEs of ModelingToolkit currently not supported.") + end + ODEFunction{iip, SciMLBase.FullSpecialize}(f.f, jac = f.jac, mass_matrix = f.mass_matrix, + initializeprobmap = f.initializeprobmap, + initializeprobpmap = f.initializeprobpmap, + update_initializeprob! = f.update_initializeprob!, + initialization_data = nothing, initializeprob = nothing) +end + mutable struct GPUTsit5Integrator{IIP, S, T, ST, P, F, TS, CB, AlgType} <: DiffEqBase.AbstractODEIntegrator{AlgType, IIP, S, T} alg::AlgType diff --git a/src/solve.jl b/src/solve.jl index c8326f73..23abbe66 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -273,7 +273,7 @@ function batch_solve_up_kernel(ensembleprob, probs, alg, ensemblealg, I, adaptiv _callback.continuous_callbacks)...) dev = ensemblealg.dev - probs = adapt(dev, probs) + probs = adapt(dev,adapt.((dev,), probs)) #Adaptive version only works with saveat if adaptive