From 04b4c2214ddaa03612f30b1b0deb633edb4b9bb4 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Sat, 29 Aug 2026 17:01:29 +0000 Subject: [PATCH 1/2] Lower computed initialization maps via symbolic tracing Trace MTK initialization state and parameter maps with Symbolics variables instead of integer index tokens. Entries copied from a source slot stay static gather recipes; literal outputs become constants; computed entries (for example ODE states that are observed variables of the torn initialization system) compile through `build_function` into isbits generated functions evaluated against the same source vector on the device. The integer-token scheme could not represent computed entries, and worse, an integer-valued computed output could alias a valid source index and initialize with wrong values while reporting success. The symbolic trace removes that aliasing by construction. Raw homotopy initialization problems now fail with a clear error, and the tutorial documents rewriting `homotopy(actual, simplified)` nodes to `actual` to obtain an equivalent supported initialization. Co-Authored-By: Chris Rackauckas --- Project.toml | 6 +- docs/src/tutorials/modelingtoolkit.md | 33 ++- ext/ModelingToolkitBaseExt.jl | 203 ++++++++++++------ src/utils.jl | 9 + .../stiff_ode/gpu_ode_modelingtoolkit_dae.jl | 60 ++++++ 5 files changed, 238 insertions(+), 73 deletions(-) diff --git a/Project.toml b/Project.toml index 4e6d6d1a..8a71f862 100644 --- a/Project.toml +++ b/Project.toml @@ -38,6 +38,8 @@ JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" ModelingToolkitBase = "7771a370-6774-4173-bd38-47e70ca0b839" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" +RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47" +Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [extensions] @@ -45,7 +47,7 @@ AMDGPUExt = ["AMDGPU"] CUDAExt = ["CUDA"] JLArraysExt = ["JLArrays"] MetalExt = ["Metal"] -ModelingToolkitBaseExt = ["ModelingToolkitBase"] +ModelingToolkitBaseExt = ["ModelingToolkitBase", "RuntimeGeneratedFunctions", "Symbolics"] OpenCLExt = ["OpenCL"] oneAPIExt = ["oneAPI"] @@ -73,6 +75,7 @@ Parameters = "0.13" PrecompileTools = "1.2.1" Random = "1" RecursiveArrayTools = "4.2" +RuntimeGeneratedFunctions = "0.5.9" SciMLBase = "3.49" SciMLPublic = "1" Setfield = "1" @@ -85,6 +88,7 @@ Test = "1" UnPack = "1" ZygoteRules = "0.2.7" SciMLTesting = "2.10" +Symbolics = "7" julia = "1.10" oneAPI = "2" pocl_jll = "7" diff --git a/docs/src/tutorials/modelingtoolkit.md b/docs/src/tutorials/modelingtoolkit.md index 50960f40..b90dfd57 100644 --- a/docs/src/tutorials/modelingtoolkit.md +++ b/docs/src/tutorials/modelingtoolkit.md @@ -170,13 +170,32 @@ The current initialization path has the following restrictions: - 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 may directly select, - reorder, and repack numeric values from the ODE and initialization problems. DiffEqGPU - traces those operations on the host and stores only static gather recipes in the - kernel. Fallback getters that evaluate derived symbolic expressions are not yet - lowered. - - Ordinary nonlinear and linear SCC initialization blocks are supported. SCC - initialization containing Modelica homotopy blocks is not supported. + - ModelingToolkit's state and parameter initialization maps are traced symbolically on + the host: entries copied from the ODE or initialization problem become static gather + recipes, and entries computed from them (for example states that are observed + variables of the torn initialization system) are compiled into device-side generated + functions. Maps that branch on runtime values cannot be traced and are not supported. + - Ordinary nonlinear and linear SCC initialization blocks are supported. Initialization + containing Modelica `homotopy(actual, simplified)` nodes is not supported: continuation + solves do not fit the kernel's static solver set. Since `homotopy(actual, simplified)` + evaluates numerically to `actual`, an equivalent supported system is obtained by + rewriting each homotopy node to its `actual` branch before `mtkcompile`: + + ```julia + using ModelingToolkit + SU = ModelingToolkit.Symbolics.SymbolicUtils + rule = SU.@rule ModelingToolkit.homotopy(~a, ~s) => ~a + strip_homotopy = SU.Rewriters.Postwalk(SU.Rewriters.PassThrough(rule)) + strip_eq(eq) = + ModelingToolkit.Symbolics.wrap(strip_homotopy(ModelingToolkit.Symbolics.unwrap(eq.lhs))) ~ + ModelingToolkit.Symbolics.wrap(strip_homotopy(ModelingToolkit.Symbolics.unwrap(eq.rhs))) + + stripped_eqs = map(strip_eq, equations_with_homotopy) + ``` + + This discards only the `simplified` starting heuristic used by continuation solvers; + the solved system is unchanged. Robustness for hard initial guesses then rests on the + kernel's `SimpleTrustRegion` solve. Structured `MTKParameters` storage is converted recursively to static storage, so this path does not require `split = false`. diff --git a/ext/ModelingToolkitBaseExt.jl b/ext/ModelingToolkitBaseExt.jl index f1b4ea2c..efd2cd53 100644 --- a/ext/ModelingToolkitBaseExt.jl +++ b/ext/ModelingToolkitBaseExt.jl @@ -1,6 +1,8 @@ module ModelingToolkitBaseExt using ModelingToolkitBase: MTKParameters, System, unknowns +using Symbolics +using RuntimeGeneratedFunctions: drop_expr using StaticArraysCore: SArray, StaticArray, SVector import DiffEqGPU import SciMLBase @@ -87,6 +89,20 @@ struct InitializationStructureRecipe{T, R} values::R end +struct InitializationConstant{V} + value::V +end + +struct InitializationScalarExpression{F} + f::F +end + +struct InitializationArrayExpression{A, F} + f::F +end + +InitializationArrayExpression{A}(f::F) where {A, F} = InitializationArrayExpression{A, F}(f) + struct InitializationStateMap{R} recipe::R end @@ -128,6 +144,14 @@ function initialization_sources(valp) end gather_initialization_value(::InitializationSourceIndex{I}, sources) where {I} = sources[I] +gather_initialization_value(recipe::InitializationConstant, sources) = recipe.value +gather_initialization_value(recipe::InitializationScalarExpression, sources) = + recipe.f(sources) +function gather_initialization_value( + recipe::InitializationArrayExpression{A}, sources + ) where {A} + return A(Tuple(recipe.f(sources))) +end function gather_initialization_value( recipe::InitializationArrayRecipe{A}, sources ) where {A} @@ -165,36 +189,50 @@ function (map::InitializationParameterMap)(prob, sol) ) end -function next_source_index!(counter) - counter[] += 1 - return counter[] +struct InitializationSourceTrace + vars::Vector{Symbolics.Num} + lookup::Dict{Any, Int} +end +InitializationSourceTrace() = InitializationSourceTrace(Symbolics.Num[], Dict{Any, Int}()) + +function next_source_token!(trace::InitializationSourceTrace) + i = length(trace.vars) + 1 + var = Symbolics.variable(:ˍdiffeqgpu_source, i) + push!(trace.vars, var) + trace.lookup[Symbolics.unwrap(var)] = i + return var end -index_numeric(::Number, counter) = next_source_index!(counter) -index_numeric(x::AbstractArray, counter) = map(Base.Fix2(index_numeric, counter), x) -index_numeric(x::Tuple, counter) = map(Base.Fix2(index_numeric, counter), x) -index_numeric(x::NamedTuple, counter) = map(Base.Fix2(index_numeric, counter), x) -function index_numeric(x, counter) - values = ntuple(i -> index_numeric(getfield(x, i), counter), fieldcount(typeof(x))) +source_index_of(trace::InitializationSourceTrace, x) = + get(trace.lookup, Symbolics.unwrap(x), nothing) + +issymbolic(x) = Symbolics.unwrap(x) isa Symbolics.BasicSymbolic + +index_numeric(::Number, trace) = next_source_token!(trace) +index_numeric(x::AbstractArray, trace) = map(Base.Fix2(index_numeric, trace), x) +index_numeric(x::Tuple, trace) = map(Base.Fix2(index_numeric, trace), x) +index_numeric(x::NamedTuple, trace) = map(Base.Fix2(index_numeric, trace), x) +function index_numeric(x, trace) + values = ntuple(i -> index_numeric(getfield(x, i), trace), fieldcount(typeof(x))) return typeof(x)(values...) end -function index_parameters(p::MTKParameters, counter) +function index_parameters(p::MTKParameters, trace) return MTKParameters( - index_numeric(p.tunable, counter), - index_numeric(p.initials, counter), - index_numeric(p.discrete, counter), - index_numeric(p.constant, counter), + index_numeric(p.tunable, trace), + index_numeric(p.initials, trace), + index_numeric(p.discrete, trace), + index_numeric(p.constant, trace), p.nonnumeric, p.caches ) end -function index_initialization_problem(prob::SciMLBase.NonlinearLeastSquaresProblem, counter) +function index_initialization_problem(prob::SciMLBase.NonlinearLeastSquaresProblem, trace) return SciMLBase.NonlinearLeastSquaresProblem{SciMLBase.isinplace(prob)}( prob.f, - index_numeric(prob.u0, counter), - index_parameters(prob.p, counter); + index_numeric(prob.u0, trace), + index_parameters(prob.p, trace); lb = prob.lb, ub = prob.ub, prob.kwargs... @@ -202,105 +240,140 @@ function index_initialization_problem(prob::SciMLBase.NonlinearLeastSquaresProbl end function index_initialization_problem( - prob::DiffEqGPU.ImmutableSCCNonlinearProblem, counter + prob::DiffEqGPU.ImmutableSCCNonlinearProblem, trace ) - return index_initialization_problem(prob.problem, counter) + return index_initialization_problem(prob.problem, trace) end function index_initialization_problem( - prob::Union{SciMLBase.NonlinearProblem, SciMLBase.ImmutableNonlinearProblem}, counter + prob::Union{SciMLBase.NonlinearProblem, SciMLBase.ImmutableNonlinearProblem}, trace ) return SciMLBase.ImmutableNonlinearProblem{SciMLBase.isinplace(prob)}( prob.f, - index_numeric(prob.u0, counter), - index_parameters(prob.p, counter), + index_numeric(prob.u0, trace), + index_parameters(prob.p, trace), prob.problem_type; prob.kwargs... ) end -function index_ode_problem(prob, counter) +function index_ode_problem(prob, trace) return SciMLBase.ImmutableODEProblem( prob.f, - index_numeric(prob.u0, counter), + index_numeric(prob.u0, trace), prob.tspan, - index_parameters(prob.p, counter), + index_parameters(prob.p, trace), prob.problem_type; prob.kwargs... ) end -function source_recipe(index::Number, source_count) - source_index = try - Int(index) - catch - nothing - end - if source_index === nothing || !isequal(index, source_index) || - !(1 <= source_index <= source_count) - error( - "ModelingToolkit initialization maps must copy numeric values from the ODE or initialization problem." +@inline function compile_sources_function(exprs, trace::InitializationSourceTrace) + built = Symbolics.build_function(exprs, trace.vars; expression = Val(false)) + f = built isa Tuple ? built[1] : built + # An RGF's stored `body::Expr` is only for re-generation; dropping it makes the + # callable isbits so the recipe can live inside GPU kernel problems. + return drop_expr(f) +end + +function scalar_source_recipe(x, trace::InitializationSourceTrace) + if issymbolic(x) + i = source_index_of(trace, x) + i === nothing || return InitializationSourceIndex{i}() + return InitializationScalarExpression( + compile_sources_function(Symbolics.unwrap(x), trace) ) end - return InitializationSourceIndex{source_index}() -end -function source_recipe(x::AbstractArray, source_count) - values = ntuple(i -> source_recipe(x[i], source_count), length(x)) - storage_type = typeof(SArray{Tuple{size(x)...}}(x)) - return InitializationArrayRecipe{storage_type}(values) + x isa Number && return InitializationConstant(x) + return error( + "ModelingToolkit initialization maps produced an unsupported value of type $(typeof(x))." + ) end -source_recipe(x::Tuple, source_count) = map(Base.Fix2(source_recipe, source_count), x) -source_recipe(x::NamedTuple, source_count) = map(Base.Fix2(source_recipe, source_count), x) -function source_recipe(x, source_count) + +# A map output entry is either a value copied verbatim from a source slot (a bare traced +# variable), a literal constant, or a computed symbolic expression. Whole arrays with any +# computed entry are compiled into one generated gather-and-compute function so the device +# recipe stays a single call. +is_direct_entry(x, trace) = !issymbolic(x) || source_index_of(trace, x) !== nothing + +function source_recipe(x, trace::InitializationSourceTrace) + (x isa Number || issymbolic(x)) && return scalar_source_recipe(x, trace) values = ntuple( - i -> source_recipe(getfield(x, i), source_count), fieldcount(typeof(x)) + i -> source_recipe(getfield(x, i), trace), fieldcount(typeof(x)) ) return InitializationStructureRecipe{typeof(x)}(values) end +function source_recipe(x::AbstractArray, trace::InitializationSourceTrace) + storage_type = SArray{Tuple{size(x)...}} + return array_source_recipe(x, trace, storage_type) +end +source_recipe(x::Tuple, trace::InitializationSourceTrace) = + map(Base.Fix2(source_recipe, trace), x) +source_recipe(x::NamedTuple, trace::InitializationSourceTrace) = + map(Base.Fix2(source_recipe, trace), x) +function source_recipe( + x::Union{Symbolics.Num, Symbolics.BasicSymbolic}, trace::InitializationSourceTrace + ) + return scalar_source_recipe(x, trace) +end -source_recipe(x, source_count, prototype) = source_recipe(x, source_count) -function source_recipe(x::AbstractArray, source_count, prototype::AbstractArray) - values = ntuple(i -> source_recipe(x[i], source_count, prototype[i]), length(x)) +function array_source_recipe(x, trace, storage_type) + if all(el -> is_direct_entry(el, trace), x) + values = ntuple(i -> scalar_source_recipe(x[i], trace), length(x)) + return InitializationArrayRecipe{storage_type}(values) + end + exprs = SVector{length(x)}(map(Symbolics.unwrap, vec(x))...) + return InitializationArrayExpression{storage_type}( + compile_sources_function(exprs, trace) + ) +end + +source_recipe(x, trace::InitializationSourceTrace, prototype) = source_recipe(x, trace) +function source_recipe( + x::AbstractArray, trace::InitializationSourceTrace, prototype::AbstractArray + ) storage_type = typeof(static_parameter_storage(prototype)) - return InitializationArrayRecipe{storage_type}(values) + return array_source_recipe(x, trace, storage_type) end -function source_recipe(x::Tuple, source_count, prototype::Tuple) +function source_recipe(x::Tuple, trace::InitializationSourceTrace, prototype::Tuple) return map( - (value, target) -> source_recipe(value, source_count, target), x, prototype + (value, target) -> source_recipe(value, trace, target), x, prototype ) end -function source_recipe(x::NamedTuple, source_count, prototype::NamedTuple) +function source_recipe( + x::NamedTuple, trace::InitializationSourceTrace, prototype::NamedTuple + ) values = map( - (value, target) -> source_recipe(value, source_count, target), x, prototype + (value, target) -> source_recipe(value, trace, target), x, prototype ) return NamedTuple{keys(x)}(values) end function make_state_map(initprob, map) map === nothing && return nothing - # Evaluate the host-only MTK map on sequential source indices, then retain its - # device-compatible gather recipe. - counter = Ref(0) - indexed_initprob = index_initialization_problem(initprob, counter) - return InitializationStateMap(source_recipe(map(indexed_initprob), counter[])) + # Evaluate the host-only MTK map on symbolically traced sources: copied slots become + # static gather recipes and computed entries compile into generated device functions. + trace = InitializationSourceTrace() + indexed_initprob = index_initialization_problem(initprob, trace) + return InitializationStateMap(source_recipe(map(indexed_initprob), trace)) end function make_parameter_map(prob, initprob, map) map === nothing && return nothing # Parameter maps may select from both the ODE problem and nonlinear solution. - counter = Ref(0) - indexed_prob = index_ode_problem(prob, counter) - indexed_initprob = index_initialization_problem(initprob, counter) + trace = InitializationSourceTrace() + indexed_prob = index_ode_problem(prob, trace) + indexed_initprob = index_initialization_problem(initprob, trace) p = map(indexed_prob, indexed_initprob) p isa MTKParameters || error( "ModelingToolkit initialization parameter maps must return `MTKParameters`." ) prototype = SciMLBase.parameter_values(prob) recipe = InitializationParameterRecipe( - source_recipe(p.tunable, counter[], prototype.tunable), - source_recipe(p.initials, counter[], prototype.initials), - source_recipe(p.discrete, counter[], prototype.discrete), - source_recipe(p.constant, counter[], prototype.constant) + source_recipe(p.tunable, trace, prototype.tunable), + source_recipe(p.initials, trace, prototype.initials), + source_recipe(p.discrete, trace, prototype.discrete), + source_recipe(p.constant, trace, prototype.constant) ) return InitializationParameterMap(recipe) end diff --git a/src/utils.jl b/src/utils.jl index 8d4c7b5d..ed7d9c3c 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -93,6 +93,15 @@ end make_initialization_maps_compatible(prob, initprob, umap, pmap, p) = (umap, pmap) lower_initialization_problem(prob) = prob +function lower_initialization_problem(prob::SciMLBase.HomotopyProblem) + throw( + ArgumentError( + "Homotopy initialization problems are not supported by EnsembleGPUKernel. \ + Rewrite `homotopy(actual, simplified)` nodes to their `actual` branch before \ + `mtkcompile`; see the ModelingToolkit tutorial in the DiffEqGPU documentation." + ) + ) +end make_static_storage(x::StaticArrays.StaticArray) = StaticArrays.SArray{Tuple{size(x)...}}(map(make_static_storage, x)) diff --git a/test/gpu_kernel_de/stiff_ode/gpu_ode_modelingtoolkit_dae.jl b/test/gpu_kernel_de/stiff_ode/gpu_ode_modelingtoolkit_dae.jl index bdf114a6..29a97863 100644 --- a/test/gpu_kernel_de/stiff_ode/gpu_ode_modelingtoolkit_dae.jl +++ b/test/gpu_kernel_de/stiff_ode/gpu_ode_modelingtoolkit_dae.jl @@ -1,6 +1,8 @@ using DiffEqGPU, StaticArrays, SciMLBase, LinearAlgebra, Test using ModelingToolkit, OrdinaryDiffEq using ModelingToolkit: t_nounits as t, D_nounits as D +using ModelingToolkit.Symbolics.SymbolicUtils: @rule +using ModelingToolkit.Symbolics.SymbolicUtils.Rewriters: Postwalk, PassThrough using KernelAbstractions: CPU const GROUP = get(ENV, "GROUP", "CUDA") @@ -267,6 +269,64 @@ end @test sol.u ≈ SA[2.0f0, 3.0f0] atol = 1.0f-5 end +@testset "Computed initialization maps" begin + @variables cmx(t) cmy(t) + # cmy is torn into an observed variable of the initialization system, so the state + # map has to evaluate `2cmx + 1` instead of copying a solved value. + @mtkcompile cmsys = ODESystem( + [D(cmx) ~ -cmx, D(cmy) ~ -cmy], t; + initialization_eqs = [cmx^3 + cmx ~ 2, cmy ~ 2cmx + 1] + ) + cmprob = ODEProblem(cmsys, [], (0.0, 1.0); guesses = [cmx => 1.0, cmy => 1.0]) + @test !SciMLBase.is_trivial_initialization(cmprob.f.initialization_data) + + compatible_prob = DiffEqGPU.make_prob_compatible(cmprob) + @test isbitstype(typeof(compatible_prob)) + + ensemble_prob = EnsembleProblem(cmprob, safetycopy = false) + sol = solve( + ensemble_prob, GPUTsit5(), EnsembleGPUKernel(backend); + trajectories = 2, dt = 0.01, adaptive = false, save_everystep = false + ) + @test length(sol.u) == 2 + @test sort(collect(sol.u[1].u[1])) ≈ [1.0, 3.0] atol = 1.0e-5 +end + +@testset "Homotopy-stripped initialization" begin + rule = @rule ModelingToolkit.homotopy(~a, ~s) => ~a + strip_homotopy = Postwalk(PassThrough(rule)) + strip_eq(eq) = + ModelingToolkit.Symbolics.wrap(strip_homotopy(ModelingToolkit.Symbolics.unwrap(eq.lhs))) ~ + ModelingToolkit.Symbolics.wrap(strip_homotopy(ModelingToolkit.Symbolics.unwrap(eq.rhs))) + + @variables hsx(t) hsy(t) + homotopy_eqs = [ + ModelingToolkit.homotopy(atan(hsx - 3) + hsx - 3, hsx - 3) ~ 0, + hsy ~ hsx, + ] + + @mtkcompile hsys = ODESystem( + [D(hsx) ~ -hsx, D(hsy) ~ -hsy], t; initialization_eqs = homotopy_eqs + ) + hprob = ODEProblem(hsys, [], (0.0, 1.0); guesses = [hsx => 2.5, hsy => 2.5]) + @test_throws ArgumentError DiffEqGPU.make_prob_compatible(hprob) + + @mtkcompile hsys2 = ODESystem( + [D(hsx) ~ -hsx, D(hsy) ~ -hsy], t; + initialization_eqs = map(strip_eq, homotopy_eqs) + ) + hprob2 = ODEProblem(hsys2, [], (0.0, 1.0); guesses = [hsx => 2.5, hsy => 2.5]) + compatible_prob = DiffEqGPU.make_prob_compatible(hprob2) + @test isbitstype(typeof(compatible_prob)) + + ensemble_prob = EnsembleProblem(hprob2, safetycopy = false) + sol = solve( + ensemble_prob, GPUTsit5(), EnsembleGPUKernel(backend); + trajectories = 2, dt = 0.01, adaptive = false, save_everystep = false + ) + @test sol.u[1].u[1] ≈ [3.0, 3.0] atol = 1.0e-4 +end + # ============================================================================ # Test 4: Host symbolic setters and trivial initialization # ============================================================================ From 1ddb4ca5049efb03652df09239fe26fd23333bfb Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Sat, 29 Aug 2026 21:53:20 +0000 Subject: [PATCH 2/2] Gate the homotopy rejection test on HomotopyProblem being emitted Older ModelingToolkit versions in the downgrade CI lower homotopy nodes without producing a HomotopyProblem, so there is no rejection to assert. Co-Authored-By: Chris Rackauckas --- test/gpu_kernel_de/stiff_ode/gpu_ode_modelingtoolkit_dae.jl | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/test/gpu_kernel_de/stiff_ode/gpu_ode_modelingtoolkit_dae.jl b/test/gpu_kernel_de/stiff_ode/gpu_ode_modelingtoolkit_dae.jl index 29a97863..1b0e516b 100644 --- a/test/gpu_kernel_de/stiff_ode/gpu_ode_modelingtoolkit_dae.jl +++ b/test/gpu_kernel_de/stiff_ode/gpu_ode_modelingtoolkit_dae.jl @@ -309,7 +309,11 @@ end [D(hsx) ~ -hsx, D(hsy) ~ -hsy], t; initialization_eqs = homotopy_eqs ) hprob = ODEProblem(hsys, [], (0.0, 1.0); guesses = [hsx => 2.5, hsy => 2.5]) - @test_throws ArgumentError DiffEqGPU.make_prob_compatible(hprob) + # Older ModelingToolkit versions lower `homotopy` without emitting a + # `HomotopyProblem`, in which case there is no rejection to test. + if hprob.f.initialization_data.initializeprob isa SciMLBase.HomotopyProblem + @test_throws ArgumentError DiffEqGPU.make_prob_compatible(hprob) + end @mtkcompile hsys2 = ODESystem( [D(hsx) ~ -hsx, D(hsy) ~ -hsy], t;