Merge pull request #4431 from SciML/as/irstructure

refactor: use `IRStructure` for codegen

Author: AayushSabharwal

lib/ModelingToolkitBase/Project.toml

@@ -167,8 +167,8 @@ StaticArrays = "1.9.14"
 StochasticDelayDiffEq = "1.11"
 StochasticDiffEq = "6.82.0"
 SymbolicIndexingInterface = "0.3.39"
-SymbolicUtils = "4.20.1"
-Symbolics = "7.12"
+SymbolicUtils = "4.23.1"
+Symbolics = "7.18.1"
 UnPack = "0.1, 1.0"
 julia = "1.9"
 

lib/ModelingToolkitBase/src/ModelingToolkitBase.jl

@@ -190,6 +190,7 @@ include("systems/imperative_affect.jl")
 include("systems/callbacks.jl")
 include("systems/system.jl")
 include("systems/analysis_points.jl")
+include("systems/ir_info.jl")
 include("systems/codegen_utils.jl")
 include("problems/docs.jl")
 include("systems/codegen.jl")
@@ -382,6 +383,8 @@ function __init__()
     SU.hashcons(BVP_SOLUTION, true)
     SU.hashcons(unwrap(W_GAMMA), true)
     SU.hashcons(unwrap(ASSERTION_LOG_VARIABLE), true)
+    SU.hashcons(DDE_AT_IDX_SYM, true)
+    SU.hashcons(DDE_DELAY_SYM, true)
     return nothing
 end
 

lib/ModelingToolkitBase/src/systems/abstractsystem.jl

