Type stability issues when broadcasting

General Usage Performance
, , ,
1

I am seeing some odd type instability issues when broadcasting.

Simple functions seem to become type-unstable when broadcasting over multiple views:

f1(x::AbstractVector{T}) where T = 1
f2(x::AbstractVector{T}) where T = x .+ 1

#@code_warntype f1(rand(2))                  # Type Stable
#@code_warntype f1.(eachcol(rand(2,2)))      # Type Stable
#@code_warntype f2(rand(2))                  # Type Stable
@code_warntype f2.(eachcol(rand(2,2)))      # Not Type Stable

Output:

MethodInstance for (::var"##dotfunction#358#54")(::Base.Generator{Base.OneTo{Int64}, Base.var"#242#243"{Matrix{Float64}}})
  from (::var"##dotfunction#358#54")(x1) in Main
Arguments
  #self#::Core.Const(var"##dotfunction#358#54"())
  x1::Base.Generator{Base.OneTo{Int64}, Base.var"#242#243"{Matrix{Float64}}}
Body::Any
1 ─ %1 = Base.broadcasted(Main.f2, x1)::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(f2), Tuple{Vector{SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}}
│   %2 = Base.materialize(%1)::Any
└──      return %2

Broadcasting also appears to be type-unstable for any function that is called as member of a struct, even when the struct is parametrically typed to the function

struct MyStruct{F<:Function}
    f::F
end

f3(x) = 1
s3 = MyStruct(f3)

#@code_warntype s3.f(1)                         # Type Stable
@code_warntype s3.f.([1,2,3])                  # Not type Stable

Output

MethodInstance for (::var"##dotfunction#357#53")(::Vector{Int64})
  from (::var"##dotfunction#357#53")(x1) in Main
Arguments
  #self#::Core.Const(var"##dotfunction#357#53"())
  x1::Vector{Int64}
Body::Any
1 ─ %1 = Base.getproperty(Main.s, :f)::Any
│   %2 = Base.broadcasted(%1, x1)::Any
│   %3 = Base.materialize(%2)::Any
└──      return %3

Not sure if these two examples are related or reflect separate issues.

2 Likes
2

Hilariously, the type is properly inferred when you do

julia> f3(x::AbstractArray{T,N}) where {T,N} = x .+ 1
f3 (generic function with 1 method)

julia> @code_warntype f3.(eachcol(rand(2,2)))
MethodInstance for (::var"##dotfunction#341#7")(::Base.Generator{Base.OneTo{Int64}, Base.var"#242#243"{Matrix{Float64}}})
  from (::var"##dotfunction#341#7")(x1) in Main
Arguments
  #self#::Core.Const(var"##dotfunction#341#7"())
  x1::Base.Generator{Base.OneTo{Int64}, Base.var"#242#243"{Matrix{Float64}}}
Body::Vector{Vector{Float64}}
1 ─ %1 = Base.broadcasted(Main.f3, x1)::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(f3), Tuple{Vector{SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}}
│   %2 = Base.materialize(%1)::Vector{Vector{Float64}}
└──      return %2

I’m sure there is a very good reason I too would love to understand.

3

The second example is resolved when nested within another function

julia> fs3(s3) = s3.f.([1,2,3])
fs3 (generic function with 1 method)
julia> @code_warntype fs3(s3)
  # stable

so is likely an artifact. Note that code_warntype is not entirely trustworthy. It takes a rather naive view of how things work and often neglects certain exceptions that occur in practice (I recall Be-aware-of-when-Julia-avoids-specializing being one such example). Nesting within functions sometimes helps.

The first example is not resolved, however:

julia> ff2(x) = f2.(eachcol(x))
ff2 (generic function with 1 method)

julia> @code_warntype ff2(rand(2,2))
  # still unstable

Someone else will have to elaborate on this case. I don’t fully understand it.

4

I thought f2.(x) is equivalent to Base.materialize(Base.broadcasted(f2, x)):

julia> @code_lowered f2.(eachcol(rand(2,2)))
CodeInfo(
1 ─ %1 = Base.broadcasted(Main.f2, x1)
│   %2 = Base.materialize(%1)
└──      return %2
)

Strangely, the direct execution of this seemingly equivalent code is type-stable:

julia> @code_warntype Base.materialize(Base.broadcasted(f2, eachcol(rand(2,2))))
MethodInstance for Base.Broadcast.materialize(::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(f2), Tuple{ColumnSlices{Matrix{Float64}, Tuple{Base.OneTo{Int64}}, SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}})
  from materialize(bc::Base.Broadcast.Broadcasted) @ Base.Broadcast broadcast.jl:873
Arguments
  #self#::Core.Const(Base.Broadcast.materialize)
  bc::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(f2), Tuple{ColumnSlices{Matrix{Float64}, Tuple{Base.OneTo{Int64}}, SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}}
Body::Vector{Vector{Float64}}
1 ─      nothing
│   %2 = Base.Broadcast.instantiate(bc)::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Tuple{Base.OneTo{Int64}}, typeof(f2), Tuple{ColumnSlices{Matrix{Float64}, Tuple{Base.OneTo{Int64}}, SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}}}
│   %3 = Base.Broadcast.copy(%2)::Vector{Vector{Float64}}
└──      return %3
5

It is if you try this:

julia> ff2(x) = f2.(eachcol(x))
ff2 (generic function with 1 method)

julia> ff2(rand(2,2)); # run it at least once

julia> @code_warntype ff2(rand(2,2))
  # stable

same.