SVD not consistently working with ROCArrays

[abussy]

I have noticed that SVD does not consistently produce correct numerical results with ROCArrays as input. As a result, calculating the condition number of a matrix also fails.

Below is a reproducer. These tests fail more often than not.

using LinearAlgebra
using AMDGPU

cpu = randn(10, 10)
gpu = ROCArray(cpu)

for i in 1:5
    cpu_vals = svd(cpu).S
    gpu_vals = Vector(svd(gpu).S)
    println("Test #$i:")
    println("SVD equal? ", cpu_vals ≈ gpu_vals)
    println("Cond equal? ", cond(cpu) ≈ cond(gpu))
end

I managed to track it down to the LinearAlgebra.LAPACK.gesdd! function being called under the hood. rocSOLVER only implements the generic gesvd!. If I define the following in my code (or the above test), everything works fine:

function LinearAlgebra.LAPACK.gesdd!(jobz::Char, A::AMDGPU.ROCArray{T}) where {T}
    AMDGPU.rocSOLVER.gesvd!(jobz, jobz, A)
end

[luraess]

I am even getting a segfault testing the MWE on the CI machine if not providing your additional method.

[luraess]

We are actually missing a gesdd! (divide-and-conquer) implementation in rocSOLVER currently and svd tries to call into this as default https://github.com/JuliaLang/LinearAlgebra.jl/blob/8d79b34b0423e616dc2899e41299e36f5190bb82/src/svd.jl#L127-L132. If you specify QRIteration as algo then it works as it uses gesvd implementation:

svd(gpu, alg = LinearAlgebra.QRIteration()).S

So current fix is to use non-default algorithm and one could add a wrapper to rocSOLVER high-level for gesdd!.

[luraess]

The issue seems to be rooted slightly deeper (unless I am overlooking something obvious). It seems that in AMDGPU calling into svd calls into LinearAlgebra's svd https://github.com/JuliaLang/LinearAlgebra.jl/blob/8d79b34b0423e616dc2899e41299e36f5190bb82/src/svd.jl#L125-L132 and AMDGPU provides the the methods for the underlying LAPACK.gesvd!(c, c, A) (wrapping https://rocm.docs.amd.com/projects/rocSOLVER/en/latest/reference/lapack.html#rocsolver-type-gesvd) for input type ROCArray, but only for the non-default algorithm QRIteration.

However, the default algorithm for svd in LinearALgebra is DivideAndConquer which calls into gesdd! for which it seems not to be a rocSOLVER implementation.

Looking how e.g. CUDA handles the svd approach, it seems that they provide methods for LinearAlgebra.svd https://github.com/JuliaGPU/CUDA.jl/blob/f1c4cdda21e268028ade0b87947ad358c934e35d/lib/cusolver/linalg.jl#L229-L253 which bypasses the algorithm selection from LinearAlgebra but instead uses a different approach specific to CUDA. There, the default algorithm is JacobiAlgorithm https://github.com/JuliaGPU/CUDA.jl/blob/f1c4cdda21e268028ade0b87947ad358c934e35d/lib/cusolver/linalg.jl#L238-L242 which calls into LAPACK-like function gesvdj!. This Jacobi algorithm is currently available to rocSOLVER (https://rocm.docs.amd.com/projects/rocSOLVER/en/latest/reference/lapacklike.html#rocsolver-type-gesvdj) and wrapped in AMDGPU

AMDGPU.jl/src/solver/highlevel.jl

