refactor: several TTFX improvements

lib/ModelingToolkitBase/src/ModelingToolkitBase.jl

@@ -2,6 +2,11 @@
 $(DocStringExtensions.README)
 """
 module ModelingToolkitBase
+
+if isdefined(Base, :Experimental) && isdefined(Base.Experimental, Symbol("@max_methods"))
+    @eval Base.Experimental.@compiler_options max_methods=1
+end
+
 using PrecompileTools, Reexport
 @recompile_invalidations begin
     using StaticArrays
@@ -394,5 +399,5 @@ function __init__()
     return nothing
 end
 
-# include("precompile.jl")
+include("precompile.jl")
 end # module

lib/ModelingToolkitBase/src/problems/linearproblem.jl

@@ -119,21 +119,29 @@ function SciMLBase.LinearProblem{iip}(
     return maybe_codegen_scimlproblem(expression, LinearProblem{iip}, args; kwargs...)
 end
 
-function __make_fww(@nospecialize(f), retT, argsT)
+function __make_fww(@nospecialize(f), retT::DataType, argsT::DataType)
     fwt = (
         FunctionWrapper{retT, argsT}(f),
     )
     cs = FunctionWrappersWrappers.SingleCacheStorage()
     return FunctionWrappersWrapper{typeof(fwt), FunctionWrappersWrappers.AllowNonIsBits, typeof(cs)}(fwt, cs)
 end
 
+function __make_fww(@nospecialize(f), retT::NTuple{N, DataType}, argsT::NTuple{N, DataType}) where {N}
+    fwt = ntuple(i -> FunctionWrapper{retT[i], argsT[i]}(f), Val(N))
+    cs = FunctionWrappersWrappers.SingleCacheStorage()
+    return FunctionWrappersWrapper{typeof(fwt), FunctionWrappersWrappers.AllowNonIsBits, typeof(cs)}(fwt, cs)
+end
+
 Base.@nospecializeinfer function wrap_symbolic_linear_interface(
         symbolic_interface::SciMLBase.SymbolicLinearInterface, iip::Bool, A, b, p
     )
     # We'll never infer the type of these, so no point specializing on them
     @nospecialize symbolic_interface, A, b, p
     if iip
-        update_A! = __make_fww(SciMLBase.Void{Any}(symbolic_interface.update_A!), Nothing, typeof((A, p)))
+        elT = eltype(A)
+        pT = typeof(p)
+        update_A! = __make_fww(SciMLBase.Void{Any}(symbolic_interface.update_A!), (Nothing, Nothing, Nothing), (Tuple{Matrix{elT}, pT}, Tuple{Diagonal{elT, Vector{elT}}, pT}, Tuple{BandedMatrix{elT, Matrix{elT}, Base.OneTo{Int}}, pT}))
         update_b! = __make_fww(SciMLBase.Void{Any}(symbolic_interface.update_b!), Nothing, typeof((b, p)))
     else
         update_A! = __make_fww(symbolic_interface.update_A!, typeof(A), typeof((p,)))

lib/ModelingToolkitBase/src/systems/codegen.jl

@@ -1509,6 +1509,19 @@ function generate_update_A(
     )
     ps = reorder_parameters(sys)
 
+    regions = SU.RegionsT()
+    values = Symbolics.SArgsT()
+    N = size(A, 1)
+    push!(regions, SU.ShapeVecT((1:N, 1:size(A, 2))))
+    push!(values, SU.Fill(last(regions))(Symbolics.COMMON_ZERO))
+    for i in axes(A, 1), j in axes(A, 2)
+        v = A[i, j]
+        SU._iszero(v) && continue
+        push!(regions, SU.ShapeVecT((i:i, j:j)))
+        push!(values, Symbolics.SConst([v;;]))
+    end
+    A = SU.ArrayMaker{VartypeT}(regions, values; shape = first(regions))
+
     res = build_function_wrapper(
         sys, A, ps..., cachesyms...; p_start = 1, expression = Val{true},
         similarto = typeof(A), add_observed = false, kwargs...
@@ -1623,6 +1636,19 @@ function generate_update_b(
     )
     ps = reorder_parameters(sys)
 
+    regions = SU.RegionsT()
+    values = Symbolics.SArgsT()
+    N = length(b)
+    push!(regions, SU.ShapeVecT((1:N,)))
+    push!(values, SU.Fill(last(regions))(Symbolics.COMMON_ZERO))
+    for i in axes(b, 1)
+        v = b[i]
+        SU._iszero(v) && continue
+        push!(regions, SU.ShapeVecT((i:i,)))
+        push!(values, Symbolics.SConst([v]))
+    end
+    b = SU.ArrayMaker{VartypeT}(regions, values; shape = first(regions))
+
     res = build_function_wrapper(
         sys, b, ps..., cachesyms...; p_start = 1, expression = Val{true},
         similarto = typeof(b), add_observed = false, kwargs...

lib/ModelingToolkitBase/src/systems/problem_utils.jl

@@ -907,7 +907,7 @@ function CopyParamsByTemplate(srcsys::AbstractSystem, syms::AbstractArray{Symbol
 
     for i in eachindex(template)
         if template[i] isa Vector{SymbolicT}
-            template[i] = concrete_getu(srcsys, template[i]; wrap_as_any = true, kws...)
+            template[i] = concrete_getu(srcsys, Symbolics.SConst(template[i]); wrap_as_any = true, kws...)
             delete!(elem_types, Vector{SymbolicT})
             push!(elem_types, typeof(template[i]))
         end

src/ModelingToolkit.jl

@@ -2,6 +2,7 @@
 $(DocStringExtensions.README)
 """
 module ModelingToolkit
