`getindex` is type-unstable for `Diagonal{<:AbstractMatrix}`

[jishnub]

julia> D = Diagonal([Diagonal(1:3), Diagonal(1:4)])
2×2 Diagonal{Diagonal{Int64, UnitRange{Int64}}, Vector{Diagonal{Int64, UnitRange{Int64}}}}:
 [1 0 0; 0 2 0; 0 0 3]                      ⋅             
              ⋅                [1 0 0 0; 0 2 0 0; 0 0 3 0; 0 0 0 4]

julia> @code_warntype D[1,1]
MethodInstance for getindex(::Diagonal{Diagonal{Int64, UnitRange{Int64}}, Vector{Diagonal{Int64, UnitRange{Int64}}}}, ::Int64, ::Int64)
  from getindex(D::Diagonal, i::Int64, j::Int64) in LinearAlgebra at /home/jishnu/packages/julias/julia-latest/share/julia/stdlib/v1.9/LinearAlgebra/src/diagonal.jl:112
Arguments
  #self#::Core.Const(getindex)
  D::Diagonal{Diagonal{Int64, UnitRange{Int64}}, Vector{Diagonal{Int64, UnitRange{Int64}}}}
  i::Int64
  j::Int64
Locals
  val::Diagonal{Int64, UnitRange{Int64}}
  r::Union{Diagonal{Int64, UnitRange{Int64}}, Matrix{Int64}}
Body::Union{Diagonal{Int64, UnitRange{Int64}}, Matrix{Int64}}
1 ─       nothing
│         Core.NewvarNode(:(val))
│         Core.NewvarNode(:(r))
└──       goto JuliaLang/julia#3 if not $(Expr(:boundscheck))
2 ─       LinearAlgebra.checkbounds(D, i, j)
3 ┄ %6  = (i == j)::Bool
└──       goto JuliaLang/julia#5 if not %6
4 ─       nothing
│   %9  = Base.getproperty(D, :diag)::Vector{Diagonal{Int64, UnitRange{Int64}}}
│   %10 = Base.getindex(%9, i)::Diagonal{Int64, UnitRange{Int64}}
│         (r = %10)
│         (val = %10)
│         nothing
│         val
└──       goto JuliaLang/julia#6
5 ─       (r = LinearAlgebra.diagzero(D, i, j))
6 ┄       return r

One solution to this might be to coerce the types of the diagonal elements to that of diagzero, or vice versa. In this former case, one may preserve the eltype using diag(D)[i] instead of D[i,i], although this involves one (small) allocation.

julia> D = Diagonal([Diagonal(zeros(10000)), Diagonal(zeros(20000))]);

julia> @btime diag($D);
  44.823 ns (1 allocation: 80 bytes)

On the other hand, this might have performance implications, as the diagonal elements will be copied on each access. The latter case might be trickier if the conversions aren't defined, although it's arguably a better approach.

Lines of code:
https://github.com/JuliaLang/julia/blob/d84f8901aefbb5d7ae0a4d591812e351af6866ca/stdlib/LinearAlgebra/src/diagonal.jl#L112-L122

[dkarrasch]

I don't think this can be reasonably resolved. Reading a diagonal element should just read D.diag. Reading an off-diagonal element may not be representable by the same type, as in your case. But, OTOH, this is not a very bad type-instability, it's just a union of two concrete types, which should be handled well by union splitting?

[N5N3]

Union splitting seems OK. But eltype should be Union{Diagonal{Int64, UnitRange{Int64}}, Matrix{Int64}}? (On master it's Diagonal{Int64, UnitRange{Int64}})

[KristofferC]

There could also be a completely different function, e.g. getdiagelement or something which should be inferrable.

So

getdiagelement(A::AbstractMatrix, i::Integer) = A[i, i]
getdiagelement(A::Diagonal, i::Integer) = ....