DifferentialEquations.jl求解ODE异常内存分配问题咨询

问题现象 #

使用DifferentialEquations.jl求解标准常微分方程（ODE）、在指定时间点计算解值时，观测到异常大额内存分配：首次基准测试内存开销极高、垃圾回收（GC）耗时占比大，但仅在首次测试后重新编译积分函数时，异常分配才会消失，性能大幅提升。

复现代码 #

依赖引入 #

using BenchmarkTools
import DifferentialEquations: Vern9, ODEProblem, solve

二体动力学积分函数 #

@inbounds function dyn_twobody!(du::Vector{Float64}, u::Vector{Float64}, 
                                 GM::Float64, t::Float64)

    r = (u[1]*u[1] + u[2]*u[2] + u[3]*u[3])^1.5

    du[1] = u[4]
    du[2] = u[5]
    du[3] = u[6]
    du[4] = -GM*u[1]/r
    du[5] = -GM*u[2]/r
    du[6] = -GM*u[3]/r
    
    nothing 
end

指定时间点求解函数 #

function test(x::Vector{Float64}, tms::Vector{Float64})

    GM_MOON = 4.902780137400001e3;
    user_prob = ODEProblem(dyn_twobody!, x, [0., 86400], GM_MOON)

    user_sol = solve(user_prob, Vern9(), abstol=1e-10, reltol=1e-9, saveat=tms, save_everystep=false).u

    return user_sol 
end

测试调用代码 #

x =  [50500., 232., 32321., -1.2, 0.01, 0.3];
tms = collect(LinRange(0., 86400., 86400));

@benchmark test($x, $tms)

测试结果对比 #

• 首次测试结果：内存估算304.61 MiB，分配次数432505次，中位耗时151.338ms，GC耗时中位占比14.82%，最高达74.42%
• 重新编译dyn_twobody!后测试结果：内存估算17.22 MiB，分配次数173311次，中位耗时59.441ms，GC耗时中位占比0%，性能提升显著

问题诱因 #

该现象本质是Julia JIT按需编译机制下，首次运行类型信息不足导致非特化代码触发动态派发，具体逻辑：

1. Julia的方法特化为按需触发：第一次调用solve求解该二体问题时，编译器尚未拿到dyn_twobody!在积分循环内部所有调用点的完整类型信息。DifferentialEquations.jl内部会对用户传入的积分函数、参数、初值做多层封装，首次运行时封装层的类型未被完全推断为具体类型，编译器只能生成通用动态派发代码，积分每一步都要做类型检查、生成临时堆对象，这就是300余MiB异常分配的核心来源，高占比GC时间就是在回收这些临时对象。
2. 重新编译dyn_twobody!时，Julia会自动回溯该函数的所有调用点，用上一次运行时采集到的完整、精确的类型信息，重新生成全特化的静态派发代码。此时积分循环无多余动态类型检查，临时分配大幅减少，GC占比直接降至0，性能自然大幅提升。
3. 开销放大因素：代码中将ODEProblem构造放在test函数内部，每次调用test都会重新执行问题对象构造和类型适配，进一步拉长了首次运行的额外开销。

解决方案 #

• 提前预热编译：在正式运行计算或基准测试前，先执行一次极小步长的求解调用，触发全链路的方法特化。比如正式测试前加一行test(x, [0., 1.])，跑通一次短时长积分，让编译器完成所有所需特化方法的编译，后续正式运行就不会出现异常分配。
• 将问题构造移出热路径：不要每次调用求解函数都重新构造ODEProblem，把引力常数、问题对象的定义移到test函数外部，全局构造一次问题对象，test函数内部直接调用solve求解即可，减少重复的对象构造和类型推导开销。
• 优化状态存储类型：对于这类固定维度的小尺寸ODE系统（本次为6维二体问题），使用StaticArrays.jl提供的固定长度静态数组SVector{6, Float64}代替普通Vector{Float64}存储状态和导数，可彻底消除积分循环里的堆分配，性能还能进一步提升。
• 显式封装ODE函数：如需进一步固定类型，可使用ODEFunction(dyn_twobody!)提前封装积分函数，显式指定参数、状态的类型，帮助编译器在首次调用前就完成类型推断，避免动态派发。

内容的提问来源于stack exchange，提问作者7vn_