+
 using PrecompileTools, Reexport
 @recompile_invalidations begin
     using StaticArrays
@@ -68,6 +69,7 @@ import Symbolics: rename, get_variables!, _solve, hessian_sparsity,
     jacobian, hessian, derivative, sparsejacobian, sparsehessian,
     scalarize, hasderiv
 import ModelingToolkitBase as MTKBase
+import SimpleNonlinearSolve
 
 import DiffEqBase: @add_kwonly
 @reexport using Symbolics
@@ -167,7 +169,7 @@ function FMIComponent end
 @public linearize_symbolic, reorder_unknowns
 @public similarity_transform
 
-include(pkgdir(ModelingToolkitBase, "src", "precompile.jl"))
+include("precompile.jl")
 
 function __init__()
     SU.hashcons(StructuralTransformations.NOTHING_EQ.lhs, true)

src/precompile.jl

@@ -0,0 +1,41 @@
+include(pkgdir(ModelingToolkitBase, "src", "precompile.jl"))
+
+PrecompileTools.@compile_workload begin
+    t = ModelingToolkitBase.t_nounits
+    D = ModelingToolkitBase.D_nounits
+
+    function f!(du, u, p)
+        du[1] = cos(u[2]) - u[1]
+        du[2] = sin(u[1] + u[2]) + u[2]
+        du[3] = 2u[4] + u[3] + 1.0
+        du[4] = u[5]^2 + u[4]
+        du[5] = u[3]^2 + u[5]
+        du[6] = u[1] + u[2] + u[3] + u[4] + u[5] + 2.0u[6] + 2.5u[7] + 1.5u[8]
+        du[7] = u[1] + u[2] + u[3] + 2.0u[4] + u[5] + 4.0u[6] - 1.5u[7] + 1.5u[8]
+        du[8] = u[1] + 2.0u[2] + 3.0u[3] + 5.0u[4] + 6.0u[5] + u[6] - u[7] - u[8]
+
+        du[9] = u[4] + u[5] + u[6] + u[7] + u[8] + 5.3u[9] + 5.8u[10] + 4.8u[11]
+        du[10] = u[4] + u[5] + u[6] + 5.3u[7] + u[8] + 7.3u[9] - 4.8u[10] + 4.8u[11]
+        du[11] = u[4] + 5.3u[5] + 6.3u[6] + 8.3u[7] + 9.3u[8] + u[9] - u[10] - u[11]
+
+        du[12] = u[7] + u[8] + u[9] + u[10] + u[11] + 8.6u[12] + 8.11u[13] + 7.11u[14]
+        du[13] = u[7] + u[8] + u[9] + 8.6u[10] + u[11] + 10.6u[12] - 7.11u[13] + 7.11u[14]
+        du[14] = u[7] + 8.6u[8] + 9.6u[9] + 11.6u[10] + 12.6u[11] + u[12] - u[13] - u[14]
+
+        du[15] = u[10] + u[11] + u[12] + u[13] + u[14] + 11.9u[15] + 11.14u[16] + 10.14u[17]
+        du[16] = u[10] + u[11] + u[12] + 11.9u[13] + u[14] + 13.9u[15] - 10.14u[16] + 10.14u[17]
+        du[17] = u[10] + 11.9u[11] + 12.9u[12] + 14.9u[13] + 15.9u[14] + u[15] - u[16] - u[17]
+
+        du[18] = u[13] + u[14] + u[15] + u[16] + u[17] + 14.12u[18] + 14.17u[19] + 13.17u[20]
+        du[19] = u[13] + u[14] + u[15] + 14.12u[16] + u[17] + 16.12u[18] - 13.17u[19] + 13.17u[20]
+        du[20] = u[13] + 14.12u[14] + 15.12u[15] + 17.12u[16] + 18.12u[17] + u[18] - u[19] - u[20]
+    end
+    @variables u[1:20] = rand(20) [irreducible = true]
+    eqs = Num[0 for _ in 1:20]
+    f!(eqs, u, nothing)
+    eqs = 0 .~ eqs
+    @mtkcompile model = System(eqs)
+    sccprob = SCCNonlinearProblem(model, nothing)
+    solve(sccprob, SimpleNonlinearSolve.SimpleNewtonRaphson())
+end
+