@@ -55,10 +55,10 @@ function generate_custom_function(
     if expression == Val{true}
         return fnexpr
     end
-    if fnexpr isa Tuple
+    if SU.is_array_shape(SU.shape(unwrap(exprs)))
         return eval_or_rgf.(fnexpr; eval_expression, eval_module)
     else
-        return eval_or_rgf(fnexpr; eval_expression, eval_module)
+        return eval_or_rgf(fnexpr[1]; eval_expression, eval_module)
     end
 end
 
@@ -941,6 +941,7 @@ const SYS_PROPS = [
     :state_priorities
     :irreducibles
     :maybe_zeros
+    :irstructure
     :assertions
     :ignored_connections
     :parent

lib/ModelingToolkitBase/src/systems/callbacks.jl

@@ -171,12 +171,11 @@ function AffectSystem(
         p in pre_params && continue
         push!(sys_params, p)
     end
-    discretes = map(tovar, discrete_parameters)
+    discretes = discrete_parameters
     dvs = collect(dvs)
     _dvs = map(default_toterm, dvs)
 
-    rev_map = Dict{SymbolicT, SymbolicT}(zip(discrete_parameters, discretes))
-    subs = merge(rev_map, Dict{SymbolicT, SymbolicT}(zip(dvs, _dvs)))
+    subs = Dict{SymbolicT, SymbolicT}(zip(dvs, _dvs))
     affect = substitute(affect, subs)
 
     ps = collect(gather_array_params(union(pre_params, sys_params)))
@@ -985,9 +984,6 @@ Base.@nospecializeinfer function compile_condition(
     fs = build_function_wrapper(
         sys, condit, u, p..., t; kwargs..., cse = false
     )
-    if is_discrete(cbs)
-        fs = (fs, nothing)
-    end
     fs = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(
         Val{false}, fs...; eval_expression, eval_module
     )
@@ -1384,7 +1380,13 @@ Base.@nospecializeinfer function compile_explicit_affect(
     dvs_to_update = setdiff(unknowns(aff), getfield.(observed(sys), :lhs))
 
     _affsys = unhack_system(affsys)
-    obseqs = observed(_affsys)
+    aff_ir_info = get_ir_info(affsys)
+    aff_ir = get_irstructure(affsys)
+    obseqs = Equation[]
+    sizehint!(obseqs, length(observed(_affsys)))
+    for (i, eq) in enumerate(observed(_affsys))
+        push!(obseqs, eq.lhs ~ aff_ir_info.obs_subber(eq.rhs))
+    end
 
     update_eqs = substitute(
         obseqs, Dict([p => unPre(p) for p in parameters(affsys)])
@@ -1416,18 +1418,19 @@ Base.@nospecializeinfer function compile_explicit_affect(
     u_up,
         u_up! = build_function_wrapper(
         sys, (@view rhss[is_u]), dvs, _ps..., t;
-        wrap_code = add_integrator_header(sys, integ, :u), expression = Val{false},
-        outputidxs = u_idxs, wrap_mtkparameters, cse = false, eval_expression,
-        eval_module
+        wrap_code = add_integrator_header(sys, integ, :u),
+        outputidxs = u_idxs, wrap_mtkparameters, iip_config = (false, true)
     )
     p_up,
         p_up! = build_function_wrapper(
         sys, (@view rhss[is_p]), dvs, _ps..., t;
-        wrap_code = add_integrator_header(sys, integ, :p), expression = Val{false},
-        outputidxs = p_idxs, wrap_mtkparameters, cse = false, eval_expression,
-        eval_module
+        wrap_code = add_integrator_header(sys, integ, :p),
+        outputidxs = p_idxs, wrap_mtkparameters, iip_config = (false, true)
     )
 
+    u_up! = eval_or_rgf(u_up!; eval_expression, eval_module)
+    p_up! = eval_or_rgf(p_up!; eval_expression, eval_module)
+
     return ExplicitAffect(dvs_to_update, ps_to_update, reset_jumps, u_up!, p_up!)
 end
 

lib/ModelingToolkitBase/src/systems/codegen.jl

@@ -698,11 +698,10 @@ function generate_cost(
     res = build_function_wrapper(
         sys, obj, args...; expression = Val{true}, p_start, p_end, wrap_delays,
         histfn = (p, t) -> BVP_SOLUTION(t), histfn_symbolic = BVP_SOLUTION, kwargs...
-    )
+    )[1]
     if expression == Val{true}
         return res
     end
-    f_oop = eval_or_rgf(res; eval_expression, eval_module)
     return maybe_compile_function(
         expression, wrap_gfw, (2, nargs, is_split(sys)), res; eval_expression, eval_module
     )
@@ -750,7 +749,7 @@ function generate_bvp_cost(
         histfn = (p, t) -> BVP_SOLUTION(t),
         histfn_symbolic = BVP_SOLUTION,
         cse, checkbounds, kwargs...
-    )
+    )[1]
 
     # (2, 2, is_split) means: 2 args out-of-place, 2 original args, split status
     return maybe_compile_function(
@@ -1061,8 +1060,9 @@ function generate_rate_function(js::System, rate)
     return build_function_wrapper(
         js, rate, unknowns(js), p...,
         get_iv(js),
-        expression = Val{true}
-    )
+        expression = Val{true},
+        iip_config = (true, false),
+    )[1]
 end
 
 function generate_affect_function(js::System, affect; kwargs...)
@@ -1178,7 +1178,7 @@ Generates a function that computes the observed value(s) `ts` in the system `sys
 - `ts`: The symbolic observed values whose value should be computed
 
 ## Keywords
-- `return_inplace = false`: If true and the observed value is a vector, then return both the in place and out of place methods.
+- `return_inplace = Val(false)`: If true and the observed value is a vector, then return both the in place and out of place methods. Can take boolean `true` or `false` values, but `Val(true)` or `Val(false)` is preferred.
 - `expression = false`: Generates a Julia `Expr`` computing the observed value if `expression` is true
 - `eval_expression = false`: If true and `expression = false`, evaluates the returned function in the module `eval_module`
 - `output_type = Array` the type of the array generated by a out-of-place vector-valued function
@@ -1228,7 +1228,7 @@ function build_explicit_observed_function(
         output_type = Array,
         checkbounds = true,
         ps = parameters(sys; initial_parameters = true),
-        return_inplace = false,
+        return_inplace = Val(false),
         param_only = false,
         throw = true,
         cse = true,
@@ -1320,10 +1320,6 @@ function build_explicit_observed_function(
             SU.query(pred, x)
         end
     end
-    extra_assignments = Assignment[]
-    for var in pred.present_dervars
-        push!(extra_assignments, var ← dervals[var])
-    end
     ts = substitute(ts, namespace_subs)
 
     dvs = if param_only
@@ -1371,35 +1367,22 @@ function build_explicit_observed_function(
     fns = build_function_wrapper(
         sys, ts, args...; p_start, p_end,
         output_type, mkarray, try_namespaced = true, expression = Val{true}, cse,
-        wrap_delays, extra_assignments, kwargs...
-    )
-    if fns isa Tuple
-        if expression == true || expression === Val{true}
-            return return_inplace ? fns : fns[1]
-        end
-        oop, iip = eval_or_rgf.(fns; eval_expression, eval_module)
-        f = GeneratedFunctionWrapper{
-            (
-                p_start + wrap_delays, length(args) - length(rps) + 1 + wrap_delays, is_split(sys),
-            ),
-        }(
+        wrap_delays, kwargs...
+    )::NTuple{2, Expr}
+    if expression === true || expression === Val{true}
+        return (return_inplace isa Val{true} || return_inplace isa Bool && return_inplace) ? fns : fns[1]
+    end
+
+    oop = eval_or_rgf(fns[1]; eval_expression, eval_module)
+    iip = eval_or_rgf(fns[2]; eval_expression, eval_module)
+    f = GeneratedFunctionWrapper{
+        (
+            p_start + wrap_delays, length(args) - length(rps) + 1 + wrap_delays, is_split(sys),
+        ),
+    }(
             oop, iip
         )
-        return return_inplace ? (f, f) : f
-    else
-        if expression == true || expression === Val{true}
-            return fns
-        end
-        f = eval_or_rgf(fns; eval_expression, eval_module)
-        f = GeneratedFunctionWrapper{
-            (
-                p_start + wrap_delays, length(args) - length(rps) + 1 + wrap_delays, is_split(sys),
-            ),
-        }(
-            f, nothing
-        )
-        return f
-    end
+    return (return_inplace isa Val{true} || return_inplace isa Bool && return_inplace) ? (f, f) : f
 end
 
 struct CachedLinearAb