Extra allocations performing dot product in ODE callback

albertomercurio
Thu, 22 Jun 2023 06:03:57 +0000
Performance

Hello,

I’m experiencing some extra memory allocations when calculating the dot product inside a ODE callback. Here I write a minimal example

A = sprand(ComplexF64, 100, 100, 0.05)
u0 = rand(ComplexF64, 100)
y = rand(ComplexF64, 100)
p = [Dict(:A => A, :y => y)]

dudt!(du, u, p, t) = mul!(du, p[1][:A], u)

Now I check that dudt! is inplace

du = similar(u0)
@btime dudt!($du, $u0, $p, 0.0);
533.511 ns (0 allocations: 0 bytes)

Now I check the ODEproblem solution without any callback

prob = ODEProblem(dudt!, u0, (0.0, 10.0), p, save_on=false)
@btime solve(prob, Tsit5(), saveat=[10.0]);

193.700 μs (48 allocations: 32.59 KiB)

And now I introduce a simple callback without doing anything

function myaffect!(integrator)
    internal_params = integrator.p[1]
    y = internal_params[:y]

    # dot(y, integrator.u)
end

tlist = range(0, 10, 1000)
cb = PresetTimeCallback(tlist, myaffect!)
prob = ODEProblem(dudt!, u0, (0.0, 10.0), p, save_on=false, callback=cb)
@btime solve(prob, Tsit5(), saveat=[10.0]);

5.394 ms (58 allocations: 75.09 KiB)

Where I think that the extra 10 allocations are for creating the callback, which should be fine.
But if I uncomment the dot function inside the callback function:

5.685 ms (1058 allocations: 106.34 KiB)

Which seems that is allocates the memory every time this function is called.

By benchmarking the dot function it returns zero allocations of course

integrator = init(prob, Tsit5())
@btime dot($y, $(integrator.u))

39.152 ns (0 allocations: 0 bytes)
53.40957206916622 - 1.5882023409748207im

while benchmarking the callback function I get one memory allocation

@btime myaffect!($integrator)

 56.135 ns (1 allocation: 32 bytes)
53.40957206916622 - 1.5882023409748207im

I tried to calculate dot(integrator.u, integrator.u) and there is no extra memory allocation, so it comes from the dot product involving y.

how can I perform such operation without allocating extra memory?

4 posts - 3 participants