add LQViaTransposedQR, CUDA LQ and tests

Author: Jutho

Project.toml

@@ -18,6 +18,7 @@ MatrixAlgebraKitCUDAExt = "CUDA"
 Aqua = "0.6, 0.7, 0.8"
 ChainRulesCore = "1"
 ChainRulesTestUtils = "1"
+CUDA = "5"
 JET = "0.9"
 LinearAlgebra = "1"
 SafeTestsets = "0.1"

ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl

@@ -4,11 +4,13 @@ using MatrixAlgebraKit
 using MatrixAlgebraKit: @algdef, Algorithm, check_input
 using MatrixAlgebraKit: one!, zero!, uppertriangular!, lowertriangular!
 using MatrixAlgebraKit: diagview, sign_safe
+using MatrixAlgebraKit: LQViaTransposedQR
 using CUDA
 using LinearAlgebra
 using LinearAlgebra: BlasFloat
 
 include("yacusolver.jl")
 include("implementations/qr.jl")
+include("implementations/lq.jl")
 
 end

ext/MatrixAlgebraKitCUDAExt/implementations/lq.jl

@@ -0,0 +1,4 @@
+function MatrixAlgebraKit.default_lq_algorithm(A::CuMatrix{<:BlasFloat}; kwargs...)
+    qr_alg = CUSOLVER_HouseholderQR(; kwargs...)
+    return LQViaTransposedQR(qr_alg)
+end

ext/MatrixAlgebraKitCUDAExt/implementations/qr.jl

@@ -11,31 +11,6 @@ function MatrixAlgebraKit.default_qr_algorithm(A::CuMatrix{<:BlasFloat}; kwargs.
     return CUSOLVER_HouseholderQR(; kwargs...)
 end
 
-# Outputs
-# -------
-function MatrixAlgebraKit.initialize_output(::typeof(qr_full!), A::AbstractMatrix,
-                                            ::CUSOLVER_HouseholderQR)
-    m, n = size(A)
-    Q = similar(A, (m, m))
-    R = similar(A, (m, n))
-    return (Q, R)
-end
-function MatrixAlgebraKit.initialize_output(::typeof(qr_compact!), A::AbstractMatrix,
-                                            ::CUSOLVER_HouseholderQR)
-    m, n = size(A)
-    minmn = min(m, n)
-    Q = similar(A, (m, minmn))
-    R = similar(A, (minmn, n))
-    return (Q, R)
-end
-function MatrixAlgebraKit.initialize_output(::typeof(qr_null!), A::AbstractMatrix,
-                                            ::CUSOLVER_HouseholderQR)
-    m, n = size(A)
-    minmn = min(m, n)
-    N = similar(A, (m, m - minmn))
-    return N
-end
-
 # Implementation
 # --------------
 # actual implementation
@@ -80,23 +55,23 @@ function _cusolver_qr!(A::AbstractMatrix, Q::AbstractMatrix, R::AbstractMatrix;
     if positive # already fix Q even if we do not need R
         # TODO: report that `lmul!` and `rmul!` with `Diagonal` don't work with CUDA
         τ .= sign_safe.(diagview(A))
-        Q .= Q .* transpose(τ)
+        Qf = view(Q, 1:m, 1:minmn) # first minmn columns of Q
+        Qf .= Qf .* transpose(τ)
     end
 
     if computeR
         R̃ = uppertriangular!(view(A, axes(R)...))
         if positive
-            R̃ .= conj.(τ) .* R̃
+            R̃f = view(R̃, 1:minmn, 1:n) # first minmn rows of R
+            R̃f .= conj.(τ) .* R̃f
         end
         copyto!(R, R̃)
     end
     return Q, R
 end
 
 function _cusolver_qr_null!(A::AbstractMatrix, N::AbstractMatrix;
-                            positive=false,
-                            pivoted=false,
-                            blocksize=1)
+                            positive=false, blocksize=1)
     blocksize > 1 &&
         throw(ArgumentError("CUSOLVER does not provide a blocked implementation for a QR decomposition"))
     m, n = size(A)

src/MatrixAlgebraKit.jl

@@ -2,12 +2,12 @@ module MatrixAlgebraKit
 
 using LinearAlgebra: LinearAlgebra
 using LinearAlgebra: norm # TODO: eleminate if we use VectorInterface.jl?
-using LinearAlgebra: mul!, rmul!, lmul!
+using LinearAlgebra: mul!, rmul!, lmul!, adjoint!, rdiv!, ldiv!
 using LinearAlgebra: sylvester
 using LinearAlgebra: isposdef, ishermitian
 using LinearAlgebra: Diagonal, diag, diagind
 using LinearAlgebra: UpperTriangular, LowerTriangular
-using LinearAlgebra: BlasFloat, BlasReal, BlasComplex, BlasInt, triu!, tril!, rdiv!, ldiv!
+using LinearAlgebra: BlasFloat, BlasReal, BlasComplex, BlasInt
 
 export qr_compact, qr_full, qr_null, lq_compact, lq_full, lq_null
 export qr_compact!, qr_full!, qr_null!, lq_compact!, lq_full!, lq_null!

src/algorithms.jl

@@ -98,7 +98,7 @@ macro algdef(name)
                 return $name{typeof(kw)}(kw)
             end
             function Base.show(io::IO, alg::$name)
-                return _show_alg(io, alg)
+                return ($_show_alg)(io, alg)
             end
 
             Core.@__doc__ $name

src/implementations/eig.jl

@@ -30,14 +30,14 @@ end
 
 # Outputs
 # -------
-function initialize_output(::typeof(eig_full!), A::AbstractMatrix, ::LAPACK_EigAlgorithm)
+function initialize_output(::typeof(eig_full!), A::AbstractMatrix, ::AbstractAlgorithm)
     n = size(A, 1) # square check will happen later
     Tc = complex(eltype(A))
     D = Diagonal(similar(A, Tc, n))
     V = similar(A, Tc, (n, n))
     return (D, V)
 end
-function initialize_output(::typeof(eig_vals!), A::AbstractMatrix, ::LAPACK_EigAlgorithm)
+function initialize_output(::typeof(eig_vals!), A::AbstractMatrix, ::AbstractAlgorithm)
     n = size(A, 1) # square check will happen later
     Tc = complex(eltype(A))
     D = similar(A, Tc, n)

src/implementations/eigh.jl

@@ -29,13 +29,13 @@ end
 
 # Outputs
 # -------
-function initialize_output(::typeof(eigh_full!), A::AbstractMatrix, ::LAPACK_EighAlgorithm)
+function initialize_output(::typeof(eigh_full!), A::AbstractMatrix, ::AbstractAlgorithm)
     n = size(A, 1) # square check will happen later
     D = Diagonal(similar(A, real(eltype(A)), n))
     V = similar(A, (n, n))
     return (D, V)
 end
-function initialize_output(::typeof(eigh_vals!), A::AbstractMatrix, ::LAPACK_EighAlgorithm)
+function initialize_output(::typeof(eigh_vals!), A::AbstractMatrix, ::AbstractAlgorithm)
     n = size(A, 1) # square check will happen later
     D = similar(A, real(eltype(A)), n)
     return D

src/implementations/lq.jl

@@ -42,20 +42,20 @@ end
 
 # Outputs
 # -------
-function initialize_output(::typeof(lq_full!), A::AbstractMatrix, ::LAPACK_HouseholderLQ)
+function initialize_output(::typeof(lq_full!), A::AbstractMatrix, ::AbstractAlgorithm)
     m, n = size(A)
     L = similar(A, (m, n))
     Q = similar(A, (n, n))
     return (L, Q)
 end
-function initialize_output(::typeof(lq_compact!), A::AbstractMatrix, ::LAPACK_HouseholderLQ)
+function initialize_output(::typeof(lq_compact!), A::AbstractMatrix, ::AbstractAlgorithm)
     m, n = size(A)
     minmn = min(m, n)
     L = similar(A, (m, minmn))
     Q = similar(A, (minmn, n))
     return (L, Q)
 end
-function initialize_output(::typeof(lq_null!), A::AbstractMatrix, ::LAPACK_HouseholderLQ)
+function initialize_output(::typeof(lq_null!), A::AbstractMatrix, ::AbstractAlgorithm)
     m, n = size(A)
     minmn = min(m, n)
     Nᴴ = similar(A, (n - minmn, n))
@@ -71,17 +71,34 @@ function lq_full!(A::AbstractMatrix, LQ, alg::LAPACK_HouseholderLQ)
     _lapack_lq!(A, L, Q; alg.kwargs...)
     return L, Q
 end
+function lq_full!(A::AbstractMatrix, LQ, alg::LQViaTransposedQR)
+    check_input(lq_full!, A, LQ)
+    L, Q = LQ
+    lq_via_qr!(A, L, Q, alg.qr_alg)
+    return L, Q
+end
 function lq_compact!(A::AbstractMatrix, LQ, alg::LAPACK_HouseholderLQ)
     check_input(lq_compact!, A, LQ)
     L, Q = LQ
     _lapack_lq!(A, L, Q; alg.kwargs...)
     return L, Q
 end
+function lq_compact!(A::AbstractMatrix, LQ, alg::LQViaTransposedQR)
+    check_input(lq_compact!, A, LQ)
+    L, Q = LQ
+    lq_via_qr!(A, L, Q, alg.qr_alg)
+    return L, Q
+end
 function lq_null!(A::AbstractMatrix, Nᴴ, alg::LAPACK_HouseholderLQ)
     check_input(lq_null!, A, Nᴴ)
     _lapack_lq_null!(A, Nᴴ; alg.kwargs...)
     return Nᴴ
 end
+function lq_null!(A::AbstractMatrix, Nᴴ, alg::LQViaTransposedQR)
+    check_input(lq_null!, A, Nᴴ)
+    lq_null_via_qr!(A, Nᴴ, alg.qr_alg)
+    return Nᴴ
+end
 
 function _lapack_lq!(A::AbstractMatrix, L::AbstractMatrix, Q::AbstractMatrix;
                      positive=false,
@@ -158,3 +175,31 @@ function _lapack_lq_null!(A::AbstractMatrix, Nᴴ::AbstractMatrix;
     end
     return Nᴴ
 end
+
+# LQ via transposition and QR
+function lq_via_qr!(A::AbstractMatrix, L::AbstractMatrix, Q::AbstractMatrix,
+                    qr_alg::AbstractAlgorithm)
+    m, n = size(A)
+    minmn = min(m, n)
+    At = adjoint!(similar(A'), A)
+    Qt = (A === Q) ? At : similar(Q')
+    Lt = similar(L')
+    if size(Q) == (n, n)
+        Qt, Lt = qr_full!(At, (Qt, Lt), qr_alg)
+    else
+        Qt, Lt = qr_compact!(At, (Qt, Lt), qr_alg)
+    end
+    adjoint!(Q, Qt)
+    !isempty(L) && adjoint!(L, Lt)
+    return L, Q
+end
+
+function lq_null_via_qr!(A::AbstractMatrix, N::AbstractMatrix, qr_alg::AbstractAlgorithm)
+    m, n = size(A)
+    minmn = min(m, n)
+    At = adjoint!(similar(A'), A)
+    Nt = similar(N')
+    Nt = qr_null!(At, Nt, qr_alg)
+    !isempty(N) && adjoint!(N, Nt)
+    return N
+end

src/implementations/polar.jl

@@ -30,13 +30,13 @@ end
 
 # Outputs
 # -------
-function initialize_output(::typeof(left_polar!), A::AbstractMatrix, ::PolarViaSVD)
+function initialize_output(::typeof(left_polar!), A::AbstractMatrix, ::AbstractAlgorithm)
     m, n = size(A)
     W = similar(A)
     P = similar(A, (n, n))
     return (W, P)
 end
-function initialize_output(::typeof(right_polar!), A::AbstractMatrix, ::PolarViaSVD)
+function initialize_output(::typeof(right_polar!), A::AbstractMatrix, ::AbstractAlgorithm)
     m, n = size(A)
     P = similar(A, (m, m))
     Wᴴ = similar(A)