safe_ldiv builds a fresh LinearCache per call: 2.6 kB and 39 allocations per RHS evaluation of an inline-linear-SCC model, 1 MB per 10 ms control tick

[baggepinnen]

What I am trying to do

Run a nonlinear MPC on hardware at a 10 ms tick with a multibody model as its prediction model. The controller (acados via MPCComponents, ERK with 2 stages over a 60-interval horizon) evaluates the model's right-hand side and its AD Jacobian 120 times each per tick, so anything the right-hand side allocates per call is multiplied by 120 and lands inside a hard real-time budget.

The model is a 4-state Furuta pendulum compiled with MultibodyComponents.multibody (i.e. DefaultReassembleAlgorithm(inline_linear_sccs = true)) and evaluated numerically through generate_control_function(io_sys, ins; simplify = false) with split = true.

What I measured

The right-hand side allocates on every call, in steady state — not a warm-up or cache-growth effect:

f! (generate_control_function's in-place rhs) : 2624 B / 39 allocations per call
J! (ForwardDiff of that f!, 5 directions)     : 6896 B per call

At 120 evaluations a tick that is 1.0 MB per 10 ms tick. Consequences on a 64-core Xeon, measured over 14 400 ticks of the closed loop: 9 % of tick time in garbage collection, and 46 ticks over the 10 ms period with the worst at 40–150 ms. The median tick is 0.97 ms, so the budget is not otherwise tight — the GC pauses are the only thing that misses it.

Profile.Allocs at sample_rate = 1.0 over 2000 f! calls attributes essentially all of it to LinearSolve setup, roughly two solves' worth per call:

   1168000 B    4000 allocs  :-1
    896000 B   16000 allocs  inner_ctors.jl:259 in array_literal
    512000 B    2000 allocs  LinearSolve/default.jl:15 in DefaultLinearSolverInit
    512000 B    2000 allocs  LinearSolve/common.jl:282 in LinearCache
    288000 B    8000 allocs  RuntimeGeneratedFunctions.jl:243 in macro expansion
    192000 B    4000 allocs  LinearSolve/common.jl:833 in #__init#27
    192000 B    4000 allocs  LinearSolve/factorization.jl:76 in _typed_copy
    160000 B    4000 allocs  MTKBase basic_transformations.jl:1177 in DiffCacheAllocatorAPIWrapper
    144000 B    6000 allocs  LinearSolve/appleaccelerate.jl:327 in init_cacheval
    128000 B    4000 allocs  LinearSolve/common.jl:585 in __init_u0_from_Ab
    128000 B    4000 allocs  LinearSolve/factorization.jl:778 in init_cacheval
    128000 B    4000 allocs  LinearSolve/default.jl:662 in macro expansion
     96000 B    2000 allocs  LinearSolve/factorization.jl:364 in _GenericLUFactorizationCache
     96000 B    2000 allocs  ArrayInterface.jl:509 in lu_instance

That is a fresh LinearCache, cacheval and LU instance per solve. The call site is ModelingToolkitTearing.safe_ldiv (src/reassemble.jl:715), which the generated code calls for the inlined linear SCC:

return CommonSolve.solve(LinearProblem(A, b)).u

Minimal reproducer, no model involved

using ModelingToolkit, LinearAlgebra, Printf
const MTKT = Base.require(Base.PkgId(Base.UUID("6bb917b9-1269-42b9-9f7c-b0dca72083ab"),
                                     "ModelingToolkitTearing"))
const safe_ldiv = MTKT.safe_ldiv

function bench(f, args...; n = 2000)
    f(args...)
    g0 = Base.gc_num()
    for _ in 1:n; f(args...); end
    d = Base.GC_Diff(Base.gc_num(), g0)
    (bytes = d.allocd / n, allocs = (d.malloc + d.poolalloc + d.bigalloc) / n)
end
reference(y, F, A, b) = (copyto!(F.factors, A); ldiv!(y, lu!(F.factors), b))

for m in (2, 3, 6)
    A = rand(m, m) + m * I; b = rand(m)
    s = bench(safe_ldiv, A, b)
    y = similar(b); F = lu(copy(A))
    r = bench(reference, y, F, A, b)
    @printf("m = %d:  safe_ldiv %5.0f B / %3.0f allocs   lu!+ldiv! %4.0f B / %2.0f allocs\n",
            m, s.bytes, s.allocs, r.bytes, r.allocs)
end

m = 2:  safe_ldiv  1344 B /  23 allocs   lu!+ldiv!   80 B /  2 allocs
m = 3:  safe_ldiv  1424 B /  23 allocs   lu!+ldiv!   80 B /  2 allocs
m = 6:  safe_ldiv  2032 B /  23 allocs   lu!+ldiv!  112 B /  2 allocs

23 allocations per call regardless of size, so it is the setup rather than the arithmetic. Wall time is 310–380 ns against 170–320 ns for the preallocated reference, i.e. the time is a secondary concern; the allocation is the problem.

The question

Is there an intended way to get an allocation-free (or at least allocation-bounded) right-hand side out of a model compiled with inlined linear SCCs? I could not find a knob for it from the outside — the A and b of the inlined SCC are built inside the generated function, so I have no handle on them, and nothing I pass to mtkcompile or generate_control_function seemed to change the call.

Two things I would like your read on rather than guessing:

1. Is the per-call LinearProblem construction something you consider fixable in place (the shape of A and b is fixed at compile time, so in principle a cache could live alongside the diffcaches already in MTKParameters), or is it load-bearing in a way that is not obvious from the outside — aliasing, thread safety, the dual-number widening of A under AD, or the rank-deficiency handling discussed in Inline linear SCCs: 296×296 dense runtime solve for an ~18-unknown core; non-rank-tolerant; symbolic shrinkage intractable — solve sparse+rank-tolerant instead #95?
2. Is the intended route for a workload like this simply not to inline the linear SCCs — take the semi-explicit index-1 DAE form instead, which my MPC layer does support — and is that expected to be allocation-free where this is not? I have not tried it, and would rather ask than build the whole controller a second way and find out that path allocates too.

Happy to test whatever you would like measured; the loop above is a two-minute experiment and the real closed loop is a ten-minute one.

Versions

ModelingToolkit 11.41.0, ModelingToolkitTearing 1.20.6, ModelingToolkitBase (nEoFQ), Julia 1.12.7, Linux x86_64.

Related: #95 (the size and rank tolerance of this same emitted runtime solve; this issue is about its per-call allocation, which is orthogonal — a 2×2 SCC allocates the same 23 times).

[ChrisRackauckas]

This shouldn't be necessary with the latest LinearSolve.jl? @DhairyaLGandhi . @AJ0070 can look into this?

[DhairyaLGandhi]

Yes. With a cache setup to reuse existing buffers, we shouldn't need to allocate a new one every time.

[ChrisRackauckas]

So is this then due to using an old version? Or is that caching part already done?

[DhairyaLGandhi]

The caching is not done, it is still constructing the LinearProblem at every call.

StateSelection.jl/lib/ModelingToolkitTearing/src/reassemble.jl

Line 715 in 0fdca6f

return CommonSolve.solve(LinearProblem(A, b)).u

That needs to get split up/ lowered as init + setting the new matrices/ vectors.

This would be a good piece for @AJ0070 indeed!

[AJ0070]

Assign me this, I'll look into it.