refactor + deprecate CUSOLVER randomized SVD

Author: lkdvos

ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl

@@ -4,11 +4,12 @@ using MatrixAlgebraKit
 using MatrixAlgebraKit: @algdef, Algorithm, check_input
 using MatrixAlgebraKit: one!, zero!, uppertriangular!, lowertriangular!
 using MatrixAlgebraKit: diagview, sign_safe
-using MatrixAlgebraKit: CUSOLVER, LQViaTransposedQR, TruncationByValue, AbstractAlgorithm
+using MatrixAlgebraKit: CUSOLVER, LQViaTransposedQR, TruncationByValue, TruncationByOrder, AbstractAlgorithm
+using MatrixAlgebraKit: GaussianSketching, SketchingStrategy, SketchedAlgorithm
 using MatrixAlgebraKit: default_qr_algorithm, default_lq_algorithm, default_svd_algorithm, default_eig_algorithm, default_eigh_algorithm
 import MatrixAlgebraKit: geqrf!, ungqr!, unmqr!, gesvd!, gesvdp!, gesvdr!, gesvdj!
 import MatrixAlgebraKit: heevj!, heevd!, geev!
-import MatrixAlgebraKit: _gpu_Xgesvdr!, _sylvester, svd_rank
+import MatrixAlgebraKit: _sylvester, svd_rank
 using CUDA, CUDA.cuBLAS
 using CUDA: i32
 using LinearAlgebra
@@ -17,6 +18,7 @@ using LinearAlgebra: BlasFloat
 include("yacusolver.jl")
 
 MatrixAlgebraKit.default_driver(::Type{TA}) where {TA <: StridedCuVecOrMat{<:BlasFloat}} = CUSOLVER()
+MatrixAlgebraKit.default_driver(::Type{<:SketchedAlgorithm}, ::Type{TA}) where {TA <: StridedCuVecOrMat{<:BlasFloat}} = CUSOLVER()
 
 function MatrixAlgebraKit.default_svd_algorithm(::Type{T}; kwargs...) where {T <: StridedCuVecOrMat{<:BlasFloat}}
     return QRIteration(; kwargs...)
@@ -50,8 +52,31 @@ end
 gesvdp!(::CUSOLVER, A::StridedCuMatrix, S::StridedCuVector, U::StridedCuMatrix, Vᴴ::StridedCuMatrix; kwargs...) =
     YACUSOLVER.gesvdp!(A, S, U, Vᴴ; kwargs...)
 
-_gpu_Xgesvdr!(A::StridedCuMatrix, S::StridedCuVector, U::StridedCuMatrix, Vᴴ::StridedCuMatrix; kwargs...) =
-    YACUSOLVER.gesvdr!(A, S, U, Vᴴ; kwargs...)
+# Sketched SVD via cuSOLVER's gesvdr kernel
+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
+end
 
 geev!(::CUSOLVER, A::StridedCuMatrix, Dd::StridedCuVector, V::StridedCuMatrix) =
     YACUSOLVER.Xgeev!(A, Dd, V)

src/algorithms.jl

