From c54c80cb7ee90b4cb367871cb01c5bb71cf72cd3 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Sat, 22 Aug 2026 21:04:02 +0200 Subject: [PATCH 1/6] Start on batched SVD support --- Project.toml | 3 + ext/TensorKitGPUArraysExt.jl | 2 +- src/factorizations/batched.jl | 85 ++++++++++++++++++ src/factorizations/factorizations.jl | 1 + src/factorizations/matrixalgebrakit.jl | 52 ++++++++++- test/amd/batched_svd.jl | 118 +++++++++++++++++++++++++ test/cuda/batched_svd.jl | 109 +++++++++++++++++++++++ 7 files changed, 367 insertions(+), 3 deletions(-) create mode 100644 src/factorizations/batched.jl create mode 100644 test/amd/batched_svd.jl create mode 100644 test/cuda/batched_svd.jl diff --git a/Project.toml b/Project.toml index d1e0c834d..9fafc118f 100644 --- a/Project.toml +++ b/Project.toml @@ -45,6 +45,9 @@ TensorKitFiniteDifferencesExt = "FiniteDifferences" TensorKitGPUArraysExt = "GPUArrays" TensorKitMooncakeExt = "Mooncake" +[sources] +MatrixAlgebraKit = {url = "https://github.com/QuantumKitHub/MatrixAlgebraKit.jl", rev = "main"} + [compat] AMDGPU = "2" Adapt = "4" diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 7843460c3..43857e3ae 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -411,7 +411,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. diff --git a/src/factorizations/batched.jl b/src/factorizations/batched.jl new file mode 100644 index 000000000..3124045ac --- /dev/null +++ b/src/factorizations/batched.jl @@ -0,0 +1,85 @@ +# Figure out which sectors are even worth batching, and if some share a batch size +# `uniform = true` additionally demands that every block already has that exact size, i.e. +# that no padding is needed. Full decompositions require this, compact decompositions only +# read back the leading `min(m, n)` columns, which padding out with zero doesn't affect. +function _batchable(t::AbstractTensorMap, alg::AbstractAlgorithm, uniform::Bool = false) + cs = collect(blocksectors(t)) + isempty(cs) && return cs, (0, 0) + szs = [size(block(t, c)) for c in cs] + m, n = maximum(first, szs), maximum(last, szs) + lim = MAK.max_batched_blocksize(alg, storagetype(t)) + (length(cs) < MAK.BATCHED_SVD_THRESHOLD || m > lim || n > lim) && return empty(cs), (m, n) + # The *padded* batch is (m, n) even if *individual* blocks are tall + # so check the padded shape rather than the blocks'. + (MAK.requires_tall(alg) && m < n) && return empty(cs), (m, n) + (uniform && !all(isequal((m, n)), szs)) && return empty(cs), (m, n) + return cs, (m, n) +end + +# The batched solvers work on 3D arrays: input `(m, n, nbatch)`, and correspondingly +# `U = (m, minmn, nbatch)`, `S = (minmn, nbatch)`, `Vᴴ = (minmn, n, nbatch)`. +function _pack(t::AbstractTensorMap, cs, m, n) + b1 = block(t, first(cs)) + A = similar(b1, m, n, length(cs)) + fill!(A, zero(scalartype(A))) + for (i, c) in enumerate(cs) + b = block(t, c) + copyto!(view(A, axes(b, 1), axes(b, 2), i), b) + end + return A +end + +for (bf!, f!) in ((:batched_svd_compact!, :svd_compact!), (:batched_svd_full!, :svd_full!)) + full = bf! === :batched_svd_full! + @eval function MAK.$bf!(t::AbstractTensorMap, F, alg) + U, S, Vᴴ = F + cs, (m, n) = _batchable(t, alg, $full) + if isempty(cs) # not worth batching, or the library doesn't support these sizes + foreachblock(t, U, S, Vᴴ) do _, (b, u, sv, v) + MAK.$f!(b, (u, sv, v), alg) + return nothing + end + return F + end + nb, minmn = length(cs), min(m, n) + A = _pack(t, cs, m, n) + Ub = similar(A, m, $full ? m : minmn, nb) + rT = real(scalartype(t)) + Sb = $full ? similar(A, rT, m, n, nb) : + similar(diagview(block(S, first(cs))), minmn, nb) + Vb = similar(A, $full ? n : minmn, n, nb) + $full && fill!(Sb, zero(rT)) + MAK.$bf!(A, (Ub, Sb, Vb), alg) + for (i, c) in enumerate(cs) + u, sv, v = block(U, c), block(S, c), block(Vᴴ, c) + copyto!(u, view(Ub, axes(u, 1), axes(u, 2), i)) + if $full + copyto!(sv, view(Sb, axes(sv, 1), axes(sv, 2), i)) + else + copyto!(diagview(sv), view(Sb, axes(diagview(sv), 1), i)) + end + copyto!(v, view(Vb, axes(v, 1), axes(v, 2), i)) + end + return F + end +end + +function MAK.batched_svd_vals!(t::AbstractTensorMap, S, alg) + cs, (M, N) = _batchable(t, alg) + if isempty(cs) + foreachblock(t, S) do _, (b, sv) + MAK.svd_vals!(b, sv, alg) + return nothing + end + return S + end + nb, K = length(cs), min(M, N) + A = _pack(t, cs, M, N) + Sb = similar(block(S, first(cs)), K, nb) + MAK.batched_svd_vals!(A, Sb, alg) + for (i, c) in enumerate(cs) + sv = block(S, c) + copyto!(sv, view(Sb, axes(sv, 1), i)) + end + return S +end diff --git a/src/factorizations/factorizations.jl b/src/factorizations/factorizations.jl index e49cc8b29..0cdf64048 100644 --- a/src/factorizations/factorizations.jl +++ b/src/factorizations/factorizations.jl @@ -28,6 +28,7 @@ using MatrixAlgebraKit: diagview include("utility.jl") include("matrixalgebrakit.jl") +include("batched.jl") include("truncation.jl") include("adjoint.jl") include("diagonal.jl") diff --git a/src/factorizations/matrixalgebrakit.jl b/src/factorizations/matrixalgebrakit.jl index 5974a60de..52f3fd0f4 100644 --- a/src/factorizations/matrixalgebrakit.jl +++ b/src/factorizations/matrixalgebrakit.jl @@ -1,8 +1,20 @@ # Algorithm selection # ------------------- + +""" + _tensor_algorithm(f!, ::Type{<:AbstractTensorMap}; kwargs...) + +Algorithm a `TensorMap` factorization defaults to. Blocks are decomposed one at a time, so +the default is whatever the block type would use. For algorithms that have a batched version, +like `QRIteration` or `Jacobi`, this can be overridden to point to the batching version. +""" +function _tensor_algorithm(f!, ::Type{T}; kwargs...) where {T <: AbstractTensorMap} + return MAK.default_algorithm(f!, blocktype(T); kwargs...) +end for f in [ :svd_compact, :svd_full, :svd_vals, + :batched_svd_compact, :batched_svd_full, :batched_svd_vals, :qr_compact, :qr_full, :qr_null, :lq_compact, :lq_full, :lq_null, :eig_full, :eig_vals, :eigh_full, :eigh_vals, @@ -12,7 +24,7 @@ for f in ] f! = Symbol(f, :!) @eval function MAK.default_algorithm(::typeof($f!), ::Type{T}; kwargs...) where {T <: AbstractTensorMap} - return MAK.default_algorithm($f!, blocktype(T); kwargs...) + return _tensor_algorithm($f!, T; kwargs...) end @eval function MAK.copy_input(::typeof($f), t::AbstractTensorMap) return @timeit_debug GLOBAL_TIMER "alloc: copy_input" copy_oftype( @@ -58,7 +70,7 @@ end # Handle these separately because single output instead of tuple for f! in ( - :qr_null!, :lq_null!, + :qr_null!, :lq_null!, :batched_svd_vals, :svd_vals!, :eig_vals!, :eigh_vals!, :project_hermitian!, :project_antihermitian!, :project_isometric!, :exponential!, @@ -107,6 +119,7 @@ for f in [ :eig_full, :eig_vals, :eig_trunc, :eig_trunc_no_error, :eigh_full, :eigh_vals, :eigh_trunc, :eigh_trunc_no_error, :svd_full, :svd_compact, :svd_trunc, :svd_trunc_no_error, :svd_vals, + :batched_svd_full, :batched_svd_compact, :batched_svd_vals, :left_polar, :right_polar, :left_orth, :right_orth, :left_null, :right_null, :project_hermitian, :project_antihermitian, :project_isometric, @@ -169,6 +182,41 @@ function MAK.initialize_output(::typeof(svd_vals!), t::AbstractTensorMap, alg::A end end +# Batched singular value decomposition +# ------------------------------------ +# We should be able to just reuse the +# above and handle the actual batched +# object filling deeper down +function MAK.initialize_output(::typeof(batched_svd_full!), t::AbstractTensorMap, ::AbstractAlgorithm) + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_cod = fuse(codomain(t)) + V_dom = fuse(domain(t)) + U = similar(t, codomain(t) ← V_cod) + S = similar(t, real(scalartype(t)), V_cod ← V_dom) + Vᴴ = similar(t, V_dom ← domain(t)) + return U, S, Vᴴ + end +end + +function MAK.initialize_output(::typeof(batched_svd_compact!), t::AbstractTensorMap, ::AbstractAlgorithm) + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_cod = V_dom = infimum(fuse(codomain(t)), fuse(domain(t))) + U = similar(t, codomain(t) ← V_cod) + S = similar_diagonal(t, real(scalartype(t)), V_cod) + Vᴴ = similar(t, V_dom ← domain(t)) + return U, S, Vᴴ + end +end + +function MAK.initialize_output(::typeof(batched_svd_vals!), t::AbstractTensorMap, alg::AbstractAlgorithm) + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_cod = infimum(fuse(codomain(t)), fuse(domain(t))) + T = real(scalartype(t)) + A = similarstoragetype(t, T) + return SectorVector{T, sectortype(t), A}(undef, V_cod) + end +end + # Eigenvalue decomposition # ------------------------ function MAK.initialize_output(::typeof(eigh_full!), t::AbstractTensorMap, ::AbstractAlgorithm) diff --git a/test/amd/batched_svd.jl b/test/amd/batched_svd.jl new file mode 100644 index 000000000..d0eb9ab78 --- /dev/null +++ b/test/amd/batched_svd.jl @@ -0,0 +1,118 @@ +using Adapt, AMDGPU +using Test, TestExtras +using TensorKit +using LinearAlgebra: LinearAlgebra +using MatrixAlgebraKit: MatrixAlgebraKit, DivideAndConquer, QRIteration, Jacobi, + svd_compact, svd_compact!, svd_vals, svd_vals!, svd_full, svd_full!, + batched_svd_compact, batched_svd_compact!, batched_svd_vals, + batched_svd_vals!, batched_svd_full, batched_svd_full! + +const Factorizations = TensorKit.Factorizations + +# SU(2) space with enough sectors to exceed the batching threshold +Vsu2 = Vect[SU2Irrep](0 => 2, 1 // 2 => 2, 1 => 2, 3 // 2 => 1) + +# spectra are returned either +# as `Diagonal` blocks (`svd_compact`) or +# as plain vectors (`svd_vals`) +_spec(b) = Array(b isa AbstractMatrix ? TensorKit.diagview(b) : b) + +function specdiff(S1, S2) + d = 0.0 + for (c, b) in TensorKit.blocks(S1) + d = max(d, maximum(abs, _spec(b) .- _spec(TensorKit.block(S2, c)))) + end + return d +end + +_id(t, V) = TensorKit.id(TensorKit.storagetype(t), V) + +@timedtestset "batched SVD on ROCArray" verbose = true begin + @testset "many blocks: $T" for T in (Float64, ComplexF64) + t_cpu = randn(T, Vsu2 ⊗ Vsu2 ← Vsu2) + nblocks = length(TensorKit.blocksectors(t_cpu)) + @test nblocks >= MatrixAlgebraKit.BATCHED_SVD_THRESHOLD + t = adapt(ROCArray, t_cpu) + + U, S, Vᴴ = batched_svd_compact(t; alg = DivideAndConquer()) + Ur, Sr, Vr = svd_compact(t; alg = DivideAndConquer()) + + # singular values agree with the unbatched path + @test specdiff(S, Sr) < 1.0e-10 + @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 + @test norm(U' * U - _id(U, domain(U))) < 1.0e-10 + @test norm(Vᴴ * Vᴴ' - _id(Vᴴ, codomain(Vᴴ))) < 1.0e-10 + + Sv = batched_svd_vals(t; alg = DivideAndConquer()) + @test specdiff(Sv, Sr) < 1.0e-10 + end + + @testset "few blocks fall back: $T" for T in (Float64, ComplexF64) + # trivial sector -> a single block, below the threshold + V = ComplexSpace(6) + t = adapt(ROCArray, randn(T, V ⊗ V ← V)) + @test length(TensorKit.blocksectors(t)) < MatrixAlgebraKit.BATCHED_SVD_THRESHOLD + U, S, Vᴴ = batched_svd_compact(t; alg = DivideAndConquer()) + Ur, Sr, Vr = svd_compact(t; alg = DivideAndConquer()) + @test specdiff(S, Sr) < 1.0e-10 + @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 + end + + @testset "other batched algorithms" for alg in (QRIteration(), Jacobi()) + t_cpu = randn(Float64, Vsu2 ⊗ Vsu2 ← Vsu2) + t = adapt(ROCArray, t_cpu) + U, S, Vᴴ = batched_svd_compact(t; alg) + _, Sr, _ = svd_compact(t; alg = DivideAndConquer()) + @test specdiff(S, Sr) < 1.0e-8 + @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-8 + end + + # `svd_full!` returns U (m, m) and Vᴴ (n, n) rather than the compact (m, minmn) / (minmn, n) + @testset "svd_full: uniform blocks batch: $T" for T in (Float64, ComplexF64) + Vu = Vect[Z4Irrep](0 => 3, 1 => 3, 2 => 3, 3 => 3) + t_cpu = randn(T, Vu ← Vu) + t = adapt(ROCArray, t_cpu) + szs = [size(TensorKit.block(t, c)) for c in TensorKit.blocksectors(t)] + @test length(szs) >= MatrixAlgebraKit.BATCHED_SVD_THRESHOLD + @test all(isequal(first(szs)), szs) + cs, _ = Factorizations._batchable(t, QRIteration(), true) + @test !isempty(cs) + + U, S, Vᴴ = batched_svd_full(t; alg = QRIteration()) + Ur, Sr, Vr = svd_full(t; alg = QRIteration()) + @test specdiff(S, Sr) < 1.0e-10 + @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 + @test norm(U' * U - _id(U, domain(U))) < 1.0e-10 + @test norm(U * U' - _id(U, codomain(U))) < 1.0e-10 + @test norm(Vᴴ' * Vᴴ - _id(Vᴴ, domain(Vᴴ))) < 1.0e-10 + end + + @testset "svd_full: ragged blocks fall back: $T" for T in (Float64,) + t = adapt(ROCArray, randn(T, Vsu2 ⊗ Vsu2 ← Vsu2)) + szs = [size(TensorKit.block(t, c)) for c in TensorKit.blocksectors(t)] + @test !all(isequal(first(szs)), szs) + cs, _ = Factorizations._batchable(t, QRIteration(), true) + @test isempty(cs) + U, S, Vᴴ = batched_svd_full(t; alg = QRIteration()) + Ur, Sr, Vr = svd_full(t; alg = QRIteration()) + @test specdiff(S, Sr) < 1.0e-10 + @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 + end + + # tall uniform blocks: hits the full mode where U (m, m) and Vᴴ (n, n) have different sizes + @testset "svd_full: tall uniform blocks: $T" for T in (Float64,) + Vbig = Vect[Z4Irrep](0 => 4, 1 => 4, 2 => 4, 3 => 4) + Vsml = Vect[Z4Irrep](0 => 2, 1 => 2, 2 => 2, 3 => 2) + t = adapt(ROCArray, randn(T, Vbig ← Vsml)) + szs = [size(TensorKit.block(t, c)) for c in TensorKit.blocksectors(t)] + @test all(isequal((4, 2)), szs) + cs, _ = Factorizations._batchable(t, QRIteration(), true) + @test !isempty(cs) + U, S, Vᴴ = batched_svd_full(t; alg = QRIteration()) + Ur, Sr, Vr = svd_full(t; alg = QRIteration()) + @test specdiff(S, Sr) < 1.0e-10 + @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 + @test norm(U' * U - _id(U, domain(U))) < 1.0e-10 + @test norm(Vᴴ * Vᴴ' - _id(Vᴴ, codomain(Vᴴ))) < 1.0e-10 + end +end diff --git a/test/cuda/batched_svd.jl b/test/cuda/batched_svd.jl new file mode 100644 index 000000000..966bfa687 --- /dev/null +++ b/test/cuda/batched_svd.jl @@ -0,0 +1,109 @@ +using Adapt, CUDA +using Test, TestExtras +using TensorKit +using LinearAlgebra: LinearAlgebra +using MatrixAlgebraKit: MatrixAlgebraKit, Jacobi, QRIteration, + svd_compact, svd_compact!, svd_vals, svd_vals!, svd_full, svd_full!, + batched_svd_compact, batched_svd_compact!, batched_svd_vals, + batched_svd_vals!, batched_svd_full, batched_svd_full! + +const Factorizations = TensorKit.Factorizations + +# SU(2) space with enough sectors to exceed the batching threshold +Vsu2 = Vect[SU2Irrep](0 => 2, 1 // 2 => 2, 1 => 2, 3 // 2 => 1) + +# spectra are returned either +# as `Diagonal` blocks (`svd_compact`) or +# as plain vectors (`svd_vals`) +_spec(b) = Array(b isa AbstractMatrix ? TensorKit.diagview(b) : b) + +function specdiff(S1, S2) + d = 0.0 + for (c, b) in TensorKit.blocks(S1) + d = max(d, maximum(abs, _spec(b) .- _spec(TensorKit.block(S2, c)))) + end + return d +end + +_id(t, V) = TensorKit.id(TensorKit.storagetype(t), V) + +@timedtestset "batched SVD on CuArray" verbose = true begin + @testset "many blocks: $T" for T in (Float64, ComplexF64) + t_cpu = randn(T, Vsu2 ⊗ Vsu2 ← Vsu2) + nblocks = length(TensorKit.blocksectors(t_cpu)) + @test nblocks >= MatrixAlgebraKit.BATCHED_SVD_THRESHOLD + t = adapt(CuArray, t_cpu) + + U, S, Vᴴ = batched_svd_compact(t; alg = Jacobi()) + Ur, Sr, Vr = svd_compact(t; alg = Jacobi()) + + # singular values agree with the unbatched path + @test specdiff(S, Sr) < 1.0e-10 + @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 + @test norm(U' * U - _id(U, domain(U))) < 1.0e-10 + @test norm(Vᴴ * Vᴴ' - _id(Vᴴ, codomain(Vᴴ))) < 1.0e-10 + + Sv = batched_svd_vals(t; alg = Jacobi()) + @test specdiff(Sv, Sr) < 1.0e-10 + end + + @testset "few blocks fall back: $T" for T in (Float64, ComplexF64) + # trivial sector -> a single block, below the threshold + V = ComplexSpace(6) + t = adapt(CuArray, randn(T, V ⊗ V ← V)) + @test length(TensorKit.blocksectors(t)) < MatrixAlgebraKit.BATCHED_SVD_THRESHOLD + U, S, Vᴴ = batched_svd_compact(t; alg = Jacobi()) + Ur, Sr, Vr = svd_compact(t; alg = Jacobi()) + @test specdiff(S, Sr) < 1.0e-10 + @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 + end + + # `svd_full!` returns U (m, m) and Vᴴ (n, n) rather than the compact (m, minmn) / (minmn, n) + @testset "svd_full: uniform blocks batch: $T" for T in (Float64, ComplexF64) + Vu = Vect[Z4Irrep](0 => 3, 1 => 3, 2 => 3, 3 => 3) + t_cpu = randn(T, Vu ← Vu) + t = adapt(CuArray, t_cpu) + szs = [size(TensorKit.block(t, c)) for c in TensorKit.blocksectors(t)] + @test length(szs) >= MatrixAlgebraKit.BATCHED_SVD_THRESHOLD + @test all(isequal(first(szs)), szs) + cs, _ = Factorizations._batchable(t, Jacobi(), true) + @test !isempty(cs) + + U, S, Vᴴ = batched_svd_full(t; alg = Jacobi()) + Ur, Sr, Vr = svd_full(t; alg = Jacobi()) + @test specdiff(S, Sr) < 1.0e-10 + @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 + @test norm(U' * U - _id(U, domain(U))) < 1.0e-10 + @test norm(U * U' - _id(U, codomain(U))) < 1.0e-10 + @test norm(Vᴴ' * Vᴴ - _id(Vᴴ, domain(Vᴴ))) < 1.0e-10 + end + + @testset "svd_full: ragged blocks fall back: $T" for T in (Float64,) + t = adapt(CuArray, randn(T, Vsu2 ⊗ Vsu2 ← Vsu2)) + szs = [size(TensorKit.block(t, c)) for c in TensorKit.blocksectors(t)] + @test !all(isequal(first(szs)), szs) + cs, _ = Factorizations._batchable(t, Jacobi(), true) + @test isempty(cs) + U, S, Vᴴ = batched_svd_full(t; alg = Jacobi()) + Ur, Sr, Vr = svd_full(t; alg = Jacobi()) + @test specdiff(S, Sr) < 1.0e-10 + @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 + end + + # tall uniform blocks: hits the full mode where U (m, m) and Vᴴ (n, n) have different sizes + @testset "svd_full: tall uniform blocks: $T" for T in (Float64,) + Vbig = Vect[Z4Irrep](0 => 4, 1 => 4, 2 => 4, 3 => 4) + Vsml = Vect[Z4Irrep](0 => 2, 1 => 2, 2 => 2, 3 => 2) + t = adapt(CuArray, randn(T, Vbig ← Vsml)) + szs = [size(TensorKit.block(t, c)) for c in TensorKit.blocksectors(t)] + @test all(isequal((4, 2)), szs) + cs, _ = Factorizations._batchable(t, Jacobi(), true) + @test !isempty(cs) + U, S, Vᴴ = batched_svd_full(t; alg = Jacobi()) + Ur, Sr, Vr = svd_full(t; alg = Jacobi()) + @test specdiff(S, Sr) < 1.0e-10 + @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 + @test norm(U' * U - _id(U, domain(U))) < 1.0e-10 + @test norm(Vᴴ * Vᴴ' - _id(Vᴴ, codomain(Vᴴ))) < 1.0e-10 + end +end From 245964861351be75751b1bc0ab7395daa7ac90c6 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Mon, 5 Oct 2026 16:31:38 +0200 Subject: [PATCH 2/6] Offload a lot to MAK --- Project.toml | 2 +- src/factorizations/batched.jl | 85 ------------------ src/factorizations/factorizations.jl | 1 - src/factorizations/matrixalgebrakit.jl | 100 ++++++++++++--------- test/amd/batched_svd.jl | 118 ------------------------- test/cuda/batched_svd.jl | 109 ----------------------- 6 files changed, 60 insertions(+), 355 deletions(-) delete mode 100644 src/factorizations/batched.jl delete mode 100644 test/amd/batched_svd.jl delete mode 100644 test/cuda/batched_svd.jl diff --git a/Project.toml b/Project.toml index 9fafc118f..9dfe6350e 100644 --- a/Project.toml +++ b/Project.toml @@ -46,7 +46,7 @@ TensorKitGPUArraysExt = "GPUArrays" TensorKitMooncakeExt = "Mooncake" [sources] -MatrixAlgebraKit = {url = "https://github.com/QuantumKitHub/MatrixAlgebraKit.jl", rev = "main"} +MatrixAlgebraKit = {url = "https://github.com/QuantumKitHub/MatrixAlgebraKit.jl", rev = "ksh/supports_ragged_batch"} [compat] AMDGPU = "2" diff --git a/src/factorizations/batched.jl b/src/factorizations/batched.jl deleted file mode 100644 index 3124045ac..000000000 --- a/src/factorizations/batched.jl +++ /dev/null @@ -1,85 +0,0 @@ -# Figure out which sectors are even worth batching, and if some share a batch size -# `uniform = true` additionally demands that every block already has that exact size, i.e. -# that no padding is needed. Full decompositions require this, compact decompositions only -# read back the leading `min(m, n)` columns, which padding out with zero doesn't affect. -function _batchable(t::AbstractTensorMap, alg::AbstractAlgorithm, uniform::Bool = false) - cs = collect(blocksectors(t)) - isempty(cs) && return cs, (0, 0) - szs = [size(block(t, c)) for c in cs] - m, n = maximum(first, szs), maximum(last, szs) - lim = MAK.max_batched_blocksize(alg, storagetype(t)) - (length(cs) < MAK.BATCHED_SVD_THRESHOLD || m > lim || n > lim) && return empty(cs), (m, n) - # The *padded* batch is (m, n) even if *individual* blocks are tall - # so check the padded shape rather than the blocks'. - (MAK.requires_tall(alg) && m < n) && return empty(cs), (m, n) - (uniform && !all(isequal((m, n)), szs)) && return empty(cs), (m, n) - return cs, (m, n) -end - -# The batched solvers work on 3D arrays: input `(m, n, nbatch)`, and correspondingly -# `U = (m, minmn, nbatch)`, `S = (minmn, nbatch)`, `Vᴴ = (minmn, n, nbatch)`. -function _pack(t::AbstractTensorMap, cs, m, n) - b1 = block(t, first(cs)) - A = similar(b1, m, n, length(cs)) - fill!(A, zero(scalartype(A))) - for (i, c) in enumerate(cs) - b = block(t, c) - copyto!(view(A, axes(b, 1), axes(b, 2), i), b) - end - return A -end - -for (bf!, f!) in ((:batched_svd_compact!, :svd_compact!), (:batched_svd_full!, :svd_full!)) - full = bf! === :batched_svd_full! - @eval function MAK.$bf!(t::AbstractTensorMap, F, alg) - U, S, Vᴴ = F - cs, (m, n) = _batchable(t, alg, $full) - if isempty(cs) # not worth batching, or the library doesn't support these sizes - foreachblock(t, U, S, Vᴴ) do _, (b, u, sv, v) - MAK.$f!(b, (u, sv, v), alg) - return nothing - end - return F - end - nb, minmn = length(cs), min(m, n) - A = _pack(t, cs, m, n) - Ub = similar(A, m, $full ? m : minmn, nb) - rT = real(scalartype(t)) - Sb = $full ? similar(A, rT, m, n, nb) : - similar(diagview(block(S, first(cs))), minmn, nb) - Vb = similar(A, $full ? n : minmn, n, nb) - $full && fill!(Sb, zero(rT)) - MAK.$bf!(A, (Ub, Sb, Vb), alg) - for (i, c) in enumerate(cs) - u, sv, v = block(U, c), block(S, c), block(Vᴴ, c) - copyto!(u, view(Ub, axes(u, 1), axes(u, 2), i)) - if $full - copyto!(sv, view(Sb, axes(sv, 1), axes(sv, 2), i)) - else - copyto!(diagview(sv), view(Sb, axes(diagview(sv), 1), i)) - end - copyto!(v, view(Vb, axes(v, 1), axes(v, 2), i)) - end - return F - end -end - -function MAK.batched_svd_vals!(t::AbstractTensorMap, S, alg) - cs, (M, N) = _batchable(t, alg) - if isempty(cs) - foreachblock(t, S) do _, (b, sv) - MAK.svd_vals!(b, sv, alg) - return nothing - end - return S - end - nb, K = length(cs), min(M, N) - A = _pack(t, cs, M, N) - Sb = similar(block(S, first(cs)), K, nb) - MAK.batched_svd_vals!(A, Sb, alg) - for (i, c) in enumerate(cs) - sv = block(S, c) - copyto!(sv, view(Sb, axes(sv, 1), i)) - end - return S -end diff --git a/src/factorizations/factorizations.jl b/src/factorizations/factorizations.jl index 0cdf64048..e49cc8b29 100644 --- a/src/factorizations/factorizations.jl +++ b/src/factorizations/factorizations.jl @@ -28,7 +28,6 @@ using MatrixAlgebraKit: diagview include("utility.jl") include("matrixalgebrakit.jl") -include("batched.jl") include("truncation.jl") include("adjoint.jl") include("diagonal.jl") diff --git a/src/factorizations/matrixalgebrakit.jl b/src/factorizations/matrixalgebrakit.jl index 52f3fd0f4..b1d4a29fd 100644 --- a/src/factorizations/matrixalgebrakit.jl +++ b/src/factorizations/matrixalgebrakit.jl @@ -14,7 +14,6 @@ end for f in [ :svd_compact, :svd_full, :svd_vals, - :batched_svd_compact, :batched_svd_full, :batched_svd_vals, :qr_compact, :qr_full, :qr_null, :lq_compact, :lq_full, :lq_null, :eig_full, :eig_vals, :eigh_full, :eigh_vals, @@ -47,8 +46,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))) @@ -68,10 +66,40 @@ 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()) + if MAK.supports_ragged_batch(MAK.$bf!, alg, driver, storagetype(t)) + cs = collect(blocksectors(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]) + if applicable(MAK.$bf!, As, Fs, alg) + @timeit_debug GLOBAL_TIMER "batched: MatrixAlgebraKit" MAK.$bf!(As, Fs, alg) + return F + end + 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!, :batched_svd_vals, - :svd_vals!, :eig_vals!, :eigh_vals!, + :qr_null!, :lq_null!, :eig_vals!, :eigh_vals!, :project_hermitian!, :project_antihermitian!, :project_isometric!, :exponential!, ) @@ -91,6 +119,32 @@ 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, N, alg::AbstractAlgorithm) + @timeit_debug GLOBAL_TIMER $(string(f!)) begin + driver = get(alg.kwargs, :driver, MAK.DefaultDriver()) + if MAK.supports_ragged_batch(MAK.$bf!, alg, driver, storagetype(t)) + cs = collect(blocksectors(t)) + As, Ns = [block(t, c) for c in cs], [block(N, c) for c in cs] + if applicable(MAK.$bf!, As, Ns, alg) + @timeit_debug GLOBAL_TIMER "batched: MatrixAlgebraKit" MAK.$bf!(As, Ns, alg) + return N + end + end + foreachblock(t, N) do _, (tblock, Nblock) + @timeit_debug GLOBAL_TIMER "dense: MatrixAlgebraKit" begin + Nblock′ = $f!(tblock, Nblock, alg) + # deal with the case where the output is not the same as the input + Nblock === Nblock′ || copy!(Nblock, Nblock′) + end + return nothing + end + end + return N + end +end + # Exponential with Tuple function MAK.exponential!((τ, t)::Tuple{E, T}, N, alg::AbstractAlgorithm) where {E <: Number, T <: AbstractTensorMap} LinearAlgebra.checksquare(t) @@ -119,7 +173,6 @@ for f in [ :eig_full, :eig_vals, :eig_trunc, :eig_trunc_no_error, :eigh_full, :eigh_vals, :eigh_trunc, :eigh_trunc_no_error, :svd_full, :svd_compact, :svd_trunc, :svd_trunc_no_error, :svd_vals, - :batched_svd_full, :batched_svd_compact, :batched_svd_vals, :left_polar, :right_polar, :left_orth, :right_orth, :left_null, :right_null, :project_hermitian, :project_antihermitian, :project_isometric, @@ -182,41 +235,6 @@ function MAK.initialize_output(::typeof(svd_vals!), t::AbstractTensorMap, alg::A end end -# Batched singular value decomposition -# ------------------------------------ -# We should be able to just reuse the -# above and handle the actual batched -# object filling deeper down -function MAK.initialize_output(::typeof(batched_svd_full!), t::AbstractTensorMap, ::AbstractAlgorithm) - @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin - V_cod = fuse(codomain(t)) - V_dom = fuse(domain(t)) - U = similar(t, codomain(t) ← V_cod) - S = similar(t, real(scalartype(t)), V_cod ← V_dom) - Vᴴ = similar(t, V_dom ← domain(t)) - return U, S, Vᴴ - end -end - -function MAK.initialize_output(::typeof(batched_svd_compact!), t::AbstractTensorMap, ::AbstractAlgorithm) - @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin - V_cod = V_dom = infimum(fuse(codomain(t)), fuse(domain(t))) - U = similar(t, codomain(t) ← V_cod) - S = similar_diagonal(t, real(scalartype(t)), V_cod) - Vᴴ = similar(t, V_dom ← domain(t)) - return U, S, Vᴴ - end -end - -function MAK.initialize_output(::typeof(batched_svd_vals!), t::AbstractTensorMap, alg::AbstractAlgorithm) - @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin - V_cod = infimum(fuse(codomain(t)), fuse(domain(t))) - T = real(scalartype(t)) - A = similarstoragetype(t, T) - return SectorVector{T, sectortype(t), A}(undef, V_cod) - end -end - # Eigenvalue decomposition # ------------------------ function MAK.initialize_output(::typeof(eigh_full!), t::AbstractTensorMap, ::AbstractAlgorithm) diff --git a/test/amd/batched_svd.jl b/test/amd/batched_svd.jl deleted file mode 100644 index d0eb9ab78..000000000 --- a/test/amd/batched_svd.jl +++ /dev/null @@ -1,118 +0,0 @@ -using Adapt, AMDGPU -using Test, TestExtras -using TensorKit -using LinearAlgebra: LinearAlgebra -using MatrixAlgebraKit: MatrixAlgebraKit, DivideAndConquer, QRIteration, Jacobi, - svd_compact, svd_compact!, svd_vals, svd_vals!, svd_full, svd_full!, - batched_svd_compact, batched_svd_compact!, batched_svd_vals, - batched_svd_vals!, batched_svd_full, batched_svd_full! - -const Factorizations = TensorKit.Factorizations - -# SU(2) space with enough sectors to exceed the batching threshold -Vsu2 = Vect[SU2Irrep](0 => 2, 1 // 2 => 2, 1 => 2, 3 // 2 => 1) - -# spectra are returned either -# as `Diagonal` blocks (`svd_compact`) or -# as plain vectors (`svd_vals`) -_spec(b) = Array(b isa AbstractMatrix ? TensorKit.diagview(b) : b) - -function specdiff(S1, S2) - d = 0.0 - for (c, b) in TensorKit.blocks(S1) - d = max(d, maximum(abs, _spec(b) .- _spec(TensorKit.block(S2, c)))) - end - return d -end - -_id(t, V) = TensorKit.id(TensorKit.storagetype(t), V) - -@timedtestset "batched SVD on ROCArray" verbose = true begin - @testset "many blocks: $T" for T in (Float64, ComplexF64) - t_cpu = randn(T, Vsu2 ⊗ Vsu2 ← Vsu2) - nblocks = length(TensorKit.blocksectors(t_cpu)) - @test nblocks >= MatrixAlgebraKit.BATCHED_SVD_THRESHOLD - t = adapt(ROCArray, t_cpu) - - U, S, Vᴴ = batched_svd_compact(t; alg = DivideAndConquer()) - Ur, Sr, Vr = svd_compact(t; alg = DivideAndConquer()) - - # singular values agree with the unbatched path - @test specdiff(S, Sr) < 1.0e-10 - @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 - @test norm(U' * U - _id(U, domain(U))) < 1.0e-10 - @test norm(Vᴴ * Vᴴ' - _id(Vᴴ, codomain(Vᴴ))) < 1.0e-10 - - Sv = batched_svd_vals(t; alg = DivideAndConquer()) - @test specdiff(Sv, Sr) < 1.0e-10 - end - - @testset "few blocks fall back: $T" for T in (Float64, ComplexF64) - # trivial sector -> a single block, below the threshold - V = ComplexSpace(6) - t = adapt(ROCArray, randn(T, V ⊗ V ← V)) - @test length(TensorKit.blocksectors(t)) < MatrixAlgebraKit.BATCHED_SVD_THRESHOLD - U, S, Vᴴ = batched_svd_compact(t; alg = DivideAndConquer()) - Ur, Sr, Vr = svd_compact(t; alg = DivideAndConquer()) - @test specdiff(S, Sr) < 1.0e-10 - @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 - end - - @testset "other batched algorithms" for alg in (QRIteration(), Jacobi()) - t_cpu = randn(Float64, Vsu2 ⊗ Vsu2 ← Vsu2) - t = adapt(ROCArray, t_cpu) - U, S, Vᴴ = batched_svd_compact(t; alg) - _, Sr, _ = svd_compact(t; alg = DivideAndConquer()) - @test specdiff(S, Sr) < 1.0e-8 - @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-8 - end - - # `svd_full!` returns U (m, m) and Vᴴ (n, n) rather than the compact (m, minmn) / (minmn, n) - @testset "svd_full: uniform blocks batch: $T" for T in (Float64, ComplexF64) - Vu = Vect[Z4Irrep](0 => 3, 1 => 3, 2 => 3, 3 => 3) - t_cpu = randn(T, Vu ← Vu) - t = adapt(ROCArray, t_cpu) - szs = [size(TensorKit.block(t, c)) for c in TensorKit.blocksectors(t)] - @test length(szs) >= MatrixAlgebraKit.BATCHED_SVD_THRESHOLD - @test all(isequal(first(szs)), szs) - cs, _ = Factorizations._batchable(t, QRIteration(), true) - @test !isempty(cs) - - U, S, Vᴴ = batched_svd_full(t; alg = QRIteration()) - Ur, Sr, Vr = svd_full(t; alg = QRIteration()) - @test specdiff(S, Sr) < 1.0e-10 - @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 - @test norm(U' * U - _id(U, domain(U))) < 1.0e-10 - @test norm(U * U' - _id(U, codomain(U))) < 1.0e-10 - @test norm(Vᴴ' * Vᴴ - _id(Vᴴ, domain(Vᴴ))) < 1.0e-10 - end - - @testset "svd_full: ragged blocks fall back: $T" for T in (Float64,) - t = adapt(ROCArray, randn(T, Vsu2 ⊗ Vsu2 ← Vsu2)) - szs = [size(TensorKit.block(t, c)) for c in TensorKit.blocksectors(t)] - @test !all(isequal(first(szs)), szs) - cs, _ = Factorizations._batchable(t, QRIteration(), true) - @test isempty(cs) - U, S, Vᴴ = batched_svd_full(t; alg = QRIteration()) - Ur, Sr, Vr = svd_full(t; alg = QRIteration()) - @test specdiff(S, Sr) < 1.0e-10 - @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 - end - - # tall uniform blocks: hits the full mode where U (m, m) and Vᴴ (n, n) have different sizes - @testset "svd_full: tall uniform blocks: $T" for T in (Float64,) - Vbig = Vect[Z4Irrep](0 => 4, 1 => 4, 2 => 4, 3 => 4) - Vsml = Vect[Z4Irrep](0 => 2, 1 => 2, 2 => 2, 3 => 2) - t = adapt(ROCArray, randn(T, Vbig ← Vsml)) - szs = [size(TensorKit.block(t, c)) for c in TensorKit.blocksectors(t)] - @test all(isequal((4, 2)), szs) - cs, _ = Factorizations._batchable(t, QRIteration(), true) - @test !isempty(cs) - U, S, Vᴴ = batched_svd_full(t; alg = QRIteration()) - Ur, Sr, Vr = svd_full(t; alg = QRIteration()) - @test specdiff(S, Sr) < 1.0e-10 - @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 - @test norm(U' * U - _id(U, domain(U))) < 1.0e-10 - @test norm(Vᴴ * Vᴴ' - _id(Vᴴ, codomain(Vᴴ))) < 1.0e-10 - end -end diff --git a/test/cuda/batched_svd.jl b/test/cuda/batched_svd.jl deleted file mode 100644 index 966bfa687..000000000 --- a/test/cuda/batched_svd.jl +++ /dev/null @@ -1,109 +0,0 @@ -using Adapt, CUDA -using Test, TestExtras -using TensorKit -using LinearAlgebra: LinearAlgebra -using MatrixAlgebraKit: MatrixAlgebraKit, Jacobi, QRIteration, - svd_compact, svd_compact!, svd_vals, svd_vals!, svd_full, svd_full!, - batched_svd_compact, batched_svd_compact!, batched_svd_vals, - batched_svd_vals!, batched_svd_full, batched_svd_full! - -const Factorizations = TensorKit.Factorizations - -# SU(2) space with enough sectors to exceed the batching threshold -Vsu2 = Vect[SU2Irrep](0 => 2, 1 // 2 => 2, 1 => 2, 3 // 2 => 1) - -# spectra are returned either -# as `Diagonal` blocks (`svd_compact`) or -# as plain vectors (`svd_vals`) -_spec(b) = Array(b isa AbstractMatrix ? TensorKit.diagview(b) : b) - -function specdiff(S1, S2) - d = 0.0 - for (c, b) in TensorKit.blocks(S1) - d = max(d, maximum(abs, _spec(b) .- _spec(TensorKit.block(S2, c)))) - end - return d -end - -_id(t, V) = TensorKit.id(TensorKit.storagetype(t), V) - -@timedtestset "batched SVD on CuArray" verbose = true begin - @testset "many blocks: $T" for T in (Float64, ComplexF64) - t_cpu = randn(T, Vsu2 ⊗ Vsu2 ← Vsu2) - nblocks = length(TensorKit.blocksectors(t_cpu)) - @test nblocks >= MatrixAlgebraKit.BATCHED_SVD_THRESHOLD - t = adapt(CuArray, t_cpu) - - U, S, Vᴴ = batched_svd_compact(t; alg = Jacobi()) - Ur, Sr, Vr = svd_compact(t; alg = Jacobi()) - - # singular values agree with the unbatched path - @test specdiff(S, Sr) < 1.0e-10 - @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 - @test norm(U' * U - _id(U, domain(U))) < 1.0e-10 - @test norm(Vᴴ * Vᴴ' - _id(Vᴴ, codomain(Vᴴ))) < 1.0e-10 - - Sv = batched_svd_vals(t; alg = Jacobi()) - @test specdiff(Sv, Sr) < 1.0e-10 - end - - @testset "few blocks fall back: $T" for T in (Float64, ComplexF64) - # trivial sector -> a single block, below the threshold - V = ComplexSpace(6) - t = adapt(CuArray, randn(T, V ⊗ V ← V)) - @test length(TensorKit.blocksectors(t)) < MatrixAlgebraKit.BATCHED_SVD_THRESHOLD - U, S, Vᴴ = batched_svd_compact(t; alg = Jacobi()) - Ur, Sr, Vr = svd_compact(t; alg = Jacobi()) - @test specdiff(S, Sr) < 1.0e-10 - @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 - end - - # `svd_full!` returns U (m, m) and Vᴴ (n, n) rather than the compact (m, minmn) / (minmn, n) - @testset "svd_full: uniform blocks batch: $T" for T in (Float64, ComplexF64) - Vu = Vect[Z4Irrep](0 => 3, 1 => 3, 2 => 3, 3 => 3) - t_cpu = randn(T, Vu ← Vu) - t = adapt(CuArray, t_cpu) - szs = [size(TensorKit.block(t, c)) for c in TensorKit.blocksectors(t)] - @test length(szs) >= MatrixAlgebraKit.BATCHED_SVD_THRESHOLD - @test all(isequal(first(szs)), szs) - cs, _ = Factorizations._batchable(t, Jacobi(), true) - @test !isempty(cs) - - U, S, Vᴴ = batched_svd_full(t; alg = Jacobi()) - Ur, Sr, Vr = svd_full(t; alg = Jacobi()) - @test specdiff(S, Sr) < 1.0e-10 - @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 - @test norm(U' * U - _id(U, domain(U))) < 1.0e-10 - @test norm(U * U' - _id(U, codomain(U))) < 1.0e-10 - @test norm(Vᴴ' * Vᴴ - _id(Vᴴ, domain(Vᴴ))) < 1.0e-10 - end - - @testset "svd_full: ragged blocks fall back: $T" for T in (Float64,) - t = adapt(CuArray, randn(T, Vsu2 ⊗ Vsu2 ← Vsu2)) - szs = [size(TensorKit.block(t, c)) for c in TensorKit.blocksectors(t)] - @test !all(isequal(first(szs)), szs) - cs, _ = Factorizations._batchable(t, Jacobi(), true) - @test isempty(cs) - U, S, Vᴴ = batched_svd_full(t; alg = Jacobi()) - Ur, Sr, Vr = svd_full(t; alg = Jacobi()) - @test specdiff(S, Sr) < 1.0e-10 - @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 - end - - # tall uniform blocks: hits the full mode where U (m, m) and Vᴴ (n, n) have different sizes - @testset "svd_full: tall uniform blocks: $T" for T in (Float64,) - Vbig = Vect[Z4Irrep](0 => 4, 1 => 4, 2 => 4, 3 => 4) - Vsml = Vect[Z4Irrep](0 => 2, 1 => 2, 2 => 2, 3 => 2) - t = adapt(CuArray, randn(T, Vbig ← Vsml)) - szs = [size(TensorKit.block(t, c)) for c in TensorKit.blocksectors(t)] - @test all(isequal((4, 2)), szs) - cs, _ = Factorizations._batchable(t, Jacobi(), true) - @test !isempty(cs) - U, S, Vᴴ = batched_svd_full(t; alg = Jacobi()) - Ur, Sr, Vr = svd_full(t; alg = Jacobi()) - @test specdiff(S, Sr) < 1.0e-10 - @test norm(U * S * Vᴴ - t) / norm(t) < 1.0e-10 - @test norm(U' * U - _id(U, domain(U))) < 1.0e-10 - @test norm(Vᴴ * Vᴴ' - _id(Vᴴ, codomain(Vᴴ))) < 1.0e-10 - end -end From 84051136e4601ebafb0c0e8e112964d8891512b5 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 7 Oct 2026 07:09:26 +0200 Subject: [PATCH 3/6] Update Project.toml --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 9dfe6350e..9fafc118f 100644 --- a/Project.toml +++ b/Project.toml @@ -46,7 +46,7 @@ TensorKitGPUArraysExt = "GPUArrays" TensorKitMooncakeExt = "Mooncake" [sources] -MatrixAlgebraKit = {url = "https://github.com/QuantumKitHub/MatrixAlgebraKit.jl", rev = "ksh/supports_ragged_batch"} +MatrixAlgebraKit = {url = "https://github.com/QuantumKitHub/MatrixAlgebraKit.jl", rev = "main"} [compat] AMDGPU = "2" From 6125746e241ebe98caffc390d23a2f7351d83aa1 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 7 Oct 2026 14:49:52 +0200 Subject: [PATCH 4/6] Get rid of unneeded bypass --- src/factorizations/matrixalgebrakit.jl | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/src/factorizations/matrixalgebrakit.jl b/src/factorizations/matrixalgebrakit.jl index b1d4a29fd..8bc67a71f 100644 --- a/src/factorizations/matrixalgebrakit.jl +++ b/src/factorizations/matrixalgebrakit.jl @@ -1,16 +1,5 @@ # Algorithm selection # ------------------- - -""" - _tensor_algorithm(f!, ::Type{<:AbstractTensorMap}; kwargs...) - -Algorithm a `TensorMap` factorization defaults to. Blocks are decomposed one at a time, so -the default is whatever the block type would use. For algorithms that have a batched version, -like `QRIteration` or `Jacobi`, this can be overridden to point to the batching version. -""" -function _tensor_algorithm(f!, ::Type{T}; kwargs...) where {T <: AbstractTensorMap} - return MAK.default_algorithm(f!, blocktype(T); kwargs...) -end for f in [ :svd_compact, :svd_full, :svd_vals, @@ -23,7 +12,7 @@ for f in ] f! = Symbol(f, :!) @eval function MAK.default_algorithm(::typeof($f!), ::Type{T}; kwargs...) where {T <: AbstractTensorMap} - return _tensor_algorithm($f!, T; kwargs...) + return MAK.default_algorithm($f!, blocktype(T); kwargs...) end @eval function MAK.copy_input(::typeof($f), t::AbstractTensorMap) return @timeit_debug GLOBAL_TIMER "alloc: copy_input" copy_oftype( From f6ba1cdb6fd20c07bfc374beeb888749dc57b734 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 8 Oct 2026 10:29:31 +0200 Subject: [PATCH 5/6] Force the batched path to actually run and use QRIteration for AMD + isometric --- src/factorizations/matrixalgebrakit.jl | 4 ++-- test/amd/factorizations.jl | 13 ++++++++----- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/src/factorizations/matrixalgebrakit.jl b/src/factorizations/matrixalgebrakit.jl index 8bc67a71f..6acf906e4 100644 --- a/src/factorizations/matrixalgebrakit.jl +++ b/src/factorizations/matrixalgebrakit.jl @@ -62,7 +62,7 @@ for (f!, bf!) in ((:svd_compact!, :batched_svd_compact!), (:svd_full!, :batched_ @timeit_debug GLOBAL_TIMER $(string(f!)) begin U, S, Vᴴ = F driver = get(alg.kwargs, :driver, MAK.DefaultDriver()) - if MAK.supports_ragged_batch(MAK.$bf!, alg, driver, storagetype(t)) + if MAK.supports_ragged_batch(MAK.$f!, alg, driver, storagetype(t)) cs = collect(blocksectors(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]) @@ -113,7 +113,7 @@ for (f!, bf!) in ((:svd_vals!, :batched_svd_vals!),) @eval function MAK.$f!(t::AbstractTensorMap, N, alg::AbstractAlgorithm) @timeit_debug GLOBAL_TIMER $(string(f!)) begin driver = get(alg.kwargs, :driver, MAK.DefaultDriver()) - if MAK.supports_ragged_batch(MAK.$bf!, alg, driver, storagetype(t)) + if MAK.supports_ragged_batch(MAK.$f!, alg, driver, storagetype(t)) cs = collect(blocksectors(t)) As, Ns = [block(t, c) for c in cs], [block(N, c) for c in cs] if applicable(MAK.$bf!, As, Ns, alg) diff --git a/test/amd/factorizations.jl b/test/amd/factorizations.jl index a4c907d63..a8887cd47 100644 --- a/test/amd/factorizations.jl +++ b/test/amd/factorizations.jl @@ -2,7 +2,7 @@ using Adapt, AMDGPU using Test, TestExtras using TensorKit using LinearAlgebra: LinearAlgebra -using MatrixAlgebraKit: DefaultAlgorithm, diagview +using MatrixAlgebraKit: DefaultAlgorithm, PolarViaSVD, QRIteration, diagview const AMDGPUExt = Base.get_extension(TensorKit, :TensorKitAMDGPUExt) @assert !isnothing(AMDGPUExt) "Failed to load TensorKit - AMDGPU extension" const ROCTensorMap = getglobal(AMDGPUExt, :ROCTensorMap) @@ -588,6 +588,9 @@ for V in spacelist end @testset "Isometric projections" begin + # rocSOLVER's gesvdx (Bisection) returns wrong singular vectors for degenerate + # singular values, see https://github.com/ROCm/rocm-libraries/issues/12188 + alg = PolarViaSVD(QRIteration()) for T in eltypes, t in ( AMDGPU.randn(T, W, W), @@ -595,25 +598,25 @@ for V in spacelist AMDGPU.randn(T, (V1 ⊗ V2 ⊗ V3), (V4 ⊗ V5)'), AMDGPU.randn(T, (V1 ⊗ V2)', (V3 ⊗ V4 ⊗ V5))', ) - t2 = project_isometric(t) + t2 = project_isometric(t, alg) @test isisometric(t2) t2′ = @constinferred project_isometric(t, DefaultAlgorithm()) @test isisometric(t2′) @test t2′ * ((t2′)' * t) ≈ t - t3 = project_isometric(t2) + t3 = project_isometric(t2, alg) @test t3 ≈ t2 # stability of the projection @test t2 * (t2' * t) ≈ t tc = similar(t) - t3 = @constinferred project_isometric!(copy!(tc, t), t2) + t3 = @constinferred project_isometric!(copy!(tc, t), t2, alg) @test t3 === t2 @test isisometric(t2) # test that t2 is closer to A then any other isometry for k in 1:10 δt = AMDGPU.randn!(similar(t)) - t3 = project_isometric(t + δt / 100) + t3 = project_isometric(t + δt / 100, alg) @test norm(t - t3) > norm(t - t2) end end From 3eb9f67d4f73712a5debff63837b40e0ec1f8ee3 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 8 Oct 2026 15:19:33 +0200 Subject: [PATCH 6/6] Update Project.toml --- Project.toml | 3 --- 1 file changed, 3 deletions(-) diff --git a/Project.toml b/Project.toml index 9fafc118f..d1e0c834d 100644 --- a/Project.toml +++ b/Project.toml @@ -45,9 +45,6 @@ TensorKitFiniteDifferencesExt = "FiniteDifferences" TensorKitGPUArraysExt = "GPUArrays" TensorKitMooncakeExt = "Mooncake" -[sources] -MatrixAlgebraKit = {url = "https://github.com/QuantumKitHub/MatrixAlgebraKit.jl", rev = "main"} - [compat] AMDGPU = "2" Adapt = "4"