Variable access feature is very slow

[stelmo]

I am using ModelingToolkit to simulate a system and noticed that there seems to be a serious performance cost (~3 orders of magnitude) associated with using the nice observed variables functionality of ModelingToolkit. Here is a minimal working example of what I mean.

Problem setup:

using DifferentialEquations, ModelingToolkit, BenchmarkTools

@parameters t σ ρ β
@variables x(t) y(t) z(t) v_x(t) v_y(t) v_z(t)
D = Differential(t) 

eqs = [
    v_x ~ σ *(y - x),
    v_y ~ x * (ρ - z) - y,
    v_z ~ x*y - β * z,
    D(x) ~ v_x,
    D(y) ~ v_y,
    D(z) ~ v_z
]

sys = ODESystem(eqs)
sys = structural_simplify(sys) # get rid of v_x, v_y, v_z

ps = Dict(σ => 10.0, ρ => 28.0, β => 2.66)
u0 = Dict(x => 1.0, y => 0.5, z => 0.8)

tspan = (0.0,1.0)
prob = ODEProblem(sys, u0, tspan, ps)
sol = solve(prob, saveat=[0.1]) # more saveat locations makes everything worse

Then when using the observed variables:

f(sol) = sol[x][end] * sol[v_x][end] + sol[y][end] * sol[v_y][end] + sol[z][end] * sol[v_z][end]

julia> @benchmark f(sol)
BechmarkTools.Trial: 10000 samples with 1 evaluations.
 Range (min … max):  38.800 μs …  10.927 ms  ┊ GC (min … max): 0.00% … 98.40%
 Time  (median):     46.200 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   54.137 μs ± 149.940 μs  ┊ GC (mean ± σ):  3.79% ±  1.40%

  ▆▆██▇▇▆▆▅▄▄▃▃▄▃▂▂▁        ▁▁▁                                ▂
  ████████████████████▇▇█████████▇▇▇▇▇▇▇█▇▇▇▇▅▇▇▆▆▆▅▆▆▅▅▅▅▄▅▃▆ █
  38.8 μs       Histogram: log(frequency) by time       133 μs <

 Memory estimate: 17.22 KiB, allocs estimate: 235.

Versus when not using them:

vec = [sol[x][end], sol[y][end], sol[z][end], sol[v_x][end], sol[v_y][end], sol[v_z][end]]
g(vec) = vec[1] * vec[4] + vec[2] * vec[5] + vec[3] * vec[6]

julia> @benchmark g(vec)
BechmarkTools.Trial: 10000 samples with 979 evaluations.
 Range (min … max):  27.681 ns …  12.044 μs  ┊ GC (min … max): 0.00% … 99.43%
 Time  (median):     39.122 ns               ┊ GC (median):    0.00%
 Time  (mean ± σ):   48.855 ns ± 123.743 ns  ┊ GC (mean ± σ):  2.45% ±  0.99%

  ▇▆▃▅██▅▄▃▂▃▅▄▃▂▁▁▂▂▂▂▂▁    ▁▁▁                               ▂
  █████████████████████████████████▇▇▇█▇▇▇▇▇▆▆▆▆▆▅▆▅▆▆▄▂▄▅▆▄▆▅ █
  27.7 ns       Histogram: log(frequency) by time       164 ns <

 Memory estimate: 16 bytes, allocs estimate: 1.

It seems that all the overhead is coming from using the variable access features, e.g.:

julia> @benchmark begin
           sol[x][end] + sol[y][end] + sol[z][end]
           end
BechmarkTools.Trial: 10000 samples with 1 evaluations.
 Range (min … max):  12.700 μs … 299.100 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     14.100 μs               ┊ GC (median):    0.00%        
 Time  (mean ± σ):   17.402 μs ±   8.769 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

  ▇█▆▃▂▁ ▁▆▆▄▃▂▁   ▄▄▂ ▁                                       ▂
  ███████████████▆████████▆▇█▇▇▇▅▆▄▆▆▇▆▅▅▆▆▆▆▇▅▆▅▆▅▆▆▆▅▅▆▅▄▅▅▅ █
  12.7 μs       Histogram: log(frequency) by time      47.2 μs <

 Memory estimate: 3.16 KiB, allocs estimate: 40.