@@ -356,17 +356,25 @@ struct TruncatedAlgorithm{A <: AbstractAlgorithm, T <: TruncationStrategy} <: Ab
 end
 
 """
-    SketchedAlgorithm(alg::AbstractAlgorithm, sketch::SketchingStrategy, trunc::TruncationStrategy)
+    SketchedAlgorithm(;
+        alg::AbstractAlgorithm, sketch::SketchingStrategy,
+        trunc::TruncationStrategy, driver::Driver = DefaultDriver()
+    )
 
 Generic wrapper type for self-truncating algorithms that produce an approximate low-rank
 factorization by first applying a sketching operation specified by `sketch`, then computing
 a small dense decomposition of the projected matrix using `alg`. The `driver` selects the
-backend (e.g. `DefaultDriver()`, `CUSOLVER()`).
-"""
-struct SketchedAlgorithm{A <: AbstractAlgorithm, S <: SketchingStrategy, T <: TruncationStrategy} <: AbstractAlgorithm
-    alg::A
+backend implementing the sketched factorization (e.g. `Native()` for the generic
+sketch-then-decompose pipeline, `CUSOLVER()` for the fused `gesvdr` kernel).
+"""
+@kwdef struct SketchedAlgorithm{
+        A <: AbstractAlgorithm, S <: SketchingStrategy,
+        T <: TruncationStrategy, D <: Driver,
+    } <: AbstractAlgorithm
+    alg::A = DefaultAlgorithm()
     sketch::S
-    trunc::T
+    trunc::T = notrunc()
+    driver::D = DefaultDriver()
 end
 
 # utility conversion constructor

src/implementations/svd.jl

@@ -308,35 +308,47 @@ check_input(::typeof(svd_trunc!), A::AbstractMatrix, USVᴴ, alg::SketchedAlgori
 
 function svd_trunc_no_error!(A::AbstractMatrix, (U, S, Vᴴ), alg::SketchedAlgorithm)
     check_input(svd_trunc_no_error!, A, (U, S, Vᴴ), alg)
+    return gesvdr!(alg.driver, A, S, U, Vᴴ; alg.sketch, alg.alg, alg.trunc)
+end
+
+function svd_trunc!(A::AbstractMatrix, USVᴴ, alg::SketchedAlgorithm)
+    U, S, Vᴴ = svd_trunc_no_error!(A, USVᴴ, alg)
+    Na = norm(A)
+    Ns = norm(S)
+    return U, S, Vᴴ, sqrt(max(zero(Na), (Na + Ns) * (Na - Ns)))
+end
+
+# gesvdr! drivers
+# ---------------
+default_driver(::Type{<:SketchedAlgorithm}, ::Type{<:AbstractArray}) = Native()
+
+gesvdr!(::DefaultDriver, A, S, U, Vᴴ; kwargs...) =
+    gesvdr!(default_driver(SketchedAlgorithm, A), A, S, U, Vᴴ; kwargs...)
+
+function gesvdr!(
+        ::Native, A::AbstractMatrix, S, U, Vᴴ;
+        sketch::SketchingStrategy, alg::AbstractAlgorithm,
+        trunc::TruncationStrategy
+    )
     m, n = size(A)
     if m ≥ n
-        Q, B = left_sketch!(A, (U, Vᴴ), alg.sketch)
+        Q, B = left_sketch!(A, (U, Vᴴ), sketch)
         k = size(B, 1)
         U′ = similar(B, (k, k))
         Vᴴ′ = similar(B)
-        USVᴴ_inner = svd_compact!(B, (U′, S, Vᴴ′), alg.alg)
-        (Uout′, Sout, Vᴴout), _ = truncate(svd_trunc!, USVᴴ_inner, alg.trunc)
+        Uout′, Sout, Vᴴout, _ = svd_trunc!(B, (U′, S, Vᴴ′), TruncatedAlgorithm(alg, trunc))
         Uout = Q * Uout′
     else
-        B, Pᴴ = right_sketch!(A, (U, Vᴴ), alg.sketch)
+        B, Pᴴ = right_sketch!(A, (U, Vᴴ), sketch)
         k = size(B, 2)
         U′ = similar(B)
         Vᴴ′ = similar(B, (k, k))
-        USVᴴ_inner = svd_compact!(B, (U′, S, Vᴴ′), alg.alg)
-        (Uout, Sout, Vᴴout′), _ = truncate(svd_trunc!, USVᴴ_inner, alg.trunc)
+        Uout, Sout, Vᴴout′, _ = svd_trunc!(B, (U′, S, Vᴴ′), TruncatedAlgorithm(alg, trunc))
         Vᴴout = Vᴴout′ * Pᴴ
     end
-    get(alg.alg.kwargs, :fixgauge, true) && gaugefix!(svd_trunc!, Uout, Vᴴout)
     return Uout, Sout, Vᴴout
 end
 
-function svd_trunc!(A::AbstractMatrix, USVᴴ, alg::SketchedAlgorithm)
-    U, S, Vᴴ = svd_trunc_no_error!(A, USVᴴ, alg)
-    Na = norm(A)
-    Ns = norm(S)
-    return U, S, Vᴴ, sqrt(max(zero(Na), (Na + Ns) * (Na - Ns)))
-end
-
 # Deprecations
 # ------------
 for algtype in (:SafeDivideAndConquer, :DivideAndConquer, :QRIteration, :Jacobi, :Bisection)
@@ -380,6 +392,40 @@ for (algtype, newtype, drivertype) in (
     end
 end
 
+# CUSOLVER_Randomized → SketchedAlgorithm with driver = CUSOLVER()
+function _cusolver_randomized_to_sketched(alg::CUSOLVER_Randomized)
+    k = alg.kwargs.k
+    p = alg.kwargs.p
+    niters = alg.kwargs.niters
+    return SketchedAlgorithm(
+        QRIteration(),
+        GaussianSketching(k + p; numiter = niters + 1),
+        truncrank(k);
+        driver = CUSOLVER(),
+    )
+end
+
+for f! in (:svd_trunc!, :svd_trunc_no_error!)
+    @eval Base.@deprecate(
+        $f!(A::AbstractMatrix, USVᴴ, alg::CUSOLVER_Randomized),
+        $f!(A, USVᴴ, _cusolver_randomized_to_sketched(alg))
+    )
+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.",
+        :select_algorithm,
+    )
+    isempty(kwargs) ||
+        throw(ArgumentError("Additional keyword arguments are not allowed when algorithm parameters are specified."))
+    return _cusolver_randomized_to_sketched(alg)
+end
+@inline function select_algorithm(::typeof(svd_trunc_no_error!), A, alg::CUSOLVER_Randomized; kwargs...)
+    return select_algorithm(svd_trunc!, A, alg; kwargs...)
+end
+
 # GLA_QRIteration SVD deprecations (eigh methods remain in the GLA extension)
 Base.@deprecate(
     svd_compact!(A::AbstractMatrix, USVᴴ, alg::GLA_QRIteration),

src/interface/decompositions.jl

@@ -408,8 +408,6 @@ for more information.
 """
 @algdef CUSOLVER_Randomized
 
-does_truncate(::TruncatedAlgorithm{<:CUSOLVER_Randomized}) = true
-
 """
     CUSOLVER_Simple(; fixgauge = default_fixgauge())
 

test/decompositions/svd.jl

@@ -33,7 +33,7 @@ if !is_buildkite
         m, n = 54, 63
         rtol = sqrt(TestSuite.precision(T)) # extra square root
         algs = [
-            SketchedAlgorithm(DefaultAlgorithm(), GaussianSketching(m ÷ 2, numiter = 4), truncrank(m ÷ 4)),
+            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)