Suboptimal performance for tuples of small unions

[wheeheee]

For example, these two functions look like they should be equivalent, but the inferred return type of the first is a small union, as desired, while the other one isn't. Using (parametric) structs instead of tuples gives similar results.

function bar(x, n)
    return n < 0 ? (inv(x), inv(n)) : (x, n)
end

function bar2(x, n)
    y, k = n < 0 ? (inv(x), inv(n)) : (x, n)
    return (y, k)
end

With the return type of bar2 in red, bar in yellow

julia> @code_warntype bar(1, -2)
MethodInstance for bar(::Int64, ::Int64)
  from bar(x, n) @ Main REPL[1]:1
Arguments
  #self#::Core.Const(Main.bar)
  x::Int64
  n::Int64
Body::Union{Tuple{Float64, Float64}, Tuple{Int64, Int64}}
1 ─ %1 = Main.:<::Core.Const(<)
│   %2 = (%1)(n, 0)::Bool
└──      goto #3 if not %2
2 ─ %4 = Main.inv(x)::Float64
│   %5 = Main.inv(n)::Float64
│   %6 = Core.tuple(%4, %5)::Tuple{Float64, Float64}
└──      return %6
3 ─ %8 = Core.tuple(x, n)::Tuple{Int64, Int64}
└──      return %8

julia> @code_warntype bar2(1, -2)
MethodInstance for bar2(::Int64, ::Int64)
  from bar2(x, n) @ Main REPL[3]:1
Arguments
  #self#::Core.Const(Main.bar2)
  x::Int64
  n::Int64
Locals
  @_4::Int64
  k::Union{Float64, Int64}
  y::Union{Float64, Int64}
  @_7::Union{Tuple{Float64, Float64}, Tuple{Int64, Int64}}
Body::Tuple{Union{Float64, Int64}, Union{Float64, Int64}}
1 ─       Core.NewvarNode(:(@_4))
│         Core.NewvarNode(:(k))
│         Core.NewvarNode(:(y))
│   %4  = (n < 0)::Bool
└──       goto #3 if not %4
2 ─ %6  = Main.inv(x)::Float64
│   %7  = Main.inv(n)::Float64
│         (@_7 = Core.tuple(%6, %7))
└──       goto #4
3 ─       (@_7 = Core.tuple(x, n))
4 ┄ %11 = @_7::Union{Tuple{Float64, Float64}, Tuple{Int64, Int64}}
│   %12 = Base.indexed_iterate(%11, 1)::Union{Tuple{Float64, Int64}, Tuple{Int64, Int64}}
│         (y = Core.getfield(%12, 1))
│         (@_4 = Core.getfield(%12, 2))
│   %15 = @_4::Int64
│   %16 = Base.indexed_iterate(%11, 2, %15)::Union{Tuple{Float64, Int64}, Tuple{Int64, Int64}}
│         (k = Core.getfield(%16, 1))
│   %18 = y::Union{Float64, Int64}
│   %19 = k::Union{Float64, Int64}
│   %20 = Core.tuple(%18, %19)::Tuple{Union{Float64, Int64}, Union{Float64, Int64}}
└──       return %20

[raminammour]

Take a look at this thread and, in particular, Jameson's comment at the end. I guess inference could be better, but it is difficult (not saying that the issue should be closed :)).

[mbauman]

This is precisely the perf issue that plagued InvertedIndices for so long (cf InvertedIndices.jl#39). I see it more as a general performance optimization than a specific inference problem — it'd generally be very nice if tuples of small unions could themselves optimize as though the elements weren't in a tuple.

The trouble of course is that there can easily be a combinatorial explosion beyond what counts as a small union, but even just supporting Tuple{Union{A,B}, C} as an effective Union{Tuple{A,C}, Tuple{B,C}} would be super-valuable. And the current small-union threshold of 4 could even theoretically support splitting through bar2's Tuple{Union{Float64, Int64}, Union{Float64, Int64}} — even without learning that it can rule out the heterogeneous mixes.