Implement DiagonalAlgorithm for GPU

Author: kshyatt

ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl

@@ -27,6 +27,17 @@ end
 function MatrixAlgebraKit.default_eigh_algorithm(::Type{T}; kwargs...) where {T <: StridedROCMatrix}
     return ROCSOLVER_DivideAndConquer(; kwargs...)
 end
+for f in (:(MatrixAlgebraKit.default_lq_algorithm),
+          :(MatrixAlgebraKit.default_qr_algorithm),
+          :(MatrixAlgebraKit.default_eig_algorithm),
+          :(MatrixAlgebraKit.default_eigh_algorithm),
+          :(MatrixAlgebraKit.default_svd_algorithm),
+         )
+
+    @eval function $f(::Type{T}; kwargs...) where {S, T <: Diagonal{S, <:StridedROCVector}}
+        return ROCM_DiagonalAlgorithm(; kwargs...)
+    end
+end
 
 _gpu_geqrf!(A::StridedROCMatrix) = YArocSOLVER.geqrf!(A)
 _gpu_ungqr!(A::StridedROCMatrix, τ::StridedROCVector) = YArocSOLVER.ungqr!(A, τ)

ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl

@@ -32,6 +32,18 @@ function MatrixAlgebraKit.default_eigh_algorithm(::Type{T}; kwargs...) where {TT
     return CUSOLVER_DivideAndConquer(; kwargs...)
 end
 
+for f in (:(MatrixAlgebraKit.default_lq_algorithm),
+          :(MatrixAlgebraKit.default_qr_algorithm),
+          :(MatrixAlgebraKit.default_eig_algorithm),
+          :(MatrixAlgebraKit.default_eigh_algorithm),
+          :(MatrixAlgebraKit.default_svd_algorithm),
+         )
+
+    @eval function $f(::Type{T}; kwargs...) where {S, T <: Diagonal{S, <:StridedCuVector}}
+        return CUDA_DiagonalAlgorithm(; kwargs...)
+    end
+end
+
 # include for block sector support
 function MatrixAlgebraKit.default_qr_algorithm(::Type{Base.ReshapedArray{T, 2, SubArray{T, 1, A, Tuple{UnitRange{Int}}, true}, Tuple{}}}; kwargs...) where {T <: BlasFloat, A <: CuVecOrMat{T}}
     return CUSOLVER_HouseholderQR(; kwargs...)

src/MatrixAlgebraKit.jl

@@ -36,7 +36,7 @@ export LAPACK_HouseholderQR, LAPACK_HouseholderLQ, LAPACK_Simple, LAPACK_Expert,
 export GLA_HouseholderQR, GLA_QRIteration, GS_QRIteration
 export LQViaTransposedQR
 export PolarViaSVD, PolarNewton
-export DiagonalAlgorithm
+export DiagonalAlgorithm, GPU_DiagonalAlgorithm, CUDA_DiagonalAlgorithm, ROCM_DiagonalAlgorithm
 export NativeBlocked
 export CUSOLVER_Simple, CUSOLVER_HouseholderQR, CUSOLVER_QRIteration, CUSOLVER_SVDPolar,
     CUSOLVER_Jacobi, CUSOLVER_Randomized, CUSOLVER_DivideAndConquer

src/implementations/eig.jl

@@ -28,7 +28,7 @@ function check_input(::typeof(eig_vals!), A::AbstractMatrix, D, ::AbstractAlgori
     return nothing
 end
 
-function check_input(::typeof(eig_full!), A::AbstractMatrix, DV, ::DiagonalAlgorithm)
+function check_input(::typeof(eig_full!), A::AbstractMatrix, DV, ::DiagonalAlgorithms)
     m, n = size(A)
     @assert m == n && isdiag(A)
     D, V = DV
@@ -40,7 +40,7 @@ function check_input(::typeof(eig_full!), A::AbstractMatrix, DV, ::DiagonalAlgor
     @check_scalar(V, A)
     return nothing
 end
-function check_input(::typeof(eig_vals!), A::AbstractMatrix, D, ::DiagonalAlgorithm)
+function check_input(::typeof(eig_vals!), A::AbstractMatrix, D, ::DiagonalAlgorithms)
     m, n = size(A)
     @assert m == n && isdiag(A)
     @assert D isa AbstractVector
@@ -69,10 +69,10 @@ function initialize_output(::typeof(eig_trunc!), A, alg::TruncatedAlgorithm)
     return initialize_output(eig_full!, A, alg.alg)
 end
 
-function initialize_output(::typeof(eig_full!), A::Diagonal, ::DiagonalAlgorithm)
+function initialize_output(::typeof(eig_full!), A::Diagonal, ::DiagonalAlgorithms)
     return A, similar(A)
 end
-function initialize_output(::typeof(eig_vals!), A::Diagonal, ::DiagonalAlgorithm)
+function initialize_output(::typeof(eig_vals!), A::Diagonal, ::DiagonalAlgorithms)
     return diagview(A)
 end
 
@@ -123,14 +123,14 @@ end
 
 # Diagonal logic
 # --------------
-function eig_full!(A::Diagonal, (D, V)::Tuple{Diagonal, Diagonal}, alg::DiagonalAlgorithm)
+function eig_full!(A::Diagonal, (D, V)::Tuple{Diagonal, Diagonal}, alg::DiagonalAlgorithms)
     check_input(eig_full!, A, (D, V), alg)
     D === A || copy!(D, A)
     one!(V)
     return D, V
 end
 
-function eig_vals!(A::Diagonal, D::AbstractVector, alg::DiagonalAlgorithm)
+function eig_vals!(A::Diagonal, D::AbstractVector, alg::DiagonalAlgorithms)
     check_input(eig_vals!, A, D, alg)
     Ad = diagview(A)
     D === Ad || copy!(D, Ad)

src/implementations/eigh.jl

@@ -38,7 +38,7 @@ function check_input(::typeof(eigh_vals!), A::AbstractMatrix, D, alg::AbstractAl
     return nothing
 end
 
-function check_input(::typeof(eigh_full!), A::AbstractMatrix, DV, alg::DiagonalAlgorithm)
+function check_input(::typeof(eigh_full!), A::AbstractMatrix, DV, alg::DiagonalAlgorithms)
     check_hermitian(A, alg)
     @assert isdiag(A)
     m = size(A, 1)
@@ -51,7 +51,7 @@ function check_input(::typeof(eigh_full!), A::AbstractMatrix, DV, alg::DiagonalA
     return nothing
 end
 
-function check_input(::typeof(eigh_vals!), A::AbstractMatrix, D, alg::DiagonalAlgorithm)
+function check_input(::typeof(eigh_vals!), A::AbstractMatrix, D, alg::DiagonalAlgorithms)
     check_hermitian(A, alg)
     @assert isdiag(A)
     m = size(A, 1)
@@ -78,10 +78,10 @@ function initialize_output(::typeof(eigh_trunc!), A, alg::TruncatedAlgorithm)
     return initialize_output(eigh_full!, A, alg.alg)
 end
 
-function initialize_output(::typeof(eigh_full!), A::Diagonal, ::DiagonalAlgorithm)
+function initialize_output(::typeof(eigh_full!), A::Diagonal, ::DiagonalAlgorithms)
     return eltype(A) <: Real ? A : similar(A, real(eltype(A))), similar(A)
 end
-function initialize_output(::typeof(eigh_vals!), A::Diagonal, ::DiagonalAlgorithm)
+function initialize_output(::typeof(eigh_vals!), A::Diagonal, ::DiagonalAlgorithms)
     return eltype(A) <: Real ? diagview(A) : similar(A, real(eltype(A)), size(A, 1))
 end
 
@@ -137,15 +137,15 @@ end
 
 # Diagonal logic
 # --------------
-function eigh_full!(A::Diagonal, DV, alg::DiagonalAlgorithm)
+function eigh_full!(A::Diagonal, DV, alg::DiagonalAlgorithms)
     check_input(eigh_full!, A, DV, alg)
     D, V = DV
     D === A || (diagview(D) .= real.(diagview(A)))
     one!(V)
     return D, V
 end
 
-function eigh_vals!(A::Diagonal, D, alg::DiagonalAlgorithm)
+function eigh_vals!(A::Diagonal, D, alg::DiagonalAlgorithms)
     check_input(eigh_vals!, A, D, alg)
     Ad = diagview(A)
     D === Ad || (D .= real.(Ad))

src/implementations/lq.jl

@@ -36,7 +36,7 @@ function check_input(::typeof(lq_null!), A::AbstractMatrix, Nᴴ, ::AbstractAlgo
     return nothing
 end
 
-function check_input(::typeof(lq_full!), A::AbstractMatrix, (L, Q), ::DiagonalAlgorithm)
+function check_input(::typeof(lq_full!), A::AbstractMatrix, (L, Q), ::DiagonalAlgorithms)
     m, n = size(A)
     @assert m == n && isdiag(A)
     @assert Q isa Diagonal && L isa Diagonal
@@ -46,10 +46,10 @@ function check_input(::typeof(lq_full!), A::AbstractMatrix, (L, Q), ::DiagonalAl
     @check_scalar(Q, A)
     return nothing
 end
-function check_input(::typeof(lq_compact!), A::AbstractMatrix, LQ, alg::DiagonalAlgorithm)
+function check_input(::typeof(lq_compact!), A::AbstractMatrix, LQ, alg::DiagonalAlgorithms)
     return check_input(lq_full!, A, LQ, alg)
 end
-function check_input(::typeof(lq_null!), A::AbstractMatrix, N, ::DiagonalAlgorithm)
+function check_input(::typeof(lq_null!), A::AbstractMatrix, N, ::DiagonalAlgorithms)
     m, n = size(A)
     @assert m == n && isdiag(A)
     @assert N isa AbstractMatrix
@@ -81,7 +81,7 @@ function initialize_output(::typeof(lq_null!), A::AbstractMatrix, ::AbstractAlgo
 end
 
 for f! in (:lq_full!, :lq_compact!)
-    @eval function initialize_output(::typeof($f!), A::AbstractMatrix, ::DiagonalAlgorithm)
+    @eval function initialize_output(::typeof($f!), A::AbstractMatrix, ::DiagonalAlgorithms)
         return similar(A), A
     end
 end
@@ -124,19 +124,19 @@ function lq_null!(A::AbstractMatrix, Nᴴ, alg::LQViaTransposedQR)
     return Nᴴ
 end
 
-function lq_full!(A::AbstractMatrix, LQ, alg::DiagonalAlgorithm)
+function lq_full!(A::AbstractMatrix, LQ, alg::DiagonalAlgorithms)
     check_input(lq_full!, A, LQ, alg)
     L, Q = LQ
     _diagonal_lq!(A, L, Q; alg.kwargs...)
     return L, Q
 end
-function lq_compact!(A::AbstractMatrix, LQ, alg::DiagonalAlgorithm)
+function lq_compact!(A::AbstractMatrix, LQ, alg::DiagonalAlgorithms)
     check_input(lq_compact!, A, LQ, alg)
     L, Q = LQ
     _diagonal_lq!(A, L, Q; alg.kwargs...)
     return L, Q
 end
-function lq_null!(A::AbstractMatrix, N, alg::DiagonalAlgorithm)
+function lq_null!(A::AbstractMatrix, N, alg::DiagonalAlgorithms)
     check_input(lq_null!, A, N, alg)
     return _diagonal_lq_null!(A, N; alg.kwargs...)
 end

src/implementations/qr.jl

@@ -36,7 +36,7 @@ function check_input(::typeof(qr_null!), A::AbstractMatrix, N, ::AbstractAlgorit
     return nothing
 end
 
-function check_input(::typeof(qr_full!), A::AbstractMatrix, (Q, R), alg::DiagonalAlgorithm)
+function check_input(::typeof(qr_full!), A::AbstractMatrix, (Q, R), alg::DiagonalAlgorithms)
     m, n = size(A)
     @assert m == n && isdiag(A)
     @assert Q isa Diagonal && R isa Diagonal
@@ -46,10 +46,10 @@ function check_input(::typeof(qr_full!), A::AbstractMatrix, (Q, R), alg::Diagona
     @check_scalar(R, A)
     return nothing
 end
-function check_input(::typeof(qr_compact!), A::AbstractMatrix, QR, alg::DiagonalAlgorithm)
+function check_input(::typeof(qr_compact!), A::AbstractMatrix, QR, alg::DiagonalAlgorithms)
     return check_input(qr_full!, A, QR, alg)
 end
-function check_input(::typeof(qr_null!), A::AbstractMatrix, N, ::DiagonalAlgorithm)
+function check_input(::typeof(qr_null!), A::AbstractMatrix, N, ::DiagonalAlgorithms)
     m, n = size(A)
     @assert m == n && isdiag(A)
     @assert N isa AbstractMatrix
@@ -81,7 +81,7 @@ function initialize_output(::typeof(qr_null!), A::AbstractMatrix, ::AbstractAlgo
 end
 
 for f! in (:qr_full!, :qr_compact!)
-    @eval function initialize_output(::typeof($f!), A::AbstractMatrix, ::DiagonalAlgorithm)
+    @eval function initialize_output(::typeof($f!), A::AbstractMatrix, ::DiagonalAlgorithms)
         return A, similar(A)
     end
 end
@@ -107,19 +107,19 @@ function qr_null!(A::AbstractMatrix, N, alg::LAPACK_HouseholderQR)
     return N
 end
 
-function qr_full!(A::AbstractMatrix, QR, alg::DiagonalAlgorithm)
+function qr_full!(A::AbstractMatrix, QR, alg::DiagonalAlgorithms)
     check_input(qr_full!, A, QR, alg)
     Q, R = QR
     _diagonal_qr!(A, Q, R; alg.kwargs...)
     return Q, R
 end
-function qr_compact!(A::AbstractMatrix, QR, alg::DiagonalAlgorithm)
+function qr_compact!(A::AbstractMatrix, QR, alg::DiagonalAlgorithms)
     check_input(qr_compact!, A, QR, alg)
     Q, R = QR
     _diagonal_qr!(A, Q, R; alg.kwargs...)
     return Q, R
 end
-function qr_null!(A::AbstractMatrix, N, alg::DiagonalAlgorithm)
+function qr_null!(A::AbstractMatrix, N, alg::DiagonalAlgorithms)
     check_input(qr_null!, A, N, alg)
     _diagonal_qr_null!(A, N; alg.kwargs...)
     return N

src/implementations/svd.jl

@@ -42,7 +42,7 @@ function check_input(::typeof(svd_vals!), A::AbstractMatrix, S, ::AbstractAlgori
     return nothing
 end
 
-function check_input(::typeof(svd_full!), A::AbstractMatrix, USVᴴ, ::DiagonalAlgorithm)
+function check_input(::typeof(svd_full!), A::AbstractMatrix, USVᴴ, ::DiagonalAlgorithms)
     m, n = size(A)
     @assert m == n && isdiag(A)
     U, S, Vᴴ = USVᴴ
@@ -56,11 +56,11 @@ function check_input(::typeof(svd_full!), A::AbstractMatrix, USVᴴ, ::DiagonalA
     return nothing
 end
 function check_input(
-        ::typeof(svd_compact!), A::AbstractMatrix, USVᴴ, alg::DiagonalAlgorithm
+        ::typeof(svd_compact!), A::AbstractMatrix, USVᴴ, alg::DiagonalAlgorithms
     )
     return check_input(svd_full!, A, USVᴴ, alg)
 end
-function check_input(::typeof(svd_vals!), A::AbstractMatrix, S, ::DiagonalAlgorithm)
+function check_input(::typeof(svd_vals!), A::AbstractMatrix, S, ::DiagonalAlgorithms)
     m, n = size(A)
     @assert m == n && isdiag(A)
     @assert S isa AbstractVector
@@ -93,15 +93,15 @@ function initialize_output(::typeof(svd_trunc!), A, alg::TruncatedAlgorithm)
     return initialize_output(svd_compact!, A, alg.alg)
 end
 
-function initialize_output(::typeof(svd_full!), A::Diagonal, ::DiagonalAlgorithm)
+function initialize_output(::typeof(svd_full!), A::Diagonal, ::DiagonalAlgorithms)
     TA = eltype(A)
     TUV = Base.promote_op(sign_safe, TA)
     return similar(A, TUV, size(A)), similar(A, real(TA)), similar(A, TUV, size(A))
 end
-function initialize_output(::typeof(svd_compact!), A::Diagonal, alg::DiagonalAlgorithm)
+function initialize_output(::typeof(svd_compact!), A::Diagonal, alg::DiagonalAlgorithms)
     return initialize_output(svd_full!, A, alg)
 end
-function initialize_output(::typeof(svd_vals!), A::Diagonal, ::DiagonalAlgorithm)
+function initialize_output(::typeof(svd_vals!), A::Diagonal, ::DiagonalAlgorithms)
     return eltype(A) <: Real ? diagview(A) : similar(A, real(eltype(A)), size(A, 1))
 end
 
@@ -214,11 +214,16 @@ end
 
 # Diagonal logic
 # --------------
-function svd_full!(A::AbstractMatrix, USVᴴ, alg::DiagonalAlgorithm)
+function svd_full!(A::AbstractMatrix, USVᴴ, alg::DiagonalAlgorithms)
     check_input(svd_full!, A, USVᴴ, alg)
     Ad = diagview(A)
     U, S, Vᴴ = USVᴴ
-    p = sortperm(Ad; by = abs, rev = true)
+    p = if isempty(Ad)
+            Int[]
+        else
+            sortperm(Ad; by = abs, rev = true)
+        end
+       
     zero!(U)
     zero!(Vᴴ)
     n = size(A, 1)
@@ -239,10 +244,10 @@ function svd_full!(A::AbstractMatrix, USVᴴ, alg::DiagonalAlgorithm)
 
     return U, S, Vᴴ
 end
-function svd_compact!(A, USVᴴ, alg::DiagonalAlgorithm)
+function svd_compact!(A, USVᴴ, alg::DiagonalAlgorithms)
     return svd_full!(A, USVᴴ, alg)
 end
-function svd_vals!(A::AbstractMatrix, S, alg::DiagonalAlgorithm)
+function svd_vals!(A::AbstractMatrix, S, alg::DiagonalAlgorithms)
     check_input(svd_vals!, A, S, alg)
     Ad = diagview(A)
     S .= abs.(Ad)

src/interface/decompositions.jl

@@ -193,6 +193,14 @@ the diagonal structure of the input and outputs.
 """
 @algdef DiagonalAlgorithm
 
+"""
+    GPUDiagonalAlgorithm(; kwargs...)
+
+Algorithm type to denote a native *on-GPU* Julia implementation of the decompositions making use of
+the diagonal structure of the input and outputs.
+"""
+@algdef GPUDiagonalAlgorithm
+
 """
     LQViaTransposedQR(qr_alg)
 
@@ -296,6 +304,14 @@ const CUSOLVER_SVDAlgorithm = Union{
     CUSOLVER_QRIteration, CUSOLVER_SVDPolar, CUSOLVER_Jacobi, CUSOLVER_Randomized,
 }
 
+"""
+    CUDA_DiagonalAlgorithm(; kwargs...)
+
+Algorithm type to denote a native Julia implementation of the decompositions making use of
+the diagonal structure of the input and outputs, specialized for CUDA.
+"""
+@algdef CUDA_DiagonalAlgorithm
+
 # =========================
 # ROCSOLVER ALGORITHMS
 # =========================
@@ -353,6 +369,14 @@ singular vectors, see also [`gaugefix!`](@ref).
 
 const ROCSOLVER_SVDAlgorithm = Union{ROCSOLVER_QRIteration, ROCSOLVER_Jacobi}
 
+"""
+    ROCM_DiagonalAlgorithm(; kwargs...)
+
+Algorithm type to denote a native Julia implementation of the decompositions making use of
+the diagonal structure of the input and outputs, specialized for ROCM.
+"""
+@algdef ROCM_DiagonalAlgorithm
+
 # Various consts and unions
 # -------------------------
 
@@ -370,11 +394,15 @@ const GPU_SVDAlgorithm = Union{CUSOLVER_SVDAlgorithm, ROCSOLVER_SVDAlgorithm}
 const GPU_SVDPolar = Union{CUSOLVER_SVDPolar}
 const GPU_Randomized = Union{CUSOLVER_Randomized}
 
+const GPU_DiagonalAlgorithm = Union{CUDA_DiagonalAlgorithm, ROCM_DiagonalAlgorithm}
+
 const QRAlgorithms = Union{LAPACK_HouseholderQR, CUSOLVER_HouseholderQR, ROCSOLVER_HouseholderQR}
 const LQAlgorithms = Union{LAPACK_HouseholderLQ, LQViaTransposedQR}
 const SVDAlgorithms = Union{LAPACK_SVDAlgorithm, GPU_SVDAlgorithm}
 const PolarAlgorithms = Union{PolarViaSVD, PolarNewton}
 
+const DiagonalAlgorithms = Union{DiagonalAlgorithm, GPU_DiagonalAlgorithm}
+
 # ================================
 # ORTHOGONALIZATION ALGORITHMS
 # ================================