Batched SVD support for ROCSOLVER and CUSOLVER

[kshyatt]

Basically what it says on the tin. For CTMRG and other algorithms, we're getting absolutely slaughtered on GPU performance for TensorMaps with sectors because we have to spin up huge numbers of very small SVDs. I'm wrapping the batched SVDs each library provides to try to address this. Extremely open to comments but I wanted to get this rolling so I can unblock others.

[kshyatt]

TODOs here:

• Finish wrapping the BisectionBatched logic for AMD
• Add the checks for CUSOLVER gesvdj_batched (blocks may not be larger than 32 x 32)
• Finish the tests for svd_trunc!
• Pullback/pushforward rules

Since I'll be out for 3 weeks everyone should feel free to just push to this.

[lkdvos]

Overall looks great, I think the main thing I am wondering about is about the interface decision around wether we implement this as a set of BatchedAlg versions, or rather as a set of batched_f(...) functions. I definitely like using dispatch for switching between the strided and non-strided inputs, but I am wondering if there might be a benefit to really having a batched_svd_compact etc function.

This is also partially since that allows us to have a CPU version for this as well, so we can just offload all of this from TensorKit to here, (and possibly play with multithreading?).

[Jutho]

I mostly agree with Lukas here, I think I would prefer

• separate batched_f methods instead of separate algorithms (with possible options provided in the driver rather than the algorithm if we need those)
• support for varying sizes, which we will anyway need to implement, so we can do it on the MAK side I would think

[kshyatt]

support for varying sizes, which we will anyway need to implement, so we can do it on the MAK side I would think

This is already implemented in a separate branch I have over at TensorKit, I can move it over here of course!

[kshyatt]

Latest commit gets rid of the JacobiBatched and friends algos, and adds instead batched_svd_compact etc. Remaining TODOs:

• Add tests for batches with varying sizes of input matrix (done!)
• Add padding support in here, rather than in TensorKit (done!)
• Add support for svd_via_adjoint! for batched arrays (done!)

[codecov]

Codecov Report

❌ Patch coverage is 80.75916% with 147 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
ext/MatrixAlgebraKitAMDGPUExt/yarocsolver.jl	65.18%	133 Missing ⚠️
...ixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl	78.12%	7 Missing ⚠️
...MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl	83.33%	2 Missing ⚠️
src/implementations/batched_svd.jl	99.32%	2 Missing ⚠️
src/interface/batched_svd.jl	0.00%	2 Missing ⚠️
ext/MatrixAlgebraKitCUDAExt/yacusolver.jl	97.22%	1 Missing ⚠️

Files with missing lines	Coverage Δ
src/MatrixAlgebraKit.jl	100.00% <ø> (ø)
src/implementations/svd.jl	95.47% <100.00%> (+0.12%)	⬆️
ext/MatrixAlgebraKitCUDAExt/yacusolver.jl	96.25% <97.22%> (+0.13%)	⬆️
...MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl	81.08% <83.33%> (+0.12%)	⬆️
src/implementations/batched_svd.jl	99.32% <99.32%> (ø)
src/interface/batched_svd.jl	0.00% <0.00%> (ø)
...ixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl	66.27% <78.12%> (+3.12%)	⬆️
ext/MatrixAlgebraKitAMDGPUExt/yarocsolver.jl	72.18% <65.18%> (-15.92%)	⬇️

... and 8 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[kshyatt]

One remaining thing to do here is restore the specific path that targets the AMDGPU ROCVector{<:ROCMatrix} path, otherwise I think things are looking ok. stack not only isn't supported on GPU, it doesn't work for the ragged case (unless we pad separately for each input, which I think is inefficient?).