diff --git a/src/implementations/svd.jl b/src/implementations/svd.jl index 9f2b6e8ae..6375c0e93 100644 --- a/src/implementations/svd.jl +++ b/src/implementations/svd.jl @@ -319,19 +319,33 @@ end # GPU logic (randomized SVD - CUSOLVER_Randomized has no CPU analog, kept as-is) # --------------------------------------------------------------------------------- +# A randomized SVD only computes the `k` largest singular triplets, so the outputs only need +# room for those: `U` is `m × r`, `S` is `r × r` and `Vᴴ` is `r × n` for some `k ≤ r ≤ minmn` +# (`initialize_output` uses `r = k`). Only the first `k` columns/values/rows are written. + +function _randomized_rank(alg::CUSOLVER_Randomized) + haskey(alg.kwargs, :k) || + throw(ArgumentError("$alg requires the number of singular values to compute, `k`")) + return alg.kwargs.k::Int +end function check_input( ::Union{typeof(svd_trunc!), typeof(svd_trunc_no_error!)}, A::AbstractMatrix, USVᴴ, alg::CUSOLVER_Randomized ) m, n = size(A) minmn = min(m, n) + k = _randomized_rank(alg) U, S, Vᴴ = USVᴴ @assert U isa AbstractMatrix && S isa Diagonal && Vᴴ isa AbstractMatrix - @check_size(U, (m, m)) + r = size(S, 1) + k ≤ r ≤ minmn || + throw(DimensionMismatch("size of S ($r) must be between k ($k) and min(size(A)...) ($minmn)")) + (size(U, 1) == m && k ≤ size(U, 2) ≤ m) || + throw(DimensionMismatch("size of U $(size(U)) must be ($m, r) with r between k ($k) and $m")) + (size(Vᴴ, 2) == n && k ≤ size(Vᴴ, 1) ≤ n) || + throw(DimensionMismatch("size of Vᴴ $(size(Vᴴ)) must be (r, $n) with r between k ($k) and $n")) @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 @@ -340,10 +354,10 @@ function initialize_output( ::Union{typeof(svd_trunc!), typeof(svd_trunc_no_error!)}, A::AbstractMatrix, alg::TruncatedAlgorithm{<:CUSOLVER_Randomized} ) m, n = size(A) - minmn = min(m, n) - U = similar(A, (m, m)) - S = Diagonal(similar(A, real(eltype(A)), (minmn,))) - Vᴴ = similar(A, (n, n)) + k = _randomized_rank(alg.alg) + U = similar(A, (m, k)) + S = Diagonal(similar(A, real(eltype(A)), (k,))) + Vᴴ = similar(A, (k, n)) return (U, S, Vᴴ) end @@ -356,13 +370,18 @@ end function svd_trunc_no_error!(A::AbstractMatrix, USVᴴ, alg::TruncatedAlgorithm{<:GPU_Randomized}) U, S, Vᴴ = USVᴴ check_input(svd_trunc_no_error!, A, (U, S, Vᴴ), alg.alg) - _gpu_Xgesvdr!(A, diagview(S), U, Vᴴ; alg.alg.kwargs...) + k = _randomized_rank(alg.alg) + # only the first `k` values are computed: restrict to those, so that the truncation never + # sees the (uninitialized) remainder of larger outputs + Uk = size(U, 2) == k ? U : view(U, :, 1:k) + Sk = size(S, 1) == k ? S : Diagonal(view(diagview(S), 1:k)) + Vᴴk = size(Vᴴ, 1) == k ? Vᴴ : view(Vᴴ, 1:k, :) + kwargs = Base.structdiff(alg.alg.kwargs, NamedTuple{(:fixgauge,)}) + _gpu_Xgesvdr!(A, diagview(Sk), Uk, Vᴴk; kwargs...) - # TODO: make sure that truncation is based on maxrank, otherwise this might be wrong - (Utr, Str, Vᴴtr), _ = truncate(svd_trunc!, (U, S, Vᴴ), alg.trunc) + (Utr, Str, Vᴴtr), _ = truncate(svd_trunc!, (Uk, Sk, Vᴴk), alg.trunc) do_gauge_fix = get(alg.alg.kwargs, :fixgauge, default_fixgauge())::Bool - # the output matrices here are the same size as for svd_full! do_gauge_fix && gaugefix!(svd_trunc!, Utr, Vᴴtr) return Utr, Str, Vᴴtr diff --git a/test/testsuite/decompositions/svd.jl b/test/testsuite/decompositions/svd.jl index 1baeebbec..6cdbae0eb 100644 --- a/test/testsuite/decompositions/svd.jl +++ b/test/testsuite/decompositions/svd.jl @@ -712,8 +712,39 @@ function test_randomized_svd(T::Type, sz, algs; kwargs...) Ac = deepcopy(A) m, n = size(A) minmn = min(m, n) + k = alg.alg.k S₀ = collect(svd_vals(A)) + atol = sqrt(eps(real(eltype(T)))) * S₀[1] + + # outputs only hold the `k` computed triplets U1, S1, V1ᴴ, ϵ1 = @testinferred svd_trunc(A; alg) - @test collect(diagview(S1))[1:alg.alg.k] ≈ S₀[1:alg.alg.k] + @test A == Ac + @test size(U1) == (m, k) && size(S1) == (k, k) && size(V1ᴴ) == (k, n) + @test collect(diagview(S1)) ≈ S₀[1:k] + @test isisometric(U1) + @test isisometric(V1ᴴ; side = :right) + @test A * V1ᴴ' ≈ U1 * S1 + @test ϵ1 ≈ norm(S₀[(k + 1):end]) atol = atol + + # further truncation only looks at the `k` computed values + alg2 = MatrixAlgebraKit.TruncatedAlgorithm(alg.alg, truncrank(k - 2) & trunctol(; atol = S₀[k] / 2)) + U2, S2, V2ᴴ, ϵ2 = svd_trunc(A; alg = alg2) + @test size(S2) == (k - 2, k - 2) + @test collect(diagview(S2)) ≈ S₀[1:(k - 2)] + @test ϵ2 ≈ norm(S₀[(k - 1):end]) atol = atol + + # larger (e.g. full-size) outputs are accepted, only the first `k` are used + USVᴴ = (similar(A, (m, m)), Diagonal(similar(A, real(eltype(A)), (minmn,))), similar(A, (n, n))) + U3, S3, V3ᴴ, ϵ3 = svd_trunc!(copy(A), USVᴴ, alg) + @test collect(diagview(S3)) ≈ S₀[1:k] + @test ϵ3 ≈ ϵ1 atol = atol + + # gauge fixing is not passed on to the driver + alg4 = MatrixAlgebraKit.TruncatedAlgorithm(CUSOLVER_Randomized(; alg.alg.kwargs..., fixgauge = true), alg.trunc) + _, S4, _, _ = svd_trunc(A; alg = alg4) + @test collect(diagview(S4)) ≈ S₀[1:k] + + # `k` is required + @test_throws ArgumentError svd_trunc(A; alg = MatrixAlgebraKit.TruncatedAlgorithm(CUSOLVER_Randomized(), alg.trunc)) end end