Merge pull request #3514 from AayushSabharwal/as/cse2

feat: enable CSE in `build_function_wrapper`

Author: ChrisRackauckas

Project.toml

@@ -150,7 +150,7 @@ StaticArrays = "0.10, 0.11, 0.12, 1.0"
 StochasticDelayDiffEq = "1.8.1"
 StochasticDiffEq = "6.72.1"
 SymbolicIndexingInterface = "0.3.37"
-SymbolicUtils = "3.14"
+SymbolicUtils = "3.25.1"
 Symbolics = "6.36"
 URIs = "1"
 UnPack = "0.1, 1.0"

src/systems/abstractsystem.jl

@@ -535,7 +535,8 @@ end
 SymbolicIndexingInterface.supports_tuple_observed(::AbstractSystem) = true
 
 function SymbolicIndexingInterface.observed(
-        sys::AbstractSystem, sym; eval_expression = false, eval_module = @__MODULE__, checkbounds = true)
+        sys::AbstractSystem, sym; eval_expression = false, eval_module = @__MODULE__,
+        checkbounds = true, cse = true)
     if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
         if sym isa Symbol
             _sym = get(ic.symbol_to_variable, sym, nothing)
@@ -559,7 +560,7 @@ function SymbolicIndexingInterface.observed(
         end
     end
     return build_explicit_observed_function(
-        sys, sym; eval_expression, eval_module, checkbounds)
+        sys, sym; eval_expression, eval_module, checkbounds, cse)
 end
 
 function SymbolicIndexingInterface.default_values(sys::AbstractSystem)
@@ -1774,13 +1775,14 @@ struct ObservedFunctionCache{S}
     eval_expression::Bool
     eval_module::Module
     checkbounds::Bool
+    cse::Bool
 end
 
 function ObservedFunctionCache(
         sys; steady_state = false, eval_expression = false,
-        eval_module = @__MODULE__, checkbounds = true)
+        eval_module = @__MODULE__, checkbounds = true, cse = true)
     return ObservedFunctionCache(
-        sys, Dict(), steady_state, eval_expression, eval_module, checkbounds)
+        sys, Dict(), steady_state, eval_expression, eval_module, checkbounds, cse)
 end
 
 # This is hit because ensemble problems do a deepcopy
@@ -1791,8 +1793,9 @@ function Base.deepcopy_internal(ofc::ObservedFunctionCache, stackdict::IdDict)
     eval_expression = ofc.eval_expression
     eval_module = ofc.eval_module
     checkbounds = ofc.checkbounds
+    cse = ofc.cse
     newofc = ObservedFunctionCache(
-        sys, dict, steady_state, eval_expression, eval_module, checkbounds)
+        sys, dict, steady_state, eval_expression, eval_module, checkbounds, cse)
     stackdict[ofc] = newofc
     return newofc
 end
@@ -1801,7 +1804,7 @@ function (ofc::ObservedFunctionCache)(obsvar, args...)
     obs = get!(ofc.dict, value(obsvar)) do
         SymbolicIndexingInterface.observed(
             ofc.sys, obsvar; eval_expression = ofc.eval_expression,
-            eval_module = ofc.eval_module, checkbounds = ofc.checkbounds)
+            eval_module = ofc.eval_module, checkbounds = ofc.checkbounds, cse = ofc.cse)
     end
     if ofc.steady_state
         obs = let fn = obs

src/systems/callbacks.jl

@@ -699,7 +699,7 @@ function compile_affect(eqs::Vector{Equation}, cb, sys, dvs, ps; outputidxs = no
                         add_integrator_header(sys, integ, outvar),
             outputidxs = update_inds,
             create_bindings = false,
-            kwargs...)
+            kwargs..., cse = false)
         # applied user-provided function to the generated expression
         if postprocess_affect_expr! !== nothing
             postprocess_affect_expr!(rf_ip, integ)
@@ -729,7 +729,7 @@ function generate_single_rootfinding_callback(
     end
 
     rf_oop, rf_ip = generate_custom_function(
-        sys, [eq.rhs], dvs, ps; expression = Val{false}, kwargs...)
+        sys, [eq.rhs], dvs, ps; expression = Val{false}, kwargs..., cse = false)
     affect_function = compile_affect_fn(cb, sys, dvs, ps, kwargs)
     cond = function (u, t, integ)
         if DiffEqBase.isinplace(integ.sol.prob)
@@ -780,7 +780,7 @@ function generate_vector_rootfinding_callback(
 
     rhss = map(x -> x.rhs, eqs)
     _, rf_ip = generate_custom_function(
-        sys, rhss, dvs, ps; expression = Val{false}, kwargs...)
+        sys, rhss, dvs, ps; expression = Val{false}, kwargs..., cse = false)
 
     affect_functions = @NamedTuple{
         affect::Function,

src/systems/codegen_utils.jl

@@ -132,7 +132,7 @@ function build_function_wrapper(sys::AbstractSystem, expr, args...; p_start = 2,
         wrap_delays = is_dde(sys), wrap_code = identity,
         add_observed = true, filter_observed = Returns(true),
         create_bindings = false, output_type = nothing, mkarray = nothing,
-        wrap_mtkparameters = true, extra_assignments = Assignment[], kwargs...)
+        wrap_mtkparameters = true, extra_assignments = Assignment[], cse = true, kwargs...)
     isscalar = !(expr isa AbstractArray || symbolic_type(expr) == ArraySymbolic())
     # filter observed equations
     obs = filter(filter_observed, observed(sys))
