refactor: improve precompile-friendliness

[AayushSabharwal]

This enables greatly improved precompilation of many common MTK workflows. On 1.11:

System constructor:

  0.000670 seconds (217 allocations: 10.641 KiB)
elapsed time (ns):  669875.0
gc time (ns):       0
bytes allocated:    10896
pool allocs:        217
non-pool GC allocs: 0
minor collections:  0
full collections:   0

This is for a particular signature. Different combinations of kwargs and different input types will cause compilation and we can't foreseeably compile every signature. This signature is System(::Vector{Equation}, ::Num, ::Vector{Num}, ::Vector{Num}; initial_conditions::Dict{SymbolicT, SymbolicT}, guesses::Dict{SymbolicT, SymbolicT}, name::Symbol). With some discussion, we can arrive at a common signature to precompile. It will likely involve unwrapping the Nums.

complete(sys):

  0.001191 seconds (1.08 k allocations: 2.554 MiB)
elapsed time (ns):  1.190625e6
gc time (ns):       0
bytes allocated:    2678088
pool allocs:        1077
non-pool GC allocs: 0
malloc() calls:     3
free() calls:       0
minor collections:  0
full collections:   0

This is for a system that does not have callbacks. The callback handling code is pretty big and needs a lot of work to precompile, so it is left for a future PR.

MTKTearing.TearingState(sys)

  0.014592 seconds (3.96 k allocations: 4.456 MiB, 92.30% compilation time)
elapsed time (ns):  1.4592042e7
gc time (ns):       0
bytes allocated:    4672936
pool allocs:        3955
non-pool GC allocs: 0
malloc() calls:     4
free() calls:       0
minor collections:  0
full collections:   0

The compilation here is all related to printing the equations to lexicographically sort them. The printing code in SU needs an overhaul to support precompilation (specifically, src/ordering.jl)

mtkcompile(sys):

  0.018629 seconds (20.74 k allocations: 932.062 KiB, 89.89% compilation time)
elapsed time (ns):  1.8628542e7
gc time (ns):       0
bytes allocated:    954432
pool allocs:        20727
non-pool GC allocs: 0
malloc() calls:     13
free() calls:       0
minor collections:  0
full collections:   0

This is the version of mtkcompile with ModelingToolkit.jl loaded. The printing from TearingState(sys) is already compiled here, so that is not part of this timing. The signatures reported by --trace-compile for this workload are

precompile(Tuple{typeof(Base.getindex), Array{Array{Int64, 1}, 1}, Int64})
precompile(Tuple{typeof(Base.setindex!), GenericMemory{:not_atomic, Number, Core.AddrSpace{Core}(0x00)}, Int64, Int64})
precompile(Tuple{typeof(Base.push!), Array{Symbolics.Equation, 1}, Symbolics.Equation})

And despite significant digging with Cthulhu I haven't been able to find where they are hit.

Script to reproduce these results:

import REPL
using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D
using Symbolics: SymbolicT

@variables x(t) = 1 y(t)
@parameters k f = 2.3
eqs = [D(x) ~ x * k + y, y ~ 2x + t]
dvs = Num[x, y]
ps = Num[k, f]
ics = Dict{SymbolicT, SymbolicT}()
ics[y] = 2.3
print(Val(:z))
println("System")
@timev System(eqs, t, dvs, ps; initial_conditions = ics, guesses = ics, name = :sys)
# print(Val(:a))
sys = System(eqs, t, dvs, ps; initial_conditions = ics, guesses = ics, name = :sys)
print(Val(:a))
println("Complete")
@timev complete(sys)
print(Val(:b))
println("TearingState")
@timev TearingState(sys)
print(Val(:c))
println("mtkcompile")
@timev mtkcompile(sys)

Which should be run as

julia --check-bounds=yes -O 3 --project=@. --trace-compile=bad.jl workload.jl

The print(Val(..)) are to have unique markers in the precompile statements in bad.jl so that it is possible to identify which sections of the script triggers the compilations. They typically show up as

precompile(Tuple{typeof(Base.print), Base.Val{:z}})
precompile(Tuple{typeof(Base.print), Base.TTY, Base.Val{:z}})

This requires JuliaSymbolics/SymbolicUtils.jl#824, JuliaSymbolics/Symbolics.jl#1711 and JuliaComputing/StateSelection.jl#11

[AayushSabharwal]

System and complete also precompile with just MTKBase loaded.