Some GPU fixes

Author: lkdvos

ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl

@@ -52,30 +52,27 @@ end
 gesvdp!(::CUSOLVER, A::StridedCuMatrix, S::StridedCuVector, U::StridedCuMatrix, Vᴴ::StridedCuMatrix; kwargs...) =
     YACUSOLVER.gesvdp!(A, S, U, Vᴴ; kwargs...)
 
-# Sketched SVD via cuSOLVER's gesvdr kernel
+# Sketched SVD via cuSOLVER's gesvdr kernel.
+# The full m×m / n×n shapes of U / Vᴴ allow YACUSOLVER.gesvdr! to reuse them as cuSOLVER workspace.
+# `alg` is accepted but unused: cuSOLVER's gesvdr fuses the inner SVD itself.
 function gesvdr!(
         ::CUSOLVER, A::StridedCuMatrix, S, U::StridedCuMatrix, Vᴴ::StridedCuMatrix;
         sketch::GaussianSketching, trunc::TruncationByOrder, alg::AbstractAlgorithm = DefaultAlgorithm()
     )
     isempty(A) && return U, S, Vᴴ
-    m, n = size(A); minmn = min(m, n)
-    k = trunc.howmany
-    1 ≤ k ≤ minmn ||
-        throw(ArgumentError("trunc.howmany=$k must satisfy 1 ≤ k ≤ min(size(A))=$minmn"))
-    p = sketch.howmany - k
-    p ≥ 0 || throw(
-        ArgumentError(
-            "sketch.howmany=$(sketch.howmany) must be ≥ trunc.howmany=$k"
-        )
-    )
-    p = min(p, minmn - k - 1)
-    niters = sketch.numiter - 1
-
-    Uk = view(U, :, 1:k)
-    Vᴴk = view(Vᴴ, 1:k, :)
-    Sk = view(diagview(S), 1:k)
-    YACUSOLVER.gesvdr!(A, Sk, Uk, Vᴴk; k, p, niters)
-    return Uk, S, Vᴴk
+    m, n = size(A)
+    sketch_amount = min(sketch.howmany, m, n)
+    k = min(trunc.howmany, m, n)
+    p = max(sketch_amount - k, 0)
+    numiter = sketch.numiter
+
+    V = Vᴴ # gesvdr returns V, but this has to be the same size so we will use this as workspace
+
+    YACUSOLVER.gesvdr!(A, diagview(S), U, V; k, p, numiter)
+
+    # Truncate requires Vᴴ, so we adjoint here
+    USVᴴtrunc, _ = MatrixAlgebraKit.truncate(MatrixAlgebraKit.svd_trunc!, (U, S, V'), trunc)
+    return USVᴴtrunc
 end
 
 geev!(::CUSOLVER, A::StridedCuMatrix, Dd::StridedCuVector, V::StridedCuMatrix) =

ext/MatrixAlgebraKitCUDAExt/yacusolver.jl

@@ -266,41 +266,53 @@ for (bname, fname, elty, relty) in
     end
 end
 
-# Wrapper for randomized SVD
+# Wrapper for randomized SVD.
+# Caller must supply full-size buffers: U is (m, m) and Vᴴ is (n, n); both are reused
+# directly as cuSOLVER's workspace, and Vᴴ is converted in place from V to Vᴴ on the
+# leading k rows after cuSOLVER returns.
+# !!! Warning: this function takes in/returns V instead of Vᴴ
 function gesvdr!(
         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));
+        U::StridedCuMatrix{T} = similar(A, T, size(A, 1), size(A, 1)),
+        V::StridedCuMatrix{T} = similar(A, T, size(A, 2), size(A, 2));
         k::Int = length(S),
         p::Int = min(size(A)...) - k - 1,
-        niters::Int = 1
+        numiter::Int = 1,
     ) where {T <: BlasFloat}
-    chkstride1(A, U, S, Vᴴ)
+    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'
     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))
+    length(S) == minmn ||
+        throw(DimensionMismatch("length of S ($(length(S))) must equal min(size(A)) = $minmn"))
+    size(U) == (m, m) ||
+        throw(DimensionMismatch("U must have shape (m, m) = ($m, $m); got $(size(U))"))
+    size(V) == (n, n) ||
+        throw(DimensionMismatch("V must have shape (n, n) = ($n, $n); got $(size(V))"))
+    k < minmn ||
+        throw(DimensionMismatch("rank k ($k) must be less than min(size(A)) = $minmn"))
+    k + p < minmn ||
+        throw(DimensionMismatch("k + p ($(k + p)) must be less than min(size(A)) = $minmn"))
+
+    isempty(A) && return S, U, V
+
+    jobu = 'S'
+    jobv = 'S'
     lda = max(1, stride(A, 2))
