Stray allocation when stepping stochastic integrator

[JesseSlim]

Problem
I'm hunting down unwanted allocations in my scripts, and it seems that one of them actually comes from within StochasticDiffEq.jl. It is triggered when a I step a stochastic integrator (I use SOSRA2() at the moment, but I've seen it for other ones as well). I'm using in-place, non-allocating derivative functions, and indeed, when stepping an ODE integrator like Tsit5(), I see no allocations.

In the simulations that I want to run, I expect to take many many timesteps (~1e8), so a stray allocation that happens every step takes up a lot of memory and gc time. With my limited Julia experience (less than three weeks), I've tried my best to investigate the issue! My findings are included below.

I'm more than happy to try and fix the issue myself, but my Julia-fu is not up to spec yet, so any pointers are welcome :).

Example to reproduce
I have a script called test_alloc.jl with the following contents:

using DifferentialEquations

# (S)DE system
# dx/dt = v(t)
# dv/dt = -Omega^2 x(t) - Gamma v(t)
#
# x(t) --> u[1] 
# v(t) --> u[2] 
function f_mech!(du,u,p,t)
    du[1] = u[2]
    du[2] = -1.0*u[1] - 0.001*u[2]
end

# noise function
# acts only on v(t)
function g!(g_out,u,p,t)
    g_out[1] = 0.0
    g_out[2] = 1.0
end

# initial condition
r0 = [1.0, 0.0]

# solution time span
tspan = (0.0, 1e4)
dt = 2pi/100
num = (tspan[2]-tspan[1]) / dt
maxiters = 10*num

# ODE problem (without noise driving)
ode_problem = ODEProblem(
    f_mech!, r0, tspan,
    
    adaptive=false, dt=dt, dense=false, maxiters=maxiters,
    save_everystep=false, 
    saveat=[0.0, tspan[2]]
)
ode_integrator = init(ode_problem, Tsit5())

# SDE problem (with noise driving)
sde_problem = SDEProblem(
    f_mech!, g!, r0, tspan, 
    alg_hints=[:additive], save_noise=false,
    
    adaptive=false, dt=dt, dense=false, maxiters=maxiters,
    save_everystep=false, 
    saveat=[0.0, tspan[2]]
)
sde_integrator = init(sde_problem, SOSRA2())

;

I run the following commands in the REPL to show the issue:

julia> include("test_alloc.jl")

julia> @time step!(ode_integrator)
  2.999299 seconds (8.16 M allocations: 377.232 MiB, 13.43% gc time)

julia> @time step!(ode_integrator)
  0.000016 seconds

julia> @time step!(sde_integrator)
  4.846923 seconds (9.64 M allocations: 396.540 MiB, 3.31% gc time)

julia> @time step!(sde_integrator)
  0.000014 seconds (1 allocation: 16 bytes)

The problem also pops up when I compare solveing the ODE vs solveing the SDE:

julia> @time solve(ode_problem, Tsit5())
  0.032882 seconds (46 allocations: 4.812 KiB)
retcode: Success
Interpolation: 1st order linear
t: 2-element Array{Float64,1}:
     0.0
 10000.0
u: 2-element Array{Array{Float64,1},1}:
 [1.0, 0.0]
 [-0.006419164047044267, 0.002051192765260945]

julia> @time solve(sde_problem, SOSRA2())
  0.050645 seconds (159.23 k allocations: 2.435 MiB)
retcode: Success
Interpolation: 1st order linear
t: 2-element Array{Float64,1}:
     0.0
 10000.0
u: 2-element Array{Array{Float64,1},1}:
 [1.0, 0.0]
 [8.572151454186196, 14.606344527329984]

julia> num
159155

We see that the number of expected iterations (adaptive timestepping is off) is very close to the number of allocations in the SDE case.

My investigation
To pinpoint the problem, I've been tracking memory allocations with julia --track-allocation=user. After analysis with Coverage.analyze_malloc on my package directory, it seems that the only allocation is made in src/iterator_interface.jl:11:

StochasticDiffEq.jl/src/iterator_interface.jl

Line 11 in b6015ae

@inbounds perform_step!(integrator,integrator.cache)

The .mem file that is left after doing a single step on sde_integrator with allocation tracking on looks like this:

- @inline function step!(integrator::SDEIntegrator)
0   if integrator.opts.advance_to_tstop
0     @inbounds while integrator.tdir * integrator.t < first(integrator.opts.tstops)
0       loopheader!(integrator)
0       check_error!(integrator) != :Success && return
0       perform_step!(integrator,integrator.cache)
0       loopfooter!(integrator)
-     end
-   else
0     @inbounds loopheader!(integrator)
16     @inbounds perform_step!(integrator,integrator.cache)
0     @inbounds loopfooter!(integrator)
0     @inbounds while !integrator.accept_step
0       loopheader!(integrator)
0       check_error!(integrator) != :Success && return
0       perform_step!(integrator,integrator.cache)
0       loopfooter!(integrator)
-     end
-   end
0   @inbounds handle_tstop!(integrator)
- end

The .mem file that corresponds to src/perform_step/sra.jl, where perform_step! is located, contains only 0's in the allocation column.

Version information
Julia version information:

julia> versioninfo()
Julia Version 1.5.3
Commit 788b2c77c1 (2020-11-09 13:37 UTC)
Platform Info:
  OS: macOS (x86_64-apple-darwin18.7.0)
  CPU: Intel(R) Core(TM) i5-5287U CPU @ 2.90GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-9.0.1 (ORCJIT, broadwell)

StochasticDiffEq: version 6.32.1

[ChrisRackauckas]

The return from the perform_step! is unused, but it's not type-stable. I wonder if the compiler isn't optimizing it away. See if #401 helps.