Merge pull request #4044 from SciML/as/precompile

refactor: improve precompile-friendliness

Author: AayushSabharwal

Project.toml

@@ -109,7 +109,7 @@ LinearSolve = "3.19.2"
 Logging = "1"
 ModelingToolkitBase = "1"
 ModelingToolkitStandardLibrary = "2.20"
-ModelingToolkitTearing = "1"
+ModelingToolkitTearing = "1.0.2"
 Moshi = "0.3"
 NonlinearSolve = "4.3"
 OffsetArrays = "1"
@@ -130,12 +130,12 @@ Serialization = "1"
 Setfield = "0.7, 0.8, 1"
 SimpleNonlinearSolve = "0.1.0, 1, 2"
 SparseArrays = "1"
-StateSelection = "1"
+StateSelection = "1.1.1"
 StaticArrays = "1.9.14"
 StochasticDelayDiffEq = "1.11"
 StochasticDiffEq = "6.82.0"
 SymbolicIndexingInterface = "0.3.39"
-SymbolicUtils = "4.5.1"
+SymbolicUtils = "4.7.1"
 Symbolics = "7"
 UnPack = "0.1, 1.0"
 julia = "1.9"

lib/ModelingToolkitBase/src/ModelingToolkitBase.jl

@@ -10,6 +10,9 @@ using PrecompileTools, Reexport
     using JumpProcesses
     # ONLY here for the invalidations
     import REPL
+    using OffsetArrays: Origin
+    import BlockArrays: BlockArray, BlockedArray, Block, blocksize, blocksizes, blockpush!,
+                        undef_blocks, blocks
 end
 
 import SymbolicUtils
@@ -60,9 +63,6 @@ using Moshi.Data: @data
 using Reexport
 using RecursiveArrayTools
 import Graphs: SimpleDiGraph, add_edge!, incidence_matrix
-import BlockArrays: BlockArray, BlockedArray, Block, blocksize, blocksizes, blockpush!,
-                    undef_blocks, blocks
-using OffsetArrays: Origin
 import CommonSolve
 import EnumX
 import ReadOnlyDicts: ReadOnlyDict

lib/ModelingToolkitBase/src/discretedomain.jl

@@ -23,7 +23,7 @@ struct Shift <: Operator
     """Fixed Shift"""
     t::Union{Nothing, SymbolicT}
     steps::Int
-    Shift(t, steps = 1) = new(value(t), steps)
+    Shift(t, steps = 1) = new(unwrap(t), steps)
 end
 Shift(steps::Int) = new(nothing, steps)
 normalize_to_differential(s::Shift) = Differential(s.t)^s.steps

lib/ModelingToolkitBase/src/precompile.jl

@@ -12,6 +12,7 @@ PrecompileTools.@compile_workload begin
         x ^ 5
         6 ^ x
         x - y
+        x * x * q[1]
         -y
         2y
         z = 2
@@ -81,11 +82,22 @@ PrecompileTools.@compile_workload begin
         q[1]
         q'q
      using ModelingToolkitBase
-    @variables x(ModelingToolkitBase.t_nounits) y(ModelingToolkitBase.t_nounits)
-    isequal(ModelingToolkitBase.D_nounits.x, ModelingToolkitBase.t_nounits)
+    @parameters g
+    @variables x(ModelingToolkitBase.t_nounits)
+    @variables y(ModelingToolkitBase.t_nounits) [state_priority = 10]
+    @variables λ(ModelingToolkitBase.t_nounits)
+    eqs = [
+        ModelingToolkitBase.D_nounits(ModelingToolkitBase.D_nounits(x)) ~ λ * x
+        ModelingToolkitBase.D_nounits(ModelingToolkitBase.D_nounits(y)) ~ λ * y - g
+        x^2 + y^2 ~ 1
+    ]
+    dvs = Num[x, y, λ]
+    ps = Num[g]
     ics = Dict{SymbolicT, SymbolicT}()
