From 32a24603613f9040d6d81a092604288d88a7a97f Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 21 Aug 2026 16:41:22 +0200 Subject: [PATCH 1/2] Batched SVD support for ROCSOLVER and CUSOLVER --- .../MatrixAlgebraKitAMDGPUExt.jl | 43 +- ext/MatrixAlgebraKitAMDGPUExt/yarocsolver.jl | 654 ++++++++++++++++++ .../MatrixAlgebraKitCUDAExt.jl | 9 +- ext/MatrixAlgebraKitCUDAExt/yacusolver.jl | 86 +++ src/MatrixAlgebraKit.jl | 7 +- src/implementations/svd.jl | 224 +++++- src/interface/decompositions.jl | 104 ++- test/decompositions/svd.jl | 17 +- test/testsuite/TestSuite.jl | 14 + test/testsuite/decompositions/svd.jl | 179 +++++ 10 files changed, 1321 insertions(+), 16 deletions(-) diff --git a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl index 0bdb10497..020c43dff 100644 --- a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl +++ b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl @@ -6,7 +6,8 @@ using MatrixAlgebraKit: one!, zero!, uppertriangular!, lowertriangular! using MatrixAlgebraKit: diagview, sign_safe using MatrixAlgebraKit: ROCSOLVER, LQViaTransposedQR, TruncationStrategy, NoTruncation, TruncationByValue, AbstractAlgorithm using MatrixAlgebraKit: default_qr_algorithm, default_lq_algorithm, default_svd_algorithm, default_eigh_algorithm -import MatrixAlgebraKit: geqrf!, ungqr!, unmqr!, gesvd!, gesvdj! +import MatrixAlgebraKit: geqrf!, ungqr!, unmqr!, gesvd!, gesdd!, gesvdx!, gesvdj! +import MatrixAlgebraKit: gesvdj_batched!, gesvdx_batched!, gesdd_batched!, gesvd_batched! import MatrixAlgebraKit: heevj!, heevd!, heev!, heevx! import MatrixAlgebraKit: _sylvester, svd_rank, svd_pullback! using AMDGPU @@ -15,10 +16,17 @@ using LinearAlgebra: BlasFloat include("yarocsolver.jl") -MatrixAlgebraKit.default_driver(::Type{TA}) where {TA <: StridedROCVecOrMat{<:BlasFloat}} = ROCSOLVER() +MatrixAlgebraKit.default_driver(::Type{TA}) where {TA <: StridedROCArray{<:BlasFloat}} = ROCSOLVER() +MatrixAlgebraKit.default_driver(::Type{TA}) where {TA <: AbstractVector{<:StridedROCMatrix{<:BlasFloat}}} = ROCSOLVER() -function MatrixAlgebraKit.default_svd_algorithm(::Type{T}; kwargs...) where {T <: StridedROCVecOrMat{<:BlasFloat}} - return QRIteration(; kwargs...) +function MatrixAlgebraKit.default_svd_algorithm(::Type{T}; kwargs...) where {T <: StridedROCMatrix{<:BlasFloat}} + return DivideAndConquer(; kwargs...) +end +function MatrixAlgebraKit.default_svd_algorithm(::Type{T}; kwargs...) where {T <: StridedROCArray{<:BlasFloat, 3}} + return DivideAndConquerBatched(; kwargs...) +end +function MatrixAlgebraKit.default_svd_algorithm(::Type{T}; kwargs...) where {T <: AbstractVector{<:StridedROCMatrix{<:BlasFloat}}} + return DivideAndConquerBatched(; kwargs...) end function MatrixAlgebraKit.default_eigh_algorithm(::Type{T}; kwargs...) where {T <: StridedROCVecOrMat{<:BlasFloat}} return DivideAndConquer(; kwargs...) @@ -28,7 +36,7 @@ for f in (:geqrf!, :ungqr!, :unmqr!) @eval $f(::ROCSOLVER, args...) = YArocSOLVER.$f(args...) end -MatrixAlgebraKit.supports_svd_full(::ROCSOLVER, f::Symbol) = f in (:qr_iteration, :jacobi) +MatrixAlgebraKit.supports_svd_full(::ROCSOLVER, f::Symbol) = f in (:qr_iteration, :jacobi, :divide_and_conquer, :qr_iteration_batched, :jacobi_batched, :divide_and_conquer_batched) function gesvd!(::ROCSOLVER, A::StridedROCMatrix, S::StridedROCVector, U::StridedROCMatrix, Vᴴ::StridedROCMatrix; kwargs...) m, n = size(A) @@ -42,6 +50,31 @@ function gesvdj!(::ROCSOLVER, A::StridedROCMatrix, S::StridedROCVector, U::Strid return MatrixAlgebraKit.svd_via_adjoint!(gesvdj!, ROCSOLVER(), A, S, U, Vᴴ; kwargs...) end +gesvd_batched!(::ROCSOLVER, As::Vector{<:StridedROCMatrix}, Ss::StridedROCMatrix, Us::StridedROCArray{T, 3}, Vᴴs::StridedROCArray{T, 3}; kwargs...) where {T <: BlasFloat} = + YArocSOLVER.gesvd_batched!(As, Ss, Us, Vᴴs; kwargs...) +gesvd_batched!(::ROCSOLVER, As::StridedROCArray{T, 3}, Ss::StridedROCMatrix, Us::StridedROCArray{T, 3}, Vᴴs::StridedROCArray{T, 3}; kwargs...) where {T <: BlasFloat} = + YArocSOLVER.gesvd_strided_batched!(As, Ss, Us, Vᴴs; kwargs...) + +gesdd_batched!(::ROCSOLVER, As::Vector{<:StridedROCMatrix}, Ss::StridedROCMatrix, Us::StridedROCArray{T, 3}, Vᴴs::StridedROCArray{T, 3}; kwargs...) where {T <: BlasFloat} = + YArocSOLVER.gesdd_batched!(As, Ss, Us, Vᴴs; kwargs...) +gesdd_batched!(::ROCSOLVER, As::StridedROCArray{T, 3}, Ss::StridedROCMatrix, Us::StridedROCArray{T, 3}, Vᴴs::StridedROCArray{T, 3}; kwargs...) where {T <: BlasFloat} = + YArocSOLVER.gesdd_strided_batched!(As, Ss, Us, Vᴴs; kwargs...) + +gesvdj_batched!(::ROCSOLVER, As::Vector{<:StridedROCMatrix}, Ss::StridedROCMatrix, Us::StridedROCArray{T, 3}, Vᴴs::StridedROCArray{T, 3}; kwargs...) where {T <: BlasFloat} = + YArocSOLVER.gesvdj_batched!(As, Ss, Us, Vᴴs; kwargs...) +gesvdj_batched!(::ROCSOLVER, As::StridedROCArray{T, 3}, Ss::StridedROCMatrix, Us::StridedROCArray{T, 3}, Vᴴs::StridedROCArray{T, 3}; kwargs...) where {T <: BlasFloat} = + YArocSOLVER.gesvdj_strided_batched!(As, Ss, Us, Vᴴs; kwargs...) + +gesvdx_batched!(::ROCSOLVER, As::Vector{<:StridedROCMatrix}, Ss::StridedROCMatrix, Us::StridedROCArray{T, 3}, Vᴴs::StridedROCArray{T, 3}; kwargs...) where {T <: BlasFloat} = + YArocSOLVER.gesvdx_batched!(As, Ss, Us, Vᴴs; kwargs...) +gesvdx_batched!(::ROCSOLVER, As::StridedROCArray{T, 3}, Ss::StridedROCMatrix, Us::StridedROCArray{T, 3}, Vᴴs::StridedROCArray{T, 3}; kwargs...) where {T <: BlasFloat} = + YArocSOLVER.gesvdx_strided_batched!(As, Ss, Us, Vᴴs; kwargs...) + +gesvdx!(::ROCSOLVER, A::StridedROCMatrix, S::StridedROCVector, U::StridedROCMatrix, Vᴴ::StridedROCMatrix; kwargs...) = + YArocSOLVER.gesvdx!(A, S, U, Vᴴ; kwargs...) +gesdd!(::ROCSOLVER, A::StridedROCMatrix, S::StridedROCVector, U::StridedROCMatrix, Vᴴ::StridedROCMatrix; kwargs...) = + YArocSOLVER.gesdd!(A, S, U, Vᴴ; kwargs...) + heevj!(::ROCSOLVER, A::StridedROCMatrix, Dd::StridedROCVector, V::StridedROCMatrix; kwargs...) = YArocSOLVER.heevj!(A, Dd, V; kwargs...) heevd!(::ROCSOLVER, A::StridedROCMatrix, Dd::StridedROCVector, V::StridedROCMatrix; kwargs...) = diff --git a/ext/MatrixAlgebraKitAMDGPUExt/yarocsolver.jl b/ext/MatrixAlgebraKitAMDGPUExt/yarocsolver.jl index e0c5f084d..2842e0a99 100644 --- a/ext/MatrixAlgebraKitAMDGPUExt/yarocsolver.jl +++ b/ext/MatrixAlgebraKitAMDGPUExt/yarocsolver.jl @@ -93,6 +93,405 @@ for (fname, elty, relty) in end end +# Wrappers for batched SVD via QR Iteration +for (fname, elty, relty) in + ( + (:rocsolver_sgesvd_batched, :Float32, :Float32), + (:rocsolver_dgesvd_batched, :Float64, :Float64), + (:rocsolver_cgesvd_batched, :ComplexF32, :Float32), + (:rocsolver_zgesvd_batched, :ComplexF64, :Float64), + ) + @eval begin + function gesvd_batched!( + A::StridedROCVector{<:StridedROCMatrix{$elty}}, + S::StridedROCMatrix{$relty} = similar(first(A), $relty, (min(size(first(A))...), length(A))), + U::StridedROCArray{$elty, 3} = similar(first(A), $elty, size(first(A), 1), min(size(first(A))...), length(A)), + Vᴴ::StridedROCArray{$elty, 3} = similar(first(A), $elty, min(size(first(A))...), size(first(A), 2), length(A)), + ) + for A_ in A + chkstride1(A_, U, Vᴴ, S) + end + m, n = size(first(A)) + (m < n) && throw(ArgumentError("rocSOLVER's gesvd_batched requires m ≥ n")) + minmn = min(m, n) + if length(U) == 0 + jobu = rocSOLVER.rocblas_svect_none + else + size(U, 1) == m || + throw(DimensionMismatch("row size mismatch between A and U")) + if size(U, 2) == minmn + if U === A # seems impossible? + jobu = rocSOLVER.rocblas_svect_overwrite + else + jobu = rocSOLVER.rocblas_svect_singular + end + elseif size(U, 2) == m + jobu = rocSOLVER.rocblas_svect_all + else + throw(DimensionMismatch("invalid column size of U")) + end + end + if length(Vᴴ) == 0 + jobvt = rocSOLVER.rocblas_svect_none + else + size(Vᴴ, 2) == n || + throw(DimensionMismatch("column size mismatch between A and Vᴴ")) + if size(Vᴴ, 1) == minmn + if Vᴴ === A # seems impossible? + jobvt = rocSOLVER.rocblas_svect_overwrite + else + jobvt = rocSOLVER.rocblas_svect_singular + end + elseif size(Vᴴ, 1) == n + jobvt = rocSOLVER.rocblas_svect_all + else + throw(DimensionMismatch("invalid row size of Vᴴ")) + end + end + length(S) == minmn * length(A) || + throw(DimensionMismatch("length mismatch between A and S")) + + lda = max(1, stride(first(A), 2)) + ldu = max(1, stride(U, 2)) + strideU = ldu * size(U, 2) + ldv = max(1, stride(Vᴴ, 2)) + strideV = ldv * n + strideS = minmn + + strideE = minmn - 1 + E = ROCArray{$relty}(undef, length(A) * strideE) + dh = rocBLAS.handle() + dev_info = ROCVector{Cint}(undef, length(A)) + pA = map(pointer, A) + rocSOLVER.$fname( + dh, jobu, jobvt, m, n, pA, lda, + S, strideS, U, ldu, strideU, Vᴴ, ldv, strideV, + E, strideE, convert(rocSOLVER.rocblas_workmode, 'I'), + dev_info, length(A) + ) + AMDGPU.unsafe_free!(pA) + AMDGPU.unsafe_free!(E) + + rocSOLVER.chkargsok.(BlasInt.(collect(dev_info))) + + return (S, U, Vᴴ) + end + end +end + +for (fname, elty, relty) in + ( + (:rocsolver_sgesvd_strided_batched, :Float32, :Float32), + (:rocsolver_dgesvd_strided_batched, :Float64, :Float64), + (:rocsolver_cgesvd_strided_batched, :ComplexF32, :Float32), + (:rocsolver_zgesvd_strided_batched, :ComplexF64, :Float64), + ) + @eval begin + function gesvd_strided_batched!( + A::StridedROCArray{$elty, 3}, + S::StridedROCMatrix{$relty} = similar(A, $relty, min(size(A, 1, size(A, 2))), size(A, 3)), + U::StridedROCArray{$elty, 3} = similar(A, $elty, size(A, 1), min(size(A, 1), size(A, 2)), size(A, 3)), + Vᴴ::StridedROCArray{$elty, 3} = similar(A, $elty, min(size(A, 1), size(A, 2)), size(A, 2), size(A, 3)), + ) + chkstride1(A, U, Vᴴ, S) + m, n, batch_size = size(A) + (m < n) && throw(ArgumentError("rocSOLVER's gesvd_strided_batched requires m ≥ n")) + minmn = min(m, n) + if length(U) == 0 + jobu = rocSOLVER.rocblas_svect_none + else + size(U, 1) == m || + throw(DimensionMismatch("row size mismatch between A and U")) + if size(U, 2) == minmn + if U === A + jobu = rocSOLVER.rocblas_svect_overwrite + else + jobu = rocSOLVER.rocblas_svect_singular + end + elseif size(U, 2) == m + jobu = rocSOLVER.rocblas_svect_all + else + throw(DimensionMismatch("invalid column size of U")) + end + end + if length(Vᴴ) == 0 + jobvt = rocSOLVER.rocblas_svect_none + else + size(Vᴴ, 2) == n || + throw(DimensionMismatch("column size mismatch between A and Vᴴ")) + if size(Vᴴ, 1) == minmn + if Vᴴ === A + jobvt = rocSOLVER.rocblas_svect_overwrite + else + jobvt = rocSOLVER.rocblas_svect_singular + end + elseif size(Vᴴ, 1) == n + jobvt = rocSOLVER.rocblas_svect_all + else + throw(DimensionMismatch("invalid row size of Vᴴ")) + end + end + length(S) == minmn * batch_size || + throw(DimensionMismatch("length mismatch between A and S")) + + lda = max(1, stride(A, 2)) + strideA = lda * n + ldu = max(1, stride(U, 2)) + strideU = ldu * size(U, 2) + ldv = max(1, stride(Vᴴ, 2)) + strideV = ldv * n + strideS = minmn + + strideE = minmn - 1 + E = ROCArray{$relty}(undef, batch_size * strideE) + dh = rocBLAS.handle() + dev_info = ROCVector{Cint}(undef, batch_size) + rocSOLVER.$fname( + dh, jobu, jobvt, m, n, A, lda, strideA, + S, strideS, U, ldu, strideU, Vᴴ, ldv, strideV, + E, strideE, convert(rocSOLVER.rocblas_workmode, 'I'), + dev_info, batch_size + ) + AMDGPU.unsafe_free!(E) + + rocSOLVER.chkargsok.(BlasInt.(collect(dev_info))) + + return (S, U, Vᴴ) + end + end +end + +# Wrapper for SVD via DivideAndConquer +for (fname, elty, relty) in + ( + (:rocsolver_sgesdd, :Float32, :Float32), + (:rocsolver_dgesdd, :Float64, :Float64), + (:rocsolver_cgesdd, :ComplexF32, :Float32), + (:rocsolver_zgesdd, :ComplexF64, :Float64), + ) + @eval begin + function gesdd!( + A::StridedROCMatrix{$elty}, + S::StridedROCVector{$relty} = similar(A, $relty, min(size(A)...)), + U::StridedROCMatrix{$elty} = similar(A, $elty, size(A, 1), min(size(A)...)), + Vᴴ::StridedROCMatrix{$elty} = similar(A, $elty, min(size(A)...), size(A, 2)) + ) + chkstride1(A, U, Vᴴ, S) + m, n = size(A) + minmn = min(m, n) + if length(U) == 0 + jobu = rocSOLVER.rocblas_svect_none + else + size(U, 1) == m || + throw(DimensionMismatch("row size mismatch between A and U")) + if size(U, 2) == minmn + if U === A + jobu = rocSOLVER.rocblas_svect_overwrite + else + jobu = rocSOLVER.rocblas_svect_singular + end + elseif size(U, 2) == m + jobu = rocSOLVER.rocblas_svect_all + else + throw(DimensionMismatch("invalid column size of U")) + end + end + if length(Vᴴ) == 0 + jobvt = rocSOLVER.rocblas_svect_none + else + size(Vᴴ, 2) == n || + throw(DimensionMismatch("column size mismatch between A and Vᴴ")) + if size(Vᴴ, 1) == minmn + if Vᴴ === A + jobvt = rocSOLVER.rocblas_svect_overwrite + else + jobvt = rocSOLVER.rocblas_svect_singular + end + elseif size(Vᴴ, 1) == n + jobvt = rocSOLVER.rocblas_svect_all + else + throw(DimensionMismatch("invalid row size of Vᴴ")) + end + end + length(S) == minmn || + throw(DimensionMismatch("length mismatch between A and S")) + + lda = max(1, stride(A, 2)) + ldu = max(1, stride(U, 2)) + ldv = max(1, stride(Vᴴ, 2)) + + dh = rocBLAS.handle() + dev_info = ROCVector{Cint}(undef, 1) + rocSOLVER.$fname( + dh, jobu, jobvt, m, n, + A, lda, S, U, ldu, Vᴴ, ldv, + dev_info + ) + + info = @allowscalar dev_info[1] + rocSOLVER.chkargsok(BlasInt(info)) + + return (S, U, Vᴴ) + end + end +end + +# Wrapper for batched SVD via DivideAndConquer +for (fname, elty, relty) in + ( + (:rocsolver_sgesdd_batched, :Float32, :Float32), + (:rocsolver_dgesdd_batched, :Float64, :Float64), + (:rocsolver_cgesdd_batched, :ComplexF32, :Float32), + (:rocsolver_zgesdd_batched, :ComplexF64, :Float64), + ) + @eval begin + function gesdd_batched!( + A::StridedROCVector{<:StridedROCMatrix{$elty}}, + S::StridedROCMatrix{$relty} = similar(first(A), $relty, (min(size(first(A))...), length(A))), + U::StridedROCArray{$elty, 3} = similar(first(A), $elty, size(first(A), 1), min(size(first(A))...), length(A)), + Vᴴ::StridedROCArray{$elty, 3} = similar(first(A), $elty, min(size(first(A))...), size(first(A), 2), length(A)), + ) + for A_ in A + chkstride1(A_, U, Vᴴ, S) + end + m, n = size(first(A)) + minmn = min(m, n) + if length(U) == 0 + jobu = rocSOLVER.rocblas_svect_none + else + size(U, 1) == m || + throw(DimensionMismatch("row size mismatch between A and U")) + if size(U, 2) == minmn + if U === A # seems impossible? + jobu = rocSOLVER.rocblas_svect_overwrite + else + jobu = rocSOLVER.rocblas_svect_singular + end + elseif size(U, 2) == m + jobu = rocSOLVER.rocblas_svect_all + else + throw(DimensionMismatch("invalid column size of U")) + end + end + if length(Vᴴ) == 0 + jobvt = rocSOLVER.rocblas_svect_none + else + size(Vᴴ, 2) == n || + throw(DimensionMismatch("column size mismatch between A and Vᴴ")) + if size(Vᴴ, 1) == minmn + if Vᴴ === A # seems impossible? + jobvt = rocSOLVER.rocblas_svect_overwrite + else + jobvt = rocSOLVER.rocblas_svect_singular + end + elseif size(Vᴴ, 1) == n + jobvt = rocSOLVER.rocblas_svect_all + else + throw(DimensionMismatch("invalid row size of Vᴴ")) + end + end + length(S) == minmn * length(A) || + throw(DimensionMismatch("length mismatch between A and S")) + + lda = max(1, stride(first(A), 2)) + ldu = max(1, stride(U, 2)) + strideU = ldu * size(U, 2) + ldv = max(1, stride(Vᴴ, 2)) + strideV = ldv * n + strideS = minmn + + dh = rocBLAS.handle() + dev_info = ROCVector{Cint}(undef, length(A)) + pA = map(pointer, A) + rocSOLVER.$fname( + dh, jobu, jobvt, m, n, pA, lda, + S, strideS, U, ldu, strideU, Vᴴ, ldv, strideV, + dev_info, length(A) + ) + AMDGPU.unsafe_free!(pA) + rocSOLVER.chkargsok.(BlasInt.(collect(dev_info))) + + return (S, U, Vᴴ) + end + end +end + +for (fname, elty, relty) in + ( + (:rocsolver_sgesdd_strided_batched, :Float32, :Float32), + (:rocsolver_dgesdd_strided_batched, :Float64, :Float64), + (:rocsolver_cgesdd_strided_batched, :ComplexF32, :Float32), + (:rocsolver_zgesdd_strided_batched, :ComplexF64, :Float64), + ) + @eval begin + function gesdd_strided_batched!( + A::StridedROCArray{$elty, 3}, + S::StridedROCMatrix{$relty} = similar(first(A), $relty, (min(size(first(A))...), length(A))), + U::StridedROCArray{$elty, 3} = similar(first(A), $elty, size(first(A), 1), min(size(first(A))...), length(A)), + Vᴴ::StridedROCArray{$elty, 3} = similar(first(A), $elty, min(size(first(A))...), size(first(A), 2), length(A)), + ) + chkstride1(A, U, Vᴴ, S) + m, n, batch_size = size(A) + minmn = min(m, n) + if length(U) == 0 + jobu = rocSOLVER.rocblas_svect_none + else + size(U, 1) == m || + throw(DimensionMismatch("row size mismatch between A and U")) + if size(U, 2) == minmn + if U === A + jobu = rocSOLVER.rocblas_svect_overwrite + else + jobu = rocSOLVER.rocblas_svect_singular + end + elseif size(U, 2) == m + jobu = rocSOLVER.rocblas_svect_all + else + throw(DimensionMismatch("invalid column size of U")) + end + end + if length(Vᴴ) == 0 + jobvt = rocSOLVER.rocblas_svect_none + else + size(Vᴴ, 2) == n || + throw(DimensionMismatch("column size mismatch between A and Vᴴ")) + if size(Vᴴ, 1) == minmn + if Vᴴ === A + jobvt = rocSOLVER.rocblas_svect_overwrite + else + jobvt = rocSOLVER.rocblas_svect_singular + end + elseif size(Vᴴ, 1) == n + jobvt = rocSOLVER.rocblas_svect_all + else + throw(DimensionMismatch("invalid row size of Vᴴ")) + end + end + length(S) == minmn * batch_size || + throw(DimensionMismatch("length mismatch between A and S")) + + lda = max(1, stride(A, 2)) + strideA = lda * n + ldu = max(1, stride(U, 2)) + strideU = ldu * size(U, 2) + ldv = max(1, stride(Vᴴ, 2)) + strideV = ldv * n + strideS = minmn + + dh = rocBLAS.handle() + dev_info = ROCVector{Cint}(undef, batch_size) + rocSOLVER.$fname( + dh, jobu, jobvt, m, n, A, lda, strideA, + S, strideS, U, ldu, strideU, Vᴴ, ldv, strideV, + dev_info, batch_size + ) + rocSOLVER.chkargsok.(BlasInt.(collect(dev_info))) + + return (S, U, Vᴴ) + end + end +end + # Wrapper for SVD via Jacobi for (fname, elty, relty) in ( @@ -175,6 +574,261 @@ for (fname, elty, relty) in end end +# Wrapper for batched SVD via Jacobi +for (fname, elty, relty) in + ( + (:rocsolver_sgesvdj_batched, :Float32, :Float32), + (:rocsolver_dgesvdj_batched, :Float64, :Float64), + (:rocsolver_cgesvdj_batched, :ComplexF32, :Float32), + (:rocsolver_zgesvdj_batched, :ComplexF64, :Float64), + ) + @eval begin + function gesvdj_batched!( + A::StridedROCVector{<:StridedROCMatrix{$elty}}, + S::StridedROCMatrix{$relty} = similar(first(A), $relty, (min(size(first(A))...), length(A))), + U::StridedROCArray{$elty, 3} = similar(first(A), $elty, size(first(A), 1), min(size(first(A))...), length(A)), + Vᴴ::StridedROCArray{$elty, 3} = similar(first(A), $elty, min(size(first(A))...), size(first(A), 2), length(A)), + tol::$relty = eps($relty), + max_sweeps::Int = 100, + ) + for A_ in A + chkstride1(A_, U, Vᴴ, S) + end + m, n = size(first(A)) + minmn = min(m, n) + + if length(U) == 0 + jobu = rocSOLVER.rocblas_svect_none + else + size(U, 1) == m || + throw(DimensionMismatch("row size mismatch between A and U")) + if size(U, 2) == minmn + if U === A + throw(ArgumentError("overwrite mode is not supported for gesvdj")) + else + jobu = rocSOLVER.rocblas_svect_singular + end + elseif size(U, 2) == m + jobu = rocSOLVER.rocblas_svect_all + else + throw(DimensionMismatch("invalid column size of U")) + end + end + if length(Vᴴ) == 0 + jobvt = rocSOLVER.rocblas_svect_none + else + size(Vᴴ, 2) == n || + throw(DimensionMismatch("column size mismatch between A and Vᴴ")) + if size(Vᴴ, 1) == minmn + if Vᴴ === A + throw(ArgumentError("overwrite mode is not supported for gesvdj")) + else + jobvt = rocSOLVER.rocblas_svect_singular + end + elseif size(Vᴴ, 1) == n + jobvt = rocSOLVER.rocblas_svect_all + else + throw(DimensionMismatch("invalid row size of Vᴴ")) + end + end + length(S) == minmn * length(A) || + throw(DimensionMismatch("length mismatch between A and S")) + + lda = max(1, stride(first(A), 2)) + ldu = max(1, stride(U, 2)) + strideU = ldu * size(U, 2) + ldv = max(1, stride(Vᴴ, 2)) + strideV = ldv * n + strideS = minmn + dev_info = ROCVector{Cint}(undef, length(A)) + dev_residual = ROCVector{$relty}(undef, length(A)) + dev_n_sweeps = ROCVector{Cint}(undef, length(A)) + + dh = rocBLAS.handle() + pA = map(pointer, A) + rocSOLVER.$fname( + dh, jobu, jobvt, m, n, pA, lda, tol, + dev_residual, max_sweeps, dev_n_sweeps, + S, strideS, U, ldu, strideU, Vᴴ, ldv, strideV, + dev_info, length(A) + ) + + rocSOLVER.chkargsok.(BlasInt.(collect(dev_info))) + + AMDGPU.unsafe_free!(pA) + AMDGPU.unsafe_free!(dev_residual) + AMDGPU.unsafe_free!(dev_n_sweeps) + return (S, U, Vᴴ) + end + end +end + +for (fname, elty, relty) in + ( + (:rocsolver_sgesvdj_strided_batched, :Float32, :Float32), + (:rocsolver_dgesvdj_strided_batched, :Float64, :Float64), + (:rocsolver_cgesvdj_strided_batched, :ComplexF32, :Float32), + (:rocsolver_zgesvdj_strided_batched, :ComplexF64, :Float64), + ) + @eval begin + function gesvdj_strided_batched!( + A::StridedROCArray{$elty, 3}, + S::StridedROCMatrix{$relty} = similar(A, $relty, min(size(A, 1, size(A, 2))), size(A, 3)), + U::StridedROCArray{$elty, 3} = similar(A, $elty, size(A, 1), min(size(A, 1), size(A, 2)), size(A, 3)), + Vᴴ::StridedROCArray{$elty, 3} = similar(A, $elty, min(size(A, 1), size(A, 2)), size(A, 2), size(A, 3)); + tol::$relty = eps($relty), + max_sweeps::Int = 100, + ) + chkstride1(A, U, Vᴴ, S) + m, n, batch_size = size(A) + minmn = min(m, n) + + if length(U) == 0 + jobu = rocSOLVER.rocblas_svect_none + else + size(U, 1) == m || + throw(DimensionMismatch("row size mismatch between A and U")) + if size(U, 2) == minmn + if U === A + throw(ArgumentError("overwrite mode is not supported for gesvdj")) + else + jobu = rocSOLVER.rocblas_svect_singular + end + elseif size(U, 2) == m + jobu = rocSOLVER.rocblas_svect_all + else + throw(DimensionMismatch("invalid column size of U")) + end + end + if length(Vᴴ) == 0 + jobvt = rocSOLVER.rocblas_svect_none + else + size(Vᴴ, 2) == n || + throw(DimensionMismatch("column size mismatch between A and Vᴴ")) + if size(Vᴴ, 1) == minmn + if Vᴴ === A + throw(ArgumentError("overwrite mode is not supported for gesvdj")) + else + jobvt = rocSOLVER.rocblas_svect_singular + end + elseif size(Vᴴ, 1) == n + jobvt = rocSOLVER.rocblas_svect_all + else + throw(DimensionMismatch("invalid row size of Vᴴ")) + end + end + length(S) == minmn * batch_size || + throw(DimensionMismatch("length mismatch between A and S")) + + lda = max(1, stride(A, 2)) + strideA = lda * n + ldu = max(1, stride(U, 2)) + strideU = ldu * size(U, 2) + ldv = max(1, stride(Vᴴ, 2)) + strideV = ldv * n + strideS = minmn + dev_info = ROCVector{Cint}(undef, batch_size) + dev_residual = ROCVector{$relty}(undef, batch_size) + dev_n_sweeps = ROCVector{Cint}(undef, batch_size) + + dh = rocBLAS.handle() + rocSOLVER.$fname( + dh, jobu, jobvt, m, n, A, lda, strideA, tol, + dev_residual, max_sweeps, dev_n_sweeps, + S, strideS, U, ldu, strideU, Vᴴ, ldv, strideV, + dev_info, batch_size + ) + + rocSOLVER.chkargsok.(BlasInt.(collect(dev_info))) + + AMDGPU.unsafe_free!(dev_residual) + AMDGPU.unsafe_free!(dev_n_sweeps) + return (S, U, Vᴴ) + end + end +end + +# Wrapper for SVD via Bisection +for (fname, elty, relty) in + ( + (:rocsolver_sgesvdx, :Float32, :Float32), + (:rocsolver_dgesvdx, :Float64, :Float64), + (:rocsolver_cgesvdx, :ComplexF32, :Float32), + (:rocsolver_zgesvdx, :ComplexF64, :Float64), + ) + @eval begin + function gesvdx!( + A::StridedROCMatrix{$elty}, + S::StridedROCVector{$relty} = similar(A, $relty, min(size(A)...)), + U::StridedROCMatrix{$elty} = similar(A, $elty, size(A, 1), min(size(A)...)), + Vᴴ::StridedROCMatrix{$elty} = similar(A, $elty, min(size(A)...), size(A, 2)); + kwargs... + ) + #! format: on + chkstride1(A, U, Vᴴ, S) + m, n = size(A) + minmn = min(m, n) + if haskey(kwargs, :irange) + irange = convert(UnitRange{Int}, kwargs[:irange]) + il = first(irange) + iu = last(irange) + vl = vu = zero($relty) + range = 'I' + elseif haskey(kwargs, :vl) || haskey(kwargs, :vu) + vl = convert($relty, get(kwargs, :vl, -Inf)) + vu = convert($relty, get(kwargs, :vu, +Inf)) + il = iu = 0 + range = 'V' + else + il = iu = 0 + vl = vu = zero($relty) + range = 'A' + end + + if length(U) == 0 + jobu = 'N' + else + size(U, 1) == m || + throw(DimensionMismatch("row size mismatch between A ($m) and U ($(size(U, 1)))")) + size(U, 2) >= (range == 'I' ? iu - il + 1 : minmn) || + throw(DimensionMismatch("invalid column size of U")) + jobu = 'V' + end + if length(Vᴴ) == 0 + jobvt = 'N' + else + size(Vᴴ, 2) == n || + throw(DimensionMismatch("column size mismatch between A ($n) and Vᴴ ($(size(Vᴴ, 2)))")) + size(Vᴴ, 1) >= (range == 'I' ? iu - il + 1 : minmn) || + throw(DimensionMismatch("invalid row size of Vᴴ")) + jobvt = 'V' + end + length(S) == minmn || + throw(DimensionMismatch("length mismatch between A ($minmn) and S ($(length(S)))")) + + lda = max(1, stride(A, 2)) + ldu = max(1, stride(U, 2)) + ldv = max(1, stride(Vᴴ, 2)) + ifail = ROCVector{Cint}(undef, minmn) + ns = ROCVector{Cint}(undef, 1) + dh = rocBLAS.handle() + dev_info = ROCVector{Cint}(undef, 1) + rocSOLVER.$fname( + dh, jobu, jobvt, m, n, + A, lda, vl, vu, il, iu, nsv, + S, U, ldu, Vᴴ, ldv, ifail, + dev_info + ) + info = @allowscalar dev_info[1] + rocSOLVER.chkargsok(BlasInt(info)) + + AMDGPU.unsafe_free!(ns) + AMDGPU.unsafe_free!(ifail) + return (S, U, Vᴴ) + end + end +end + # for (jname, bname, fname, elty, relty) in # ((:sygvd!, :rocsolverDnSsygvd_bufferSize, :rocsolverDnSsygvd, :Float32, :Float32), # (:sygvd!, :rocsolverDnDsygvd_bufferSize, :rocsolverDnDsygvd, :Float64, :Float64), diff --git a/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl b/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl index 9fa58750f..6251c1a98 100644 --- a/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl +++ b/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl @@ -7,6 +7,7 @@ using MatrixAlgebraKit: diagview, sign_safe using MatrixAlgebraKit: CUSOLVER, LQViaTransposedQR, TruncationByValue, AbstractAlgorithm 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: gesvdj_batched! import MatrixAlgebraKit: heevj!, heevd!, geev! import MatrixAlgebraKit: _gpu_Xgesvdr!, _sylvester, svd_rank, svd_pullback!, eigh_pullback!, eig_pullback!, svd_pushforward! using CUDA, CUDA.cuBLAS @@ -21,6 +22,9 @@ MatrixAlgebraKit.default_driver(::Type{TA}) where {TA <: StridedCuVecOrMat{<:Bla function MatrixAlgebraKit.default_svd_algorithm(::Type{T}; kwargs...) where {T <: StridedCuVecOrMat{<:BlasFloat}} return QRIteration(; kwargs...) end +function MatrixAlgebraKit.default_svd_algorithm(::Type{T}; kwargs...) where {T <: StridedCuArray{<:BlasFloat, 3}} + return JacobiBatched(; kwargs...) +end function MatrixAlgebraKit.default_eig_algorithm(::Type{T}; kwargs...) where {T <: StridedCuVecOrMat{<:BlasFloat}} return QRIteration(; kwargs...) end @@ -35,7 +39,7 @@ end MatrixAlgebraKit.prefers_ungqr(::CUSOLVER) = true -MatrixAlgebraKit.supports_svd_full(::CUSOLVER, f::Symbol) = f in (:qr_iteration, :jacobi, :svd_polar) +MatrixAlgebraKit.supports_svd_full(::CUSOLVER, f::Symbol) = f in (:qr_iteration, :jacobi, :svd_polar, :jacobi_batched) function gesvd!(::CUSOLVER, A::StridedCuMatrix, S::StridedCuVector, U::StridedCuMatrix, Vᴴ::StridedCuMatrix; kwargs...) m, n = size(A) @@ -49,6 +53,9 @@ function gesvdj!(::CUSOLVER, A::StridedCuMatrix, S::StridedCuVector, U::StridedC return MatrixAlgebraKit.svd_via_adjoint!(gesvdj!, CUSOLVER(), A, S, U, Vᴴ; kwargs...) end +gesvdj_batched!(::CUSOLVER, As::StridedCuArray{T, 3}, Ss::StridedCuMatrix, Us::StridedCuArray{T, 3}, Vᴴs::StridedCuArray{T, 3}; kwargs...) where {T <: BlasFloat} = + YACUSOLVER.gesvdj_strided_batched!(As, Ss, Us, Vᴴs; kwargs...) + gesvdp!(::CUSOLVER, A::StridedCuMatrix, S::StridedCuVector, U::StridedCuMatrix, Vᴴ::StridedCuMatrix; kwargs...) = YACUSOLVER.gesvdp!(A, S, U, Vᴴ; kwargs...) diff --git a/ext/MatrixAlgebraKitCUDAExt/yacusolver.jl b/ext/MatrixAlgebraKitCUDAExt/yacusolver.jl index 0cfa64c3c..14dc84d90 100644 --- a/ext/MatrixAlgebraKitCUDAExt/yacusolver.jl +++ b/ext/MatrixAlgebraKitCUDAExt/yacusolver.jl @@ -266,6 +266,92 @@ for (bname, fname, elty, relty) in end end +# Wrapper for batched SVD via Jacobi +for (bname, fname, elty, relty) in + ( + (:cusolverDnSgesvdjBatched_bufferSize, :cusolverDnSgesvdjBatched, :Float32, :Float32), + (:cusolverDnDgesvdjBatched_bufferSize, :cusolverDnDgesvdjBatched, :Float64, :Float64), + (:cusolverDnCgesvdjBatched_bufferSize, :cusolverDnCgesvdjBatched, :ComplexF32, :Float32), + (:cusolverDnZgesvdjBatched_bufferSize, :cusolverDnZgesvdjBatched, :ComplexF64, :Float64), + ) + @eval begin + #! format: off + function gesvdj_batched!( + A::StridedCuArray{$elty, 3}, + S::StridedCuMatrix{$relty} = similar(A, $relty, min(size(A, 1), size(A, 2)), size(A, 3)), + U::StridedCuArray{$elty, 3} = similar(A, $elty, size(A, 1), min(size(A, 1), size(A, 2)), size(A, 3)), + Vᴴ::StridedCuArray{$elty, 3} = similar(A, $elty, min(size(A, 1), size(A, 2)), size(A, 2), size(A, 3)); + tol::$relty = eps($relty), + max_sweeps::Int = 100, + kwargs... + ) + #! format: on + chkstride1(A, U, Vᴴ, S) + m, n, batch_size = size(A) + minmn = min(m, n) + + if length(U) == 0 && length(Vᴴ) == 0 + jobz = 'N' + econ = 0 + else + jobz = 'V' + size(U, 1) == m || + throw(DimensionMismatch("row size mismatch between A and U")) + size(Vᴴ, 2) == n || + throw(DimensionMismatch("column size mismatch between A and Vᴴ")) + if size(U, 2) == size(Vᴴ, 1) == minmn + econ = 1 + elseif size(U, 2) == m && size(Vᴴ, 1) == n + econ = 0 + else + throw(DimensionMismatch("invalid column size of U or row size of Vᴴ")) + end + end + length(S) == minmn || + throw(DimensionMismatch("length mismatch between A and S")) + + Ṽ = (jobz == 'V') ? similar(Vᴴ') : similar(Vᴴ, (n, minmn)) + Ũ = (jobz == 'V') ? U : similar(U, (m, minmn)) + lda = max(1, stride(A, 2)) + ldu = max(1, stride(Ũ, 2)) + ldv = max(1, stride(Ṽ, 2)) + + params = Ref{cuSOLVER.gesvdjInfo_t}(C_NULL) + cuSOLVER.cusolverDnCreateGesvdjInfo(params) + cuSOLVER.cusolverDnXgesvdjSetTolerance(params[], tol) + cuSOLVER.cusolverDnXgesvdjSetMaxSweeps(params[], max_sweeps) + dh = cuSOLVER.dense_handle() + + function bufferSize() + out = Ref{Cint}(0) + cuSOLVER.$bname( + dh, jobz, econ, m, n, A, lda, S, Ũ, ldu, Ṽ, ldv, + out, params[] + ) + return out[] * sizeof($elty) + end + + cuSOLVER.with_workspace(dh.workspace_gpu, bufferSize) do buffer + return cuSOLVER.$fname( + dh, jobz, econ, m, n, A, lda, S, Ũ, ldu, Ṽ, ldv, + buffer, sizeof(buffer) ÷ sizeof($elty), dh.info, + params[], batch_size + ) + end + + info = collect(dh.info) + cuSOLVER.chkargsok.(BlasInt.(info)) + + cuSOLVER.cusolverDnDestroyGesvdjInfo(params[]) + + if jobz == 'V' + adjoint!(Vᴴ, Ṽ) + end + return S, U, Vᴴ + end + end +end + # Wrapper for randomized SVD function gesvdr!( A::StridedCuMatrix{T}, diff --git a/src/MatrixAlgebraKit.jl b/src/MatrixAlgebraKit.jl index 4d4e0084e..db660e771 100644 --- a/src/MatrixAlgebraKit.jl +++ b/src/MatrixAlgebraKit.jl @@ -34,7 +34,7 @@ export exponential, exponential! export Householder, Native_HouseholderQR, Native_HouseholderLQ export DivideAndConquer, SafeDivideAndConquer, QRIteration, Bisection, Jacobi, SVDViaPolar -export RobustRepresentations +export RobustRepresentations, DivideAndConquerBatched, QRIterationBatched, BisectionBatched, JacobiBatched export LAPACK_HouseholderQR, LAPACK_HouseholderLQ, LAPACK_Simple, LAPACK_Expert, LAPACK_QRIteration, LAPACK_Bisection, LAPACK_MultipleRelativelyRobustRepresentations, LAPACK_DivideAndConquer, LAPACK_Jacobi, LAPACK_SafeDivideAndConquer @@ -46,9 +46,10 @@ export DefaultAlgorithm export DiagonalAlgorithm export NativeBlocked export CUSOLVER_Simple, CUSOLVER_HouseholderQR, CUSOLVER_QRIteration, CUSOLVER_SVDPolar, - CUSOLVER_Jacobi, CUSOLVER_Randomized, CUSOLVER_DivideAndConquer + CUSOLVER_Jacobi, CUSOLVER_Randomized, CUSOLVER_DivideAndConquer, CUSOLVER_JacobiBatched export ROCSOLVER_HouseholderQR, ROCSOLVER_QRIteration, ROCSOLVER_Jacobi, - ROCSOLVER_DivideAndConquer, ROCSOLVER_Bisection + ROCSOLVER_DivideAndConquer, ROCSOLVER_Bisection, ROCSOLVER_QRIterationBatched, ROCSOLVER_JacobiBatched, + ROCSOLVER_DivideAndConquerBatched, ROCSOLVER_BisectionBatched export notrunc, truncrank, trunctol, truncerror, truncfilter diff --git a/src/implementations/svd.jl b/src/implementations/svd.jl index 3d20e96d4..9d358cac7 100644 --- a/src/implementations/svd.jl +++ b/src/implementations/svd.jl @@ -1,5 +1,7 @@ # Input # ------ +copy_input(::typeof(svd_full), As::AbstractVector{<:AbstractMatrix}) = map(A -> copy!(similar(A, float(eltype(A))), A), As) +copy_input(::typeof(svd_full), A::AbstractArray{T, 3}) where {T} = copy!(similar(A, float(T)), A) copy_input(::typeof(svd_full), A::AbstractMatrix) = copy!(similar(A, float(eltype(A))), A) copy_input(::typeof(svd_compact), A) = copy_input(svd_full, A) copy_input(::typeof(svd_vals), A) = copy_input(svd_full, A) @@ -42,6 +44,80 @@ function check_input(::typeof(svd_vals!), A::AbstractMatrix, S, ::AbstractAlgori return nothing end +# batched varieties +function check_input(::typeof(svd_full!), A::AbstractVector{<:AbstractMatrix}, USVᴴ, ::AbstractAlgorithm) + @assert all(==(size(first(A))), size.(A)) + m, n = size(first(A)) + batch_size = length(A) + U, S, Vᴴ = USVᴴ + @assert U isa AbstractArray && S isa AbstractArray && Vᴴ isa AbstractArray + @check_size(U, (m, m, batch_size)) + @check_scalar(U, first(A)) + @check_size(S, (m, n, batch_size)) + @check_scalar(S, first(A), real) + @check_size(Vᴴ, (n, n, batch_size)) + @check_scalar(Vᴴ, first(A)) + return nothing +end +function check_input(::typeof(svd_compact!), A::AbstractVector{<:AbstractMatrix}, USVᴴ, ::AbstractAlgorithm) + @assert all(==(size(first(A))), size.(A)) + m, n = size(first(A)) + batch_size = length(A) + minmn = min(m, n) + U, S, Vᴴ = USVᴴ + @assert U isa AbstractArray && S isa AbstractArray && Vᴴ isa AbstractArray + @check_size(U, (m, minmn, batch_size)) + @check_scalar(U, first(A)) + @check_size(S, (minmn, batch_size)) + @check_scalar(S, first(A), real) + @check_size(Vᴴ, (minmn, n, batch_size)) + @check_scalar(Vᴴ, first(A)) + return nothing +end +function check_input(::typeof(svd_vals!), A::AbstractVector{<:AbstractMatrix}, S, ::AbstractAlgorithm) + @assert all(==(size(first(A))), size.(A)) + m, n = size(first(A)) + batch_size = length(A) + minmn = min(m, n) + @assert S isa AbstractMatrix + @check_size(S, (minmn, batch_size)) + @check_scalar(S, first(A), real) + return nothing +end +function check_input(::typeof(svd_full!), A::AbstractArray{T, 3}, USVᴴ, ::AbstractAlgorithm) where {T} + m, n, batch_size = size(A) + U, S, Vᴴ = USVᴴ + @assert U isa AbstractArray && S isa AbstractArray && Vᴴ isa AbstractArray + @check_size(U, (m, m, batch_size)) + @check_scalar(U, A) + @check_size(S, (m, n, batch_size)) + @check_scalar(S, A, real) + @check_size(Vᴴ, (n, n, batch_size)) + @check_scalar(Vᴴ, A) + return nothing +end +function check_input(::typeof(svd_compact!), A::AbstractArray{T, 3}, USVᴴ, ::AbstractAlgorithm) where {T} + m, n, batch_size = size(A) + minmn = min(m, n) + U, S, Vᴴ = USVᴴ + @assert U isa AbstractArray && S isa AbstractArray && Vᴴ isa AbstractArray + @check_size(U, (m, minmn, batch_size)) + @check_scalar(U, A) + @check_size(S, (minmn, batch_size)) + @check_scalar(S, A, real) + @check_size(Vᴴ, (minmn, n, batch_size)) + @check_scalar(Vᴴ, A) + return nothing +end +function check_input(::typeof(svd_vals!), A::AbstractArray{T, 3}, S, ::AbstractAlgorithm) where {T} + m, n, batch_size = size(A) + minmn = min(m, n) + @assert S isa AbstractMatrix + @check_size(S, (minmn, batch_size)) + @check_scalar(S, A, real) + return nothing +end + function check_input(::typeof(svd_full!), A::AbstractMatrix, USVᴴ, ::DiagonalAlgorithm) m, n = size(A) @assert m == n && isdiag(A) @@ -92,6 +168,47 @@ end function initialize_output(::Union{typeof(svd_trunc!), typeof(svd_trunc_no_error!)}, A, alg::TruncatedAlgorithm) return initialize_output(svd_compact!, A, alg.alg) end +# batched versions +function initialize_output(::typeof(svd_full!), A::AbstractVector{<:AbstractMatrix}, ::AbstractAlgorithm) + m, n = size(first(A)) + U = similar(first(A), (m, m, length(A))) + S = similar(first(A), real(eltype(first(A))), (m, n, length(A))) # TODO: Rectangular diagonal type? + Vᴴ = similar(first(A), (n, n, length(A))) + return (U, S, Vᴴ) +end +function initialize_output(::typeof(svd_full!), A::AbstractArray{T, 3}, ::AbstractAlgorithm) where {T} + m, n, batch_size = size(A) + U = similar(A, (m, m, batch_size)) + S = similar(A, real(eltype(A)), (m, n, batch_size)) + Vᴴ = similar(A, (n, n, batch_size)) + return (U, S, Vᴴ) +end +function initialize_output(::typeof(svd_compact!), A::AbstractVector{<:AbstractMatrix}, ::AbstractAlgorithm) + @assert all(==(size(first(A))), size.(A)) + m, n = size(first(A)) + minmn = min(m, n) + U = similar(first(A), (m, minmn, length(A))) + S = similar(first(A), real(eltype(first(A))), minmn, length(A)) + Vᴴ = similar(first(A), (minmn, n, length(A))) + return (U, S, Vᴴ) +end +function initialize_output(::typeof(svd_compact!), A::AbstractArray{T, 3}, ::AbstractAlgorithm) where {T} + m, n, batch_size = size(A) + minmn = min(m, n) + U = similar(A, (m, minmn, batch_size)) + S = similar(A, real(eltype(A)), (minmn, batch_size)) + Vᴴ = similar(A, (minmn, n, batch_size)) + return (U, S, Vᴴ) +end +function initialize_output(::typeof(svd_vals!), A::AbstractVector{<:AbstractMatrix}, ::AbstractAlgorithm) + @assert all(==(size(first(A))), size.(A)) + m, n = size(first(A)) + return similar(first(A), real(eltype(first(A))), (min(m, n), length(A))) +end +function initialize_output(::typeof(svd_vals!), A::AbstractArray{T, 3}, ::AbstractAlgorithm) where {T} + m, n, batch_size = size(A) + return similar(A, real(eltype(A)), (min(m, n), batch_size)) +end function initialize_output(::typeof(svd_full!), A::Diagonal, ::DiagonalAlgorithm) TA = eltype(A) @@ -120,10 +237,17 @@ end # IMPLEMENTATIONS # ========================== -for f! in (:gesdd!, :gesvd!, :gesvdj!, :gesvdp!, :gesvdx!, :gesvdr!, :gesdvd!) +for f! in (:gesdd!, :gesvd!, :gesvdj!, :gesvdp!, :gesvdx!, :gesvdr!, :gesdvd!, :gesdd_batched!, :gesvd_batched!, :gesvdj_batched!, :gesvdx_batched!) @eval $f!(driver::Driver, args...) = throw(ArgumentError("$driver does not provide $($(f!))")) end +# declare these as dummies so the GPU extensions can import them safely +function gesvd_batched! end +function gesdd_batched! end +function gesvdj_batched! end +function gesvdx_batched! end + + """ svd_via_adjoint!(f!, driver, A, S, U, Vᴴ; kwargs...) @@ -220,6 +344,104 @@ for (f, f_lapack!, Alg) in ( end end +# batched varieties +for (f, f_lapack!, Alg) in ( + (:divide_and_conquer_batched, :gesdd_batched!, :DivideAndConquerBatched), + (:qr_iteration_batched, :gesvd_batched!, :QRIterationBatched), + (:bisection_batched, :gesvdx_batched!, :BisectionBatched), + (:jacobi_batched, :gesvdj_batched!, :JacobiBatched), + ) + svd_compact_f! = Symbol(:svd_compact_, f, :!) + svd_full_f! = Symbol(:svd_full_, f, :!) + svd_vals_f! = Symbol(:svd_vals_, f, :!) + + # MatrixAlgebraKit wrappers + @eval begin + function svd_compact!(A::AbstractVector{<:AbstractMatrix}, USVᴴ, alg::$Alg) + check_input(svd_compact!, A, USVᴴ, alg) + return $svd_compact_f!(A, USVᴴ...; alg.kwargs...) + end + function svd_compact!(A::AbstractArray{T, 3}, USVᴴ, alg::$Alg) where {T} + check_input(svd_compact!, A, USVᴴ, alg) + return $svd_compact_f!(A, USVᴴ...; alg.kwargs...) + end + function svd_full!(A::AbstractVector{<:AbstractMatrix}, USVᴴ, alg::$Alg) + check_input(svd_full!, A, USVᴴ, alg) + return $svd_full_f!(A, USVᴴ...; alg.kwargs...) + end + function svd_full!(A::AbstractArray{T, 3}, USVᴴ, alg::$Alg) where {T} + check_input(svd_full!, A, USVᴴ, alg) + return $svd_full_f!(A, USVᴴ...; alg.kwargs...) + end + function svd_vals!(A::AbstractVector{<:AbstractMatrix}, S, alg::$Alg) + check_input(svd_vals!, A, S, alg) + return $svd_vals_f!(A, S; alg.kwargs...) + end + function svd_vals!(A::AbstractArray{T, 3}, S, alg::$Alg) where {T} + check_input(svd_vals!, A, S, alg) + return $svd_vals_f!(A, S; alg.kwargs...) + end + end + + # driver + @eval begin + @inline $svd_compact_f!(A, U, S, Vᴴ; driver::Driver = DefaultDriver(), kwargs...) = $svd_compact_f!(driver, A, U, S, Vᴴ; kwargs...) + @inline $svd_full_f!(A, U, S, Vᴴ; driver::Driver = DefaultDriver(), kwargs...) = $svd_full_f!(driver, A, U, S, Vᴴ; kwargs...) + @inline $svd_vals_f!(A, S; driver::Driver = DefaultDriver(), kwargs...) = $svd_vals_f!(driver, A, S; kwargs...) + @inline $svd_compact_f!(::DefaultDriver, A::AbstractVector{<:AbstractMatrix}, U, S, Vᴴ; kwargs...) = $svd_compact_f!(default_driver($Alg, A), A, U, S, Vᴴ; kwargs...) + @inline $svd_compact_f!(::DefaultDriver, A::AbstractArray{<:Any, 3}, U, S, Vᴴ; kwargs...) = $svd_compact_f!(default_driver($Alg, A), A, U, S, Vᴴ; kwargs...) + @inline $svd_full_f!(::DefaultDriver, A::AbstractVector{<:AbstractMatrix}, U, S, Vᴴ; kwargs...) = $svd_full_f!(default_driver($Alg, A), A, U, S, Vᴴ; kwargs...) + @inline $svd_full_f!(::DefaultDriver, A::AbstractArray{<:Any, 3}, U, S, Vᴴ; kwargs...) = $svd_full_f!(default_driver($Alg, A), A, U, S, Vᴴ; kwargs...) + @inline $svd_vals_f!(::DefaultDriver, A::AbstractVector{<:AbstractMatrix}, S::AbstractMatrix; kwargs...) = $svd_vals_f!(default_driver($Alg, A), A, S; kwargs...) + @inline $svd_vals_f!(::DefaultDriver, A::AbstractArray{<:Any, 3}, S::AbstractMatrix; kwargs...) = $svd_vals_f!(default_driver($Alg, A), A, S; kwargs...) + end + + # Implementation + @eval begin + function $svd_compact_f!(driver::Driver, A, U, S, Vᴴ; fixgauge::Bool = true, kwargs...) + isempty(A) && return one!(U), zero!(S), one!(Vᴴ) + $f_lapack!(driver, A, S, U, Vᴴ; kwargs...) + if fixgauge + for (u, vᴴ) in zip(eachslice(U, dims = 3), eachslice(Vᴴ, dims = 3)) + gaugefix!(svd_compact!, u, vᴴ) + end + end + return U, S, Vᴴ + end + function $svd_full_f!(driver::Driver, A, U, S, Vᴴ; fixgauge::Bool = true, kwargs...) + supports_svd_full(driver, $(QuoteNode(f))) || + throw(ArgumentError(LazyString("driver ", driver, " does not provide `$($(QuoteNode(f_lapack!)))`"))) + isempty(A) && return one!(U), zero!(S), one!(Vᴴ) + zero!(S) + m, n, batch_size = size(S) + minmn = min(m, n) + Sd = similar(S, (minmn, batch_size)) + $f_lapack!(driver, A, Sd, U, Vᴴ; kwargs...) + for (s, sd) in zip(eachslice(S, dims = 3), eachslice(Sd, dims = 2)) + diagview(s) .= sd + end + if fixgauge + for (u, vᴴ) in zip(eachslice(U, dims = 3), eachslice(Vᴴ, dims = 3)) + gaugefix!(svd_full!, u, vᴴ) + end + end + return U, S, Vᴴ + end + function $svd_vals_f!(driver::Driver, A::AbstractArray{T, 3}, S::AbstractMatrix; fixgauge::Bool = true, kwargs...) where {T} + isempty(A) && return zero!(S) + U, Vᴴ = similar(A, (0, 0, 0)), similar(A, (0, 0, 0)) + $f_lapack!(driver, A, S, U, Vᴴ; kwargs...) + return S + end + function $svd_vals_f!(driver::Driver, A::AbstractVector{<:AbstractMatrix}, S::AbstractMatrix; fixgauge::Bool = true, kwargs...) + isempty(A) && return zero!(S) + U, Vᴴ = similar(first(A), (0, 0, 0)), similar(first(A), (0, 0, 0)) + $f_lapack!(driver, A, S, U, Vᴴ; kwargs...) + return S + end + end +end + supports_svd_full(::Driver, ::Symbol) = false supports_svd_full(::LAPACK, f::Symbol) = f in (:safe_divide_and_conquer, :divide_and_conquer, :qr_iteration) diff --git a/src/interface/decompositions.jl b/src/interface/decompositions.jl index 1d423564a..90999fb3c 100644 --- a/src/interface/decompositions.jl +++ b/src/interface/decompositions.jl @@ -100,6 +100,17 @@ The optional `driver` keyword can be used to choose between different implementa """ @algdef DivideAndConquer +""" + DivideAndConquerBatched(; [driver], fixgauge = default_fixgauge()) + +Algorithm type for computing the *batched* eigenvalue decompositions of a set of Hermitian matrices, +or the singular value decompositions of a set of general matrices using the divide-and-conquer algorithm. + +$_fixgauge_docs +The optional `driver` keyword can be used to choose between different implementations of this algorithm. +""" +@algdef DivideAndConquerBatched + """ SafeDivideAndConquer(; [driver], fixgauge = default_fixgauge()) @@ -162,6 +173,49 @@ The optional `driver` keyword can be used to choose between different implementa """ @algdef Jacobi +""" + QRIterationBatched(; [driver], fixgauge = default_fixgauge(), kwargs...) + +Algorithm type for computing the *batched* eigenvalue, Schur or singular value decompositions of a set of matrices via QR iteration. + +## Keyword arguments + +Various customizations are available, depending on the type of decomposition this algorithm is used for. + +Schur decompositions are not yet supported. + +For non-Hermitian eigenvalue decompositions there is `permute = true` and `scale = true` to control whether +or not to balance the input matrix before starting the QR iterations. + +For the singular value and eigenvalue decompositions, there is residual freedom in the outputs that can be resolved. +$_fixgauge_docs + +In all cases, the optional `driver` keyword can be used to choose between different implementations of this algorithm. +""" +@algdef QRIterationBatched + +""" + BisectionBatched(; [driver], fixgauge = default_fixgauge()) + +Algorithm type for computing the *batched* eigenvalue decompositions of a set of Hermitian matrices +via the bisection algorithm, or the singular value decompositions of a set of general matrices. + +$_fixgauge_docs +The optional `driver` keyword can be used to choose between different implementations of this algorithm. +""" +@algdef BisectionBatched + +""" + JacobiBatched(; [driver], fixgauge = default_fixgauge()) + +Algorithm type for computing the *batched* eigenvalue decompositions of a set of Hermitian matrices, +or the singular value decompositions of a set of general matrices using the Jacobi algorithm. + +$_fixgauge_docs +The optional `driver` keyword can be used to choose between different implementations of this algorithm. +""" +@algdef JacobiBatched + """ RobustRepresentations(; [driver], fixgauge = default_fixgauge()) @@ -397,6 +451,15 @@ $_fixgauge_docs """ @algdef CUSOLVER_Jacobi +""" + CUSOLVER_JacobiBatched(; fixgauge = default_fixgauge()) + +Algorithm type to denote the CUSOLVER driver for computing the *batched singular value decompositions +of a set of general matrices using the Jacobi algorithm. +$_fixgauge_docs +""" +@algdef CUSOLVER_JacobiBatched + """ CUSOLVER_Randomized(; k, p, niters) @@ -485,6 +548,45 @@ $_fixgauge_docs """ @algdef ROCSOLVER_DivideAndConquer +""" + ROCSOLVER_QRIterationBatched(; fixgauge = default_fixgauge()) + +Algorithm type to denote the ROCSOLVER driver for computing the *batched* eigenvalue decompositions of a +set of Hermitian matrices, or the singular value decompositions of a set of general matrices using the +QR Iteration algorithm. +$_fixgauge_docs +""" +@algdef ROCSOLVER_QRIterationBatched + +""" + ROCSOLVER_JacobiBatched(; fixgauge = default_fixgauge()) + +Algorithm type to denote the ROCSOLVER driver for computing the *batched* singular value decompositions of +a set of general matrices using the Jacobi algorithm. +$_fixgauge_docs +""" +@algdef ROCSOLVER_JacobiBatched + +""" + ROCSOLVER_BisectionBatched(; fixgauge = default_fixgauge()) + +Algorithm type to denote the ROCSOLVER driver for computing the *batched* eigenvalue decomposition of a +set of Hermitian matrices, or the singular value decompositions of a set of general matrices using the +Bisection algorithm. +$_fixgauge_docs +""" +@algdef ROCSOLVER_BisectionBatched + +""" + ROCSOLVER_DivideAndConquerBatched(; fixgauge = default_fixgauge()) + +Algorithm type to denote the ROCSOLVER driver for computing the *batched* eigenvalue decomposition of a +set of Hermitian matrices, or the singular value decompositions of a set of general matrices using the +Divide and Conquer algorithm. +$_fixgauge_docs +""" +@algdef ROCSOLVER_DivideAndConquerBatched + # Various consts and unions # ------------------------- @@ -509,7 +611,6 @@ const CUSOLVER_SVDAlgorithm = Union{ CUSOLVER_QRIteration, CUSOLVER_SVDPolar, CUSOLVER_Jacobi, CUSOLVER_Randomized, } const GPU_SVDAlgorithm = Union{CUSOLVER_SVDAlgorithm, ROCSOLVER_SVDAlgorithm} - const LAPACK_EighAlgorithm = Union{ LAPACK_QRIteration, LAPACK_Bisection, @@ -524,7 +625,6 @@ const LAPACK_EigAlgorithm = Union{LAPACK_Simple, LAPACK_Expert} const CUSOLVER_EigAlgorithm = Union{CUSOLVER_Simple} const GPU_EigAlgorithm = Union{GPU_Simple} - # List of available algorithms - for docs and convenience purposes const SVDAlgorithms = Union{ SafeDivideAndConquer, diff --git a/test/decompositions/svd.jl b/test/decompositions/svd.jl index c69ed3a0e..8c89744e9 100644 --- a/test/decompositions/svd.jl +++ b/test/decompositions/svd.jl @@ -45,15 +45,21 @@ if !is_buildkite end end +batch_size = 16 + # CUDA tests # ------------ if CUDA.functional() - # LAPACK algorithms: + # CUSOLVER algorithms: for T in BLASFloats, m in (0, 23), n in (0, 17, m, 27) TestSuite.seed_rng!(123) TestSuite.test_svd(CuMatrix{T}, (m, n)) CUDA_SVD_ALGS = (QRIteration(), SVDViaPolar(), Jacobi()) TestSuite.test_svd_algs(CuMatrix{T}, (m, n), CUDA_SVD_ALGS) + + TestSuite.test_svd_batched(CuMatrix{T}, (m, n), batch_size) + CUDA_SVD_ALGS = (JacobiBatched(),) + TestSuite.test_svd_batched_algs(CuMatrix{T}, (m, n), batch_size, CUDA_SVD_ALGS) end # Randomized SVD: @@ -77,12 +83,15 @@ end # AMDGPU tests # ------------ if AMDGPU.functional() - # LAPACK algorithms: - for T in BLASFloats, m in (0, 23), n in (0, 17, m, 27) + # ROCSOLVER algorithms: + for T in BLASFloats, m in (23,), n in (17, m) TestSuite.seed_rng!(123) TestSuite.test_svd(ROCMatrix{T}, (m, n)) - AMD_SVD_ALGS = (QRIteration(), Jacobi()) + AMD_SVD_ALGS = (QRIteration(), Jacobi(), DivideAndConquer()) TestSuite.test_svd_algs(ROCMatrix{T}, (m, n), AMD_SVD_ALGS) + TestSuite.test_svd_batched(ROCMatrix{T}, (m, n), batch_size) + AMD_SVD_ALGS = (QRIterationBatched(), JacobiBatched(), DivideAndConquerBatched()) + TestSuite.test_svd_batched_algs(ROCMatrix{T}, (m, n), batch_size, AMD_SVD_ALGS) end # Diagonal: diff --git a/test/testsuite/TestSuite.jl b/test/testsuite/TestSuite.jl index 36ae68304..653d99f80 100644 --- a/test/testsuite/TestSuite.jl +++ b/test/testsuite/TestSuite.jl @@ -42,6 +42,20 @@ instantiate_matrix(::Type{AT}, size) where {AT <: Diagonal} = Diagonal(randn(rng instantiate_matrix(::Type{AT}, size) where {T, AT <: Diagonal{T, <:CuVector}} = Diagonal(CuArray(randn(rng, eltype(AT), size))) instantiate_matrix(::Type{AT}, size) where {T, AT <: Diagonal{T, <:ROCVector}} = Diagonal(ROCArray(randn(rng, eltype(AT), size))) +# Collect a batch of separately-allocated matrices into a single contiguous batch, transferring +# the array data itself rather than a list of pointers into it. For the GPU eltypes the copies +# stay device-to-device. +function _collect_batch(As::AbstractVector{<:AbstractMatrix}) + B = similar(first(As), (size(first(As))..., length(As))) + for (i, A) in enumerate(As) + copyto!(view(B, :, :, i), A) + end + return B +end +device_batch(As::AbstractVector{<:Array}) = _collect_batch(As) +device_batch(As::AbstractVector{<:CuArray}) = _collect_batch(As) +device_batch(As::AbstractVector{<:ROCArray}) = _collect_batch(As) + precision(::Type{T}) where {T <: Number} = sqrt(eps(real(T))) precision(::Type{T}) where {T} = precision(eltype(T)) diff --git a/test/testsuite/decompositions/svd.jl b/test/testsuite/decompositions/svd.jl index 4b89d4973..1b020de86 100644 --- a/test/testsuite/decompositions/svd.jl +++ b/test/testsuite/decompositions/svd.jl @@ -11,6 +11,16 @@ function test_svd(T::Type, sz; test_compact::Bool = true, test_full::Bool = true end end +function test_svd_batched(T::Type, sz, batch_size::Int; test_compact::Bool = true, test_full::Bool = true, test_trunc::Bool = true, kwargs...) + summary_str = testargs_summary(T, sz) + return @testset "svd batched $summary_str batch_size $batch_size" begin + test_compact && test_svd_compact_batched(T, sz, batch_size; kwargs...) + test_full && test_svd_full_batched(T, sz, batch_size; kwargs...) + # TODO + #test_trunc && test_svd_trunc(T, sz; kwargs...) + end +end + function test_svd_algs(T::Type, sz, algs; test_compact::Bool = true, test_full::Bool = true, test_trunc::Bool = true, kwargs...) summary_str = testargs_summary(T, sz) return @testset "svd algorithms $summary_str" begin @@ -20,6 +30,16 @@ function test_svd_algs(T::Type, sz, algs; test_compact::Bool = true, test_full:: end end +function test_svd_batched_algs(T::Type, sz, batch_size::Int, algs; test_compact::Bool = true, test_full::Bool = true, test_trunc::Bool = true, kwargs...) + summary_str = testargs_summary(T, sz) + return @testset "svd batched algorithms $summary_str batch_size $batch_size" begin + test_compact && test_svd_compact_algs_batched(T, sz, algs, batch_size; kwargs...) + test_full && test_svd_full_algs_batched(T, sz, algs, batch_size; kwargs...) + # TODO + #test_trunc && test_svd_trunc_algs(T, sz, algs; kwargs...) + end +end + function test_svd_compact( T::Type, sz; atol::Real = 0, rtol::Real = precision(eltype(T)), @@ -54,6 +74,47 @@ function test_svd_compact( end end +function test_svd_compact_batched( + T::Type, sz, batch_size::Int; + atol::Real = 0, rtol::Real = precision(eltype(T)), + test_vals::Bool = true, kwargs... + ) + summary_str = testargs_summary(T, sz) + return @testset "svd_compact! $summary_str batch_size $batch_size" begin + As = [instantiate_matrix(T, sz) for bi in 1:batch_size] + Ad = device_batch(As) + Ac = deepcopy(Ad) + m, n = size(first(As)) + minmn = min(m, n) + U, S, Vᴴ = @testinferred svd_compact(Ad) + @test size(U) == (m, minmn, batch_size) + @test S isa AbstractMatrix{real(eltype(T))} && size(S) == (minmn, batch_size) + @test size(Vᴴ) == (minmn, n, batch_size) + for (a, u, s, vᴴ) in zip(As, eachslice(U, dims = 3), eachslice(S, dims = 2), eachslice(Vᴴ, dims = 3)) + @test u * Diagonal(s) * vᴴ ≈ a + @test isisometric(u) + @test isisometric(vᴴ; side = :right) + @test isposdef(Diagonal(s)) + end + + Sc = similar(diagview(S)) + U2, S2, V2ᴴ = @testinferred svd_compact!(Ac, (U, S, Vᴴ)) + for (a, u, s, vᴴ) in zip(As, eachslice(U2, dims = 3), eachslice(S2, dims = 2), eachslice(V2ᴴ, dims = 3)) + @test u * Diagonal(s) * vᴴ ≈ a + @test isisometric(u) + @test isisometric(vᴴ; side = :right) + @test isposdef(Diagonal(s)) + end + + if test_vals + Sd = @testinferred svd_vals(Ad) + for (s, sd) in zip(eachslice(S, dims = 2), eachslice(Sd, dims = 2)) + @test s ≈ sd + end + end + end +end + function test_svd_compact_algs( T::Type, sz, algs; atol::Real = 0, rtol::Real = precision(eltype(T)), @@ -87,6 +148,46 @@ function test_svd_compact_algs( end end +function test_svd_compact_algs_batched( + T::Type, sz, algs, batch_size::Int; + atol::Real = 0, rtol::Real = precision(eltype(T)), + test_vals::Bool = true, kwargs... + ) + summary_str = testargs_summary(T, sz) + return @testset "svd_compact! algorithm $alg $summary_str batch_size $batch_size" for alg in algs + As = [instantiate_matrix(T, sz) for bi in 1:batch_size] + Ad = device_batch(As) + Ac = deepcopy(Ad) + m, n = size(first(As)) + minmn = min(m, n) + U, S, Vᴴ = @testinferred svd_compact(Ad; alg) + @test size(U) == (m, minmn, batch_size) + @test S isa AbstractMatrix{real(eltype(T))} && size(S) == (minmn, batch_size) + @test size(Vᴴ) == (minmn, n, batch_size) + for (a, u, s, vᴴ) in zip(As, eachslice(U, dims = 3), eachslice(S, dims = 2), eachslice(Vᴴ, dims = 3)) + @test u * Diagonal(s) * vᴴ ≈ a + @test isisometric(u) + @test isisometric(vᴴ; side = :right) + @test isposdef(Diagonal(s)) + end + + U2, S2, V2ᴴ = @testinferred svd_compact!(Ac, (U, S, Vᴴ); alg) + for (a, u, s, vᴴ) in zip(As, eachslice(U2, dims = 3), eachslice(S2, dims = 2), eachslice(V2ᴴ, dims = 3)) + @test u * Diagonal(s) * vᴴ ≈ a + @test isisometric(u) + @test isisometric(vᴴ; side = :right) + @test isposdef(Diagonal(s)) + end + + if test_vals + Sd = @testinferred svd_vals(Ad; alg) + for (s, sd) in zip(eachslice(S, dims = 2), eachslice(Sd, dims = 2)) + @test s ≈ sd + end + end + end +end + function test_svd_full( T::Type, sz; atol::Real = 0, rtol::Real = precision(eltype(T)), @@ -120,6 +221,45 @@ function test_svd_full( end end +function test_svd_full_batched( + T::Type, sz, batch_size::Int; + atol::Real = 0, rtol::Real = precision(eltype(T)), + kwargs... + ) + summary_str = testargs_summary(T, sz) + return @testset "svd_full! $summary_str batch_size $batch_size" begin + As = [instantiate_matrix(T, sz) for bi in 1:batch_size] + Ad = device_batch(As) + Ac = deepcopy(Ad) + m, n = size(first(As)) + minmn = min(m, n) + U, S, Vᴴ = @testinferred svd_full(Ad) + @test size(U) == (m, m, batch_size) + @test S isa AbstractArray{real(eltype(T)), 3} && size(S) == (m, n, batch_size) + @test size(Vᴴ) == (n, n, batch_size) + for (a, u, s, vᴴ) in zip(As, eachslice(U, dims = 3), eachslice(S, dims = 3), eachslice(Vᴴ, dims = 3)) + @test u * s * vᴴ ≈ a + @test isunitary(u) + @test isunitary(vᴴ) + @test all(isposdef, diagview(s)) + end + + U2, S2, V2ᴴ = @testinferred svd_full!(Ac, (U, S, Vᴴ)) + for (a, u, s, vᴴ) in zip(As, eachslice(U2, dims = 3), eachslice(S2, dims = 3), eachslice(V2ᴴ, dims = 3)) + @test u * s * vᴴ ≈ a + @test isunitary(u) + @test isunitary(vᴴ) + @test all(isposdef, diagview(s)) + end + + Sc = similar(first(As), real(eltype(T)), min(m, n), batch_size) + Sc2 = @testinferred svd_vals!(copy!(Ac, Ad), Sc) + for (s, s2) in zip(eachslice(S, dims = 3), eachslice(Sc, dims = 2)) + @test collect(diagview(s)) ≈ collect(s2) + end + end +end + function test_svd_full_algs( T::Type, sz, algs; atol::Real = 0, rtol::Real = precision(eltype(T)), @@ -153,6 +293,45 @@ function test_svd_full_algs( end end +function test_svd_full_algs_batched( + T::Type, sz, algs, batch_size::Int; + atol::Real = 0, rtol::Real = precision(eltype(T)), + kwargs... + ) + summary_str = testargs_summary(T, sz) + return @testset "svd_full! algorithm $alg $summary_str batch_size $batch_size" for alg in algs + As = [instantiate_matrix(T, sz) for bi in 1:batch_size] + Ad = device_batch(As) + Ac = deepcopy(Ad) + m, n = size(first(As)) + minmn = min(m, n) + U, S, Vᴴ = @testinferred svd_full(Ad; alg) + @test size(U) == (m, m, batch_size) + @test S isa AbstractArray{real(eltype(T)), 3} && size(S) == (m, n, batch_size) + @test size(Vᴴ) == (n, n, batch_size) + for (a, u, s, vᴴ) in zip(As, eachslice(U, dims = 3), eachslice(S, dims = 3), eachslice(Vᴴ, dims = 3)) + @test u * s * vᴴ ≈ a + @test isunitary(u) + @test isunitary(vᴴ) + @test all(isposdef, diagview(s)) + end + + U2, S2, V2ᴴ = @testinferred svd_full!(Ac, (U, S, Vᴴ); alg) + for (a, u, s, vᴴ) in zip(As, eachslice(U2, dims = 3), eachslice(S2, dims = 3), eachslice(V2ᴴ, dims = 3)) + @test u * s * vᴴ ≈ a + @test isunitary(u) + @test isunitary(vᴴ) + @test all(isposdef, diagview(s)) + end + + Sc = similar(first(As), real(eltype(T)), min(m, n), batch_size) + Sc2 = @testinferred svd_vals!(copy!(Ac, Ad), Sc; alg) + for (s, s2) in zip(eachslice(S, dims = 3), eachslice(Sc, dims = 2)) + @test collect(diagview(s)) ≈ collect(s2) + end + end +end + function test_svd_trunc( T::Type, sz; atol::Real = 0, rtol::Real = precision(eltype(T)), From 2cb322f42f05d6b850bc47d948c3d81c266c101b Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Sat, 22 Aug 2026 20:43:16 +0200 Subject: [PATCH 2/2] Go back to QRIteration default --- ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl | 6 +++--- test/decompositions/svd.jl | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl index 020c43dff..7cdeaff31 100644 --- a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl +++ b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl @@ -20,13 +20,13 @@ MatrixAlgebraKit.default_driver(::Type{TA}) where {TA <: StridedROCArray{<:BlasF MatrixAlgebraKit.default_driver(::Type{TA}) where {TA <: AbstractVector{<:StridedROCMatrix{<:BlasFloat}}} = ROCSOLVER() function MatrixAlgebraKit.default_svd_algorithm(::Type{T}; kwargs...) where {T <: StridedROCMatrix{<:BlasFloat}} - return DivideAndConquer(; kwargs...) + return QRIteration(; kwargs...) end function MatrixAlgebraKit.default_svd_algorithm(::Type{T}; kwargs...) where {T <: StridedROCArray{<:BlasFloat, 3}} - return DivideAndConquerBatched(; kwargs...) + return QRIterationBatched(; kwargs...) end function MatrixAlgebraKit.default_svd_algorithm(::Type{T}; kwargs...) where {T <: AbstractVector{<:StridedROCMatrix{<:BlasFloat}}} - return DivideAndConquerBatched(; kwargs...) + return QRIterationBatched(; kwargs...) end function MatrixAlgebraKit.default_eigh_algorithm(::Type{T}; kwargs...) where {T <: StridedROCVecOrMat{<:BlasFloat}} return DivideAndConquer(; kwargs...) diff --git a/test/decompositions/svd.jl b/test/decompositions/svd.jl index 8c89744e9..d3b6a4425 100644 --- a/test/decompositions/svd.jl +++ b/test/decompositions/svd.jl @@ -84,7 +84,7 @@ end # ------------ if AMDGPU.functional() # ROCSOLVER algorithms: - for T in BLASFloats, m in (23,), n in (17, m) + for T in BLASFloats, m in (0, 23), n in (0, 17, m, 27) TestSuite.seed_rng!(123) TestSuite.test_svd(ROCMatrix{T}, (m, n)) AMD_SVD_ALGS = (QRIteration(), Jacobi(), DivideAndConquer())