add testsuite

Author: lkdvos

Project.toml

@@ -33,10 +33,11 @@ GenericSchur = "0.5.6"
 JET = "0.9, 0.10"
 LinearAlgebra = "1"
 Mooncake = "0.4.174"
+Random = "1"
 SafeTestsets = "0.1"
 StableRNGs = "1"
 Test = "1"
-TestExtras = "0.2,0.3"
+TestExtras = "0.3.2"
 Zygote = "0.7"
 julia = "1.10"
 
@@ -47,11 +48,12 @@ ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "JET", "SafeTestsets", "Test", "TestExtras", "ChainRulesCore", "ChainRulesTestUtils", "StableRNGs", "Zygote", "CUDA", "AMDGPU", "GenericLinearAlgebra", "GenericSchur", "Mooncake"]
+test = ["Aqua", "JET", "SafeTestsets", "Test", "TestExtras", "ChainRulesCore", "ChainRulesTestUtils", "StableRNGs", "Zygote", "CUDA", "AMDGPU", "GenericLinearAlgebra", "GenericSchur", "Random", "Mooncake"]

ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl

@@ -159,43 +159,4 @@ function MatrixAlgebraKit._avgdiff!(A::StridedROCMatrix, B::StridedROCMatrix)
     return A, B
 end
 
-function MatrixAlgebraKit.truncate(
-        ::typeof(left_null!), US::Tuple{TU, TS}, strategy::TruncationStrategy
-    ) where {TU <: ROCMatrix, TS}
-    # TODO: avoid allocation?
-    U, S = US
-    extended_S = vcat(diagview(S), zeros(eltype(S), max(0, size(S, 1) - size(S, 2))))
-    ind = MatrixAlgebraKit.findtruncated(extended_S, strategy)
-    trunc_cols = collect(1:size(U, 2))[ind]
-    Utrunc = U[:, trunc_cols]
-    return Utrunc, ind
-end
-function MatrixAlgebraKit.truncate(
-        ::typeof(right_null!), SVᴴ::Tuple{TS, TVᴴ}, strategy::TruncationStrategy
-    ) where {TS, TVᴴ <: ROCMatrix}
-    # TODO: avoid allocation?
-    S, Vᴴ = SVᴴ
-    extended_S = vcat(diagview(S), zeros(eltype(S), max(0, size(S, 2) - size(S, 1))))
-    ind = MatrixAlgebraKit.findtruncated(extended_S, strategy)
-    trunc_rows = collect(1:size(Vᴴ, 1))[ind]
-    Vᴴtrunc = Vᴴ[trunc_rows, :]
-    return Vᴴtrunc, ind
-end
-
-# disambiguate:
-function MatrixAlgebraKit.truncate(
-        ::typeof(left_null!), (U, S)::Tuple{TU, TS}, ::NoTruncation
-    ) where {TU <: ROCMatrix, TS}
-    m, n = size(S)
-    ind = (n + 1):m
-    return U[:, ind], ind
-end
-function MatrixAlgebraKit.truncate(
-        ::typeof(right_null!), (S, Vᴴ)::Tuple{TS, TVᴴ}, ::NoTruncation
-    ) where {TS, TVᴴ <: ROCMatrix}
-    m, n = size(S)
-    ind = (m + 1):n
-    return Vᴴ[ind, :], ind
-end
-
 end

src/implementations/qr.jl

@@ -270,10 +270,12 @@ function _gpu_unmqr!(
 end
 
 function _gpu_qr!(
-        A::AbstractMatrix, Q::AbstractMatrix, R::AbstractMatrix; positive = false, blocksize = 1
+        A::AbstractMatrix, Q::AbstractMatrix, R::AbstractMatrix; pivoted = false, positive = false, blocksize = 1
     )
     blocksize > 1 &&
         throw(ArgumentError("CUSOLVER/ROCSOLVER does not provide a blocked implementation for a QR decomposition"))
+    pivoted &&
+        throw(ArgumentError("CUSOLVER/ROCSOLVER does not provide a pivoted implementation for a QR decomposition"))
     m, n = size(A)
     minmn = min(m, n)
     computeR = length(R) > 0
@@ -309,10 +311,12 @@ function _gpu_qr!(
 end
 
 function _gpu_qr_null!(
-        A::AbstractMatrix, N::AbstractMatrix; positive = false, blocksize = 1
+        A::AbstractMatrix, N::AbstractMatrix; positive = false, blocksize = 1, pivoted = false
     )
     blocksize > 1 &&
         throw(ArgumentError("CUSOLVER/ROCSOLVER does not provide a blocked implementation for a QR decomposition"))
+    pivoted &&
+        throw(ArgumentError("CUSOLVER/ROCSOLVER does not provide a pivoted implementation for a QR decomposition"))
     m, n = size(A)
     minmn = min(m, n)
     fill!(N, zero(eltype(N)))

src/implementations/schur.jl

@@ -26,6 +26,29 @@ function check_input(::typeof(schur_vals!), A::AbstractMatrix, vals, ::AbstractA
     return nothing
 end
 
+function check_input(::typeof(schur_full!), A::AbstractMatrix, TZv, ::DiagonalAlgorithm)
+    m, n = size(A)
+    @assert m == n && isdiag(A)
+    T, Z, vals = TZv
+    @assert vals isa AbstractVector && Z isa Diagonal
+    @check_scalar(T, A)
+    @check_size(Z, (m, m))
+    @check_scalar(Z, A)
+    @check_size(vals, (n,))
+    # Diagonal doesn't need to promote to complex scalartype since we know it is diagonalizable
+    @check_scalar(vals, A)
+    return nothing
+end
+function check_input(::typeof(schur_vals!), A::AbstractMatrix, vals, ::DiagonalAlgorithm)
+    m, n = size(A)
+    @assert m == n && isdiag(A)
+    @assert vals isa AbstractVector
+    @check_size(vals, (n,))
+    # Diagonal doesn't need to promote to complex scalartype since we know it is diagonalizable
+    @check_scalar(vals, A)
+    return nothing
+end
+
 # Outputs
 # -------
 function initialize_output(::typeof(schur_full!), A::AbstractMatrix, ::AbstractAlgorithm)
@@ -39,6 +62,17 @@ function initialize_output(::typeof(schur_vals!), A::AbstractMatrix, ::AbstractA
     vals = similar(A, complex(eltype(A)), n)
     return vals
 end
+function initialize_output(::typeof(schur_full!), A::Diagonal, ::DiagonalAlgorithm)
+    n = size(A, 1)
+    Z = similar(A)
+    vals = similar(A, eltype(A), n)
+    return (A, Z, vals)
+end
+function initialize_output(::typeof(schur_vals!), A::Diagonal, ::DiagonalAlgorithm)
+    n = size(A, 1)
+    vals = similar(A, eltype(A), n)
+    return vals
+end
 
 # Implementation
 # --------------
@@ -72,3 +106,20 @@ function schur_vals!(A::AbstractMatrix, vals, alg::LAPACK_EigAlgorithm)
     end
     return vals
 end
+
+# Diagonal logic
+# --------------
+function schur_full!(A::Diagonal, (T, Z, vals)::Tuple{Diagonal, Diagonal, <:AbstractVector}, alg::DiagonalAlgorithm)
+    check_input(schur_full!, A, (T, Z, vals), alg)
+    copy!(vals, diagview(A))
+    one!(Z)
+    T === A || copy!(T, A)
+    return T, Z, vals
+end
+
+function schur_vals!(A::Diagonal, vals::AbstractVector, alg::DiagonalAlgorithm)
+    check_input(schur_vals!, A, vals, alg)
+    Ad = diagview(A)
+    vals === Ad || copy!(vals, Ad)
+    return vals
+end

src/implementations/svd.jl

@@ -152,6 +152,12 @@ end
 function svd_compact!(A::AbstractMatrix, USVᴴ, alg::LAPACK_SVDAlgorithm)
     check_input(svd_compact!, A, USVᴴ, alg)
     U, S, Vᴴ = USVᴴ
+    if length(A) == 0
+        one!(U)
+        zero!(S)
+        one!(Vᴴ)
+        return USVᴴ
+    end
 
     do_gauge_fix = get(alg.kwargs, :fixgauge, default_fixgauge())::Bool
     alg_kwargs = Base.structdiff(alg.kwargs, NamedTuple{(:fixgauge,)})
@@ -382,6 +388,12 @@ end
 function svd_compact!(A::AbstractMatrix, USVᴴ, alg::GPU_SVDAlgorithm)
     check_input(svd_compact!, A, USVᴴ, alg)
     U, S, Vᴴ = USVᴴ
+    if length(A) == 0
+        one!(U)
+        zero!(S)
+        one!(Vᴴ)
+        return USVᴴ
+    end
 
     do_gauge_fix = get(alg.kwargs, :fixgauge, default_fixgauge())::Bool
     alg_kwargs = Base.structdiff(alg.kwargs, NamedTuple{(:fixgauge,)})
@@ -406,6 +418,10 @@ _largest(x, y) = abs(x) < abs(y) ? y : x
 
 function svd_vals!(A::AbstractMatrix, S, alg::GPU_SVDAlgorithm)
     check_input(svd_vals!, A, S, alg)
+    if length(A) == 0
+        zero!(S)
+        return S
+    end
     U, Vᴴ = similar(A, (0, 0)), similar(A, (0, 0))
 
     alg_kwargs = Base.structdiff(alg.kwargs, NamedTuple{(:fixgauge,)})