-    ics[x] = 2.3
-    sys = System([ModelingToolkitBase.D_nounits(x) ~ x * y, y ~ 2x], ModelingToolkitBase.t_nounits, [x, y], Num[]; initial_conditions = ics, guesses = ics, name = :sys)
+    ics[y] = -1.0
+    ics[ModelingToolkitBase.D_nounits(x)] = 0.5
+    isequal(ModelingToolkitBase.D_nounits.x, ModelingToolkitBase.t_nounits)
+    sys = System(eqs, ModelingToolkitBase.t_nounits, dvs, ps; initial_conditions = ics, guesses = ics, name = :sys)
     complete(sys)
     @static if @isdefined(ModelingToolkit)
         TearingState(sys)

lib/ModelingToolkitBase/src/problems/jumpproblem.jl

@@ -172,23 +172,22 @@ end
 
 ##### MTK dispatches for Symbolic jumps #####
 eqtype_supports_collect_vars(j::MassActionJump) = true
-function collect_vars!(unknowns::OrderedSet{SymbolicT}, parameters::OrderedSet{SymbolicT}, j::MassActionJump, iv::Union{SymbolicT, Nothing}; depth = 0,
-        op = Differential)
-    collect_vars!(unknowns, parameters, j.scaled_rates, iv; depth, op)
+function collect_vars!(unknowns::OrderedSet{SymbolicT}, parameters::OrderedSet{SymbolicT}, j::MassActionJump, iv::Union{SymbolicT, Nothing}, ::Type{op} = Differential; depth = 0) where {op}
+    collect_vars!(unknowns, parameters, j.scaled_rates, iv, op; depth)
     for field in (j.reactant_stoch, j.net_stoch)
         for el in field
-            collect_vars!(unknowns, parameters, el, iv; depth, op)
+            collect_vars!(unknowns, parameters, el, iv, op; depth)
         end
     end
     return nothing
 end
 
 eqtype_supports_collect_vars(j::Union{ConstantRateJump, VariableRateJump}) = true
 function collect_vars!(unknowns::OrderedSet{SymbolicT}, parameters::OrderedSet{SymbolicT}, j::Union{ConstantRateJump, VariableRateJump},
-        iv::Union{SymbolicT, Nothing}; depth = 0, op = Differential)
-    collect_vars!(unknowns, parameters, j.rate, iv; depth, op)
+        iv::Union{SymbolicT, Nothing}, ::Type{op} = Differential; depth = 0) where {op}
+    collect_vars!(unknowns, parameters, j.rate, iv, op; depth)
     for eq in j.affect!
-        (eq isa Equation) && collect_vars!(unknowns, parameters, eq, iv; depth, op)
+        (eq isa Equation) && collect_vars!(unknowns, parameters, eq, iv, op; depth)
     end
     return nothing
 end

lib/ModelingToolkitBase/src/systems/abstractsystem.jl

@@ -1460,6 +1460,14 @@ function unknowns_toplevel(sys::AbstractSystem)
     return get_unknowns(sys)
 end
 
+function __no_initial_params_pred(x::SymbolicT)
+    arr, _ = split_indexed_var(x)
+    Moshi.Match.@match arr begin
+        BSImpl.Term(; f) && if f isa Initial end => false
+        _ => true
+    end
+end
+
 """
 $(TYPEDSIGNATURES)
 
@@ -1482,11 +1490,7 @@ function parameters(sys::AbstractSystem; initial_parameters = false)
     end
     result = collect(result)
     if !initial_parameters && !is_initializesystem(sys)
-        filter!(result) do sym
-            return !(isoperator(sym, Initial) ||
-                     iscall(sym) && operation(sym) === getindex &&
-                     isoperator(arguments(sym)[1], Initial))
-        end
+        filter!(__no_initial_params_pred, result)
     end
     return result
 end
@@ -1699,7 +1703,7 @@ Recursively substitute `dict` into `expr`. Use `Symbolics.simplify` on the expre
 if `simplify == true`.
 """
 function substitute_and_simplify(expr, dict::AbstractDict, simplify::Bool)
