Broadcasting not type stable (but list comprehensions & map are)

[kaandocal]

The following code (thanks @ysfoo) shows that broadcasting can fail to be type stable, even when equivalent map calls/list comprehensions are:

using LinearAlgebra

function chol_func1(mats::Vector{<:Matrix})
    return [getproperty(cholesky(mat), :L) for mat in mats]    
end
    
function chol_func2(mats::Vector{<:Matrix})
    choleskies = cholesky.(mats)
    return getproperty.(choleskies, :L) 
end

function chol_func3(mats::Vector{<:Matrix})
    choleskies = cholesky.(mats)
    return map(getproperty, choleskies, :L) 
end

bmats = [ [ 1 0; 0 1 ], [ 2 0; 0 2 ] ]

@code_warntype chol_func1(bmats)
@code_warntype chol_func2(bmats)    # NOT TYPE STABLE
@code_warntype chol_func3(bmats)

Output:

MethodInstance for chol_func1(::Vector{Matrix{Int64}})
  from chol_func1(mats::Vector{<:Matrix})
Arguments
  #self#::Core.Const(chol_func1)
  mats::Vector{Matrix{Int64}}
Locals
  #36::var"#36#37"
Body::Vector{LowerTriangular{Float64, Matrix{Float64}}}
1 ─      (#36 = %new(Main.:(var"#36#37")))
│   %2 = #36::Core.Const(var"#36#37"())
│   %3 = Base.Generator(%2, mats)::Base.Generator{Vector{Matrix{Int64}}, var"#36#37"}
│   %4 = Base.collect(%3)::Vector{LowerTriangular{Float64, Matrix{Float64}}}
└──      return %4

MethodInstance for chol_func2(::Vector{Matrix{Int64}})
  from chol_func2(mats::Vector{<:Matrix})
Arguments
  #self#::Core.Const(chol_func2)
  mats::Vector{Matrix{Int64}}
Locals
  choleskies::Vector{Cholesky{Float64, Matrix{Float64}}}
Body::AbstractVector
1 ─ %1 = Base.broadcasted(Main.cholesky, mats)::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(cholesky), Tuple{Vector{Matrix{Int64}}}}
│        (choleskies = Base.materialize(%1))
│   %3 = Base.broadcasted(Main.getproperty, choleskies, :L)::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(getproperty), Tuple{Vector{Cholesky{Float64, Matrix{Float64}}}, Base.RefValue{Symbol}}}
│   %4 = Base.materialize(%3)::AbstractVector
└──      return %4                               ######### NOT TYPE STABLE

MethodInstance for chol_func3(::Vector{Matrix{Int64}})
  from chol_func3(mats::Vector{<:Matrix})
Arguments
  #self#::Core.Const(chol_func3)
  mats::Vector{Matrix{Int64}}
Locals
  choleskies::Vector{Cholesky{Float64, Matrix{Float64}}}
Body::Union{}
1 ─ %1 = Base.broadcasted(Main.cholesky, mats)::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(cholesky), Tuple{Vector{Matrix{Int64}}}}
│        (choleskies = Base.materialize(%1))
│        Main.map(Main.getproperty, choleskies, :L)
└──      Core.Const(:(return %3))

Version info:

Julia Version 1.10.3
Commit 0b4590a5507 (2024-04-30 10:59 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 72 × Intel(R) Xeon(R) Gold 6254 CPU @ 3.10GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, cascadelake)
Threads: 8 default, 0 interactive, 4 GC (on 72 virtual cores)
Environment:
  LD_LIBRARY_PATH = (...)
  JULIA_NUM_THREADS = 8

Changing :L to Ref(:L) doesn't make a difference for the first two examples. For the third (with map), using Ref(:L) loses type stability:

MethodInstance for chol_func3_ref(::Vector{Matrix{Int64}})
  from chol_func3(mats::Vector{<:Matrix})
Arguments
  #self#::Core.Const(chol_func3)
  mats::Vector{Matrix{Int64}}
Locals
  choleskies::Vector{Cholesky{Float64, Matrix{Float64}}}
Body::Vector
1 ─ %1 = Base.broadcasted(Main.cholesky, mats)::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(cholesky), Tuple{Vector{Matrix{Int64}}}}
│        (choleskies = Base.materialize(%1))
│   %3 = Main.getproperty::Core.Const(getproperty)
│   %4 = choleskies::Vector{Cholesky{Float64, Matrix{Float64}}}
│   %5 = Main.Ref(:L)::Base.RefValue{Symbol}
│   %6 = Main.map(%3, %4, %5)::Vector
└──      return %6                         ### NOT TYPE STABLE

[kaandocal]

@ysfoo just pointed out to me that chol_func3 with map doesn't actually run without Ref.

[mbauman]

I believe this is a duplicate of #43333

[mbauman]

Yeah, this is all about constant propagation through higher order functions. A simple workaround is to place the constant inside a function — "closer" to the getproperty call. These flavors are type stable:

julia> function chol_func2a(mats::Vector{<:Matrix})
           choleskies = cholesky.(mats)
           return (x->getproperty(x, :L)).(choleskies)
       end
chol_func2a (generic function with 2 methods)

julia> function chol_func3a(mats::Vector{<:Matrix})
           choleskies = cholesky.(mats)
           return map(x->getproperty(x,:L), choleskies)
       end
chol_func3a (generic function with 3 methods)