specialize copyto! for Q from qr

This fixes some performance problems reported in #38972 and #37102 (multiplying `Q` by a `Diagonal` or `UniformScaling`).
ranocha · Feb 5, 2021 · c8e8b20 · c8e8b20
1 parent 0ea19e4
commit c8e8b20
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/stdlib/LinearAlgebra/src/qr.jl b/stdlib/LinearAlgebra/src/qr.jl
@@ -533,6 +533,12 @@ function getindex(Q::AbstractQ, i::Integer, j::Integer)
     return dot(x, lmul!(Q, y))
 end
 
+# specialization avoiding the fallback using slow `getindex`
+function copyto!(dest::AbstractMatrix, src::LinearAlgebra.AbstractQ)
+    copyto!(dest, I)
+    lmul!(src, dest)
+end
+
 ## Multiplication by Q
 ### QB
 lmul!(A::QRCompactWYQ{T,S}, B::StridedVecOrMat{T}) where {T<:BlasFloat, S<:StridedMatrix} =

diff --git a/stdlib/LinearAlgebra/test/qr.jl b/stdlib/LinearAlgebra/test/qr.jl
@@ -322,4 +322,13 @@ end
     end
 end
 
+@testset "QR factorization of Q" begin
+    for T in (Float32, Float64, ComplexF32, ComplexF64)
+        Q1, R1 = qr(randn(T,5,5))
+        Q2, R2 = qr(Q1)
+        @test Q1 ≈ Q2
+        @test R2 ≈ I
+    end
+end
+
 end # module TestQR