and

julia> @benchmark begin
           sol[v_x][end] + sol[v_y][end] + sol[v_z][end]
       end
BechmarkTools.Trial: 10000 samples with 1 evaluations.
 Range (min … max):  63.500 μs …  20.670 ms  ┊ GC (min … max): 0.00% … 98.80%
 Time  (median):     74.400 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   88.097 μs ± 208.651 μs  ┊ GC (mean ± σ):  2.32% ±  0.99%

  ▇█▆▇▆▆▅▅▄▄▅▄▃▃▂▁▁▁▁ ▁ ▁   ▁ ▁                                ▂
  ██████████████████████████████▇▇██▇█▆▇███▇██▇██▇▇▇▇▆▆▆▅▆▄▆▅▆ █
  63.5 μs       Histogram: log(frequency) by time       216 μs <

 Memory estimate: 10.84 KiB, allocs estimate: 187.

Is there a way to fix this? This becomes especially problematic when trying to evaluate the results of ensemble simulations etc.

[ChrisRackauckas]

@sharanry or @shashi could you take a look? Some of this looks like it will be fixed when issymbollike is fixed.

[sharanry]

Yes, it seems like the current issymbollike hack is the culprit:

using Pkg; Pkg.activate(@__DIR__)
using RecursiveArrayTools, ModelingToolkit, BenchmarkTools

@parameters t
@variables α(t)

display(@benchmark RecursiveArrayTools.issymbollike(α))

 Range (min … max):  5.354 μs … 377.278 μs  ┊ GC (min … max): 0.00% … 97.94%
 Time  (median):     5.472 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   5.568 μs ±   5.253 μs  ┊ GC (mean ± σ):  1.33% ±  1.39%

           █   ▃  ▁                                            
  ▂▂▁▂▂▂▆▆███▇▇█▆▆█▄▄▆▃▃▄▃▃▂▃▂▂▃▃▂▃▂▂▃▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂ ▃
  5.35 μs         Histogram: frequency by time        5.89 μs <

 Memory estimate: 1.53 KiB, allocs estimate: 35.

using Symbolics
display(@benchmark Symbolics.SciMLBase.issymbollike(α))

 Range (min … max):  4.500 ns … 18.459 ns  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     4.583 ns              ┊ GC (median):    0.00%
 Time  (mean ± σ):   4.590 ns ±  0.223 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

  ▂     ▇     █      ▃     ▂      ▂     ▃      ▂             ▁
  █▁▁▁▁▁█▁▁▁▁▁█▁▁▁▁▁▁█▁▁▁▁▁█▁▁▁▁▁▁█▁▁▁▁▁█▁▁▁▁▁▁█▁▁▁▁▁█▁▁▁▁▁▃ █
  4.5 ns       Histogram: log(frequency) by time     4.88 ns <

 Memory estimate: 0 bytes, allocs estimate: 0.%   

[ChrisRackauckas]

Ouch. We could mark it Base.@pure...

[ChrisRackauckas]

Oh god, if Jameson ever saw what we were doing, and then put @pure on it, he would just close our whole organization. We would be gone.

[sharanry]

Ouch. We could mark it Base.@pure...

Doesn't seem to help.

Base.@pure __parameterless_type(T) = Base.typename(T).wrapper
Base.@pure function issymbollike_new(x)
    x == Symbol ||
    # x isa AllObserved ||
    Symbol(__parameterless_type(x)) == :Operation || Symbol(__parameterless_type(x)) == Symbol("Symbolics.Operation") ||
    Symbol(__parameterless_type(x)) == :Variable || Symbol(__parameterless_type(x)) == Symbol("Symbolics.Variable") ||
    Symbol(__parameterless_type(x)) == :Sym || Symbol(__parameterless_type(x)) == Symbol("Symbolics.Sym") ||
    Symbol(__parameterless_type(x)) == :Num || Symbol(__parameterless_type(x)) == Symbol("Symbolics.Num") ||
    Symbol(__parameterless_type(x)) == :Term || Symbol(__parameterless_type(typeof(x))) == Symbol("Symbolics.Term")       
