diff --git a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl index 976c1dc04..43b6601f1 100644 --- a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl +++ b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl @@ -41,10 +41,14 @@ MatrixAlgebraKit.supports_svd_full(::ROCSOLVER, f::Symbol) = f in (:qr_iteration # rocSOLVER's `gesvd*_batched` functions take the batch as an array # of device pointers, so a group of equally sized matrices # doesn't need a copy into a 3D ROCArray. -MatrixAlgebraKit.supports_pointer_batch(::AbstractAlgorithm, ::ROCSOLVER, ::Type{<:StridedROCMatrix{<:BlasFloat}}) = true -MatrixAlgebraKit.supports_ragged_batch(::typeof(svd_full!), ::AbstractAlgorithm, ::ROCSOLVER, ::Type{<:AnyROCArray}) = true -MatrixAlgebraKit.supports_ragged_batch(::typeof(svd_compact!), ::AbstractAlgorithm, ::ROCSOLVER, ::Type{<:AnyROCArray}) = true -MatrixAlgebraKit.supports_ragged_batch(::typeof(svd_vals!), ::AbstractAlgorithm, ::ROCSOLVER, ::Type{<:AnyROCArray}) = true +# opt-in per algorithm to avoid picking up `DiagonalAlgorithm` +# which doesn't support batching +for alg in (:Bisection, :Jacobi, :QRIteration, :DivideAndConquer) + @eval MatrixAlgebraKit.supports_pointer_batch(::$alg, ::ROCSOLVER, ::Type{<:StridedROCMatrix{<:BlasFloat}}) = true + @eval MatrixAlgebraKit.supports_ragged_batch(::typeof(svd_full!), ::$alg, ::ROCSOLVER, ::Type{<:AnyROCArray}) = true + @eval MatrixAlgebraKit.supports_ragged_batch(::typeof(svd_compact!), ::$alg, ::ROCSOLVER, ::Type{<:AnyROCArray}) = true + @eval MatrixAlgebraKit.supports_ragged_batch(::typeof(svd_vals!), ::$alg, ::ROCSOLVER, ::Type{<:AnyROCArray}) = true +end function gesvd!(::ROCSOLVER, A::StridedROCMatrix, S::StridedROCVector, U::StridedROCMatrix, Vá´´::StridedROCMatrix; kwargs...) m, n = size(A) diff --git a/test/decompositions/svd.jl b/test/decompositions/svd.jl index b7a4a6ef3..8c0c671f8 100644 --- a/test/decompositions/svd.jl +++ b/test/decompositions/svd.jl @@ -68,6 +68,7 @@ if CUDA.functional() TestSuite.test_svd_algs_batched_oversized(CuMatrix{T}, (Jacobi(),), batch_size) TestSuite.test_svd_algs_batched_ragged_support(CuMatrix{T}, (Jacobi(),), batch_size) TestSuite.test_svd_algs_batched_ragged_support(CuMatrix{T}, (QRIteration(),), batch_size; supported = false) + TestSuite.test_svd_algs_batched_ragged_support(Diagonal{T, CuVector{T}}, (DiagonalAlgorithm(),), batch_size; supported = false) end # Randomized SVD: @@ -104,6 +105,7 @@ if AMDGPU.functional() TestSuite.seed_rng!(123) AMD_SVD_ALGS = (QRIteration(), Jacobi(), DivideAndConquer(), Bisection()) TestSuite.test_svd_algs_batched_ragged_support(ROCMatrix{T}, AMD_SVD_ALGS, batch_size) + TestSuite.test_svd_algs_batched_ragged_support(Diagonal{T, ROCVector{T}}, (DiagonalAlgorithm(),), batch_size; supported = false) end @testset "Bisection with min(m, n) == 1 $(testargs_summary(T, sz))" for T in BLASFloats,