broadcast of short-circuiting boolean operator || generates runtime dispatch and per-element allocations

[jcphill]

Creating a BitVector from a slightly non-trivial expression using the short-circuiting boolean operator .|| results in runtime dispatch, per-element allocations, and worse performance than the equivalent expression using the bitwise operator .|. The arrays are both in a named tuple and of SVector element type, and LinearAlgebra.norm is called on the right-hand side.

julia> versioninfo()
Julia Version 1.10.2
Commit bd47eca2c8a (2024-03-01 10:14 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (x86_64-apple-darwin22.4.0)
  CPU: 8 × Intel(R) Core(TM) i5-1038NG7 CPU @ 2.00GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, icelake-client)
Threads: 1 default, 0 interactive, 1 GC (on 8 virtual cores)

julia> using StaticArrays

julia> using LinearAlgebra

julia> ab = (a=zeros(SVector{3,Float64}, 1000000),
             b=zeros(SVector{3,Float64}, 1000000));

julia> f_bitwise(ab) = (first.(ab.a).>0) .| (norm.(ab.b) .!= 0)
f_bitwise (generic function with 1 method)

julia> f_boolean(ab) = (first.(ab.a).>0) .|| (norm.(ab.b) .!= 0)
f_boolean (generic function with 1 method)

julia> f_bitwise(ab) == f_boolean(ab)
true

julia> @time f_bitwise(ab);
  0.007586 seconds (4 allocations: 126.391 KiB)

julia> @time f_boolean(ab);
  0.117142 seconds (2.00 M allocations: 61.159 MiB, 69.11% gc time)

julia> using JET

julia> @report_opt f_bitwise(ab)
No errors detected