end
t = typeof(α)
display(@benchmark issymbollike_new(t))

BechmarkTools.Trial: 10000 samples with 6 evaluations.
 Range (min … max):  5.500 μs … 414.479 μs  ┊ GC (min … max): 0.00% … 98.14%
 Time  (median):     5.736 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   5.890 μs ±   5.655 μs  ┊ GC (mean ± σ):  1.35% ±  1.39%

         █▄█▂                                                  
  ▁▁▂▂▃▃█████▇▇▄▄▂▂▂▂▂▂▂▃▂▂▂▂▂▁▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ ▂
  5.5 μs          Histogram: frequency by time         6.9 μs <

 Memory estimate: 1.53 KiB, allocs estimate: 35.%  

[ChrisRackauckas]

String construction allocates so that must disable it.

[ChrisRackauckas]

Construct the symbols in const outside of the function.

[sharanry]

Construct the symbols in const outside of the function.

@ChrisRackauckas
I am guessing that would still not be @pure as it depends on something other than the inputs.

Same result as before:

Base.@pure __parameterless_type(T) = Base.typename(T).wrapper
const _Symbolics_Operation = Symbol("Symbolics.Operation")
const _Symbolics_Variable = Symbol("Symbolics.Variable")
const _Symbolics_Sym = Symbol("Symbolics.Sym")
const _Symbolics_Num = Symbol("Symbolics.Num")
const _Symbolics_Term = Symbol("Symbolics.Term")
Base.@pure function issymbollike_new(x)
    return typeof(x) == Symbol ||
    # typeof(x) == AllObserved ||
    Symbol(__parameterless_type(typeof(x))) == :Operation || Symbol(__parameterless_type(typeof(x))) == _Symbolics_Operation  ||
    Symbol(__parameterless_type(typeof(x))) == :Variable || Symbol(__parameterless_type(typeof(x))) == _Symbolics_Variable  ||
    Symbol(__parameterless_type(typeof(x))) == :Sym || Symbol(__parameterless_type(typeof(x))) == _Symbolics_Sym  ||
    Symbol(__parameterless_type(typeof(x))) == :Num || Symbol(__parameterless_type(typeof(x))) == _Symbolics_Num  ||
    Symbol(__parameterless_type(typeof(x))) == :Term || Symbol(__parameterless_type(typeof(x))) == _Symbolics_Term        
end

BechmarkTools.Trial: 10000 samples with 6 evaluations.
 Range (min … max):  5.319 μs … 389.854 μs  ┊ GC (min … max): 0.00% … 98.06%
 Time  (median):     5.431 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   5.519 μs ±   5.403 μs  ┊ GC (mean ± σ):  1.38% ±  1.39%

    ▂▁▂█  ▂ ▂  █                                               
  ▃▃█████▇█▆█▇██▇▆█▅█▅▅▇▄▄▄▃▄▃▂▃▂▂▂▂▂▂▂▂▂▂▂▂▂▁▂▂▂▂▂▂▂▁▂▁▁▂▂▂▂ ▃
  5.32 μs         Histogram: frequency by time        5.88 μs <

 Memory estimate: 1.53 KiB, allocs estimate: 35.% 

[ChrisRackauckas]

===? Assuming const this only depends on the type of the input.

[sharanry]

===?

Didn't help either 😕. Same result as before.

[baggepinnen]

With a function barrier, the performance can be increased some 10x

Base.@pure __parameterless_type(T) = Base.typename(T).wrapper
const _Symbolics_Operation = Symbol("Symbolics.Operation")
const _Symbolics_Variable = Symbol("Symbolics.Variable")
const _Symbolics_Sym = Symbol("Symbolics.Sym")
const _Symbolics_Num = Symbol("Symbolics.Num")
const _Symbolics_Term = Symbol("Symbolics.Term")
Base.@pure function issymbollike_new(x)
    return typeof(x) == Symbol ||
    # typeof(x) == AllObserved ||
    Symbol(__parameterless_type(typeof(x))) == :Operation || Symbol(__parameterless_type(typeof(x))) == _Symbolics_Operation  ||
    Symbol(__parameterless_type(typeof(x))) == :Variable || Symbol(__parameterless_type(typeof(x))) == _Symbolics_Variable  ||
    Symbol(__parameterless_type(typeof(x))) == :Sym || Symbol(__parameterless_type(typeof(x))) == _Symbolics_Sym  ||
    Symbol(__parameterless_type(typeof(x))) == :Num || Symbol(__parameterless_type(typeof(x))) == _Symbolics_Num  ||
    Symbol(__parameterless_type(typeof(x))) == :Term || Symbol(__parameterless_type(typeof(x))) == _Symbolics_Term        
end

@btime issymbollike_new($Ref(1)[]) #  6.744 μs (50 allocations: 2.19 KiB)

issymbollike_new2(x::Symbol) = true
function issymbollike_new2(x)
    s = Symbol(__parameterless_type(typeof(x)))
    _inner(s)
end
function _inner(s)
    s === :Operation || s === _Symbolics_Operation  ||
    s === :Variable || s === _Symbolics_Variable  ||
    s === :Sym || s === _Symbolics_Sym  ||
    s === :Num || s === _Symbolics_Num  ||
    s === :Term || s === _Symbolics_Term        
end
@btime issymbollike_new2($Ref(1)[]) # 666.180 ns (5 allocations: 224 bytes)

---

Edit: Actually, just saving the result of Symbol(__parameterless_type(typeof(x))) in a variable and using === gets you there

issymbollike_new3(x::Symbol) = true
function issymbollike_new2(x)
    s = Symbol(__parameterless_type(typeof(x)))
    s === :Operation || s === _Symbolics_Operation  ||
    s === :Variable || s === _Symbolics_Variable  ||
    s === :Sym || s === _Symbolics_Sym  ||
    s === :Num || s === _Symbolics_Num  ||
    s === :Term || s === _Symbolics_Term   
end
@btime issymbollike_new2($Ref(1)[])

[sharanry]

@baggepinnen This is great!

Old:

BechmarkTools.Trial: 10000 samples with 6 evaluations.
 Range (min … max):  5.215 μs … 467.285 μs  ┊ GC (min … max): 0.00% … 98.36%
 Time  (median):     5.382 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   5.521 μs ±   6.023 μs  ┊ GC (mean ± σ):  1.53% ±  1.39%

  ▁▃▅██▇▅▄▂▁                                                  ▂
  █████████████▇██▇▇▇▆▄▆▃▆▅▃▃▄▁▅▄▄▃▄▁▃▃▁▅▁▃▃▃▁▃▄▃▄▁▄▄▅▄▁▄▃▃▄▃ █
  5.22 μs      Histogram: log(frequency) by time      7.47 μs <

 Memory estimate: 1.53 KiB, allocs estimate: 35.

New:

BechmarkTools.Trial: 10000 samples with 109 evaluations.
 Range (min … max):  769.495 ns …  20.644 μs  ┊ GC (min … max): 0.00% … 95.35%
 Time  (median):     794.725 ns               ┊ GC (median):    0.00%
 Time  (mean ± σ):   805.460 ns ± 382.583 ns  ┊ GC (mean ± σ):  0.94% ±  1.90%

        ▂▄▅██▁                                                   
  ▂▃▃▄▄▅██████▄▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ ▂
  769 ns           Histogram: frequency by time          939 ns <

 Memory estimate: 224 bytes, allocs estimate: 5.% 

[YingboMa]

Those are a bit more complicated than it should be.

julia> using SciMLBase, ModelingToolkit, BenchmarkTools

julia> @generated function issymbollike(x)
           x <: Symbol && return true
           ss = ["Operation", "Variable", "Sym", "Num", "Term"]
           s = string(Symbol(SciMLBase.__parameterless_type(x)))
           any(x->occursin(x, s), ss)
       end
issymbollike (generic function with 1 method)

julia> @btime issymbollike(Ref(1)[])
  0.029 ns (0 allocations: 0 bytes)
false

julia> @variables x;

julia> @btime issymbollike($x)
  0.029 ns (0 allocations: 0 bytes)
true

[YingboMa]

We can also optimize sym_to_index in SciMLBase.