Unexpected memory allocation inside ODE function with EnsembleProblem

Specific Domains Modelling & Simulations
1

I am unsure if this is an issue or I just don’t understand how memory allocation works with ensemble problems.

I start Julia with --track-allocation=user and run the following small script

using DifferentialEquations
using Profile

function lorenz!(du,u,p,t)
    du[1] = 10.0*(u[2]-u[1])
    du[2] = u[1]*(28.0-u[3]) - u[2]
    du[3] = u[1]*u[2] - (8/3)*u[3]
end

u0 = [1.0;0.0;0.0]
tspan = (0.0, 100.0)
prob = ODEProblem(lorenz!, u0, tspan)

function prob_func(prob,i,repeat)
    remake(prob, tspan=(rand(), 100.0))
end

ensemble_prob = EnsembleProblem(prob,prob_func=prob_func)

sim = solve(ensemble_prob, Tsit5(), EnsembleThreads(), trajectories=100)
Profile.clear_malloc_data()
sim = solve(ensemble_prob, Tsit5(), EnsembleThreads(), trajectories=100)

The function Lorenz!(du, u, p, t) should not allocate any memory.

However, looking at the output of the .mem file, I get the following

        - using DifferentialEquations
        - using Profile
        - 
        - function lorenz!(du,u,p,t)
  1284544     du[1] = 10.0*(u[2]-u[1])
  1220016     du[2] = u[1]*(28.0-u[3]) - u[2]
  1456704     du[3] = u[1]*u[2] - (8/3)*u[3]
        - end
        - 
        - u0 = [1.0;0.0;0.0]
        - tspan = (0.0, 100.0)
        - prob = ODEProblem(lorenz!, u0, tspan)
        - 
        - function prob_func(prob,i,repeat)
        -     remake(prob, tspan=(rand(), 100.0))
        - end
        - 
        - ensemble_prob = EnsembleProblem(prob,prob_func=prob_func)
        - 
        - sim = solve(ensemble_prob, Tsit5(), EnsembleThreads(), trajectories=100)
        - Profile.clear_malloc_data()
        - sim = solve(ensemble_prob, Tsit5(), EnsembleThreads(), trajectories=100)

There are memory allocations inside the lorenz! function. If I change the ensemblealg to EnsembleSerial then I don’t observe any memory allocation in the .mem file.

2

Are you using Julia 1.4.1 or something else? Because when i run your script with track-allocation=user I get:

        - using DifferentialEquations
        - using Profile
        - 
        - function lorenz!(du,u,p,t)
        0     du[1] = 10.0*(u[2]-u[1])
        0     du[2] = u[1]*(28.0-u[3]) - u[2]
        0     du[3] = u[1]*u[2] - (8/3)*u[3]
        - end
        - 
        - u0 = [1.0;0.0;0.0]
        - tspan = (0.0, 100.0)
        - prob = ODEProblem(lorenz!, u0, tspan)
        - 
        - function prob_func(prob,i,repeat)
        -     remake(prob, tspan=(rand(), 100.0))
        - end
        - 
        - ensemble_prob = EnsembleProblem(prob,prob_func=prob_func)
        - 
        - sim = solve(ensemble_prob, Tsit5(), EnsembleThreads(), trajectories=100)
        - Profile.clear_malloc_data()
        - sim = solve(ensemble_prob, Tsit5(), EnsembleThreads(), trajectories=100)
        -
3

Yes, I’m using Julia 1.4.1 on macOS

4

Also, @pixel27 how many threads do you have? Threads.nthreads()

5

Ahh, that did it. I was only running with 1 thread. When I ran it with 4 i get the memory allocation count.

6

That’s definitely not good. I might need to call in the big guns here because that’s somewhat mysterious.

7

I’m at a loss…even looking at the generated code doesn’t really help. My fall back position would be that the memory tracker doesn’t do well with multi-threaded and that is causing some sort of false positive.