Type instability in pinv() for non-square matrix

[dcabecinhas]

Using a non-square matrix leads to a type instability in pinv. The operation is type-stable for standard arrays.

A = @SMatrix rand(3,4)

julia> @code_warntype pinv(A)
Variables
  #self#::Core.Compiler.Const(LinearAlgebra.pinv, false)
  A::SArray{Tuple{3,4},Float64,2,12}
  atol::Float64
  rtol::Float64

Body::Any
1 ─       (atol = 0.0)
│   %2  = LinearAlgebra.one($(Expr(:static_parameter, 1)))::Core.Compiler.Const(1.0, false)
│   %3  = LinearAlgebra.float(%2)::Core.Compiler.Const(1.0, false)
│   %4  = LinearAlgebra.real(%3)::Core.Compiler.Const(1.0, false)
│   %5  = LinearAlgebra.eps(%4)::Core.Compiler.Const(2.220446049250313e-16, false)
│   %6  = LinearAlgebra.size(A)::Core.Compiler.Const((3, 4), false)
│   %7  = Core._apply_iterate(Base.iterate, LinearAlgebra.min, %6)::Core.Compiler.Const(3, false)
│   %8  = (%5 * %7)::Core.Compiler.Const(6.661338147750939e-16, false)
│   %9  = LinearAlgebra.iszero(atol::Core.Compiler.Const(0.0, false))::Core.Compiler.Const(true, false)
│         (rtol = %8 * %9)
│   %11 = LinearAlgebra.:(var"#pinv#37")(atol::Core.Compiler.Const(0.0, false), rtol::Core.Compiler.Const(6.661338147750939e-16, false), #self#, A)::Any
└──       return %11

julia> @code_warntype pinv(Array(A))
Variables
  #self#::Core.Compiler.Const(LinearAlgebra.pinv, false)
  A::Array{Float64,2}
  atol::Float64
  rtol::Float64

Body::Array{Float64,2}
1 ─       (atol = 0.0)
│   %2  = LinearAlgebra.one($(Expr(:static_parameter, 1)))::Core.Compiler.Const(1.0, false)
│   %3  = LinearAlgebra.float(%2)::Core.Compiler.Const(1.0, false)
│   %4  = LinearAlgebra.real(%3)::Core.Compiler.Const(1.0, false)
│   %5  = LinearAlgebra.eps(%4)::Core.Compiler.Const(2.220446049250313e-16, false)
│   %6  = LinearAlgebra.size(A)::Tuple{Int64,Int64}
│   %7  = Core._apply_iterate(Base.iterate, LinearAlgebra.min, %6)::Int64
│   %8  = (%5 * %7)::Float64
│   %9  = LinearAlgebra.iszero(atol::Core.Compiler.Const(0.0, false))::Core.Compiler.Const(true, false)
│         (rtol = %8 * %9)
│   %11 = LinearAlgebra.:(var"#pinv#37")(atol::Core.Compiler.Const(0.0, false), rtol, #self#, A)::Array{Float64,2}
└──       return %11

[c42f]

Looking at the implementation of LinearAlgebra.pinv (the function called var"#pinv#37" above), I believe this is due to the compiler failing to constant propagate the keyword argument full=false in the call to SVD when A is non-square.

Here's a more minimal reproduction of that issue:

julia> using StaticArrays, LinearAlgebra

julia> foo(A) = svd(A, full=false)
foo (generic function with 1 method)

julia> @code_warntype foo(SA[1 2; 3 4; 5 6])
Variables
  #self#::Core.Compiler.Const(foo, false)
  A::SArray{Tuple{3,2},Int64,2,6}

Body::Union{StaticArrays.SVD{Float64,SArray{Tuple{3,2},Float64,2,6},SArray{Tuple{2},Float64,1,2},SArray{Tuple{2,2},Float64,2,4}}, StaticArrays.SVD{Float64,SArray{Tuple{3,3},Float64,2,9},SArray{Tuple{2},Float64,1,2},SArray{Tuple{2,2},Float64,2,4}}}
1 ─ %1 = (:full,)::Core.Compiler.Const((:full,), false)
│   %2 = Core.apply_type(Core.NamedTuple, %1)::Core.Compiler.Const(NamedTuple{(:full,),T} where T<:Tuple, false)
│   %3 = Core.tuple(false)::Core.Compiler.Const((false,), false)
│   %4 = (%2)(%3)::Core.Compiler.Const((full = false,), false)
│   %5 = Core.kwfunc(Main.svd)::Core.Compiler.Const(LinearAlgebra.var"#svd##kw"(), false)
│   %6 = (%5)(%4, Main.svd, A)::Union{StaticArrays.SVD{Float64,SArray{Tuple{3,2},Float64,2,6},SArray{Tuple{2},Float64,1,2},SArray{Tuple{2,2},Float64,2,4}}, StaticArrays.SVD{Float64,SArray{Tuple{3,3},Float64,2,9},SArray{Tuple{2},Float64,1,2},SArray{Tuple{2,2},Float64,2,4}}}
└──      return %6

This union is the killer and it propagates through the rest of the pinv code causing all sorts of type instability.

On the other hand, if full=Val(false) is used, we get no type instability

julia> bar(A) = svd(A, full=Val(false))

[c42f]

Options for StaticArrays include:

• Copy the Base implementation of pinv into our code, with a note as to why we're copying it (ugh! But would fix the problem immediately.)
• Fix this upstream, by investigating why constant propagation and inlining don't work in this case.

[c42f]

Fix this upstream

Specifically, this problem can be seen in the above code_warntype — the kwfunc is not inlined even though it would expose significant optimization opportunities.

Actually, if I mark one of the internal _svd explicitly as @inline this seems to fix the problem on julia-1.5-rc1 but not on earlier versions.

[c42f]

Ok, #805 should fix this on Julia 1.5, though pinv still won't return a StaticMatrix - that's a separate issue though.

Will that work for you once 1.5 is out, or do we need a more comprehensive workaround?

[andyferris]

It would be really helpful if lowering made the extra functions it makes for keyword arguments always @propagate_inbounds. You’d get these constant propagation improvements and could use bounds check elision on keyword functions.

[dcabecinhas]

Ok, #805 should fix this on Julia 1.5, though pinv still won't return a StaticMatrix - that's a separate issue though.

Will that work for you once 1.5 is out, or do we need a more comprehensive workaround?

That will do, I'm on 1.5 already. Thanks for the quick turnaround.

[c42f]

Nice!