Type instability for broadcast with abstract type

General Usage Performance
1

I have encountered a type instability when defining broadcasting for an abstract type. Consider the following (silly) example:

import Base.Broadcast: Broadcasted, BroadcastStyle, broadcastable, broadcasted, instantiate
import Base: axes, copy

abstract type T end

struct S <: BroadcastStyle end

BroadcastStyle(::Type{<:T}) = S()
broadcastable(x::T) = x

Base.axes(::T) = nothing

instantiate(bc::Broadcasted{S}) = bc

copy(bc::Broadcasted{S, Nothing, typeof(+)}) = length(bc.args)

f(v::Vector{T}) = v[1] .+ v[2]

struct A <: T end

v = T[A(),A()]

Then I get

julia> f(v)
2

julia> @code_typed f(v)
CodeInfo(
1 ─ %1 = Base.arrayref(true, v, 1)::T
│   %2 = Base.arrayref(true, v, 2)::T
│   %3 = Base.broadcasted(Main.:+, %1, %2)::Broadcasted{Style, Nothing, typeof(+)} where Style<:Union{Nothing, BroadcastStyle}
│   %4 = Base.Broadcast.instantiate(%3)::Any
│   %5 = Base.Broadcast.copy(%4)::Any
└──      return %5
) => Any

To me, the problem seems to be that Julia thinks there are two relevant methods for BroadcastStyle:

julia> methods(BroadcastStyle, (Type{<:T},))
# 2 methods for type constructor:
 [1] BroadcastStyle(::Type{Union{}}, slurp...)
     @ broadcast.jl:38
 [2] BroadcastStyle(::Type{<:T})
     @ Main REPL[5]:1

julia> code_typed(BroadcastStyle, (Type{<:T},))
2-element Vector{Any}:
 CodeInfo(
1 ─     return $(QuoteNode(Base.Broadcast.Unknown()))
) => Base.Broadcast.Unknown
 CodeInfo(
1 ─     return $(QuoteNode(S()))
) => S

However, the first method cannot apply to any concrete subtype of T. I wish Julia would ignore it when determining the return type of f. Is this a bug, or am I missing something here?

2

Isn’t the type instability coming from the vector itself? If you use v = [A(), A()] there is no type instability.

Perhaps these two examples are simpler:

g(v::Vector{Number}) = v[1] .+ v[2]

v = Number[1, 2]

g(v)
3

@code_warntype g(v)
MethodInstance for g(::Vector{Number})
  from g(v::Vector{Number}) @ Main REPL[55]:1
Arguments
  #self#::Core.Const(g)
  v::Vector{Number}
Body::Any
1 ─ %1 = Main.:+::Core.Const(+)
│   %2 = Base.getindex(v, 1)::Number
│   %3 = Base.getindex(v, 2)::Number
│   %4 = Base.broadcasted(%1, %2, %3)::Broadcasted{Style, Nothing, typeof(+)} where Style<:Union{Nothing, BroadcastStyle}
│   %5 = Base.materialize(%4)::Any
└──      return %5

The exact type of the arguments to broadcasted is not known inside the function, since the container eltype is abstract. Apparently we cannot figure out that a proper method exists for all possible subtypes of the abstract type :thinking:

The same behavior appears without broadcasting as well:

h(v::Vector{Number}) = v[1] + v[2]

h(v)

@code_warntype h(v)
MethodInstance for h(::Vector{Number})
  from h(v::Vector{Number}) @ Main REPL[58]:1
Arguments
  #self#::Core.Const(h)
  v::Vector{Number}
Body::Any
1 ─ %1 = Base.getindex(v, 1)::Number
│   %2 = Base.getindex(v, 2)::Number
│   %3 = (%1 + %2)::Any
└──      return %3
3

Note that the return type of f(v) is always Int, no matter what the concrete types of the elements of v are. (This is because of the length(bc.args) in the definition of copy.)

If I define

g(x::T, y::T) = 2

then I get

julia> code_typed(g, (T,T))
1-element Vector{Any}:
 CodeInfo(
1 ─     return 2
) => Int64

that is, a concrete return type although I didn’t fix the concrete type of the input. Why is the same not possible for f?

4

Yes, but the compiler cannot seem to figure this out, because the vector eltype is abstract. So in the typed code, the arguments of the broadcast show up as being of an abstract type and also the broadcast object is not known to have a concrete type (it is still parametric, because the broadcast style could not be figured out). I don’t know why, but at some point the compiler gives up…

My point was just that this is (I think) not specific to broadcasting, as shown in the example with just + and Vector{Number} above. My guess would be that for sufficiently many or complicated methods (e.g. length and +), it would take too long to figure out what are the possible return types for abstract inputs.

In some cases, it is possible to figure out the right bounds at least, e.g. here

typeof example 
g(v::Vector{Number}) = eltype(first(v))

v = Number[1, 2]

@code_warntype g(v)
MethodInstance for g(::Vector{Number})
  from g(v::Vector{Number}) @ Main REPL[22]:1
Arguments
  #self#::Core.Const(g)
  v::Vector{Number}
Body::Type{<:Number}
1 ─ %1 = Main.first(v)::Number
│   %2 = Main.eltype(%1)::Type{<:Number}
└──      return %2

When the function is called, the arguments always have a concrete type. But the problem here is that the information is not yet available when the “outside” function f is called. It only sees a Vector{SomeAbstractType}, which is itself concrete. What method should be called inside the function f might have to be determined at runtime. When you define only one method which also returns a constant, like g(x::T, y::T) = 2 it should work, but for more complicated functions it generally doesn’t.

v = Number[1, 2]

# Completely inferred, since inside `+` we see that the arguments are `Int`
@code_warntype v[1] + v[2]
MethodInstance for +(::Int64, ::Int64)
  from +(x::T, y::T) where T<:Union{Int128, Int16, Int32, Int64, Int8, UInt128, UInt16, UInt32, UInt64, UInt8} @ Base int.jl:87
Static Parameters
  T = Int64
Arguments
  #self#::Core.Const(+)
  x::Int64
  y::Int64
Body::Int64
1 ─ %1 = Base.add_int(x, y)::Int64
└──      return %1

g(v::Vector{Number}) = v[1] + v[2]

# Cannot be inferred one level higher, since we only know that we have `Number`s in the vector.
# If the function (`+`) is sufficiently simple, it should work, but in general it is (presumably) impossible.
@code_warntype g(v)
MethodInstance for g(::Vector{Number})
  from g(v::Vector{Number}) @ Main REPL[143]:1
Arguments
  #self#::Core.Const(g)
  v::Vector{Number}
Body::Any
1 ─ %1 = Base.getindex(v, 1)::Number
│   %2 = Base.getindex(v, 2)::Number
│   %3 = (%1 + %2)::Any
└──      return %3

In principle it should be possible to figure out that there is only one possible method in the example from your original post, but I don’t know enough about the compiler/inference to know why it doesn’t work here.

5

It still seems to me the problem is that Julia thinks there are two applicable methods for BroadcastStyle:

julia> methods(BroadcastStyle, (Type{<:T},))
# 2 methods for type constructor:
 [1] BroadcastStyle(::Type{Union{}}, slurp...)
     @ broadcast.jl:38
 [2] BroadcastStyle(::Type{<:T})
     @ Main REPL[5]:1

I don’t know what the purpose of the first (pre-defined) method is, but here it might prevent Julia from inferring the correct type.