-    expr = Symbolics.fixpoint_sub(expr, dict; operator = Union{ModelingToolkitBase.Initial, Pre})
+    expr = Symbolics.fixpoint_sub(expr, dict, Union{Initial, Pre})
     simplify ? Symbolics.simplify(expr) : expr
 end
 

lib/ModelingToolkitBase/src/systems/callbacks.jl

@@ -89,7 +89,7 @@ function AffectSystem(affect::Vector{Equation}; discrete_parameters = SymbolicT[
         if !haspre(eq) && !(isconst(eq.lhs) && isconst(eq.rhs))
             @invokelatest warn_algebraic_equation(eq)
         end
-        collect_vars!(dvs, params, eq, iv; op = Pre)
+        collect_vars!(dvs, params, eq, iv, Pre)
         empty!(_varsbuf)
         SU.search_variables!(_varsbuf, eq; is_atomic = OperatorIsAtomic{Pre}())
         filter!(x -> iscall(x) && operation(x) === Pre(), _varsbuf)
@@ -125,20 +125,35 @@ function AffectSystem(affect::Vector{Equation}; discrete_parameters = SymbolicT[
     # This `@invokelatest` should not be necessary, but it works around the inference bug
     # in https://github.com/JuliaLang/julia/issues/59943. Remove it at your own risk, the
     # bug took weeks to reduce to an MWE.
-    affectsys = @invokelatest mtkcompile(affectsys; fully_determined = nothing)
+    affectsys = (@invokelatest mtkcompile(affectsys; fully_determined = nothing))::System
     # get accessed parameters p from Pre(p) in the callback parameters
     accessed_params = Vector{SymbolicT}(filter(isparameter, map(unPre, collect(pre_params))))
     union!(accessed_params, sys_params)
 
     # add scalarized unknowns to the map.
     _obs, _ = unhack_observed(observed(affectsys), equations(affectsys))
     _dvs = vcat(unknowns(affectsys), map(eq -> eq.lhs, _obs))
-    _dvs = reduce(vcat, map(safe_vec ∘ scalarize, _dvs), init = SymbolicT[])
-    _discs = reduce(vcat, map(safe_vec ∘ scalarize, discretes); init = SymbolicT[])
+    _dvs = __safe_scalarize_vars(_dvs)
+    _discs = __safe_scalarize_vars(discretes)
     setdiff!(_dvs, _discs)
     AffectSystem(affectsys, _dvs, accessed_params, discrete_parameters)
 end
 
+function __safe_scalarize_vars(vars::Vector{SymbolicT})
+    _vars = SymbolicT[]
+    for v in vars
+        sh = SU.shape(v)::SU.ShapeVecT
+        if isempty(sh)
+            push!(_vars, v)
+            continue
+        end
+        for i in SU.stable_eachindex(v)
+            push!(_vars, v[i])
+        end
+    end
+    return _vars
+end
+
 safe_vec(@nospecialize(x)) = x isa SymbolicT ? [x] : vec(x::Array{SymbolicT})
 
 system(a::AffectSystem) = a.system
@@ -1043,13 +1058,13 @@ The `SymbolicDiscreteCallback`s in the returned vector are structs with two fiel
 See also `get_discrete_events`, which only returns the events of the top-level system.
 """
 function discrete_events(sys::AbstractSystem)
-    obs = get_discrete_events(sys)
+    cbs = get_discrete_events(sys)
     systems = get_systems(sys)
-    cbs = [obs;
-           reduce(vcat,
-               (map(cb -> namespace_callback(cb, s), discrete_events(s)) for s in systems),
-               init = SymbolicDiscreteCallback[])]
-    cbs
+    cbs = copy(cbs)
+    for s in systems
+        append!(cbs, map(Base.Fix2(namespace_callback, s), discrete_events(s)))
+    end
+    return cbs
 end
 
 """
@@ -1100,15 +1115,13 @@ The `SymbolicContinuousCallback`s in the returned vector are structs with two fi
 See also `get_continuous_events`, which only returns the events of the top-level system.
 """
 function continuous_events(sys::AbstractSystem)
-    obs = get_continuous_events(sys)
-    filter(!isempty, obs)
-
+    cbs = get_continuous_events(sys)
     systems = get_systems(sys)
-    cbs = [obs;
-           reduce(vcat,
-               (map(o -> namespace_callback(o, s), continuous_events(s)) for s in systems),
-               init = SymbolicContinuousCallback[])]
-    filter(!isempty, cbs)
+    cbs = copy(cbs)
+    for s in systems
+        append!(cbs, map(Base.Fix2(namespace_callback, s), continuous_events(s)))
+    end
+    return cbs
 end
 
 """

lib/ModelingToolkitBase/src/systems/codegen_utils.jl

@@ -42,7 +42,7 @@ reconstruct array variables if they are present scalarized in `args`.
 """
 function array_variable_assignments(args...; argument_name = generated_argument_name)
     # map array symbolic to an identically sized array where each element is (buffer_idx, idx_in_buffer)
-    var_to_arridxs = Dict{BasicSymbolic, Array{Tuple{Int, Int}}}()
+    var_to_arridxs = Dict{SymbolicT, Vector{Tuple{Int, Int}}}()
     for (i, arg) in enumerate(args)
         # filter out non-arrays
         # any element of args which is not an array is assumed to not contain a
@@ -55,13 +55,12 @@ function array_variable_assignments(args...; argument_name = generated_argument_
         for (j, var) in enumerate(arg)
             var = unwrap(var)
             # filter out non-array-symbolics
-            iscall(var) || continue
-            operation(var) == getindex || continue
-            arrvar = arguments(var)[1]
+            arrvar, isarr = split_indexed_var(var)
+            isarr || continue
             # get and/or construct the buffer storing indexes
             idxbuffer = get!(
-                () -> map(Returns((0, 0)), eachindex(arrvar)), var_to_arridxs, arrvar)
-            Origin(first.(axes(arrvar))...)(idxbuffer)[arguments(var)[2:end]...] = (i, j)
+                () -> map(Returns((0, 0)), SU.stable_eachindex(arrvar)), var_to_arridxs, arrvar)
+            idxbuffer[SU.as_linear_idx(SU.shape(arrvar), get_stable_index(var))] = (i, j)
         end
     end
 

lib/ModelingToolkitBase/src/systems/system.jl

@@ -381,12 +381,12 @@ function System(eqs::Vector{Equation}, iv, dvs, ps, brownians = SymbolicT[];
         continuous_events = SymbolicContinuousCallback[], discrete_events = SymbolicDiscreteCallback[],
         connector_type = nothing, assertions = Dict{SymbolicT, String}(),
         metadata = MetadataT(), gui_metadata = nothing,
-        is_dde = nothing, tstops = [], inputs = OrderedSet{SymbolicT}(),
+        is_dde = nothing, @nospecialize(tstops = []), inputs = OrderedSet{SymbolicT}(),
         outputs = OrderedSet{SymbolicT}(), tearing_state = nothing,
         ignored_connections = nothing, parent = nothing,
         description = "", name = nothing, discover_from_metadata = true,
         initializesystem = nothing, is_initializesystem = false, is_discrete = false,
-        preface = [], checks = true, __legacy_defaults__ = nothing)
+        @nospecialize(preface = nothing), checks = true, __legacy_defaults__ = nothing)
     name === nothing && throw(NoNameError())
 
     if __legacy_defaults__ !== nothing
@@ -480,14 +480,17 @@ function System(eqs::Vector{Equation}, iv, dvs, ps, brownians = SymbolicT[];
     filter!(!(Base.Fix1(===, COMMON_NOTHING) ∘ last), guesses)
 
     if iv === nothing
-        filter!(bindings) do kvp
-            k = kvp[1]
-            if k in all_dvs
-                initial_conditions[k] = kvp[2]
-                return false
+        filterer = let initial_conditions = initial_conditions, all_dvs = all_dvs
+            function _filterer(kvp)
+                k = kvp[1]
+                if k in all_dvs
+                    initial_conditions[k] = kvp[2]
+                    return false
+                end
+                return true
             end
-            return true
         end
+        filter!(filterer, bindings)
     end
 
     check_bindings(ps, bindings)
@@ -683,19 +686,20 @@ Create a `System` with a single equation `eq`.
 System(eq::Equation, args...; kwargs...) = System([eq], args...; kwargs...)
 
 function gather_array_params(ps)
-    new_ps = OrderedSet()
+    new_ps = OrderedSet{SymbolicT}()
     for p in ps
-        if iscall(p) && operation(p) === getindex
-            par = arguments(p)[begin]
-            if symbolic_has_known_size(par) && all(par[i] in ps for i in eachindex(par))
-                push!(new_ps, par)
+        arr, isarr = split_indexed_var(p)
+        sh = SU.shape(arr)
+        if isarr
+            if !(sh isa SU.Unknown) && all(in(ps) ∘ Base.Fix1(getindex, arr), SU.stable_eachindex(arr))
+                push!(new_ps, arr)
             else
                 push!(new_ps, p)
             end
         else
-            if symbolic_type(p) == ArraySymbolic() && symbolic_has_known_size(p)
-                for i in eachindex(p)
-                    delete!(new_ps, p[i])
+            if sh isa SU.ShapeVecT && !isempty(sh)
+                for i in SU.stable_eachindex(arr)
+                    delete!(new_ps, arr[i])
                 end
             end
             push!(new_ps, p)

lib/ModelingToolkitBase/src/systems/systems.jl

@@ -10,14 +10,15 @@ function canonicalize_io(iovars, type::String)
     iobuffer = OrderedSet{SymbolicT}()
     arrsyms = AtomicArrayDict{OrderedSet{SymbolicT}}()
     for var in iovars
-        if Symbolics.isarraysymbolic(var)
-            if !symbolic_has_known_size(var)
-                throw(ArgumentError("""
-                All $(type)s must have known shape. Found $var with unknown shape.
-                """))
+        sh = SU.shape(var)
+        if SU.is_array_shape(sh)
+            if sh isa SU.ShapeVecT
+                union!(iobuffer, vec(collect(var)::Array{SymbolicT})::Vector{SymbolicT})
+                continue
             end
-            union!(iobuffer, vec(collect(var)::Array{SymbolicT})::Vector{SymbolicT})
-            continue
+            throw(ArgumentError("""
+            All $(type)s must have known shape. Found $var with unknown shape.
+            """))
         end
         arr, isarr = split_indexed_var(var)
         if isarr
@@ -41,7 +42,7 @@ function canonicalize_io(iovars, type::String)
             or simply pass $k as an $type.
             """))
         end
-        if type != "output" && !isequal(vec(collect(k))::Vector{SymbolicT}, collect(v))
+        if type != "output" && !isequal(vec(collect(k)::Array{SymbolicT})::Vector{SymbolicT}, collect(v))
             throw(ArgumentError("""
             Elements of scalarized array variables must be in sorted order in $(type)s. \
             Either pass all scalarized elements in sorted order as $(type)s \
@@ -601,16 +602,17 @@ function __num_isdiag_noise(mat)
     true
 end
 
-function __get_num_diag_noise(mat)
-    map(axes(mat, 1)) do i
+function __get_num_diag_noise(mat::Matrix{SymbolicT})
+    result = fill(Symbolics.COMMON_ZERO, size(mat, 1))
+    for i in axes(mat, 1)
         for j in axes(mat, 2)
             mij = mat[i, j]
-            if !_iszero(mij)
-                return mij
-            end
+            _iszero(mij) && continue
+            result[i] = mij
+            break
         end
-        0
     end
+    return result
 end
 
 """