Attempting to wrap randomized SVD

Author: kshyatt

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),

src/MatrixAlgebraKit.jl

@@ -33,7 +33,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,
        ROCSOLVER_HouseholderQR, ROCSOLVER_QRIteration, ROCSOLVER_Jacobi
 export truncrank, trunctol, truncabove, TruncationKeepSorted, TruncationKeepFiltered
 

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)
@@ -174,6 +175,7 @@ function svd_vals!(A::AbstractMatrix, S, alg::LAPACK_SVDAlgorithm)
     return S
 end
 
+#YACUSOLVER.Xgesvdr!(A, view(S, 1:k, 1), U, Vᴴ; alg.kwargs...)
 function svd_trunc!(A::AbstractMatrix, USVᴴ, alg::TruncatedAlgorithm)
     USVᴴ′ = svd_compact!(A, USVᴴ, alg.alg)
     return truncate!(svd_trunc!, USVᴴ′, alg.trunc)
@@ -185,17 +187,20 @@ end
 ###
 const CUSOLVER_SVDAlgorithm = Union{CUSOLVER_QRIteration,
                                     CUSOLVER_SVDPolar,
-                                    CUSOLVER_Jacobi}
+                                    CUSOLVER_Jacobi,
+                                    CUSOLVER_Randomized}
 const ROCSOLVER_SVDAlgorithm = Union{ROCSOLVER_QRIteration,
                                      ROCSOLVER_Jacobi}
 const GPU_SVDAlgorithm = Union{CUSOLVER_SVDAlgorithm, ROCSOLVER_SVDAlgorithm}
 
 const GPU_QRIteration = Union{CUSOLVER_QRIteration, ROCSOLVER_QRIteration}
 const GPU_SVDPolar = Union{CUSOLVER_SVDPolar}
 const GPU_Jacobi = Union{CUSOLVER_Jacobi, ROCSOLVER_Jacobi}
+const GPU_Randomized = Union{CUSOLVER_Randomized}
 
 _gpu_gesvd!(A::AbstractMatrix, S::AbstractVector, U::AbstractMatrix, Vᴴ::AbstractMatrix) = throw(MethodError(_gpu_gesvd!, (A, S, U, Vᴴ)))
 _gpu_Xgesvdp!(A::AbstractMatrix, S::AbstractVector, U::AbstractMatrix, Vᴴ::AbstractMatrix; kwargs...) = throw(MethodError(_gpu_Xgesvdp!, (A, S, U, Vᴴ)))
+_gpu_Xgesvdr!(A::AbstractMatrix, S::AbstractVector, U::AbstractMatrix, Vᴴ::AbstractMatrix; kwargs...) = throw(MethodError(_gpu_Xgesvdr!, (A, S, U, Vᴴ)))
 _gpu_gesvdj!(A::AbstractMatrix, S::AbstractVector, U::AbstractMatrix, Vᴴ::AbstractMatrix; kwargs...) = throw(MethodError(_gpu_gesvdj!, (A, S, U, Vᴴ)))
 
 # GPU SVD implementation

src/interface/decompositions.jl

@@ -151,6 +151,19 @@ 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
+
 # =========================
 # ROCSOLVER ALGORITHMS
 # =========================

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