Unexpected allocations after re-including unchanged file in ODE simulation

Vladislav
Thu, 07 Aug 2025 12:15:34 +0000
General Usage

Hello!

I’m running a large number of parallel Monte Carlo simulations, so I’m working to minimize memory allocations in my code. While doing this, I encountered some behavior I don’t fully understand.

I’m using Julia Version 1.11.6 and OrdinaryDiffEq v6.101.0.

Consider the following code:

using Parameters, OrdinaryDiffEq;
const DEIntegrator = OrdinaryDiffEq.OrdinaryDiffEqCore.ODEIntegrator;

include("functions.jl");

solver_params = generate_solver_params();
integrator = get_integrator(solver_params);

@time output1 = solve_ode!(integrator, solver_params);
reinit!(integrator);
@time output1 = solve_ode!(integrator, solver_params);

include("functions.jl");

reinit!(integrator);
@time output2 = solve_ode!(integrator, solver_params);
reinit!(integrator);
@time output2 = solve_ode!(integrator, solver_params);

On the first run, the second @time shows zero allocations, as expected for compiled code. However, after re-including functions.jl —without making any changes to the file—subsequent @time calls show unexpected allocations after compilation (see the fourth @time):

0.023522 seconds (5.08 k allocations: 262.219 KiB, 19.72% compilation time)
0.020661 seconds
0.028323 seconds (85.08 k allocations: 4.528 MiB, 15.60% compilation time)
0.025087 seconds (80.00 k allocations: 4.272 MiB)

The content of functions.jl is:

function get_initial_condition(total_dim::Int)
    x0=ones(Float64, total_dim)
    return x0
end

function generate_solver_params()
    total_dim::Int = 1000               # Dimension of the system
    t_end::Float64 = 10.0               # End time
    n_tot::Int = 10_000                 # Number of time steps
    ts_all = range(start=0.0, stop=t_end, length=n_tot+1)
    dt::Float64 = ts_all[2]-ts_all[1]   # Time step size
    decay_rate::Float64 = 1.0           # Constant decay rate
    return (; total_dim, t_end, n_tot, ts_all, dt, decay_rate)
end

function rhs!(dx::Vector{Float64}, x::Vector{Float64}, solver_params::P, t::Float64) where {P<:NamedTuple}
    @unpack decay_rate = solver_params
    dx .= -decay_rate .* x
    return nothing
end

function get_integrator(solver_params::P) where {P<:NamedTuple}
    @unpack total_dim,t_end,n_tot,ts_all,dt,decay_rate = solver_params
    prob = ODEProblem(rhs!, get_initial_condition(total_dim), (0.0, t_end), solver_params)
    integrator = init(prob, RK4();
        adaptive=false,
        dt=dt,
        dense=false,
        save_everystep=false,
        save_on=false,
        save_start=false,
        save_end=false)
    return integrator
end

function solve_ode!(integrator::I, solver_params::P) where {I<:DEIntegrator, P<:NamedTuple}
    @unpack n_tot, dt = solver_params
    for _ in 1:n_tot
        step!(integrator, dt)
    end
    return integrator.u
end

I suspect this behavior is related to Julia’s method invalidation or the way functions are recompiled when re-evaluated. Re-including functions.jl redefines rhs!, which likely triggers recompilation or invalidates previously compiled method specializations, leading to unexpected allocations.

When I explicitly wrap rhs! using ODEFunction{true}, specifying that it is in-place, the issue disappears:

prob = ODEProblem(ODEFunction{true}(rhs!), get_initial_condition(total_dim), (0.0, t_end), solver_params)

After this change, allocations remain at zero even after re-including functions.jl :

  0.487818 seconds (3.92 M allocations: 219.215 MiB, 8.36% gc time, 96.45% compilation time)
  0.017568 seconds
  0.301034 seconds (2.28 M allocations: 126.925 MiB, 2.18% gc time, 94.06% compilation time)
  0.021564 seconds

However, compilation takes significantly longer and consumes considerably more memory.

My questions are the following:

1. What exactly causes the additional allocations when the wrapper is not used after re-including the file?

2. Is the solution using ODEFunction{true} a universal and safe approach?

3. In my actual application, rhs! is implemented as a functor with externally modified mutable fields. I’m concerned about potential memory issues or unintended behavior resulting from method invalidation or recompilation. Is explicitly wrapping the rhs! functor sufficient to avoid these problems in such cases?

Thank you for your time and any insights you can share!

3 posts - 3 participants