Pinv not Type Stable?

General Usage Performance
1

Why is LinearAlgebra.pinv not type stable? Can it be implemented in a type stable way?

Example (thanks to @albheim)

using LinearAlgebra
using JET

function test(M)
    A = rand(M, M + 2)
    pinv(A)
end

test(10)
report = @report_opt test(10)
println(report)

…
││┌ pinv(A::Matrix{Float64}; atol::Float64, rtol::Float64) @ LinearAlgebra /Users/julia/.julia/scratchspaces/a66863c6-20e8-4ff4-8a62-49f30b1f605e/agent-cache/default-honeycrisp-R17H3W25T9.0/build/default-honeycrisp-R17H3W25T9-0/julialang/julia-release-1-dot-9/usr/share/julia/stdlib/v1.9/LinearAlgebra/src/dense.jl:1490
│││ runtime dispatch detected: Base.materialize!(%405::Any, %408::Base.Broadcast.Broadcasted{_A, Nothing, typeof(pinv)} where _A<:Union{Nothing, Base.Broadcast.BroadcastStyle})::Any
││└────────────────────

2

Do you have an example where it is not type stable?
My test with a normal floating point matrix seems to be stable according to @code_warntype, and I’m not doubting you, it just makes it much easier to investigate with a concrete example since then we don’t have to hunt for what specific type creates an instability.

3

Amended my question. Thanks for taking a look. Could it be a false positive by JET?

4

This is what you mean with not being “type stable” I guess. This is what’s reported by @code_warntype

julia> function test(M)
         A = rand(M, M+2)
         pinv(A)
       end
test (generic function with 1 method)

julia> @code_warntype test(10)
MethodInstance for test(::Int64)
  from test(M) @ Main REPL[10]:1
Arguments
  #self#::Core.Const(test)
  M::Int64
Locals
  A::Matrix{Float64}
Body::Matrix{Float64}
1 ─ %1 = (M + 2)::Int64
│        (A = Main.rand(M, %1))
│   %3 = Main.pinv(A)::Matrix{Float64}
└──      return %3

EDIT: the thought below turned out to be wrong :wink:

The problem with the runtime dispatch is the following: you are passing an integer which is then used to generate a matrix with dimensions based on the integer’s value. This way, the compiler cannot know in advance the method call for pinv. In other words: the method call and even return value depend on the actual numerical value of M, therefore the method call will be decided during runtime (called runtime dispatch).

5

Why not? What’s going on inside pinv to cause this, seeing as size is not part of the type?

1 Like
6

I think it’s rand(), which will create an M x M+2 matrix, so based on the value of M, the return type will be different and that is passed to pinv. But maybe I have a misunderstanding :wink:

7

I belive you do. Rand is always of type ‘Matrix{Float64}’. I also tested just rand, and that seemed type stable. The issue also crept up in other code, that decidedly didn’t use rand.

Adding to that the size of a standard Julia Array is not part of its type. Only it’s dimensionality is, which is 2 for Matrix.

1 Like
8

Yes I guess you are right, I was on the wrong path :wink: (let me add that above to avoid confusion). After so many years of Julia, I still get confused by runtime dispatch stuff :laughing:

1 Like
9

Methods with keyword arguments are implemented to call hidden helper methods that have the written body, so @code_warntype is really difficult to use with them. Well, I dug out the underlying call and am including the part where things get a little unstable, even though the internal arrays and the return value are inferred concretely.

julia> @code_warntype LinearAlgebra.:(var"#pinv#33")(0.0, (eps(real(float(1.1)))*min(3, 5))*iszero(0.0), pinv,rand(3, 5))
MethodInstance for LinearAlgebra.var"#pinv#33"(::Float64, ::Float64, ::typeof(pinv), ::Matrix{Float64})
  from var"#pinv#33"(atol::Real, rtol::Real, ::typeof(pinv), A::AbstractMatrix{T}) where T in LinearAlgebra at /nix/store/4dkh426qldnal1n4m2czlhrmi4qf9ri5-julia-bin-1.8.3/share/julia/stdlib/v1.8/LinearAlgebra/src/dense.jl:1445
...
Locals
...
  index::Any
...
  tol@_14::Core.Box
...
  tol@_21::Union{}
...
│    %40 = (rtol * maxabsA)::Float64
│    %41 = LinearAlgebra.max(%40, atol)::Float64
│          Core.setfield!(tol@_14, :contents, %41)
...
│    %55 = Base.getproperty(SVD, :S)::Vector{Float64}
│    %56 = LinearAlgebra.maximum(%55)::Float64
│    %57 = (rtol * %56)::Float64
│    %58 = LinearAlgebra.max(%57, atol)::Float64
│          Core.setfield!(tol@_14, :contents, %58)
...
│    %68 = Core.isdefined(tol@_14, :contents)::Bool
└───       goto #9 if not %68
8 ──       goto #10
9 ──       Core.NewvarNode(:(tol@_21))
└───       tol@_21
10 ┄ %73 = Core.getfield(tol@_14, :contents)::Any
│    %74 = Base.broadcasted(LinearAlgebra.:>, %67, %73)::Any
│          (index = Base.materialize(%74))
│    %76 = Base.dotview(Sinv, index)::Any
│    %77 = Base.getproperty(SVD, :S)::Vector{Float64}
│    %78 = LinearAlgebra.view(%77, index)::Any
│    %79 = Base.broadcasted(LinearAlgebra.pinv, %78)::Base.Broadcast.Broadcasted{_A, Nothing, typeof(pinv)} where _A<:Union{Nothing, Base.Broadcast.BroadcastStyle}
│          Base.materialize!(%76, %79)
...

The corresponding source code section is:

        tol = max(rtol * maxabsA, atol)
...
    tol         = max(rtol*maximum(SVD.S), atol)
...
    index       = SVD.S .> tol
    Sinv[index] .= pinv.(view(SVD.S, index))

It seems as if tol is an uninferred boxed value (Core.Box) and it’s not even certain if it is defined (%68: tol@_14 vs tol@_21). I don’t really know why because in both lines (%41, %58) the final max call is inferred concretely ::Float64 before being stored in the box.

5 Likes
10

pinv reassigns tol after it is captured by a closure, I think that’s the reason. Closures often ruin inferrability because it’s understandably hard to infer a variable shared by 2 separate functions that get compiled separately. I don’t know all the cases where it is inferred (guaranteed reassignment before closure instantiation) and where it is not (conditional reassignment before closure, any reassignment after the closure).

I @evaled a pinv2 into LinearAlgebra that replaced the closure with a function that doesn’t capture tol: B[indB] .= ((x, toll) -> abs(x) > toll ? pinv(x) : zero(x)).(dA, tol). I haven’t compared the performance of these anonymous functions in isolation, so I don’t think it’s optimal. But the box vanished:

julia> @code_warntype LinearAlgebra.:(var"#pinv2#251")(0.0, (eps(real(float(1.1)))*min(3, 5))*iszero(0.0), LinearAlgebra.pinv2,rand(3, 5))
MethodInstance for LinearAlgebra.var"#pinv2#251"(::Float64, ::Float64, ::typeof(LinearAlgebra.pinv2), ::Matrix{Float64})
  from var"#pinv2#251"(atol::Real, rtol::Real, ::typeof(LinearAlgebra.pinv2), A::AbstractMatrix{T}) where T in LinearAlgebra at REPL[16]:1
...
Locals
...
  index::BitVector
...
  tol::Float64

EDIT: the least disruptive change to accomplish the same inferability would be to keep the closure, but alter the captured variable so it’s not reassigned afterward.

        tol_diag_branch = max(rtol * maxabsA, atol)
...
        B[indB] .= (x -> abs(x) > tol_diag_branch ? pinv(x) : zero(x)).(dA)
        # do not reassign tol_diag_branch, so parser lets compiler infer it
3 Likes
11

Does this have performance implications? I assume not really since it’s doing only one dispatch per call.

12

Not really; computing the pinv — even for very small arrays — is orders of magnitude more work than the cost of this internal instability. Where it might have an impact is in static compilation and potentially being a source for triggering invalidations, so it’s still a good thing to address. @Benny you should post that change as a pull request! :slight_smile:

The TL;DR of why this is unstable is:

  • tol is assigned to twice in a slightly complicated manner — once in an if branch and once outside.
  • tol is also captured inside a closure, which theoretically has to support getting an updated value from the second assignment. In this case it’s immediately evaluated (and is in a completely independent code path), but the front-end isn’t grokking that.

So this hits the infamous issue #15276. An even simpler fix is to use two completely separate names for tol for the two different codepaths.

3 Likes
13

Thanks, I learned quite a lot in this thread. I’ll run my own implementation for now since I don’t expect myself to hand me any diagonal matrices.

1 Like
14

Yh it would be useful to get this fixed as it seems to cause pinv to not work on the GPU, Support for LinearAlgebra.pinv · Issue #2070 · JuliaGPU/CUDA.jl · GitHub.

2 Likes
15

So I ran some benchmarks with the changes @mbauman suggested, and types do indeed get inferred correctly now. The two implementations are indistinguishable from a performance perspective.

"Type Stable Pinv"

BenchmarkTools.Trial: 4059 samples with 1 evaluation.
 Range (min … max):  1.138 ms …   6.104 ms  ┊ GC (min … max): 0.00% … 79.07%
 Time  (median):     1.177 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   1.231 ms ± 356.126 μs  ┊ GC (mean ± σ):  3.26% ±  8.18%

  █▆▁                                                         ▁
  ██████▇▇▄▅▅▅▅▄▄▄▄▃▄▃▁▁▄▁▁▃▁▁▁▁▁▁▁▁▁▃▃▁▁▁▁▁▃▁▁▁▁▁▁▁▁▁▁▃▅▁▅▃▅ █
  1.14 ms      Histogram: log(frequency) by time      3.71 ms <

 Memory estimate: 645.53 KiB, allocs estimate: 19.
--------------------------------------------------------------------------
"LinearAlgebra.pinv"

BenchmarkTools.Trial: 4008 samples with 1 evaluation.
 Range (min … max):  1.143 ms …   7.397 ms  ┊ GC (min … max): 0.00% … 69.65%
 Time  (median):     1.178 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   1.246 ms ± 379.645 μs  ┊ GC (mean ± σ):  3.12% ±  8.04%

  █▅                                                           
  ███▇█▇▇▆▇▇▆▅▆▄▅▃▅▄▄▁▃▄▃▄▄▄▄▄▄▄▃▃▁▁▃▃▃▃▁▃▃▁▃▃▁▁▄▁▃▁▁▄▃▄▄▅▄▄▅ █
  1.14 ms      Histogram: log(frequency) by time      3.67 ms <

 Memory estimate: 645.81 KiB, allocs estimate: 27.
2 Likes
16

The fix has been merged into master, Remove boxing in pinv by Zentrik · Pull Request #51351 · JuliaLang/julia · GitHub.

5 Likes