Type inference fails for broadcasted addition of tuples containing multiple types

General Usage Performance
1

Hi,

I want to write addition function for tuples with unknown length and unknown type e.g. something like Tuple{a, b, c, d...} where the length and type is determined at runtime, and a, b, c, and d are custom types with addition defined.

Naively, I would try something like this

function add_tuple(a, b)
    a .+ b
end

This works great for a :: Tuple{Int, Float64} :

@btime add_tuple((1,2.0), (2,3.0))
  0.013 ns (0 allocations: 0 bytes)

@code_warntype add_tuple((1,3.0), (2,2.0))

Variables
  #self#::Core.Compiler.Const(add_tuple, false)
  a::Tuple{Int64,Float64}
  b::Tuple{Int64,Float64}

Body::Tuple{Int64,Float64}
1 ─ %1 = Base.broadcasted(Main.:+, a, b)::Base.Broadcast.Broadcasted{Base.Broadcast.Style{Tuple},Nothing,typeof(+),Tuple{Tuple{Int64,Float64},Tuple{Int64,Float64}}}
│   %2 = Base.materialize(%1)::Tuple{Int64,Float64}
└──      return %2

A bit surprisingly, type inference seems to fail for slightly more complicated examples

struct MyInt
    s :: Int
end

Base.:+(a::MyInt, b::MyInt) = MyInt(a.s + b.s)

a = MyInt(1)
@code_warntype add_tuple((1,a), (2,a))

Variables
  #self#::Core.Compiler.Const(add_tuple, false)
  a::Tuple{Int64,MyInt}
  b::Tuple{Int64,MyInt}

Body::Tuple{Union{Int64, MyInt},Union{Int64, MyInt}}
1 ─ %1 = Base.broadcasted(Main.:+, a, b)::Base.Broadcast.Broadcasted{Base.Broadcast.Style{Tuple},Nothing,typeof(+),Tuple{Tuple{Int64,MyInt},Tuple{Int64,MyInt}}}
│   %2 = Base.materialize(%1)::Tuple{Union{Int64, MyInt},Union{Int64, MyInt}}
└──      return %2

Not that type inference fails for Base.materialize, even though the compiler should have enough information to figure out the output type. I suspect this has some performance cost in more complicated cases.

My question is why this happens, and if there is anything I can do to force type inference. Thanks!

2

A possibly relevant observation: getindex is type unstable for tuple containing multiple types

a = (1, 1.0)
@code_warntype a[1]

Variables
  #self#::Core.Compiler.Const(getindex, false)
  t::Tuple{Int64,Float64}
  i::Int64

Body::Union{Float64, Int64}
1 ─      nothing
│   %2 = Base.getfield(t, i, $(Expr(:boundscheck)))::Union{Float64, Int64}
└──      return %2
3

Use

add_tuple(a, b) = map(+, a, b)

If you are interested in more complex cases, search for “lispy tuple recursion”.

1 Like
4

Thanks for the suggestion! I tested your solution

function add_tuple_map(a, b)
    return map(+, a, b)
end

function add_tuple_broadcast(a, b)
    return broadcast(+, a, b)
end

function add_tuple_dot(a, b)
    return a .+ b
end

a = (1, MyInt(3))
b = (2, MyInt(4))

@code_warntype add_tuple_map(a, b)
Variables
  #self#::Core.Compiler.Const(add_tuple_map, false)
  a::Tuple{Int64,MyInt}
  b::Tuple{Int64,MyInt}

Body::Tuple{Int64,MyInt}
1 ─ %1 = Main.map(Main.:+, a, b)::Tuple{Int64,MyInt}
└──      return %1

@code_warntype add_tuple_broadcast(a, b)
Variables
  #self#::Core.Compiler.Const(add_tuple_broadcast, false)
  a::Tuple{Int64,MyInt}
  b::Tuple{Int64,MyInt}

Body::Tuple{Int64,MyInt}
1 ─ %1 = Main.broadcast(Main.:+, a, b)::Tuple{Int64,MyInt}
└──      return %1

@code_warntype add_tuple_dot(a, b)
Variables
  #self#::Core.Compiler.Const(add_tuple_dot, false)
  a::Tuple{Int64,MyInt}
  b::Tuple{Int64,MyInt}

Body::Tuple{Union{Int64, MyInt},Union{Int64, MyInt}}
1 ─ %1 = Base.broadcasted(Main.:+, a, b)::Base.Broadcast.Broadcasted{Base.Broadcast.Style{Tuple},Nothing,typeof(+),Tuple{Tuple{Int64,MyInt},Tuple{Int64,MyInt}}}
│   %2 = Base.materialize(%1)::Tuple{Union{Int64, MyInt},Union{Int64, MyInt}}
└──      return %2

As you said, map makes it type stable. I was surprised broadcast also makes it type stable, because I thought a .+ b was equivalent to broadcast(+, a, b). Do you happen to know why this is the case?