julia> @report_opt f_boolean(ab)
═════ 6 possible errors found ═════
┌ f_boolean(ab::@NamedTuple{a::Vector{SVector{3, Float64}}, b::Vector{SVector{3, Float64}}}) @ Main ./REPL[6]:1
│┌ materialize(bc::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, Base.Broadcast.var"#7#8"{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, Base.Broadcast.var"#12#14"{Base.Broadcast.var"#18#20"{Base.Broadcast.var"#15#16"{Base.Broadcast.var"#11#13"}, Base.Broadcast.var"#15#16"{Base.Broadcast.var"#17#19"}, Base.Broadcast.var"#25#26"{Base.Broadcast.var"#27#28"}, Base.Broadcast.var"#21#22"{Base.Broadcast.var"#23#24"}, typeof(norm)}, typeof(!=)}, Tuple{Vector{SVector{3, Float64}}, Int64}}}, Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(>), Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(first), Tuple{Vector{SVector{3, Float64}}}}, Int64}}, Vector{SVector{3, Float64}}, Int64}}) @ Base.Broadcast ./broadcast.jl:903
││┌ copy(bc::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Tuple{Base.OneTo{Int64}}, Base.Broadcast.var"#7#8"{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, Base.Broadcast.var"#12#14"{Base.Broadcast.var"#18#20"{Base.Broadcast.var"#15#16"{Base.Broadcast.var"#11#13"}, Base.Broadcast.var"#15#16"{Base.Broadcast.var"#17#19"}, Base.Broadcast.var"#25#26"{Base.Broadcast.var"#27#28"}, Base.Broadcast.var"#21#22"{Base.Broadcast.var"#23#24"}, typeof(norm)}, typeof(!=)}, Tuple{Vector{SVector{3, Float64}}, Int64}}}, Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(>), Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(first), Tuple{Vector{SVector{3, Float64}}}}, Int64}}, Vector{SVector{3, Float64}}, Int64}}) @ Base.Broadcast ./broadcast.jl:925
│││┌ combine_eltypes(f::Base.Broadcast.var"#7#8"{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, Base.Broadcast.var"#12#14"{Base.Broadcast.var"#18#20"{Base.Broadcast.var"#15#16"{Base.Broadcast.var"#11#13"}, Base.Broadcast.var"#15#16"{Base.Broadcast.var"#17#19"}, Base.Broadcast.var"#25#26"{Base.Broadcast.var"#27#28"}, Base.Broadcast.var"#21#22"{Base.Broadcast.var"#23#24"}, typeof(norm)}, typeof(!=)}, Tuple{Vector{SVector{3, Float64}}, Int64}}}, args::Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(>), Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(first), Tuple{Vector{SVector{3, Float64}}}}, Int64}}, Vector{SVector{3, Float64}}, Int64}) @ Base.Broadcast ./broadcast.jl:760
││││┌ promote_typejoin_union(::Type) @ Base ./promotion.jl:188
│││││┌ promote_typejoin(a::Type, b::Type) @ Base ./promotion.jl:172
││││││┌ typejoin(a::Type, b::Type) @ Base ./promotion.jl:34
│││││││┌ typejoin(a::Any, b::Type) @ Base ./promotion.jl:36
││││││││┌ typejoin(a::Any, b::Any) @ Base ./promotion.jl:127
│││││││││ runtime dispatch detected: Base.UnionAll(%403::Any, %405::Any)::Any
││││││││└────────────────────
│││││││┌ typejoin(a::Any, b::Type) @ Base ./promotion.jl:127
││││││││ runtime dispatch detected: Base.UnionAll(%398::Any, %400::Any)::Any
│││││││└────────────────────
││││││┌ typejoin(a::Type, b::Type) @ Base ./promotion.jl:127
│││││││ runtime dispatch detected: Base.UnionAll(%393::Any, %395::Any)::Any
││││││└────────────────────
││││┌ typejoin_union_tuple(T::DataType) @ Base ./promotion.jl:216
│││││ runtime dispatch detected: Base.promote_typejoin_union(%40::Any)::Type
││││└────────────────────
│││┌ combine_eltypes(f::Base.Broadcast.var"#7#8"{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, Base.Broadcast.var"#12#14"{Base.Broadcast.var"#18#20"{Base.Broadcast.var"#15#16"{Base.Broadcast.var"#11#13"}, Base.Broadcast.var"#15#16"{Base.Broadcast.var"#17#19"}, Base.Broadcast.var"#25#26"{Base.Broadcast.var"#27#28"}, Base.Broadcast.var"#21#22"{Base.Broadcast.var"#23#24"}, typeof(norm)}, typeof(!=)}, Tuple{Vector{SVector{3, Float64}}, Int64}}}, args::Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(>), Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(first), Tuple{Vector{SVector{3, Float64}}}}, Int64}}, Vector{SVector{3, Float64}}, Int64}) @ Base.Broadcast ./broadcast.jl:760
││││ runtime dispatch detected: Base._return_type(f::Base.Broadcast.var"#7#8"{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, Base.Broadcast.var"#12#14"{Base.Broadcast.var"#18#20"{Base.Broadcast.var"#15#16"{Base.Broadcast.var"#11#13"}, Base.Broadcast.var"#15#16"{Base.Broadcast.var"#17#19"}, Base.Broadcast.var"#25#26"{Base.Broadcast.var"#27#28"}, Base.Broadcast.var"#21#22"{Base.Broadcast.var"#23#24"}, typeof(norm)}, typeof(!=)}, Tuple{Vector{SVector{3, Float64}}, Int64}}}, ::Tuple{Bool, SVector{3, Float64}, Int64})::Type
│││└────────────────────
│││┌ combine_eltypes(f::Base.Broadcast.var"#7#8"{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, Base.Broadcast.var"#12#14"{Base.Broadcast.var"#18#20"{Base.Broadcast.var"#15#16"{Base.Broadcast.var"#11#13"}, Base.Broadcast.var"#15#16"{Base.Broadcast.var"#17#19"}, Base.Broadcast.var"#25#26"{Base.Broadcast.var"#27#28"}, Base.Broadcast.var"#21#22"{Base.Broadcast.var"#23#24"}, typeof(norm)}, typeof(!=)}, Tuple{Vector{SVector{3, Float64}}, Int64}}}, args::Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(>), Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(first), Tuple{Vector{SVector{3, Float64}}}}, Int64}}, Vector{SVector{3, Float64}}, Int64}) @ Base.Broadcast ./broadcast.jl:760
││││ runtime dispatch detected: Base.Broadcast.promote_typejoin_union(%2::Type)::Type
│││└────────────────────

Here is the input code only:

using StaticArrays
using LinearAlgebra

ab = (a=zeros(SVector{3,Float64}, 1000000),
      b=zeros(SVector{3,Float64}, 1000000));

f_bitwise(ab) = (first.(ab.a).>0) .| (norm.(ab.b) .!= 0)

f_boolean(ab) = (first.(ab.a).>0) .|| (norm.(ab.b) .!= 0)

f_bitwise(ab) == f_boolean(ab)

@time f_bitwise(ab);

@time f_boolean(ab);

using JET

@report_opt f_bitwise(ab)

@report_opt f_boolean(ab)

[jcphill]

Done.

[mbauman]

Fascinating; I was about to punt this to StaticArrays, but neither f_boolean nor f_bitwise are using StaticArray implementations here. I believe this has to do with how .|| flattens the broadcasted tree before evaluating it. Here's a slightly more minimal MWE:

julia> using StaticArrays, LinearAlgebra; using .Broadcast: instantiate, broadcasted, Broadcasted, oror

julia> bc = convert(Broadcasted{Nothing}, instantiate(broadcasted(oror, false, broadcasted(>, broadcasted(norm, zeros(SVector{3, Float64}, 1000000)), 0))));

julia> dest = falses(1000000);

julia> copyto!(dest, bc); @time copyto!(dest, bc);
  0.111565 seconds (3.00 M allocations: 76.294 MiB, 4.76% gc time)

julia> bc[1]; @time bc[1]
  0.000014 seconds (3 allocations: 80 bytes)
false

Interestingly all of ||, norm and SVector are required here; replacing any with |, sum, or Vector alternatives is not sufficient. I'm still not sure where this is going sideways.

[LilithHafner]

JET is clean on the OP in 1.11.0-beta1.

It makes sense that the broadcasted .|| would be slower than the .| because for operations as fast as computing the norm of three elements it's faster to always do it than branch and maybe do it.

The allocations are unexpected and remain present even in 1.11.0-beta1. Unexpected allocations are typically worth investigating, and in a task like this, they are also likely the bottleneck.

Aside: IMO it is more idiomatic to define a complex scalar function and broadcast that than to define a complex broadcast, but this does not make the issue you found invalid.

g_bitwise(a, b) = (first(a) > 0) | (norm(b) != 0)
g_bitwise.(ab.a, ab.b)

I concur that flatten is an issue here.

I suspect specialization/compiler heuristics are hurting this:

julia> using Chairmarks, StaticArrays, LinearAlgebra; using .Broadcast: instantiate, broadcasted, Broadcasted, oror

julia> bc = Base.Broadcast.flatten(broadcasted(>, broadcasted(norm, zeros(SVector{3, Float64}, 1000000)), 0))
Broadcasted(#12, (SVector{3, Float64}[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]  …  [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], 0))

julia> inp = first.(bc.args)
([0.0, 0.0, 0.0], 0)

julia> function slow(bc)
           bcf = Base.Broadcast.flatten(bc)
           broadcasted((args...) -> bcf.f(args...), bcf.args...)
       end
slow (generic function with 1 method)

julia> function fast(bc)
           bcf = Base.Broadcast.flatten(bc)
           broadcasted(bcf.f, bcf.args...)
       end
fast (generic function with 1 method)

julia> @b slow(bc) _.f($inp...)
20.280 ns (2 allocs: 64 bytes)

julia> @b fast(bc) _.f($inp...)
2.840 ns

[LilithHafner]

This fixes the allocations and major runtime discrepancies:

julia> @eval Base.Broadcast function broadcasted(::OrOr, a, bc::Broadcasted)
           bcf = flatten(bc)
           broadcasted(((a, args::Vararg{Any, N}) where {N}) -> a || bcf.f(args...), a, bcf.args...)
       end

See also: https://docs.julialang.org/en/v1/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing