Skip to content

Unnecessary heap allocations in ODE function #4336

@1-Bart-1

Description

@1-Bart-1

Describe the bug 🐞

Unnecessary heap allocations in ODE function when array variables have a mix of structurally eliminated (algebraic) and ODE state components. Also when equations are fully scalarized.

Expected behavior

The RHS function should not allocate.

Minimal Reproducible Example 👇

using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D
using OrdinaryDiffEqBDF
using OrdinaryDiffEqCore
using Symbolics
using BenchmarkTools
using Printf

function build_and_bench(label; n_static, n_dynamic)
    N = n_static + n_dynamic

    @variables begin
        pos(t)[1:3, 1:N]
        vel(t)[1:3, 1:N]
    end

    defaults = [
        pos => 0.1 .* ones(3, N),
        vel => 0.1 .* ones(3, N),
    ]

    eqs = Equation[]
    for i in 1:N
        if i <= n_static
            # STATIC: algebraic constraints
            for j in 1:3
                push!(eqs, D(pos[j, i]) ~ 0.1)
                push!(eqs, vel[j, i] ~ 0.1)
            end
        else
            # DYNAMIC: spring-damper ODE
            for j in 1:3
                push!(eqs, D(pos[j, i]) ~ vel[j, i])
                push!(eqs, D(vel[j, i]) ~
                    -100.0 * pos[j, i] - 1.0 * vel[j, i])
            end
        end
    end

    n_eqs = length(eqs)
    println("\n  Building: $label")
    println("  N=$N ($n_static static, $n_dynamic dynamic)")
    println("  $n_eqs equations")
    for eq in eqs
        @show eq
    end

    @named sys = System(eqs, t)
    sys = mtkcompile(sys)
    prob = ODEProblem(sys, defaults, (0.0, 1.0);
        warn_initialize_determined=false)
    integ = OrdinaryDiffEqCore.init(prob, FBDF();
        save_on=false, save_everystep=false)

    f = integ.f
    u = copy(integ.u)
    pp = integ.p
    du = similar(u)
    f(du, u, pp, 0.0)  # compile

    result = @benchmark $f($du, $u, $pp, 0.0) samples=5

    @printf("  Result: %6d allocs, %8.3f μs\n",
            result.allocs,
            median(result.times) / 1e3)

    # Dump @code_warntype for cases with allocations
    if result.allocs > 0
        wt_file = "/tmp/claude/mwe_warntype_" *
            replace(label, " " => "_", ":" => "") * ".txt"
        open(wt_file, "w") do io
            code_warntype(io, f.f,
                Tuple{typeof(du), typeof(u),
                      typeof(pp), Float64})
        end
        println("  Warntype written to: $wt_file")
    end

    return result
end

# --- Run ---
println("="^60)
println("  MWE: array_literal mixed-type allocations")
println("="^60)

# Baseline: all dynamic (no constant zeros)
build_and_bench("A: all dynamic";
    n_static=0, n_dynamic=10)

# Mix: some static, some dynamic
build_and_bench("B: 6 static + 4 dynamic";
    n_static=6, n_dynamic=4)

# More static points
build_and_bench("C: 8 static + 3 dynamic";
    n_static=8, n_dynamic=3)

println("\n", "="^60)

Output

============================================================
  MWE: array_literal mixed-type allocations
============================================================

  Building: A: all dynamic
  N=10 (0 static, 10 dynamic)
  60 equations
  Result:      0 allocs,    0.018 μs

  Building: B: 6 static + 4 dynamic
  N=10 (6 static, 4 dynamic)
  60 equations
  Result:      0 allocs,    0.011 μs

  Building: C: 8 static + 3 dynamic
  N=11 (8 static, 3 dynamic)
  66 equations
  Result:     12 allocs,    0.453 μs
  Warntype written to: /tmp/claude/mwe_warntype_C_8_static_+_3_dynamic.txt

============================================================

Environment (please complete the following information):

  • Output of using Pkg; Pkg.status()
  [6e4b80f9] BenchmarkTools v1.6.3
  [961ee093] ModelingToolkit v11.11.0
  [6ad6398a] OrdinaryDiffEqBDF v1.18.0
  [bbf590c4] OrdinaryDiffEqCore v3.6.0
  [7e49a35a] RuntimeGeneratedFunctions v0.5.17
  [0c5d862f] Symbolics v7.15.1
  [9abbd945] Profile v1.11.0
  • Output of versioninfo()
Julia Version 1.12.5
Commit 5fe89b8ddc1 (2026-02-09 16:05 UTC)
Build Info:
  Official https://julialang.org release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 16 × 11th Gen Intel(R) Core(TM) i7-11850H @ 2.50GHz
  WORD_SIZE: 64
  LLVM: libLLVM-18.1.7 (ORCJIT, tigerlake)
  GC: Built with stock GC
Threads: 1 default, 1 interactive, 1 GC (on 16 virtual cores)

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions