very bad performance of iterator on array of union types

[lmiq]

This a simplified example of this thread on discourse, but which appears to have revealed an inefficient behaviour of one of the common iterators when dealing with arrays of mixed Union types. Here is the simplified example:

using BenchmarkTools, Test

struct A x::Float64 end
struct B x::Float64 end
struct C x::Float64 end
struct D x::Float64 end
struct E x::Float64 end

function sum1(v)
  s = 0.
  for x in v
    s += x.x
  end
  s
end

function sum2(v)
  s = 0.
  for (i,x) in pairs(v)
    s += x.x
  end
  s
end

t = [A,B,C,D,E]
v = Union{A,B,C,D,E}[rand(t)(rand()) for i in 1:1000]

@test sum1(v) ≈ sum2(v)

@btime sum1($v)

@btime sum2($v)

result (Julia 1.6.1):

  241.063 μs (7980 allocations: 140.31 KiB)
  992.583 ns (0 allocations: 0 bytes)
501.12183238665824

Thus, iterating with for x in v is much worse than using for (i,x) in pairs(v), which does not appear to make any sense in terms of functionality.

[KristofferC]

Looking at @code_warntype optimize=true sum1(v) there is

%12 = Core.tuple(%10, %11)::Tuple{Union{A, B, C, D, E}, Int64}
...
%14 = Base.nothing::Nothing
...
%17 = φ (#2 => %12, #3 => %14)::Union{Nothing, Tuple{Any, Any}}

which loses the type information of the Union. Theoretically, that could have been inferred to Union{Nothing, Tuple{Union{A, B, C, D, E}, Int64}}. Maybe there is some heuristic that widens it.

[melonedo]

The heuristic is a test to Core.Compiler.issimpleenoughtype at

julia/base/compiler/typelimits.jl

Line 487 in 6d2c0a7

if issimpleenoughtype(u)

Here, unioncomplexity of limits the maximum elements in a union:

julia> using Core.Compiler: issimpleenoughtype, unioncomplexity

julia> Union{Nothing, Tuple{Union{A, B, C, D, E}, Int64}}|>issimpleenoughtype
false

julia> Union{Nothing, Tuple{Union{A, B, C, D, E}, Int64}}|>unioncomplexity
5

julia> Union{Nothing, Tuple{Pair{Int64, Union{A, B, C, D, E}}, Int64}}|>issimpleenoughtype
true

julia> Union{Nothing, Tuple{Pair{Int64, Union{A, B, C, D, E}}, Int64}}|>unioncomplexity
1

Here the union is wrapped in a struct Pair, which is invisible to the heuristic. We can manually wrapped the union, for a workaround, too:

struct A x::Float64 end
struct B x::Float64 end
struct C x::Float64 end
struct D x::Float64 end
struct E x::Float64 end
function sum1(v)
  s = 0.
  for x in v
    s += x.x
  end
  s
end
function sum2(v)
  s = 0.
  for (i,x) in pairs(v)
    s += x.x
  end
  s
end
function sum3(v)
  s = 0.
  for x in v
    s += x.u.x
  end
  s
end

struct Wrapper{T}
  u::T
end
Base.convert(::Type{Wrapper{U}}, x) where U = Wrapper{U}(x)

u = Wrapper{Union{A,B,C,D}}[rand([A,B,C,D])(rand()) for i in 1:1000]
v = Union{A,B,C,D,E}[rand([A,B,C,D,E])(rand()) for i in 1:1000]
@code_typed(sum(u)) |> last # => Wrapper{Union{A, B, C, D}}
@code_typed(sum(v)) |> last # => Any
using BenchmarkTools
@btime sum1(v) # => 258.500 μs (7980 allocations: 140.31 KiB)
@btime sum2(v) # => 1.020 μs (1 allocation: 16 bytes)
@btime sum3(u) # => 1.030 μs (1 allocation: 16 bytes)

[Tortar]

I just stumbled on this issue and I see identical good performance in 1.11 and nightly (but bad on 1.10):

julia> @btime sum1($v)
  686.887 ns (0 allocations: 0 bytes)
498.54917340525105

julia> @btime sum2($v)
  686.827 ns (0 allocations: 0 bytes)
498.54917340525105

so this was fixed somehow? Maybe it can be then closed

[oscardssmith]

presumably Jameson and I fixed it with our Memory rewrite. not obvious why though...

[Tortar]

just to let you know, I see a much better performance of big Union types in micro-benchmarks in 1.11, unfortunately this doesn't translate well on macro-ones