Performance issues with fusing 12+ broadcasts

[ChrisRackauckas]

Updated OP

While working through the issue, @jebej identified that the problem is fusing 12+ broadcasts in this comment (#22255 (comment)) which contains an MWE.

Original OP

In OrdinaryDiffEq.jl, I see a 10x performance regression due to using broadcast. With the testing code:

const ζ = 0.5
const ω₀ = 10.0

using OrdinaryDiffEq, DiffEqBase

const y₀ = Float64[0., sqrt(1-ζ^2)*ω₀]
const A = 1
const ϕ = 0


function f(t::Float64)
    α = sqrt(1-ζ^2)*ω₀
    x = A*exp(-ζ*ω₀*t)*sin(α*t + ϕ)
    p = A*exp(-ζ*ω₀*t)*(-ζ*ω₀*sin(α*t + ϕ) + α*cos(α*t + ϕ))
    return [x,p]
end

function df(t::Float64, y::Vector{Float64}, dy::Vector{Float64})
    dy[1] = y[2]
    dy[2] = -2*ζ*ω₀*y[2] - ω₀^2*y[1]
    return nothing
end

const T = [0.,10.]
using BenchmarkTools
prob = ODEProblem(df,y₀,(T[1],T[2]))
@benchmark init($prob,Tsit5(),dense=false,dt=1/10)
@benchmark solve($prob,Tsit5(),dense=false,dt=1/10)

using ProfileView

@profile for i in 1:1000; solve(prob,Tsit5(),dense=false,dt=1/10); end
ProfileView.view()

I get a 10x regression by changing the inner loop from:

function perform_step!(integrator,cache::Tsit5Cache,f=integrator.f)
  @unpack t,dt,uprev,u,k = integrator
  uidx = eachindex(integrator.uprev)
  @unpack c1,c2,c3,c4,c5,c6,a21,a31,a32,a41,a42,a43,a51,a52,a53,a54,a61,a62,a63,a64,a65,a71,a72,a73,a74,a75,a76,b1,b2,b3,b4,b5,b6,b7 = cache.tab
  @unpack k1,k2,k3,k4,k5,k6,k7,utilde,tmp,atmp = cache
  a = dt*a21
  for i in uidx
    tmp[i] = @muladd uprev[i]+a*k1[i]
  end
  f(@muladd(t+c1*dt),tmp,k2)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a31*k1[i]+a32*k2[i])
  end
  f(@muladd(t+c2*dt),tmp,k3)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a41*k1[i]+a42*k2[i]+a43*k3[i])
  end
  f(@muladd(t+c3*dt),tmp,k4)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a51*k1[i]+a52*k2[i]+a53*k3[i]+a54*k4[i])
  end
  f(@muladd(t+c4*dt),tmp,k5)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a61*k1[i]+a62*k2[i]+a63*k3[i]+a64*k4[i]+a65*k5[i])
  end
  f(t+dt,tmp,k6)
  for i in uidx
    u[i] = @muladd uprev[i]+dt*(a71*k1[i]+a72*k2[i]+a73*k3[i]+a74*k4[i]+a75*k5[i]+a76*k6[i])
  end
  f(t+dt,u,k7)
  if integrator.opts.adaptive
    for i in uidx
      utilde[i] = @muladd uprev[i] + dt*(b1*k1[i] + b2*k2[i] + b3*k3[i] + b4*k4[i] + b5*k5[i] + b6*k6[i] + b7*k7[i])
      atmp[i] = ((utilde[i]-u[i])./@muladd(integrator.opts.abstol+max(abs(uprev[i]),abs(u[i])).*integrator.opts.reltol))
    end
    integrator.EEst = integrator.opts.internalnorm(atmp)
  end
  @pack integrator = t,dt,u,k
end

to:

function perform_step!(integrator,cache::Tsit5Cache,f=integrator.f)
  @unpack t,dt,uprev,u,k = integrator
  @unpack c1,c2,c3,c4,c5,c6,a21,a31,a32,a41,a42,a43,a51,a52,a53,a54,a61,a62,a63,a64,a65,a71,a72,a73,a74,a75,a76,b1,b2,b3,b4,b5,b6,b7 = cache.tab
  @unpack k1,k2,k3,k4,k5,k6,k7,utilde,tmp,atmp = cache
  a = dt*a21
  tmp .= @muladd uprev+a*k1
  f(@muladd(t+c1*dt),tmp,k2)
  tmp .= @muladd uprev+dt*(a31*k1+a32*k2)
  f(@muladd(t+c2*dt),tmp,k3)
  tmp .= @muladd uprev+dt*(a41*k1+a42*k2+a43*k3)
  f(@muladd(t+c3*dt),tmp,k4)
  tmp .= @muladd uprev+dt*(a51*k1+a52*k2+a53*k3+a54*k4)
  f(@muladd(t+c4*dt),tmp,k5)
  tmp .= @muladd uprev+dt*(a61*k1+a62*k2+a63*k3+a64*k4+a65*k5)
  f(t+dt,tmp,k6)
  u .= @muladd uprev+dt*(a71*k1+a72*k2+a73*k3+a74*k4+a75*k5+a76*k6)
  f(t+dt,u,k7)
  if integrator.opts.adaptive
    utilde .= @muladd uprev + dt*(b1*k1 + b2*k2 + b3*k3 + b4*k4 + b5*k5 + b6*k6 + b7*k7)
    atmp .= ((utilde.-u)./@muladd(integrator.opts.abstol+max.(abs.(uprev),abs.(u)).*integrator.opts.reltol))
    integrator.EEst = integrator.opts.internalnorm(atmp)
  end
  @pack integrator = t,dt,u,k
end

I.e. all that's changed are loops to broadcast. The input array is y0 which is length 2. For reference, the @muladd macro acts like:

println(macroexpand(:(u .= @muladd uprev+dt*(a71*k1+a72*k2+a73*k3+a74*k4+a75*k5+a76*k6))))
#u .= (muladd).(dt, (muladd).(a71, k1, (muladd).(a72, k2, (muladd).(a73, k3, (muladd).(a74, k4, (muladd).(a75, k5, a76 .* k6))))), uprev)

println(macroexpand(:(atmp .= ((utilde.-u)./@muladd(integrator.opts.abstol+max.(abs.(uprev),abs.(u)).*integrator.opts.reltol)))))
#atmp .= (utilde .- u) ./ (muladd).(max.(abs.(uprev), abs.(u)), integrator.opts.reltol, integrator.opts.abstol)

The profile is here: https://ufile.io/2lu0f

The benchmark results are using loops:

BenchmarkTools.Trial:
  memory estimate:  87.43 KiB
  allocs estimate:  3757
  --------------
  minimum time:     281.032 μs (0.00% GC)
  median time:      526.934 μs (0.00% GC)
  mean time:        475.180 μs (2.88% GC)
  maximum time:     5.601 ms (87.13% GC)
  --------------
  samples:          10000
  evals/sample:     1

and using broadcast:

BenchmarkTools.Trial:
  memory estimate:  854.34 KiB
  allocs estimate:  37976
  --------------
  minimum time:     3.246 ms (0.00% GC)
  median time:      6.185 ms (0.00% GC)
  mean time:        5.762 ms (2.40% GC)
  maximum time:     13.919 ms (26.41% GC)
  --------------
  samples:          867
  evals/sample:     1

Am I hitting some broadcasting splatting penalty or something?

[nalimilan]

Any chance you could make the example simpler, e.g. by removing uses of @muladd and keeping only one computation line?

[yuyichao]

AFAICT you are not doing braodcasting for the slower version. Also, f uses non-const global variables.

[yuyichao]

Ah, the @muladd does the transformation. That is extremely confusing................ It is indeed possible to hit tuple length limit with braodcast. You should check code_warntype and also reduce your example.

[ChrisRackauckas]

Give me a little bit to get it simpler. I have to restart Atom and delete my precompilation cache every time I do something due to #21969, so this will take awhile.

[ChrisRackauckas]

Here's the form that expands all of the broadcasts:

function perform_step!(integrator,cache::Tsit5Cache,f=integrator.f)
  @unpack t,dt,uprev,u,k = integrator
  @unpack c1,c2,c3,c4,c5,c6,a21,a31,a32,a41,a42,a43,a51,a52,a53,a54,a61,a62,a63,a64,a65,a71,a72,a73,a74,a75,a76,b1,b2,b3,b4,b5,b6,b7 = cache.tab
  @unpack k1,k2,k3,k4,k5,k6,k7,utilde,tmp,atmp = cache
  a = dt*a21
  tmp .= (muladd).(a, k1, uprev)
  f(muladd(c1, dt, t),tmp,k2)
  tmp .= (muladd).(dt, (muladd).(a31, k1, a32 .* k2), uprev)
  f(muladd(c2, dt, t),tmp,k3)
  tmp .= (muladd).(dt, (muladd).(a41, k1, (muladd).(a42, k2, a43 .* k3)), uprev)
  f(muladd(c3, dt, t),tmp,k4)
  tmp .= (muladd).(dt, (muladd).(a51, k1, (muladd).(a52, k2, (muladd).(a53, k3, a54 .* k4))), uprev)
  f(muladd(c4, dt, t),tmp,k5)
  tmp .= (muladd).(dt, (muladd).(a61, k1, (muladd).(a62, k2, (muladd).(a63, k3, (muladd).(a64, k4, a65 .* k5)))), uprev)
  f(t+dt,tmp,k6)
  u .= (muladd).(dt, (muladd).(a71, k1, (muladd).(a72, k2, (muladd).(a73, k3, (muladd).(a74, k4, (muladd).(a75, k5, a76 .* k6))))), uprev)
  f(t+dt,u,k7)
  if integrator.opts.adaptive
    utilde .= (muladd).(dt, (muladd).(b1, k1, (muladd).(b2, k2, (muladd).(b3, k3, (muladd).(b4, k4, (muladd).(b5, k5, (muladd).(b6, k6, b7 .* k7)))))), uprev)
    atmp .= ((utilde.-u)./(muladd).(max.(abs.(uprev), abs.(u)), integrator.opts.reltol, integrator.opts.abstol))
    integrator.EEst = integrator.opts.internalnorm(atmp)
  end
  @pack integrator = t,dt,u,k
end

adding consts around doesn't really do anything, which is just because most of the time is not spent in the function calls f but rather in the broadcasts. The @code_warntype produces a gigantic output:

integrator = init(prob,Tsit5(),dense=false,dt=1/10)
OrdinaryDiffEq.loopheader!(integrator)
@code_warntype OrdinaryDiffEq.perform_step!(integrator,integrator.cache,integrator.f)

https://gist.github.com/ChrisRackauckas/11ba482f80147dd744f5c2b87e288067

[ChrisRackauckas]

This might be an MWE:

function f(a,b,c)
  a.= b.*c
end

function g(a,b,c)
  @inbounds for ii in eachindex(a)
    a[ii] = b[ii].*c[ii]
  end
end

function f(a,b,c,d,e,h,i,j,k,l,m,n,o,p,q,r,s,t,u,v,w,x,y,z)
  a.= b.*c .+ d .* e .+ h .* i .+ j .*k .+ l .* m .+ n .* o .+ p .* q .+ r .* s .+ t .* u .+ v .* w .+ x .* y .+ z
end

function g(a,b,c,d,e,h,i,j,k,l,m,n,o,p,q,r,s,t,u,v,w,x,y,z)
  @inbounds for ii in eachindex(a)
    a[ii] = b[ii].*c[ii] .+ d[ii] .* e[ii] .+ h[ii] .* i[ii] .+ j[ii] .*k[ii] .+ l[ii] .* m[ii] .+ n[ii] .* o[ii] .+ p[ii] .* q[ii] .+ r[ii] .* s[ii] .+ t[ii] .* u[ii] .+ v[ii] .* w[ii] .+ x[ii] .* y[ii] .+ z[ii]
  end
end
a = rand(10)
b = rand(10)
c = rand(10)
d = rand(10)
e = rand(10)
h = rand(10)
i = rand(10)
j = rand(10)
k = rand(10)
l = rand(10)
m = rand(10)
n = rand(10)
o = rand(10)
p = rand(10)
q = rand(10)
r = rand(10)
s = rand(10)
t = rand(10)
u = rand(10)
v = rand(10)
w = rand(10)
x = rand(10)
y = rand(10)
z = rand(10)

@benchmark f($a,$b,$c)
@benchmark g($a,$b,$c)

@benchmark f($a,$b,$c,$d,$e,$h,$i,$j,$k,$l,$m,$n,$o,$p,$q,$r,$s,$t,$u,$v,$w,$x,$y,$z)
@benchmark g($a,$b,$c,$d,$e,$h,$i,$j,$k,$l,$m,$n,$o,$p,$q,$r,$s,$t,$u,$v,$w,$x,$y,$z)

Benchmarks on the small broadcasts:

@benchmark f($a,$b,$c)
BenchmarkTools.Trial: 
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     16.393 ns (0.00% GC)
  median time:      32.494 ns (0.00% GC)
  mean time:        28.710 ns (0.00% GC)
  maximum time:     225.996 ns (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     1000

@benchmark g($a,$b,$c)
BenchmarkTools.Trial: 
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     7.611 ns (0.00% GC)
  median time:      14.051 ns (0.00% GC)
  mean time:        12.511 ns (0.00% GC)
  maximum time:     192.038 ns (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     1000

and benchmarks on the long broadcasts:

@benchmark f($a,$b,$c,$d,$e,$h,$i,$j,$k,$l,$m,$n,$o,$p,$q,$r,$s,$t,$u,$v,$w,$x,$y,$z)
BenchmarkTools.Trial: 
  memory estimate:  17.77 KiB
  allocs estimate:  650
  --------------
  minimum time:     51.229 μs (0.00% GC)
  median time:      101.874 μs (0.00% GC)
  mean time:        122.778 μs (3.60% GC)
  maximum time:     31.670 ms (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     1

@benchmark g($a,$b,$c,$d,$e,$h,$i,$j,$k,$l,$m,$n,$o,$p,$q,$r,$s,$t,$u,$v,$w,$x,$y,$z)
BenchmarkTools.Trial: 
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     222.590 ns (0.00% GC)
  median time:      413.482 ns (0.00% GC)
  mean time:        363.430 ns (0.00% GC)
  maximum time:     2.199 μs (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     434

When it gets sufficiently long, it begins to allocate quite a bit!

[ChrisRackauckas]

Indeed get a the same kinds of differences when I use some of the broadcast expressions from the test code:

function f(a,b,c,d,e,h,i,j,k,l,m,n,o,p,q)
  a .= (muladd).(b, (muladd).(c, d, (muladd).(e, h, (muladd).(i, j, (muladd).(k, l, (muladd).(m, n, o .* p))))), q)
end

function g(a,b,c,d,e,h,i,j,k,l,m,n,o,p,q)
  @inbounds for ii in eachindex(a)
    a[ii] = (muladd)(b[ii], (muladd)(c[ii], d[ii], (muladd)(e[ii], h[ii], (muladd)(i[ii], j[ii], (muladd)(k[ii], l[ii], (muladd)(m[ii], n[ii], o[ii] .* p[ii]))))), q[ii])
  end
end

@benchmark f($a,$b,$c,$d,$e,$h,$i,$j,$k,$l,$m,$n,$o,$p,$q)
@benchmark g($a,$b,$c,$d,$e,$h,$i,$j,$k,$l,$m,$n,$o,$p,$q)

@benchmark f($a,$b,$c,$d,$e,$h,$i,$j,$k,$l,$m,$n,$o,$p,$q)

BenchmarkTools.Trial: 
  memory estimate:  1.58 KiB
  allocs estimate:  44
  --------------
  minimum time:     6.382 μs (0.00% GC)
  median time:      12.471 μs (0.00% GC)
  mean time:        11.167 μs (2.64% GC)
  maximum time:     1.154 ms (94.98% GC)
  --------------
  samples:          10000
  evals/sample:     5

@benchmark g($a,$b,$c,$d,$e,$h,$i,$j,$k,$l,$m,$n,$o,$p,$q)

BenchmarkTools.Trial: 
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     48.009 ns (0.00% GC)
  median time:      83.139 ns (0.00% GC)
  mean time:        74.315 ns (0.00% GC)
  maximum time:     1.335 μs (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     1000

Interestingly, that's only with 8 dots and 15 numbers... I would've though the magic 16 would be a limit here.

[nalimilan]

Does it still happen when reducing the number of operands?

[ChrisRackauckas]

Does it still happen when reducing the number of operands?

No. See the "benchmarks on small broadcasts". For "sufficiently few" operands, it works just fine.

[nalimilan]

OK. So where's the cutoff?

[jebej]

With the earlier example Chris had, it happens here with 12 variables, but not with 11:

function fun1!(a,b,c,d,e,f,g,h,i,j,k)
  a .= b.*c .+ d.*e .+ f.*g .+ h.*i .+ j.*k
end
function fun2!(a,b,c,d,e,f,g,h,i,j,k,l)
  a .= b.*c .+ d.*e .+ f.*g .+ h.*i .+ j.*k .+ l
end
a = rand(10); b = rand(10); c = rand(10); d = rand(10); e = rand(10);
f = rand(10); g = rand(10); h = rand(10); i = rand(10); j = rand(10);
k = rand(10); l = rand(10);
using BenchmarkTools
@benchmark fun1!($a,$b,$c,$d,$e,$f,$g,$h,$i,$j,$k)
@benchmark fun2!($a,$b,$c,$d,$e,$f,$g,$h,$i,$j,$k,$l)

julia> @benchmark fun1!($a,$b,$c,$d,$e,$f,$g,$h,$i,$j,$k)
BenchmarkTools.Trial:
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     149.380 ns (0.00% GC)
  median time:      150.616 ns (0.00% GC)
  mean time:        153.680 ns (0.00% GC)
  maximum time:     382.402 ns (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     947

julia> @benchmark fun2!($a,$b,$c,$d,$e,$f,$g,$h,$i,$j,$k,$l)
BenchmarkTools.Trial:
  memory estimate:  256 bytes
  allocs estimate:  4
  --------------
  minimum time:     509.969 ns (0.00% GC)
  median time:      514.536 ns (0.00% GC)
  mean time:        540.944 ns (3.26% GC)
  maximum time:     9.924 μs (91.70% GC)
  --------------
  samples:          10000
  evals/sample:     192

[ChrisRackauckas]

Updated the title to reflect the new status of the issue. Thanks @jebej for finding that. Updating the OP with the new information.

[ChrisRackauckas]

Does anyone know which inference option would be involved here? I could 10 .s as fine, but 11 as too much, but none of the inference options have a cutoff at 10. I want to make sure this is fixed with #22545 though.