Extra allocations performing dot product in ODE callback

General Usage Performance
, ,
1

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?

2

Your parameter is type unstable:

julia> p |> typeof
Vector{Dict{Symbol, AbstractArray{ComplexF64}}} (alias for Array{Dict{Symbol, AbstractArray{Complex{Float64}}}, 1})

This is because you use two different arrays in it (sparse and a normal array)
It would be best to create your own struct for p to avoid the allocations that are part of the type instability

1 Like
3

Did you check what @code_warntype myaffect!(integrator) says?

My guess is that the parameter vector is the problem: The Dict contains different array types, so the output of y = integrator.p[1][:y] is not determined based on the type of integrator alone.

Using a NamedTuple instead of a Dict should help:

p = [(A = A, y = y)]

(the rest of the code would probably stay the same, but I didn’t check)

2 Likes
4

Thanks, it worked!

1 Like