Start on batched SVD support

ext/TensorKitGPUArraysExt.jl

@@ -410,7 +410,7 @@ end
         st_dst, offs_dst = @inbounds structs_dst[blk.dst_offset + i + 1]
         i_dst = _linear_index(coords, st_dst, offs_dst)
 
-        # dst_i = β * dst_i + α * Σ_j U[i, j] * permute(src_j, p): each output tree is a
+        # dst_i = β * dst_i + α * Σ_j U[i, j] * permute(op(src_j), p): each output tree is a
         # linear combination of the input trees weighted by the recoupling coefficients.
         # The permutation of src_j was already done by permuting its strides before the
         # kernel launched.

src/factorizations/matrixalgebrakit.jl

@@ -35,8 +35,7 @@ end
 # -----------------------
 for f! in (
         :qr_compact!, :qr_full!, :lq_compact!, :lq_full!,
-        :eig_full!, :eigh_full!, :svd_compact!, :svd_full!,
-        :left_polar!, :right_polar!,
+        :eig_full!, :eigh_full!, :left_polar!, :right_polar!,
     )
     @eval function MAK.$f!(t::AbstractTensorMap, F, alg::AbstractAlgorithm)
         $(f! in (:eig_full!, :eigh_full!) && :(LinearAlgebra.checksquare(t)))
@@ -56,10 +55,38 @@ for f! in (
     end
 end
 
+# these have batched versions, use them if the driver supports it
+# which blocks are worth batching, and how to pack them, is decided by MatrixAlgebraKit
+for (f!, bf!) in ((:svd_compact!, :batched_svd_compact!), (:svd_full!, :batched_svd_full!))
+    @eval function MAK.$f!(t::AbstractTensorMap, F, alg::AbstractAlgorithm)
+        @timeit_debug GLOBAL_TIMER $(string(f!)) begin
+            U, S, Vᴴ = F
+            driver = get(alg.kwargs, :driver, MAK.DefaultDriver())
+            cs = collect(union(blocksectors.((t, U, S, Vᴴ))...))
+            if !isempty(cs) && MAK.supports_ragged_batch(MAK.$f!, alg, driver, storagetype(t))
+                As = [block(t, c) for c in cs]
+                Fs = ([block(U, c) for c in cs], [block(S, c) for c in cs], [block(Vᴴ, c) for c in cs])
+                @timeit_debug GLOBAL_TIMER "batched: MatrixAlgebraKit" MAK.$bf!(As, Fs, alg)
+                return F
+            end
+            foreachblock(t, U, S, Vᴴ) do _, (tblock, Fblocks...)
+                @timeit_debug GLOBAL_TIMER "dense: MatrixAlgebraKit" begin
+                    Fblocks′ = $f!(tblock, Fblocks, alg)
+                    # deal with the case where the output is not in-place
+                    for (b′, b) in zip(Fblocks′, Fblocks)
+                        b === b′ || copy!(b, b′)
+                    end
+                end
+                return nothing
+            end
+            return F
+        end
+    end
+end
+
 # Handle these separately because single output instead of tuple
 for f! in (
-        :qr_null!, :lq_null!,
-        :svd_vals!, :eig_vals!, :eigh_vals!,
+        :qr_null!, :lq_null!, :eig_vals!, :eigh_vals!,
         :project_hermitian!, :project_antihermitian!, :project_isometric!,
         :exponential!,
     )
@@ -79,6 +106,30 @@ for f! in (
     end
 end
 
+# Handle these separately because single output instead of tuple AND batching is supported
+for (f!, bf!) in ((:svd_vals!, :batched_svd_vals!),)
+    @eval function MAK.$f!(t::AbstractTensorMap, S, alg::AbstractAlgorithm)
+        @timeit_debug GLOBAL_TIMER $(string(f!)) begin
+            driver = get(alg.kwargs, :driver, MAK.DefaultDriver())
+            cs = blocksectors(t)
+            if !isempty(cs) && MAK.supports_ragged_batch(MAK.$f!, alg, driver, storagetype(t))
+                As, Ss = last.(blocks(t)), last.(blocks(S))
+                @timeit_debug GLOBAL_TIMER "batched: MatrixAlgebraKit" MAK.$bf!(As, Ss, alg)
+                return S
+            end
+            foreachblock(t, S) do _, (tblock, Sblock)
+                @timeit_debug GLOBAL_TIMER "dense: MatrixAlgebraKit" begin
+                    Sblock′ = $f!(tblock, Sblock, alg)
+                    # deal with the case where the output is not the same as the input
+                    Sblock === Sblock′ || copy!(Sblock, Sblock′)
+                end
+                return nothing
+            end
+        end
+        return S
+    end
+end
+
 # Exponential with Tuple
 function MAK.exponential!((τ, t)::Tuple{E, T}, N, alg::AbstractAlgorithm) where {E <: Number, T <: AbstractTensorMap}
     LinearAlgebra.checksquare(t)