src/problems/sccnonlinearproblem.jl

@@ -51,7 +51,25 @@ function CacheWriter(
     )
     fn = eval_or_rgf(fn; eval_expression, eval_module)
     fn = GeneratedFunctionWrapper{(3, 3, is_split(sys))}(fn, nothing)
-    return CacheWriter(fn)
+    return CacheWriter{Any}(fn)
+end
+
+# This phrasing allows us to precompile the calls
+@noinline function __explicitfun_copy_states_helper(buffer, sols)
+    offset = 0
+    for sol in sols
+        u = sol.u
+        copyto!(buffer, CartesianIndices(((offset + 1):(offset + length(u)),)), u, CartesianIndices((1:length(u),)))
+        offset += length(u)
+    end
+    return nothing
+end
+
+function __explicitfun_copy_states(@nospecialize(p), sols)
+    buffer_idx = findfirst(Base.Fix2(==, eltype(sols[1].u)) ∘ eltype, p.caches)::Int
+    buffer = p.caches[buffer_idx]::Vector{eltype(sols[1].u)}
+    __explicitfun_copy_states_helper(buffer, sols)
+    return nothing
 end
 
 """
@@ -286,12 +304,41 @@ function build_caches!(sys::System, decomposition::SCCDecomposition)
     banned_vars = Set{SymbolicT}()
     state = Dict{SymbolicT, SymbolicT}()
     ir = get_irstructure(sys)
+    previous_scc_unknowns = SymbolicT[]
     for i in eachindex(decomposition.subsystems)
+        subsys = decomposition.subsystems[i]
+        if decomposition.islinear[i]
+            store_to_mutable_cache!(subsys, CachedLinearAb, nothing)
+            # For linear systems, there is no point in a complicated explicitfun. Just copy
+            # all previous states to the cache buffer.
+
+            cachevars = SCCCacheVarsExprsElT()
+            cacheexprs = SCCCacheVarsExprsElT()
+            push!(decomposition.scc_cachevars, cachevars)
+            push!(decomposition.scc_cacheexprs, cacheexprs)
+
+            if isempty(previous_scc_unknowns)
+                append!(previous_scc_unknowns, unknowns(subsys))
+                continue
+            end
+            T = Real
+            buffer = cachevars[T] = cacheexprs[T] = SymbolicT[]
+            append!(buffer, previous_scc_unknowns)
+            idx = findfirst(isequal(T), decomposition.cachetypes)
+            if idx === nothing
+                push!(decomposition.cachetypes, T)
+                push!(decomposition.cachesizes, 0)
+                idx = lastindex(decomposition.cachetypes)
+            end
+            decomposition.cachesizes[idx] = max(decomposition.cachesizes[idx], length(buffer))
+            append!(previous_scc_unknowns, unknowns(subsys))
+            continue
+        end
         empty!(banned_vars)
         empty!(state)
 
-        subsys = decomposition.subsystems[i]
         union!(banned_vars, unknowns(subsys))
+        append!(previous_scc_unknowns, unknowns(subsys))
         for u in unknowns(subsys)
             push!(banned_vars, split_indexed_var(u)[1])
         end
@@ -496,7 +543,7 @@ function SciMLBase.SCCNonlinearProblem{iip}(
     )
     op = calculate_op_from_u0_p(sys, u0, p)
 
-    explicitfuns = []
+    explicitfuns = Union{Returns{Nothing}, CacheWriter{Any}, typeof(__explicitfun_copy_states)}[]
     nlfuns = []
     decomposition = SCCDecomposition(sys, var_sccs, eq_sccs; combine_sccs)
 
@@ -513,15 +560,15 @@ function SciMLBase.SCCNonlinearProblem{iip}(
         subsys = decomposition.subsystems[i]
         if isempty(cachevars)
             push!(explicitfuns, Returns(nothing))
+        elseif decomposition.islinear[i]
+            push!(explicitfuns, __explicitfun_copy_states)
         else
             solsyms = view.((dvs,), view(decomposition.var_sccs, 1:(i - 1)))
             push!(
                 explicitfuns,
-                SciMLBase.Void{Any}(
-                    CacheWriter(
-                        sys, decomposition.cachetypes, cacheexprs, solsyms;
-                        eval_expression, eval_module, cse
-                    )
+                CacheWriter(
+                    sys, decomposition.cachetypes, cacheexprs, solsyms;
+                    eval_expression, eval_module, cse
                 )
             )
         end
@@ -640,7 +687,7 @@ function SciMLBase.SCCNonlinearProblem{iip}(
     if length(subprobs) <= 5
         return SCCNonlinearProblem(Tuple(subprobs), Tuple(explicitfuns), p, true; sys)
     else
-        return SCCNonlinearProblem(subprobs, explicitfuns, p, true; sys)
+        return SCCNonlinearProblem(subprobs, SciMLBase.Void{Any}.(explicitfuns), p, true; sys)
     end
 end