-    ldu = max(1, stride(Ũ, 2))
-    ldv = max(1, stride(Ṽ, 2))
+    ldu = max(1, stride(U, 2))
+    ldv = max(1, stride(V, 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,
+            dh, params, jobu, jobv, m, n, k, p, numiter,
+            T, A, lda, R, S, T, U, ldu, T, V, ldv,
             T, out_gpu, out_cpu
         )
 
@@ -311,8 +323,8 @@ function gesvdr!(
         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,
+            dh, params, jobu, jobv, m, n, k, p, numiter,
+            T, A, lda, R, S, T, U, ldu, T, V, ldv,
             T, buffer_gpu, sizeof(buffer_gpu),
             buffer_cpu, sizeof(buffer_cpu),
             dh.info
@@ -321,16 +333,8 @@ function gesvdr!(
 
     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ᴴ
+    return S, U, V
 end
 
 # Wrapper for general eigensolver

src/implementations/svd.jl

@@ -311,6 +311,36 @@ function svd_trunc_no_error!(A::AbstractMatrix, (U, S, Vᴴ), alg::SketchedAlgor
     return gesvdr!(alg.driver, A, S, U, Vᴴ; alg.sketch, alg.alg, alg.trunc)
 end
 
+# CUSOLVER's gesvdr kernel requires full U and Vᴴ
+function initialize_output(
+        ::typeof(svd_trunc_no_error!), A::AbstractMatrix,
+        alg::SketchedAlgorithm{<:AbstractAlgorithm, <:SketchingStrategy, <:TruncationStrategy, CUSOLVER},
+    )
+    m, n = size(A)
+    minmn = min(m, n)
+    T = float(eltype(A))
+    U = similar(A, T, (m, m))
+    S = Diagonal(similar(A, real(T), (minmn,)))
+    Vᴴ = similar(A, T, (n, n))
+    return (U, S, Vᴴ)
+end
+
+function check_input(
+        ::typeof(svd_trunc_no_error!), A::AbstractMatrix, (U, S, Vᴴ),
+        alg::SketchedAlgorithm{<:AbstractAlgorithm, <:SketchingStrategy, <:TruncationStrategy, CUSOLVER},
+    )
+    m, n = size(A)
+    minmn = min(m, n)
+    @assert U isa AbstractMatrix && S isa Diagonal && Vᴴ isa AbstractMatrix
+    @check_size(U, (m, m))
+    @check_scalar(U, A)
+    @check_size(S, (minmn, minmn))
+    @check_scalar(S, A, real)
+    @check_size(Vᴴ, (n, n))
+    @check_scalar(Vᴴ, A)
+    return nothing
+end
+
 function svd_trunc!(A::AbstractMatrix, USVᴴ, alg::SketchedAlgorithm)
     U, S, Vᴴ = svd_trunc_no_error!(A, USVᴴ, alg)
     Na = norm(A)
@@ -399,7 +429,7 @@ function _cusolver_randomized_to_sketched(alg::CUSOLVER_Randomized)
     niters = alg.kwargs.niters
     return SketchedAlgorithm(
         QRIteration(),
-        GaussianSketching(k + p; numiter = niters + 1),
+        GaussianSketching(k + p; numiter = niters),
         truncrank(k);
         driver = CUSOLVER(),
     )
@@ -415,7 +445,7 @@ end
 @inline function select_algorithm(::typeof(svd_trunc!), A, alg::CUSOLVER_Randomized; kwargs...)
     Base.depwarn(
         "`CUSOLVER_Randomized` is deprecated; use \
-         `SketchedAlgorithm(QRIteration(), GaussianSketching(k+p; numiter=niters+1), truncrank(k); driver=CUSOLVER())` instead.",
+         `SketchedAlgorithm(QRIteration(), GaussianSketching(k+p; numiter=niters), truncrank(k); driver=CUSOLVER())` instead.",
         :select_algorithm,
     )
     isempty(kwargs) ||

src/yalapack.jl

@@ -2351,7 +2351,7 @@ for (gesvd, gesdd, gesvdx, gejsv, gesvj, elty, relty) in
                 throw(DimensionMismatch("length mismatch between A ($n) and S ($(length(S)))"))
 
             lda = max(1, stride(A, 2))
-            mv = Ref{BlasInt}() # unused
+            mv = Ref{BlasInt}(0) # unused by LAPACK when JOBV='V', but must satisfy MV ≥ 0 input check
             if jobv == 'V'
                 if U !== A
                     V = view(U, 1:n, 1:n) # use U as V storage

test/decompositions/svd.jl

@@ -36,7 +36,7 @@ if !is_buildkite
             SketchedAlgorithm(; sketch = GaussianSketching(m ÷ 2, numiter = 4), trunc = truncrank(m ÷ 4)),
         ]
         TestSuite.test_sketched_svd(T, (m, n), algs; rtol)
-        TestSuite.test_sketched_svd(T, (n, m), algs)
+        TestSuite.test_sketched_svd(T, (n, m), algs; rtol)
     end
 
     # Generic floats:
@@ -57,7 +57,7 @@ end
 
 # CUDA tests
 # ------------
-if false # CUDA.functional()
+if CUDA.functional()
     # LAPACK algorithms:
     for T in BLASFloats, m in (0, 23), n in (0, 17, m, 27)
         TestSuite.seed_rng!(123)
@@ -66,18 +66,19 @@ if false # CUDA.functional()
         TestSuite.test_svd_algs(CuMatrix{T}, (m, n), CUDA_SVD_ALGS)
     end
 
-    # Randomized SVD:
+    # Sketched SVD:
     for T in BLASFloats, m in (0, 23), n in (0, 17, m, 27)
         TestSuite.seed_rng!(123)
         k = 5
         p = min(m, n) - k - 2
         p > 0 || continue
-        cusolver_sketch = SketchedAlgorithm(
-            GaussianSketching(k; oversampling = p, niters = 20),
-            DefaultAlgorithm(),
-            MatrixAlgebraKit.CUSOLVER(),
+        rtol = sqrt(TestSuite.precision(T)) # extra square root
+        cusolver_sketch = SketchedAlgorithm(;
+            sketch = GaussianSketching(k + p; numiter = 20),
+            trunc = truncrank(k),
+            driver = MatrixAlgebraKit.CUSOLVER(),
         )
-        TestSuite.test_randomized_svd(CuMatrix{T}, (m, n), (cusolver_sketch,))
+        TestSuite.test_sketched_svd(CuMatrix{T}, (m, n), (cusolver_sketch,); rtol)
     end
 
     # Diagonal: