Build function allocating on large expressions?

[lindnemi]

This problem comes up in large complex networks:

using ModelingToolkit, SparseArrays, BenchmarkTools, LinearAlgebra

N=100

const B = vcat(ones(N+1)', hcat(ones(N), -sparse(I,N,N)))


function f(dx, x)
    dx .=  B * x
end

@variables z[1:N+1]
@variables dz[1:N+1]
f(dz,z)
dz = simplify.(dz)
fastf = eval(ModelingToolkit.build_function(dz,z)[2])

x0 = randn(N+1)
dx0 = randn(N+1)

@btime f(dx0,x0)         # 388.388 ns (1 allocation: 896 bytes)
@btime fastf(dx0, x0)    # 2.107 μs (102 allocations: 2.38 KiB)

For N=10 everything works as expected: the MTK version is 4x faster and does not allocate.

[ChrisRackauckas]

This is an instance of JuliaLang/julia#34999 . @shashi @YingboMa do we have a cheap way to break down big + expressions?

[shashi]

We can walk the tree and cut it up in build_function.

using SymbolicUtils

function nest_long_plus(x, N)
    x  = to_symbolic(x)
    rule = @rule (+(~~x) => length(~~x) > N ? +(+((~~x)[1:N]...) + +((~~x)[N+1:end]...))
    Rewrtiters.Fixpoint(Rewriters.Chain([rule]))(x) |> to_mtk
end

[ChrisRackauckas]

That could work. Is there no easy way to do this at trace time, like by defining +(x...) to build the recursed plus?

[shashi]

I think in this case it did get built recursively but simplify flattened it, it's useful to flatten for simplification (e.g. cancellations are easy).

[ChrisRackauckas]

I see, so yeah then we should cut it up in build_function