Strange allocations during broadcasting

General Usage Performance
1

Can you please explain me why in the following code the second broadcasting results in 3 allocations?

using BenchmarkTools

u = rand(100)
k1 = similar(u)
k2 = similar(u)
k3 = similar(u)
k4 = similar(u)

@btime @. $u = $u + 1 * ($k1 + $k2 + $k3)
@btime @. $u = $u + 1 * ($k1 + $k2 + $k3 + $k4)
@btime @. $u = $u + ($k1 + $k2 + $k3 + $k4)

  100.137 ns (0 allocations: 0 bytes)
  143.974 ns (3 allocations: 80 bytes)
  127.223 ns (0 allocations: 0 bytes)

What is more confusing for me is that if I swap the last two lines, there will be no allocations.
Thank you.

1 Like
2

Up.

Meanwhile, in Julia 1.4 there are already 4 allocations in the above code.
Just to highlight the problem, the following code

using BenchmarkTools
# import DifferentialEquations

function foo(u, k1, k2, k3, k4)
    for i=1:1000
        @. u = u + 1.0 * (k1 + k2 + k3 + k4)
    end
    return nothing
end

function main()
    u = rand(100)
    k1 = zero(u)
    k2 = zero(u)
    k3 = zero(u)
    k4 = zero(u)

    @btime foo($u, $k1, $k2, $k3, $k4)
end

main()

results in

155.048 μs (4000 allocations: 93.75 KiB)

Surprisingly, if I import DifferentialEquations (uncomment the second line), the result will be

122.759 μs (0 allocations: 0 bytes)

That is, there are no more allocations.

The above code is typical for various Runge-Kutta methods. So far in my realizations of Runge-Kutta methods of order higher than 3 I inevitably see these broadcasting allocations. Can you, please, help me to understand what is going on. Especially, why the import of DifferentialEquations solves the problem.

3

It seems there is some heuristic that is getting hit because removing the +k4 makes it stop allocating.

4

Do you have any ideas where I can start to dig in?

5

Doing

julia> function foo(u, k1, k2, k3, k4)
           @. u = u + 1.0 * (k1 + k2 + k3 + k4)
           return nothing
       end
foo (generic function with 2 methods)

julia> @code_warntype optimize=true foo(u, k1, k2, k3, k4)

I see a

8 ─── %24  = invoke Base.Broadcast.combine_axes(1.0::Float64, %2::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1},Nothing,typeof(+),NTuple{4,Array{Float64,1}}})::Tuple{Union{Base.OneTo{Int64}, UnitRange{Int64}},Vararg{Any,N} where N}

which is not inferred to a concrete type. That’s the place I would probably start poking at.

6

Probably due to our broadcast fixes, i.e. https://github.com/SciML/DiffEqBase.jl/blob/master/src/diffeqfastbc.jl#L18-L19

1 Like
7

Do you have a PR that upstreams those?

8

There wasn’t before, but https://github.com/JuliaLang/julia/pull/35260 is it. Someone should probably double check that those changes are sufficient to get rid of these allocations.

9

It seems this is not the case:

using BenchmarkTools

import Base.Broadcast: Broadcasted
@inline combine_axes(A, B) = broadcast_shape(broadcast_axes(A), broadcast_axes(B))
@inline check_broadcast_axes(shp, A::Union{Number, Array, Broadcasted}) = check_broadcast_shape(shp, axes(A))

function foo(u, k1, k2, k3, k4)
    for i=1:1000
        @. u = u + 1.0 * (k1 + k2 + k3 + k4)
    end
    return nothing
end

function main()
    u = rand(100)
    k1 = zero(u)
    k2 = zero(u)
    k3 = zero(u)
    k4 = zero(u)

    @btime foo($u, $k1, $k2, $k3, $k4)
end

main()


177.050 μs (4000 allocations: 93.75 KiB)
10

You’re shadowing and not overloading:

@inline Base.Broadcast.combine_axes(A, B) = broadcast_shape(broadcast_axes(A), broadcast_axes(B))
@inline Base.Broadcast.check_broadcast_axes(shp, A::Union{Number, Array, Broadcasted}) = check_broadcast_shape(shp, axes(A))
1 Like
11

Thank you. Now it works:

using BenchmarkTools

@inline Base.Broadcast.combine_axes(A, B) = Base.Broadcast.broadcast_shape(Base.Broadcast.broadcast_axes(A), Base.Broadcast.broadcast_axes(B))

function foo(u, k1, k2, k3, k4)
    for i=1:1000
        @. u = u + 1.0 * (k1 + k2 + k3 + k4)
    end
    return nothing
end

function main()
    u = rand(100)
    k1 = zero(u)
    k2 = zero(u)
    k3 = zero(u)
    k4 = zero(u)

    @btime foo($u, $k1, $k2, $k3, $k4)
end

main()

136.010 μs (0 allocations: 0 bytes)

If it is not hard, can you give a brief explanation on what this patch do?

12

Splatting can sometimes cause a tuple to allocate if everything doesn’t inline or the world doesn’t align its chakras correctly. This makes it so that way a common case doesn’t ever splat, so boom now it can’t be a problem.

1 Like
13

OK. Thank you.

14

I find it’s easiest to just do @eval Base.Broadcast begin ... end to paste in these sorts of patches interactively without editing. You can also use Revise, but that can be trickier and isn’t as obvious.

2 Likes