slow down when @. assignment expression has more than 3 terms

[daviehh]

Not sure if it's a julia issue or StaticArrays issue: Assignment expressions such as @. r = a + 5 * b + 3 * c - d slows down dramatically when the right-hand-side has 4 terms [allocate and may be dynamically dispatched according to profile]:

n = 5

a = rand(n, n)
b = rand(n, n)
c = rand(n, n)
d = rand(n, n)
r = rand(n, n)

as = MMatrix{n, n, Float64}(a)
bs = MMatrix{n, n, Float64}(b)
cs = MMatrix{n, n, Float64}(c)
ds = MMatrix{n, n, Float64}(d)
rs = MMatrix{n, n, Float64}(r)

function g!(r,a,b,c,d)
    @. r = a + 5 * b + 3 * c - d
end

with standard array:

@btime g!($r, $a, $b, $c, $d)

62.803 ns (0 allocations: 0 bytes)

with static array

@btime g!($rs, $as, $bs, $cs, $ds)

17.648 μs (750 allocations: 13.28 KiB)

however, if tested with the function g_ex! which takes one term out of the assignment,

function g_ex!(r,a,b,c,d)
    @. r = a 
    @. r += 5 * b + 3 * c - d
end

benchmark:

@btime g_ex!($rs, $as, $bs, $cs, $ds)

26.951 ns (0 allocations: 0 bytes)

so this is a ~650 slowdown (17.648 μs vs 26.951 ns). This is very noticeable when functions containing such expressions are in a deeply-nested/hot loop.

Also, it's ok when the assignment only has 3 terms in the function g3!()

function g3!(r,a,b,c)
    @. r = a + 3 * b - c
end

benchmark:

@btime g3!($rs, $as, $bs, $cs)

17.306 ns (0 allocations: 0 bytes)

[daviehh]

Just noticed it's not exactly the number of terms: changing the assignment order

function g_reordered!(r,a,b,c,d)
    @. r = 5 * b + 3 * c - d + a
end

then there's no allocation and the speed is fast:

@btime g_reordered!($rs, $as, $bs, $cs, $ds)

22.522 ns (0 allocations: 0 bytes)

Also, changing the first a to 1 * a such that @. r = 1 * a + 5 * b + 3 * c - d also works...

Profiling with Juno suggests this line is dynamically dispatched:

StaticArrays.jl/src/broadcast.jl

Line 155 in 64c64b2

@inbounds $(Expr(:block, exprs...))

profiling script here. Note you may change the loop iteration number to higher or adjust Profile.init's n or delay to see the correct profile on a different machine

[c42f]

Likely this is more or less the same issue as #682.

[daviehh]

Might be slightly different since here @inferred g!(rs, as, bs, cs, ds) is ok...

[mateuszbaran]

Small update: I've proposed a solution here: JuliaLang/julia#41090 .