From c58b67fc7b9ca0939d7ab472cd77fcee75d097ab Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Thu, 5 Jun 2025 12:27:37 +0530 Subject: [PATCH 1/2] refactor: remove CSE hack --- .../symbolics_tearing.jl | 97 +++---------------- src/systems/nonlinear/initializesystem.jl | 14 --- test/code_generation.jl | 31 ++++++ test/structural_transformation/utils.jl | 73 +------------- 4 files changed, 47 insertions(+), 168 deletions(-) diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl index c5c0757295..af0973f1dc 100644 --- a/src/structural_transformation/symbolics_tearing.jl +++ b/src/structural_transformation/symbolics_tearing.jl @@ -929,7 +929,7 @@ Update the system equations, unknowns, and observables after simplification. """ function update_simplified_system!( state::TearingState, neweqs, solved_eqs, dummy_sub, var_sccs, extra_unknowns; - cse_hack = true, array_hack = true, D = nothing, iv = nothing) + array_hack = true, D = nothing, iv = nothing) @unpack fullvars, structure = state @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure diff_to_var = invview(var_to_diff) @@ -978,8 +978,7 @@ function update_simplified_system!( end @set! sys.unknowns = unknowns - obs = cse_and_array_hacks( - sys, obs, unknowns, neweqs; cse = cse_hack, array = array_hack) + obs = tearing_hacks(sys, obs, unknowns, neweqs; array = array_hack) @set! sys.eqs = neweqs @set! sys.observed = obs @@ -1035,7 +1034,7 @@ differential variables. according to `full_var_eq_matching`. """ function tearing_reassemble(state::TearingState, var_eq_matching::Matching, - full_var_eq_matching::Matching, var_sccs::Vector{Vector{Int}}; simplify = false, mm, cse_hack = true, + full_var_eq_matching::Matching, var_sccs::Vector{Vector{Int}}; simplify = false, mm, array_hack = true, fully_determined = true) extra_eqs_vars = get_extra_eqs_vars( state, var_eq_matching, full_var_eq_matching, fully_determined) @@ -1074,7 +1073,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching::Matching, # var_eq_matching and full_var_eq_matching are now invalidated sys = update_simplified_system!(state, neweqs, solved_eqs, dummy_sub, var_sccs, - extra_unknowns; cse_hack, array_hack, iv, D) + extra_unknowns; array_hack, iv, D) @set! state.sys = sys @set! sys.tearing_state = state @@ -1223,14 +1222,7 @@ function get_extra_eqs_vars( end """ -# HACK 1 - -Since we don't support array equations, any equation of the sort `x[1:n] ~ f(...)[1:n]` -gets turned into `x[1] ~ f(...)[1], x[2] ~ f(...)[2]`. Repeatedly calling `f` gets -_very_ expensive. this hack performs a limited form of CSE specifically for this case to -avoid the unnecessary cost. This and the below hack are implemented simultaneously - -# HACK 2 +# HACK Add equations for array observed variables. If `p[i] ~ (...)` are equations, add an equation `p ~ [p[1], p[2], ...]` allow topsort to reorder them only add the new equation @@ -1238,13 +1230,7 @@ if all `p[i]` are present and the unscalarized form is used in any equation (obs not) we first count the number of times the scalarized form of each observed variable occurs in observed equations (and unknowns if it's split). """ -function cse_and_array_hacks(sys, obs, unknowns, neweqs; cse = true, array = true) - # HACK 1 - # mapping of rhs to temporary CSE variable - # `f(...) => tmpvar` in above example - rhs_to_tempvar = Dict() - - # HACK 2 +function tearing_hacks(sys, obs, unknowns, neweqs; array = true) # map of array observed variable (unscalarized) to number of its # scalarized terms that appear in observed equations arr_obs_occurrences = Dict() @@ -1252,31 +1238,6 @@ function cse_and_array_hacks(sys, obs, unknowns, neweqs; cse = true, array = tru lhs = eq.lhs rhs = eq.rhs - # HACK 1 - if cse && is_getindexed_array(rhs) - rhs_arr = arguments(rhs)[1] - iscall(rhs_arr) && operation(rhs_arr) isa Symbolics.Operator && continue - if !haskey(rhs_to_tempvar, rhs_arr) - tempvar = gensym(Symbol(lhs)) - N = length(rhs_arr) - tempvar = unwrap(Symbolics.variable( - tempvar; T = Symbolics.symtype(rhs_arr))) - tempvar = setmetadata( - tempvar, Symbolics.ArrayShapeCtx, Symbolics.shape(rhs_arr)) - tempeq = tempvar ~ rhs_arr - rhs_to_tempvar[rhs_arr] = tempvar - push!(obs, tempeq) - end - - # getindex_wrapper is used because `observed2graph` treats `x` and `x[i]` as different, - # so it doesn't find a dependency between this equation and `tempvar ~ rhs_arr` - # which fails the topological sort - neweq = lhs ~ getindex_wrapper( - rhs_to_tempvar[rhs_arr], Tuple(arguments(rhs)[2:end])) - obs[i] = neweq - end - # end HACK 1 - array || continue iscall(lhs) || continue operation(lhs) === getindex || continue @@ -1287,31 +1248,6 @@ function cse_and_array_hacks(sys, obs, unknowns, neweqs; cse = true, array = tru continue end - # Also do CSE for `equations(sys)` - if cse - for (i, eq) in enumerate(neweqs) - (; lhs, rhs) = eq - is_getindexed_array(rhs) || continue - rhs_arr = arguments(rhs)[1] - if !haskey(rhs_to_tempvar, rhs_arr) - tempvar = gensym(Symbol(lhs)) - N = length(rhs_arr) - tempvar = unwrap(Symbolics.variable( - tempvar; T = Symbolics.symtype(rhs_arr))) - tempvar = setmetadata( - tempvar, Symbolics.ArrayShapeCtx, Symbolics.shape(rhs_arr)) - tempeq = tempvar ~ rhs_arr - rhs_to_tempvar[rhs_arr] = tempvar - push!(obs, tempeq) - end - # don't need getindex_wrapper, but do it anyway to know that this - # hack took place - neweq = lhs ~ getindex_wrapper( - rhs_to_tempvar[rhs_arr], Tuple(arguments(rhs)[2:end])) - neweqs[i] = neweq - end - end - # count variables in unknowns if they are scalarized forms of variables # also present as observed. e.g. if `x[1]` is an unknown and `x[2] ~ (..)` # is an observed equation. @@ -1346,18 +1282,7 @@ function cse_and_array_hacks(sys, obs, unknowns, neweqs; cse = true, array = tru return obs end -function is_getindexed_array(rhs) - (!ModelingToolkit.isvariable(rhs) || ModelingToolkit.iscalledparameter(rhs)) && - iscall(rhs) && operation(rhs) === getindex && - Symbolics.shape(rhs) != Symbolics.Unknown() -end - -# PART OF HACK 1 -getindex_wrapper(x, i) = x[i...] - -@register_symbolic getindex_wrapper(x::AbstractArray, i::Tuple{Vararg{Int}}) - -# PART OF HACK 2 +# PART OF HACK function change_origin(origin, arr) if all(isone, Tuple(origin)) return arr @@ -1385,11 +1310,11 @@ new residual equations after tearing. End users are encouraged to call [`mtkcomp instead, which calls this function internally. """ function tearing(sys::AbstractSystem, state = TearingState(sys); mm = nothing, - simplify = false, cse_hack = true, array_hack = true, fully_determined = true, kwargs...) + simplify = false, array_hack = true, fully_determined = true, kwargs...) var_eq_matching, full_var_eq_matching, var_sccs, can_eliminate = tearing(state) invalidate_cache!(tearing_reassemble( state, var_eq_matching, full_var_eq_matching, var_sccs; mm, - simplify, cse_hack, array_hack, fully_determined)) + simplify, array_hack, fully_determined)) end """ @@ -1399,7 +1324,7 @@ Perform index reduction and use the dummy derivative technique to ensure that the system is balanced. """ function dummy_derivative(sys, state = TearingState(sys); simplify = false, - mm = nothing, cse_hack = true, array_hack = true, fully_determined = true, kwargs...) + mm = nothing, array_hack = true, fully_determined = true, kwargs...) jac = let state = state (eqs, vars) -> begin symeqs = EquationsView(state)[eqs] @@ -1425,5 +1350,5 @@ function dummy_derivative(sys, state = TearingState(sys); simplify = false, state, jac; state_priority, kwargs...) tearing_reassemble(state, var_eq_matching, full_var_eq_matching, var_sccs; - simplify, mm, cse_hack, array_hack, fully_determined) + simplify, mm, array_hack, fully_determined) end diff --git a/src/systems/nonlinear/initializesystem.jl b/src/systems/nonlinear/initializesystem.jl index b1ca187e95..df97adfdb3 100644 --- a/src/systems/nonlinear/initializesystem.jl +++ b/src/systems/nonlinear/initializesystem.jl @@ -780,20 +780,6 @@ function unhack_observed(obseqs::Vector{Equation}, eqs::Vector{Equation}) push!(rm_idxs, i) continue end - if operation(eq.rhs) == StructuralTransformations.getindex_wrapper - var, idxs = arguments(eq.rhs) - subs[eq.rhs] = var[idxs...] - push!(tempvars, var) - end - end - - for (i, eq) in enumerate(eqs) - iscall(eq.rhs) || continue - if operation(eq.rhs) == StructuralTransformations.getindex_wrapper - var, idxs = arguments(eq.rhs) - subs[eq.rhs] = var[idxs...] - push!(tempvars, var) - end end for (i, eq) in enumerate(obseqs) diff --git a/test/code_generation.jl b/test/code_generation.jl index 2d5925cbca..15de194fd8 100644 --- a/test/code_generation.jl +++ b/test/code_generation.jl @@ -79,3 +79,34 @@ end @test SciMLBase.successful_retcode(sol) end end + +@testset "scalarized array observed calling same function multiple times" begin + @variables x(t) y(t)[1:2] + @parameters foo(::Real)[1:2] + val = Ref(0) + function _tmp_fn2(x) + val[] += 1 + return [x, 2x] + end + @mtkcompile sys = System([D(x) ~ y[1] + y[2], y ~ foo(x)], t) + @test length(equations(sys)) == 1 + @test length(ModelingToolkit.observed(sys)) == 3 + prob = ODEProblem(sys, [x => 1.0, foo => _tmp_fn2], (0.0, 1.0)) + val[] = 0 + @test_nowarn prob.f(prob.u0, prob.p, 0.0) + @test val[] == 1 + + @testset "CSE in equations(sys)" begin + val[] = 0 + @variables z(t)[1:2] + @mtkcompile sys = System( + [D(y) ~ foo(x), D(x) ~ sum(y), zeros(2) ~ foo(prod(z))], t) + @test length(equations(sys)) == 5 + @test length(ModelingToolkit.observed(sys)) == 0 + prob = ODEProblem( + sys, [y => ones(2), z => 2ones(2), x => 3.0, foo => _tmp_fn2], (0.0, 1.0)) + val[] = 0 + @test_nowarn prob.f(prob.u0, prob.p, 0.0) + @test val[] == 2 + end +end diff --git a/test/structural_transformation/utils.jl b/test/structural_transformation/utils.jl index 4d5dead59f..4c1ba1e807 100644 --- a/test/structural_transformation/utils.jl +++ b/test/structural_transformation/utils.jl @@ -52,7 +52,7 @@ end @mtkcompile sys = System( [D(x) ~ z[1] + z[2] + foo(z)[1], y[1] ~ 2t, y[2] ~ 3t, z ~ foo(y)], t) @test length(equations(sys)) == 1 - @test length(observed(sys)) == 7 + @test length(observed(sys)) == 6 @test any(obs -> isequal(obs, y), observables(sys)) @test any(obs -> isequal(obs, z), observables(sys)) prob = ODEProblem(sys, [x => 1.0, foo => _tmp_fn], (0.0, 1.0)) @@ -62,60 +62,11 @@ end @test length(unknowns(isys)) == 5 @test length(equations(isys)) == 4 @test !any(equations(isys)) do eq - iscall(eq.rhs) && operation(eq.rhs) in [StructuralTransformations.getindex_wrapper, - StructuralTransformations.change_origin] + iscall(eq.rhs) && operation(eq.rhs) in [StructuralTransformations.change_origin] end end -@testset "scalarized array observed calling same function multiple times" begin - @variables x(t) y(t)[1:2] - @parameters foo(::Real)[1:2] - val = Ref(0) - function _tmp_fn2(x) - val[] += 1 - return [x, 2x] - end - @mtkcompile sys = System([D(x) ~ y[1] + y[2], y ~ foo(x)], t) - @test length(equations(sys)) == 1 - @test length(observed(sys)) == 4 - prob = ODEProblem(sys, [x => 1.0, foo => _tmp_fn2], (0.0, 1.0)) - val[] = 0 - @test_nowarn prob.f(prob.u0, prob.p, 0.0) - @test val[] == 1 - - isys = ModelingToolkit.generate_initializesystem(sys) - @test length(unknowns(isys)) == 3 - @test length(equations(isys)) == 2 - @test !any(equations(isys)) do eq - iscall(eq.rhs) && operation(eq.rhs) in [StructuralTransformations.getindex_wrapper, - StructuralTransformations.change_origin] - end - - @testset "CSE hack in equations(sys)" begin - val[] = 0 - @variables z(t)[1:2] - @mtkcompile sys = System( - [D(y) ~ foo(x), D(x) ~ sum(y), zeros(2) ~ foo(prod(z))], t) - @test length(equations(sys)) == 5 - @test length(observed(sys)) == 2 - prob = ODEProblem( - sys, [y => ones(2), z => 2ones(2), x => 3.0, foo => _tmp_fn2], (0.0, 1.0)) - val[] = 0 - @test_nowarn prob.f(prob.u0, prob.p, 0.0) - @test val[] == 2 - - isys = ModelingToolkit.generate_initializesystem(sys) - @test length(unknowns(isys)) == 5 - @test length(equations(isys)) == 2 - @test !any(equations(isys)) do eq - iscall(eq.rhs) && - operation(eq.rhs) in [StructuralTransformations.getindex_wrapper, - StructuralTransformations.change_origin] - end - end -end - -@testset "array and cse hacks can be disabled" begin +@testset "array hack can be disabled" begin @testset "fully_determined = true" begin @variables x(t) y(t)[1:2] z(t)[1:2] @parameters foo(::AbstractVector)[1:2] @@ -123,15 +74,8 @@ end @named sys = System( [D(x) ~ z[1] + z[2] + foo(z)[1], y[1] ~ 2t, y[2] ~ 3t, z ~ foo(y)], t) - sys1 = mtkcompile(sys; cse_hack = false) - @test length(observed(sys1)) == 6 - @test !any(observed(sys1)) do eq - iscall(eq.rhs) && - operation(eq.rhs) == StructuralTransformations.getindex_wrapper - end - sys2 = mtkcompile(sys; array_hack = false) - @test length(observed(sys2)) == 5 + @test length(observed(sys2)) == 4 @test !any(observed(sys2)) do eq iscall(eq.rhs) && operation(eq.rhs) == StructuralTransformations.change_origin end @@ -144,15 +88,8 @@ end @named sys = System( [D(x) ~ z[1] + z[2] + foo(z)[1] + w, y[1] ~ 2t, y[2] ~ 3t, z ~ foo(y)], t) - sys1 = mtkcompile(sys; cse_hack = false, fully_determined = false) - @test length(observed(sys1)) == 6 - @test !any(observed(sys1)) do eq - iscall(eq.rhs) && - operation(eq.rhs) == StructuralTransformations.getindex_wrapper - end - sys2 = mtkcompile(sys; array_hack = false, fully_determined = false) - @test length(observed(sys2)) == 5 + @test length(observed(sys2)) == 4 @test !any(observed(sys2)) do eq iscall(eq.rhs) && operation(eq.rhs) == StructuralTransformations.change_origin end From aff913f468f5144ca7de377ff53a569c20b0d649 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Thu, 5 Jun 2025 12:32:06 +0530 Subject: [PATCH 2/2] test: test array unknowns occurring unscalarized in initializeprobpmap --- test/initializationsystem.jl | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/test/initializationsystem.jl b/test/initializationsystem.jl index fc01e9e2e6..99e73e5b17 100644 --- a/test/initializationsystem.jl +++ b/test/initializationsystem.jl @@ -1643,3 +1643,24 @@ end @test !SciMLBase.isinplace(prob) @test !SciMLBase.isinplace(prob.f.initialization_data.initializeprob) end + +@testset "Array unknowns occurring unscalarized in initializeprobpmap" begin + @variables begin + u(t)[1:2] = 0.9ones(2) + x(t)[1:2], [guess = 0.01ones(2)] + o(t)[1:2] + end + @parameters p[1:4] = [2.0, 1.875, 2.0, 1.875] + + eqs = [D(u[1]) ~ p[1] * u[1] - p[2] * u[1] * u[2] + x[1] + 0.1 + D(u[2]) ~ p[4] * u[1] * u[2] - p[3] * u[2] - x[2] + o[1] ~ sum(p) * sum(u) + o[2] ~ sum(p) * sum(x) + x[1] ~ 0.01exp(-1) + x[2] ~ 0.01cos(t)] + + @mtkbuild sys = ODESystem(eqs, t) + prob = ODEProblem(sys, [], (0.0, 1.0)) + sol = solve(prob, Tsit5()) + @test SciMLBase.successful_retcode(sol) +end