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.

[wheeheee]

I think it's not just tuples. For example, bar2 here doesn't have a small union as its return type.

struct Foo{T,V}
    x::T
    n::V
end

bar(A) = A.n < 0 ? Foo(inv(A.x), inv(A.n)^2) : Foo(A.x, A.n^2)
baz(A) = Foo(inv(A.x), inv(A.n))
function bar2(A)
    B = A.n < 0 ? baz(A) : A
    return Foo(B.x, B.n^2)
end

[mbauman]

I think that's a very subtly different problem — in that case, Julia hit too many possibilities and gave up.

The difference between tuples and structs becomes more obvious with a simpler example, wherein only one field is unstable.

julia> struct Foo{T,V}
           x::T
           n::V
       end

julia> function bar2(x, n)
           y, k = n < 0 ? (inv(x), inv(n)) : (x, n)
           return (y, k)
       end
bar2 (generic function with 1 method)

julia> function bar3(x, n)
           y, k = n < 0 ? (inv(x), inv(n)) : (x, n)
           return Foo(y, k)
       end
bar3 (generic function with 1 method)

julia> @code_warntype bar2(1, 2.)
# ...
Body::Tuple{Union{Float64, Int64}, Float64}

julia> @code_warntype bar3(1, 2.)
# ...
Body::Union{Foo{Float64, Float64}, Foo{Int64, Float64}}

[wheeheee]

You're right, these are slightly different. I was actually dealing with structs when I encountered these difficulties, so I simplified the code to make a MWE with some tuples, not realizing the differences. It would be nice if this just worked for a small number of combinations, however, for a bigger value of small. I encountered this here

https://github.com/JuliaDSP/DSP.jl/blob/06f8776fd3b7b7944413e52aa7ca56771ae602d6/src/Filters/coefficients.jl#L329-L337

In the case that inv_f and f are the same type, the repeat code appears twice in code_typed and code_llvm, which isn't nice, but because with f = e < 0 ? inv(f) : f the return type isn't a small union if they are not the same type I couldn't just do that.