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

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions