Attempting to wrap randomized SVD

Author: kshyatt

ext/MatrixAlgebraKitCUDAExt/implementations/svd.jl

@@ -1,6 +1,7 @@
 const CUSOLVER_SVDAlgorithm = Union{CUSOLVER_QRIteration,
                                     CUSOLVER_SVDPolar,
-                                    CUSOLVER_Jacobi}
+                                    CUSOLVER_Jacobi,
+                                    CUSOLVER_Randomized}
 
 # CUSOLVER SVD implementation
 function MatrixAlgebraKit.svd_full!(A::CuMatrix, USVᴴ, alg::CUSOLVER_SVDAlgorithm)
@@ -9,6 +10,7 @@ function MatrixAlgebraKit.svd_full!(A::CuMatrix, USVᴴ, alg::CUSOLVER_SVDAlgori
     fill!(S, zero(eltype(S)))
     m, n = size(A)
     minmn = min(m, n)
+    k = alg isa CUSOLVER_Randomized ? get(alg.kwargs, :k, min(size(S)...)) : minmn
     if alg isa CUSOLVER_QRIteration
         isempty(alg.kwargs) ||
             throw(ArgumentError("LAPACK_QRIteration does not accept any keyword arguments"))
@@ -17,18 +19,18 @@ function MatrixAlgebraKit.svd_full!(A::CuMatrix, USVᴴ, alg::CUSOLVER_SVDAlgori
         YACUSOLVER.Xgesvdp!(A, view(S, 1:minmn, 1), U, Vᴴ; alg.kwargs...)
     elseif alg isa CUSOLVER_Jacobi
         YACUSOLVER.gesvdj!(A, view(S, 1:minmn, 1), U, Vᴴ; alg.kwargs...)
-        # elseif alg isa LAPACK_Bisection
-        #     throw(ArgumentError("LAPACK_Bisection is not supported for full SVD"))
-        # elseif alg isa LAPACK_Jacobi
-        #     throw(ArgumentError("LAPACK_Bisection is not supported for full SVD"))
+    elseif alg isa CUSOLVER_Randomized
+        YACUSOLVER.Xgesvdr!(A, view(S, 1:k, 1), U, Vᴴ; alg.kwargs...)
+    # elseif alg isa LAPACK_Bisection
+    #     throw(ArgumentError("LAPACK_Bisection is not supported for full SVD"))
     else
         throw(ArgumentError("Unsupported SVD algorithm"))
     end
     diagview(S) .= view(S, 1:minmn, 1)
     view(S, 2:minmn, 1) .= zero(eltype(S))
     # TODO: make this controllable using a `gaugefix` keyword argument
     for j in 1:max(m, n)
-        if j <= minmn
+        if j <= minmn 
             u = view(U, :, j)
             v = view(Vᴴ, j, :)
             s = conj(sign(_argmaxabs(u)))
@@ -58,12 +60,6 @@ function MatrixAlgebraKit.svd_compact!(A::CuMatrix, USVᴴ, alg::CUSOLVER_SVDAlg
         YACUSOLVER.Xgesvdp!(A, S.diag, U, Vᴴ; alg.kwargs...)
     elseif alg isa CUSOLVER_Jacobi
         YACUSOLVER.gesvdj!(A, S.diag, U, Vᴴ; alg.kwargs...)
-        # elseif alg isa LAPACK_DivideAndConquer
-        #     isempty(alg.kwargs) ||
-        #         throw(ArgumentError("LAPACK_DivideAndConquer does not accept any keyword arguments"))
-        #     YALAPACK.gesdd!(A, S.diag, U, Vᴴ)
-        # elseif alg isa LAPACK_Bisection
-        #     YALAPACK.gesvdx!(A, S.diag, U, Vᴴ; alg.kwargs...)
     else
         throw(ArgumentError("Unsupported SVD algorithm"))
     end
@@ -81,7 +77,8 @@ _argmaxabs(x) = reduce(_largest, x; init=zero(eltype(x)))
 _largest(x, y) = abs(x) < abs(y) ? y : x
 
 function MatrixAlgebraKit.svd_vals!(A::CuMatrix, S, alg::CUSOLVER_SVDAlgorithm)
-    check_input(svd_vals!, A, S)
+    # TODO restore this also for randomized
+    alg isa CUSOLVER_Randomized || check_input(svd_vals!, A, S)
     U, Vᴴ = similar(A, (0, 0)), similar(A, (0, 0))
     if alg isa CUSOLVER_QRIteration
         isempty(alg.kwargs) ||
@@ -91,18 +88,10 @@ function MatrixAlgebraKit.svd_vals!(A::CuMatrix, S, alg::CUSOLVER_SVDAlgorithm)
         YACUSOLVER.Xgesvdp!(A, S, U, Vᴴ; alg.kwargs...)
     elseif alg isa CUSOLVER_Jacobi
         YACUSOLVER.gesvdj!(A, S, U, Vᴴ; alg.kwargs...)
-        # elseif alg isa LAPACK_DivideAndConquer
-        #     isempty(alg.kwargs) ||
-        #         throw(ArgumentError("LAPACK_DivideAndConquer does not accept any keyword arguments"))
-        #     YALAPACK.gesdd!(A, S, U, Vᴴ)
-        # elseif alg isa LAPACK_Bisection
-        #     YALAPACK.gesvdx!(A, S, U, Vᴴ; alg.kwargs...)
-        # elseif alg isa LAPACK_Jacobi
-        #     isempty(alg.kwargs) ||
-        #         throw(ArgumentError("LAPACK_Jacobi does not accept any keyword arguments"))
-        #     YALAPACK.gesvj!(A, S, U, Vᴴ)
+    elseif alg isa CUSOLVER_Randomized
+        YACUSOLVER.Xgesvdr!(A, S, U, Vᴴ; alg.kwargs...)
     else
         throw(ArgumentError("Unsupported SVD algorithm"))
     end
     return S
-end
+end

ext/MatrixAlgebraKitCUDAExt/yacusolver.jl

@@ -247,6 +247,66 @@ for (bname, fname, elty, relty) in
     end
 end
 
+# Wrapper for randomized SVD
+function Xgesvdr!(A::StridedCuMatrix{T},
+                   S::StridedCuVector=similar(A, real(T), min(size(A)...)),
+                   U::StridedCuMatrix{T}=similar(A, T, size(A, 1), min(size(A)...)),
+                   Vᴴ::StridedCuMatrix{T}=similar(A, T, min(size(A)...), size(A, 2));
+                   k::Int=length(S),
+                   p::Int=min(size(A)...)-k-1,
+                   niters::Int=1) where {T<:BlasFloat}
+    chkstride1(A, U, S, Vᴴ)
+    m, n = size(A)
+    minmn = min(m, n)
+    jobu = length(U) == 0 ? 'N' : 'S'
+    jobv = length(Vᴴ) == 0 ? 'N' : 'S'
+    k = min(size(S)...)
+    R = eltype(S)
+    k < minmn || throw(DimensionMismatch("length of S ($k) must be less than the smaller dimension of A ($minmn)"))
+    k + p < minmn || throw(DimensionMismatch("length of S ($k) plus oversampling ($p) must be less than the smaller dimension of A ($minmn)"))
+    R == real(T) ||
+        throw(ArgumentError("S does not have the matching real `eltype` of A"))
+
+    Ṽ = similar(Vᴴ, (n, n))
+    Ũ = (size(U) == (m, m)) ? U : similar(U, (m, m))
+    lda = max(1, stride(A, 2))
+    ldu = max(1, stride(Ũ, 2))
+    ldv = max(1, stride(Ṽ, 2))
+    params = CUSOLVER.CuSolverParameters()
+    dh = CUSOLVER.dense_handle()
+
+    function bufferSize()
+        out_cpu = Ref{Csize_t}(0)
+        out_gpu = Ref{Csize_t}(0)
+        CUSOLVER.cusolverDnXgesvdr_bufferSize(dh, params, jobu, jobv, m, n, k, p, niters,
+                                              T, A, lda, R, S, T, Ũ, ldu, T, Ṽ, ldv,
+                                              T, out_gpu, out_cpu)
+
+        return out_gpu[], out_cpu[]
+    end
+    CUSOLVER.with_workspaces(dh.workspace_gpu, dh.workspace_cpu,
+                             bufferSize()...) do buffer_gpu, buffer_cpu
+        return CUSOLVER.cusolverDnXgesvdr(dh, params, jobu, jobv, m, n, k, p, niters,
+                                          T, A, lda, R, S, T, Ũ, ldu, T, Ṽ, ldv,
+                                          T, buffer_gpu, sizeof(buffer_gpu),
+                                          buffer_cpu, sizeof(buffer_cpu),
+                                          dh.info)
+    end
+
+    flag = @allowscalar dh.info[1]
+    CUSOLVER.chklapackerror(BlasInt(flag))
+    if Ũ !== U && length(U) > 0
+        U .= view(Ũ, 1:m, 1:size(U, 2))
+    end
+    if length(Vᴴ) > 0
+        Vᴴ .= view(Ṽ', 1:size(Vᴴ, 1), 1:n)
+    end
+    Ũ !== U && CUDA.unsafe_free!(Ũ)
+    CUDA.unsafe_free!(Ṽ)
+
+    return S, U, Vᴴ
+end
+
 # for (jname, bname, fname, elty, relty) in
 #     ((:sygvd!, :cusolverDnSsygvd_bufferSize, :cusolverDnSsygvd, :Float32, :Float32),
 #      (:sygvd!, :cusolverDnDsygvd_bufferSize, :cusolverDnDsygvd, :Float64, :Float64),
@@ -591,4 +651,4 @@ end
 #     end
 # end
 
-end
+end

src/MatrixAlgebraKit.jl

@@ -31,7 +31,7 @@ export LAPACK_HouseholderQR, LAPACK_HouseholderLQ,
        LAPACK_QRIteration, LAPACK_Bisection, LAPACK_MultipleRelativelyRobustRepresentations,
        LAPACK_DivideAndConquer, LAPACK_Jacobi,
        LQViaTransposedQR,
-       CUSOLVER_HouseholderQR, CUSOLVER_QRIteration, CUSOLVER_SVDPolar, CUSOLVER_Jacobi
+       CUSOLVER_HouseholderQR, CUSOLVER_QRIteration, CUSOLVER_SVDPolar, CUSOLVER_Jacobi, CUSOLVER_Randomized
 export truncrank, trunctol, truncabove, TruncationKeepSorted, TruncationKeepFiltered
 
 VERSION >= v"1.11.0-DEV.469" &&

src/implementations/svd.jl

@@ -66,6 +66,7 @@ function initialize_output(::typeof(svd_trunc!), A::AbstractMatrix, alg::Truncat
     return initialize_output(svd_compact!, A, alg.alg)
 end
 
+
 # Implementation
 # --------------
 function svd_full!(A::AbstractMatrix, USVᴴ, alg::LAPACK_SVDAlgorithm)

src/interface/decompositions.jl

@@ -150,3 +150,16 @@ Algorithm type to denote the CUSOLVER driver for computing the singular value de
 a general matrix using the Jacobi algorithm.
 """
 @algdef CUSOLVER_Jacobi
+
+"""
+    CUSOLVER_Randomized(; p, niters)
+
+Algorithm type to denote the CUSOLVER driver for computing the singular value decomposition of
+a general matrix using the randomized SVD algorithm.
+
+!!! note
+  Randomized SVD cannot compute all singular values of the input matrix `A`, only the first `k` where
+  `k < min(m, n)`. The remainder are used for oversampling. See the [CUSOLVER documentation](https://docs.nvidia.com/cuda/cusolver/index.html#cusolverdnxgesvdr)
+  for more information.
+"""
+@algdef CUSOLVER_Randomized

test/cuda/svd.jl

@@ -76,12 +76,49 @@ end
             @test isapproxone(Vᴴ' * Vᴴ)
             @test all(isposdef, diagview(S))
 
-            Sc = similar(A, real(T), min(m, n))
+            minmn = min(m, n)
+            Sc = similar(A, real(T), minmn)
             Sc2 = svd_vals!(copy!(Ac, A), Sc, alg)
             @test Sc === Sc2
             @test CuArray(diagview(S)) ≈ Sc
             # CuArray is necessary because norm of CuArray view with non-unit step is broken
         end
+        k = min(m, n) - 20
+        p = min(m, n) - k - 1
+        algs = (CUSOLVER_Randomized(; k=k, p=p, niters=100),)
+        @testset "algorithm $alg" for alg in algs
+            A = CuArray(randn(rng, T, m, n))
+            Uref, Sref, Vᴴref = svd_full(A, CUSOLVER_SVDPolar())
+            U, S, Vᴴ = svd_full(A; alg)
+            @test U isa CuMatrix{T} && size(U) == (m, m)
+            @test S isa CuMatrix{real(T)} && size(S) == (m, n)
+            @test Vᴴ isa CuMatrix{T} && size(Vᴴ) == (n, n)
+            for col in 1:k
+                @test view(collect(U), :, col) ≈ view(collect(Uref), :, col)
+                @test view(collect(Vᴴ), col, :) ≈ view(collect(Vᴴref), col, :)
+            end
+            @test all(isposdef, view(diagview(S), 1:k))
+            @test view(CuArray(diagview(S)), 1:k) ≈ view(CuArray(diagview(Sref)), 1:k)
+
+            Ac = similar(A)
+            U2, S2, V2ᴴ = @constinferred svd_full!(copy!(Ac, A), (U, S, Vᴴ), alg)
+            @test U2 === U
+            @test S2 === S
+            @test V2ᴴ === Vᴴ 
+            for col in 1:k
+                @test view(collect(U), :, col) ≈ view(collect(Uref), :, col)
+                @test view(collect(Vᴴ), col, :) ≈ view(collect(Vᴴref), col, :)
+            end
+            @test all(isposdef, view(diagview(S), 1:k))
+            @test view(CuArray(diagview(S2)), 1:k) ≈ view(CuArray(diagview(Sref)), 1:k)
+
+            Sc = similar(A, real(T), k)
+            Sc2 = svd_vals!(copy!(Ac, A), Sc, alg)
+            @test Sc === Sc2 
+            @test view(Sc, 1:k) ≈ view(CuArray(diagview(Sref)), 1:k)
+            @test view(CuArray(diagview(S)), 1:k) ≈ Sc
+            # CuArray is necessary because norm of CuArray view with non-unit step is broken
+        end
     end
 end