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:
- What exactly causes the additional allocations when the wrapper is not used after re-including the file?
- Is the solution using
ODEFunction{true}a universal and safe approach? - 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 therhs!functor sufficient to avoid these problems in such cases?
Thank you for your time and any insights you can share!