Lines 376 to 426 in c2efc88

	for (fname, matrix_elty, vector_elty) in (
	(:rocsolver_zgesvdj, :ComplexF64, :Float64),
	(:rocsolver_cgesvdj, :ComplexF32, :Float32),
	(:rocsolver_dgesvdj, :Float64, :Float64),
	(:rocsolver_sgesvdj, :Float32, :Float32),
	)
	@eval begin
	function gesvdj!(A::ROCMatrix{$matrix_elty}, abstol::$vector_elty, max_sweeps::Cint)
	m, n = size(A)
	lda = max(1, stride(A, 2))
	dev_residual = ROCVector{$vector_elty}(undef, 1)
	dev_n_sweeps = ROCVector{Cint}(undef, 1)
	S = ROCArray{$vector_elty}(undef, min(m, n))
	U = ROCMatrix{$matrix_elty}(undef, (m, min(m, n)))
	ldu = m
	@assert stride(U, 2) == ldu
	V = ROCMatrix{$matrix_elty}(undef, (min(m, n), n))
	ldv = min(m, n)
	@assert stride(V, 2) == ldv
	dev_info = ROCVector{Cint}(undef, 1)
	$fname(
	rocBLAS.handle(),
	rocblas_svect_singular,
	rocblas_svect_singular,
	m, n, A, lda,
	abstol,
	dev_residual,
	max_sweeps,
	dev_n_sweeps,
	S,
	U, ldu,
	V, ldv,
	dev_info
	)
	residual = AMDGPU.@allowscalar dev_residual[1]
	AMDGPU.unsafe_free!(dev_residual)
	n_sweeps = AMDGPU.@allowscalar dev_n_sweeps[1]
	AMDGPU.unsafe_free!(dev_n_sweeps)
	info = AMDGPU.@allowscalar dev_info[1]
	AMDGPU.unsafe_free!(dev_info)
	U, S, V', residual, n_sweeps, info
	end
	end
	end

.

According to various sources, it seems discussed that gesdd may not be the most suitable algorithm in many cases (see e.g. https://discourse.julialang.org/t/svd-better-default-to-gesvd-instead-of-gesdd/20603 which introduced the alg keyword).

The question would now be how to proceed here, given that the default option for svd fails in AMDGPU. We may actually want to have something similar to CUDA.jl, leaving the choice between QRIteration and JacobiAlgorithm, which would both call into wrapped rocSOLVER functions, but we may need to change the level a which we dispatch.

I don't have much expertise in the LinALg / rocSOLVER domain, so I'd be happy to weight in suggestions.

[evelyne-ringoot]

Leaving the choice between QR and Jacobi algorithm seems reasonable to me, and in line with CUSOLVER. In terms of default, I did a quick benchmark on an NVIDIA RTX4060 GPU in FP32 for the full svd (time in ms):

size nxn	QR	jacobi
16	1.93	0.20
64	10.7	0.32
256	39.4	0.47
1024	194	5.91
2048	659	790
4096	4,549	7,566
8192	37,788	68,269

And just the singular values:

size nxn	QR	jacobi
16	1.05	0.19
64	4.36	0.25
256	20.3	0.27
1024	100	1.38
2048	282	485
4096	939	6,060
8192	7,218	55,655
16384	56,809	619,332

So I would propose setting the default for SVD to jacobi below 2048x2048, and to QR above that.

[evelyne-ringoot]

I added a draft for svdvals in jacobi in #864 and did a quick benchmark on AMD MI300X as well, and it seems the trend does not hold and in rocSOLVER jacobi should be preferred (results in ms):
Full SVD:

size nxn	QR	jacobi
16	1.67	0.58
64	12.0	4.06
256	165	40.1
1024	2,128	273
2048	8,122	747
4096	30,508	2,989
8192	122,498	16,555

Singular values:

size nxn	QR	jacobi
16	1.23	0.41
64	8.17	3.51
256	94	37
1024	1,141	243
2048	4,375	756
4096	16,422	2,930
8192	62,402	17,444
16384	239,528	118,076

[luraess]

Thanks for investigating @evelyne-ringoot ! From your testing, it seems that Jacobi alg is could be the default in AMDGPU, leaving QRIteration as the second choice.

What if we would move forward and implement those changes in https://github.com/JuliaLang/LinearAlgebra.jl/blob/8d79b34b0423e616dc2899e41299e36f5190bb82/src/svd.jl#L91 and following lines for svd, together with the additions from #864 ?