refactor: remove CSE hack

Author: AayushSabharwal

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,60 +1222,22 @@ 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
 if all `p[i]` are present and the unscalarized form is used in any equation (observed or
 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()
     for (i, eq) in enumerate(obs)
         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

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)

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

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,76 +62,20 @@ 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]
         _tmp_fn(x) = 2x
         @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