@@ -234,7 +234,7 @@ function build_function_wrapper(sys::AbstractSystem, expr, args...; p_start = 2,
     if wrap_code isa Tuple && symbolic_type(expr) == ScalarSymbolic()
         wrap_code = wrap_code[1]
     end
-    return build_function(expr, args...; wrap_code, similarto, kwargs...)
+    return build_function(expr, args...; wrap_code, similarto, cse, kwargs...)
 end
 
 """

src/systems/diffeqs/abstractodesystem.jl

@@ -312,12 +312,13 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         analytic = nothing,
         split_idxs = nothing,
         initialization_data = nothing,
+        cse = true,
         kwargs...) where {iip, specialize}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `ODEFunction`")
     end
     f_gen = generate_function(sys, dvs, ps; expression = Val{true},
-        expression_module = eval_module, checkbounds = checkbounds,
+        expression_module = eval_module, checkbounds = checkbounds, cse,
         kwargs...)
     f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
     f = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(f_oop, f_iip)
@@ -333,7 +334,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         tgrad_gen = generate_tgrad(sys, dvs, ps;
             simplify = simplify,
             expression = Val{true},
-            expression_module = eval_module,
+            expression_module = eval_module, cse,
             checkbounds = checkbounds, kwargs...)
         tgrad_oop, tgrad_iip = eval_or_rgf.(tgrad_gen; eval_expression, eval_module)
         _tgrad = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(tgrad_oop, tgrad_iip)
@@ -345,7 +346,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         jac_gen = generate_jacobian(sys, dvs, ps;
             simplify = simplify, sparse = sparse,
             expression = Val{true},
-            expression_module = eval_module,
+            expression_module = eval_module, cse,
             checkbounds = checkbounds, kwargs...)
         jac_oop, jac_iip = eval_or_rgf.(jac_gen; eval_expression, eval_module)
 
@@ -365,7 +366,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
     end
 
     observedfun = ObservedFunctionCache(
-        sys; steady_state, eval_expression, eval_module, checkbounds)
+        sys; steady_state, eval_expression, eval_module, checkbounds, cse)
 
     jac_prototype = if sparse
         uElType = u0 === nothing ? Float64 : eltype(u0)
@@ -420,12 +421,13 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         eval_module = @__MODULE__,
         checkbounds = false,
         initialization_data = nothing,
+        cse = true,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating a `DAEFunction`")
     end
     f_gen = generate_function(sys, dvs, ps; implicit_dae = true,
-        expression = Val{true},
+        expression = Val{true}, cse,
         expression_module = eval_module, checkbounds = checkbounds,
         kwargs...)
     f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
@@ -435,7 +437,7 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         jac_gen = generate_dae_jacobian(sys, dvs, ps;
             simplify = simplify, sparse = sparse,
             expression = Val{true},
-            expression_module = eval_module,
+            expression_module = eval_module, cse,
             checkbounds = checkbounds, kwargs...)
         jac_oop, jac_iip = eval_or_rgf.(jac_gen; eval_expression, eval_module)
 
@@ -445,7 +447,7 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
     end
 
     observedfun = ObservedFunctionCache(
-        sys; eval_expression, eval_module, checkbounds = get(kwargs, :checkbounds, false))
+        sys; eval_expression, eval_module, checkbounds = get(kwargs, :checkbounds, false), cse)
 
     jac_prototype = if sparse
         uElType = u0 === nothing ? Float64 : eltype(u0)
@@ -479,14 +481,15 @@ function DiffEqBase.DDEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         eval_module = @__MODULE__,
         checkbounds = false,
         initialization_data = nothing,
+        cse = true,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `DDEFunction`")
     end
     f_gen = generate_function(sys, dvs, ps; isdde = true,
         expression = Val{true},
         expression_module = eval_module, checkbounds = checkbounds,
-        kwargs...)
+        cse, kwargs...)
     f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
     f = GeneratedFunctionWrapper{(3, 4, is_split(sys))}(f_oop, f_iip)
 
@@ -503,19 +506,20 @@ function DiffEqBase.SDDEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys
         eval_module = @__MODULE__,
         checkbounds = false,
         initialization_data = nothing,
+        cse = true,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `SDDEFunction`")
     end
     f_gen = generate_function(sys, dvs, ps; isdde = true,
         expression = Val{true},
         expression_module = eval_module, checkbounds = checkbounds,
-        kwargs...)
+        cse, kwargs...)
     f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
     f = GeneratedFunctionWrapper{(3, 4, is_split(sys))}(f_oop, f_iip)
 
     g_gen = generate_diffusion_function(sys, dvs, ps; expression = Val{true},
-        isdde = true, kwargs...)
+        isdde = true, cse, kwargs...)
     g_oop, g_iip = eval_or_rgf.(g_gen; eval_expression, eval_module)
     g = GeneratedFunctionWrapper{(3, 4, is_split(sys))}(g_oop, g_iip)
 
@@ -841,6 +845,7 @@ function SciMLBase.BVProblem{iip, specialize}(sys::AbstractODESystem, u0map = []
         warn_initialize_determined = true,
         eval_expression = false,
         eval_module = @__MODULE__,
+        cse = true,
         kwargs...) where {iip, specialize}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `BVProblem`")
@@ -864,12 +869,12 @@ function SciMLBase.BVProblem{iip, specialize}(sys::AbstractODESystem, u0map = []
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
     f, u0, p = process_SciMLProblem(ODEFunction{iip, specialize}, sys, _u0map, parammap;
         t = tspan !== nothing ? tspan[1] : tspan, guesses,
-        check_length, warn_initialize_determined, eval_expression, eval_module, kwargs...)
+        check_length, warn_initialize_determined, eval_expression, eval_module, cse, kwargs...)
 
     stidxmap = Dict([v => i for (i, v) in enumerate(sts)])
     u0_idxs = has_alg_eqs(sys) ? collect(1:length(sts)) : [stidxmap[k] for (k, v) in u0map]
 
-    fns = generate_function_bc(sys, u0, u0_idxs, tspan)
+    fns = generate_function_bc(sys, u0, u0_idxs, tspan; cse)
     bc_oop, bc_iip = eval_or_rgf.(fns; eval_expression, eval_module)
     bc(sol, p, t) = bc_oop(sol, p, t)
     bc(resid, u, p, t) = bc_iip(resid, u, p, t)
@@ -988,15 +993,16 @@ function DiffEqBase.DDEProblem{iip}(sys::AbstractODESystem, u0map = [],
         eval_expression = false,
         eval_module = @__MODULE__,
         u0_constructor = identity,
+        cse = true,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating a `DDEProblem`")
     end
     f, u0, p = process_SciMLProblem(DDEFunction{iip}, sys, u0map, parammap;
         t = tspan !== nothing ? tspan[1] : tspan,
-        symbolic_u0 = true, u0_constructor,
+        symbolic_u0 = true, u0_constructor, cse,
         check_length, eval_expression, eval_module, kwargs...)
-    h_gen = generate_history(sys, u0; expression = Val{true})
+    h_gen = generate_history(sys, u0; expression = Val{true}, cse)
     h_oop, h_iip = eval_or_rgf.(h_gen; eval_expression, eval_module)
     h = h_oop
     u0 = float.(h(p, tspan[1]))
@@ -1027,15 +1033,16 @@ function DiffEqBase.SDDEProblem{iip}(sys::AbstractODESystem, u0map = [],
         eval_expression = false,
         eval_module = @__MODULE__,
         u0_constructor = identity,
+        cse = true,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating a `SDDEProblem`")
     end
     f, u0, p = process_SciMLProblem(SDDEFunction{iip}, sys, u0map, parammap;
         t = tspan !== nothing ? tspan[1] : tspan,
         symbolic_u0 = true, eval_expression, eval_module, u0_constructor,
-        check_length, kwargs...)
-    h_gen = generate_history(sys, u0; expression = Val{true})
+        check_length, cse, kwargs...)
+    h_gen = generate_history(sys, u0; expression = Val{true}, cse)
     h_oop, h_iip = eval_or_rgf.(h_gen; eval_expression, eval_module)
     h = h_oop
     u0 = h(p, tspan[1])

src/systems/diffeqs/odesystem.jl

@@ -454,8 +454,8 @@ Generates a function that computes the observed value(s) `ts` in the system `sys
 - `checkbounds = true` checks bounds if true when destructuring parameters
 - `op = Operator` sets the recursion terminator for the walk done by `vars` to identify the variables that appear in `ts`. See the documentation for `vars` for more detail.
 - `throw = true` if true, throw an error when generating a function for `ts` that reference variables that do not exist.
-- `mkarray`; only used if the output is an array (that is, `!isscalar(ts)`  and `ts` is not a tuple, in which case the result will always be a tuple). Called as `mkarray(ts, output_type)` where `ts` are the expressions to put in 
-the array and `output_type` is the argument of the same name passed to build_explicit_observed_function.
+- `mkarray`: only used if the output is an array (that is, `!isscalar(ts)`  and `ts` is not a tuple, in which case the result will always be a tuple). Called as `mkarray(ts, output_type)` where `ts` are the expressions to put in the array and `output_type` is the argument of the same name passed to build_explicit_observed_function.
+- `cse = true`: Whether to use Common Subexpression Elimination (CSE) to generate a more efficient function.
 
 ## Returns
 
@@ -493,6 +493,7 @@ function build_explicit_observed_function(sys, ts;
         param_only = false,
         op = Operator,
         throw = true,
+        cse = true,
         mkarray = nothing)
     is_tuple = ts isa Tuple
     if is_tuple
@@ -579,7 +580,7 @@ function build_explicit_observed_function(sys, ts;
     p_end = length(dvs) + length(inputs) + length(ps)
     fns = build_function_wrapper(
         sys, ts, args...; p_start, p_end, filter_observed = obsfilter,
-        output_type, mkarray, try_namespaced = true, expression = Val{true})
+        output_type, mkarray, try_namespaced = true, expression = Val{true}, cse)
     if fns isa Tuple
         if expression
             return return_inplace ? fns : fns[1]

src/systems/diffeqs/sdesystem.jl

@@ -595,23 +595,23 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         jac = false, Wfact = false, eval_expression = false,
         eval_module = @__MODULE__,
         checkbounds = false, initialization_data = nothing,
-        kwargs...) where {iip, specialize}
+        cse = true, kwargs...) where {iip, specialize}
     if !iscomplete(sys)
         error("A completed `SDESystem` is required. Call `complete` or `structural_simplify` on the system before creating an `SDEFunction`")
     end
     dvs = scalarize.(dvs)
 
-    f_gen = generate_function(sys, dvs, ps; expression = Val{true}, kwargs...)
+    f_gen = generate_function(sys, dvs, ps; expression = Val{true}, cse, kwargs...)
     f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
     g_gen = generate_diffusion_function(sys, dvs, ps; expression = Val{true},
-        kwargs...)
+        cse, kwargs...)
     g_oop, g_iip = eval_or_rgf.(g_gen; eval_expression, eval_module)
 
     f = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(f_oop, f_iip)
     g = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(g_oop, g_iip)
 
     if tgrad
-        tgrad_gen = generate_tgrad(sys, dvs, ps; expression = Val{true},
+        tgrad_gen = generate_tgrad(sys, dvs, ps; expression = Val{true}, cse,
             kwargs...)
         tgrad_oop, tgrad_iip = eval_or_rgf.(tgrad_gen; eval_expression, eval_module)
         _tgrad = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(tgrad_oop, tgrad_iip)
@@ -621,7 +621,7 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
 
     if jac
         jac_gen = generate_jacobian(sys, dvs, ps; expression = Val{true},
-            sparse = sparse, kwargs...)
+            sparse = sparse, cse, kwargs...)
         jac_oop, jac_iip = eval_or_rgf.(jac_gen; eval_expression, eval_module)
 
         _jac = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(jac_oop, jac_iip)
@@ -631,7 +631,7 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
 
     if Wfact
         tmp_Wfact, tmp_Wfact_t = generate_factorized_W(sys, dvs, ps, true;
-            expression = Val{true}, kwargs...)
+            expression = Val{true}, cse, kwargs...)
         Wfact_oop, Wfact_iip = eval_or_rgf.(tmp_Wfact; eval_expression, eval_module)
         Wfact_oop_t, Wfact_iip_t = eval_or_rgf.(tmp_Wfact_t; eval_expression, eval_module)
 
@@ -645,7 +645,7 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
     _M = (u0 === nothing || M == I) ? M : ArrayInterface.restructure(u0 .* u0', M)
 
     observedfun = ObservedFunctionCache(
-        sys; eval_expression, eval_module, checkbounds = get(kwargs, :checkbounds, false))
+        sys; eval_expression, eval_module, checkbounds = get(kwargs, :checkbounds, false), cse)
 
     SDEFunction{iip, specialize}(f, g;
         sys = sys,

src/systems/discrete_system/discrete_system.jl

@@ -360,13 +360,13 @@ function SciMLBase.DiscreteFunction{iip, specialize}(
         t = nothing,
         eval_expression = false,
         eval_module = @__MODULE__,
-        analytic = nothing,
+        analytic = nothing, cse = true,
         kwargs...) where {iip, specialize}
     if !iscomplete(sys)
         error("A completed `DiscreteSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `DiscreteProblem`")
     end
     f_gen = generate_function(sys, dvs, ps; expression = Val{true},
-        expression_module = eval_module, kwargs...)
+        expression_module = eval_module, cse, kwargs...)
     f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
     f = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(f_oop, f_iip)
 
@@ -378,7 +378,7 @@ function SciMLBase.DiscreteFunction{iip, specialize}(
     end
 
     observedfun = ObservedFunctionCache(
-        sys; eval_expression, eval_module, checkbounds = get(kwargs, :checkbounds, false))
+        sys; eval_expression, eval_module, checkbounds = get(kwargs, :checkbounds, false), cse)
 
     DiscreteFunction{iip, specialize}(f;
         sys = sys,