From cf3c11f0d2b79fba9dfb4cc3f761af4b633ae572 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Tue, 31 Mar 2026 08:23:54 -0400 Subject: [PATCH 01/15] Updates for GPU compatibility for MPSKit --- ext/BlockTensorKitGPUArraysExt.jl | 24 +++++++++++++++++++ src/linalg/factorizations.jl | 7 ++++-- .../abstractblocktensor/abstractarray.jl | 2 +- src/tensors/indexmanipulations.jl | 16 ++++++++----- src/tensors/tensoroperations.jl | 23 ++++++++++++++---- 5 files changed, 59 insertions(+), 13 deletions(-) diff --git a/ext/BlockTensorKitGPUArraysExt.jl b/ext/BlockTensorKitGPUArraysExt.jl index ff5f217..8de4451 100644 --- a/ext/BlockTensorKitGPUArraysExt.jl +++ b/ext/BlockTensorKitGPUArraysExt.jl @@ -3,9 +3,33 @@ module BlockTensorKitGPUArraysExt using BlockTensorKit, BlockArrays, GPUArrays, Strided using Strided: StridedViews using GPUArrays: KernelAbstractions +import BlockTensorKit: _full function KernelAbstractions.get_backend(BA::BlockArrays.BlockArray{T, N, A}) where {T, N, A <: AbstractArray{<:StridedView{T, N, <:AnyGPUArray}}} return KernelAbstractions.get_backend(first(BA.blocks)) end +function BlockTensorKit._full(A::BM) where {T <: Number, TA <: AnyGPUMatrix{T}, BM <: BlockMatrix{T, Matrix{TA}}} + arr = similar(first(A.blocks), size(A)) + # TODO -- should we use Threads here to parallelize these + # transfers in streams if possible? + for block_index in Iterators.product(blockaxes(A)...) + indices = getindex.(axes(A), block_index) + arr[indices...] = @view A[block_index...] + end + return arr +end + +# awful piracy but defined here as BlockArrays doesn't support this well +function Base.copyto!(dest::BM, src::TA) where {T <: Number, TA <: AnyGPUMatrix{T}, BM <: BlockMatrix{T, Matrix{TA}}} + # TODO -- should we use Threads here to parallelize these + # transfers in streams if possible? + for block_index in Iterators.product(blockaxes(dest)...) + indices = getindex.(axes(dest), block_index) + dest_view = @view dest[block_index...] + dest_view = src[indices...] + end + return dest +end + end diff --git a/src/linalg/factorizations.jl b/src/linalg/factorizations.jl index 003711a..7db8a33 100644 --- a/src/linalg/factorizations.jl +++ b/src/linalg/factorizations.jl @@ -23,6 +23,8 @@ function MatrixAlgebraKit.one!(A::BlockBlasMat) return A end +_full(A) = Array(A) + for f in [ :svd_compact, :svd_full, :svd_vals, @@ -44,7 +46,8 @@ for f! in ( ) @eval function MAK.$f!(t::AbstractBlockTensorMap, F, alg::AbstractAlgorithm) TensorKit.foreachblock(t, F...) do _, (tblock, Fblocks...) - Fblocks′ = MAK.$f!(Array(tblock), alg) + full_block = _full(tblock) + Fblocks′ = MAK.$f!(full_block, alg) # deal with the case where the output is not in-place for (b′, b) in zip(Fblocks′, Fblocks) b === b′ || copy!(b, b′) @@ -63,7 +66,7 @@ for f! in ( ) @eval function MAK.$f!(t::AbstractBlockTensorMap, N, alg::AbstractAlgorithm) TensorKit.foreachblock(t, N) do _, (tblock, Nblock) - Nblock′ = MAK.$f!(Array(tblock), alg) + Nblock′ = MAK.$f!(_full(tblock), alg) # deal with the case where the output is not the same as the input Nblock === Nblock′ || copy!(Nblock, Nblock′) return nothing diff --git a/src/tensors/abstractblocktensor/abstractarray.jl b/src/tensors/abstractblocktensor/abstractarray.jl index 99cb4fd..1f07f99 100644 --- a/src/tensors/abstractblocktensor/abstractarray.jl +++ b/src/tensors/abstractblocktensor/abstractarray.jl @@ -328,7 +328,7 @@ Base.eltypeof(t::AbstractBlockTensorMap) = eltype(t) ) where {T <: AbstractTensorMap} catdims = Base.dims2cat(dims) V = space(Base._cat(dims, eachspace.(ts)...)) - A = similar(ts[1], T, V) + A = similar(ts[1], TK.storagetype(ts[1]), V) shape = size(A) if count(!iszero, catdims)::Int > 1 zerovector!(A) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 58e3cd7..3dc82d8 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -102,7 +102,6 @@ for f! in (:add_permute!, :add_transpose!) end dstdata = parent(tdst) srcdata = permutedims(StridedView(parent(tsrc)), (p₁..., p₂...)) - @inbounds for I in eachindex(dstdata, srcdata) dstdata[I] = TK.$f!(dstdata[I], srcdata[I], (p₁, p₂), α, β, backend...) end @@ -231,25 +230,27 @@ function TK.add_braid!( end Base.@constprop :aggressive function TK.insertleftunit( - t::AbstractBlockTensorMap, i::Int = numind(t) + 1; kwargs... + t::AbstractBlockTensorMap, i::Int = numind(t) + 1; + storage_type = TK.storagetype(t), kwargs... ) W = TK.insertleftunit(space(t), i; kwargs...) tdst = similar(t, W) for (I, v) in nonzero_pairs(t) I′ = CartesianIndex(TT.insertafter(I.I, i - 1, (1,))) - tdst[I′] = TK.insertleftunit(v, i; kwargs...) + tdst[I′] = TK.insertleftunit(v, i; storage_type, kwargs...) end return tdst end Base.@constprop :aggressive function TK.insertrightunit( - t::AbstractBlockTensorMap, i::Int = numind(t) + 1; kwargs... + t::AbstractBlockTensorMap, i::Int = numind(t) + 1; + storage_type = TK.storagetype(t), kwargs... ) W = TK.insertrightunit(space(t), i; kwargs...) tdst = similar(t, W) for (I, v) in nonzero_pairs(t) I′ = CartesianIndex(TT.insertafter(I.I, i, (1,))) - tdst[I′] = TK.insertrightunit(v, i; kwargs...) + tdst[I′] = TK.insertrightunit(v, i; storage_type, kwargs...) end return tdst end @@ -261,7 +262,10 @@ Base.@constprop :aggressive function TK.removeunit( tdst = similar(t, W) for (I, v) in nonzero_pairs(t) I′ = CartesianIndex(TT.deleteat(I.I, i)) - tdst[I′] = TK.removeunit(v, i) + # pass the storagetype to account for the BraidingTensor + # case, to ensure the output TensorMap has the correct + # type of backing array + tdst[I′] = TK.removeunit(v, i; storage_type = TK.storagetype(t)) end return tdst end diff --git a/src/tensors/tensoroperations.jl b/src/tensors/tensoroperations.jl index 80bb85a..0688936 100644 --- a/src/tensors/tensoroperations.jl +++ b/src/tensors/tensoroperations.jl @@ -21,17 +21,26 @@ end # tensoralloc_contract # -------------------- -for TTA in (:AbstractTensorMap, :AbstractBlockTensorMap), TTB in (:AbstractTensorMap, :AbstractBlockTensorMap) - TTA == TTB == :AbstractTensorMap && continue +for TTB in (:AbstractTensorMap, :AbstractBlockTensorMap) @eval function TO.tensorcontract_type( TC, - A::$TTA, ::Index2Tuple, ::Bool, + A::AbstractBlockTensorMap, ::Index2Tuple, ::Bool, B::$TTB, ::Index2Tuple, ::Bool, ::Index2Tuple{N₁, N₂}, ) where {N₁, N₂} S = TK.check_spacetype(A, B) TC′ = TK.promote_permute(TC, sectortype(S)) - M = TK.promote_storagetype(TK.similarstoragetype(A, TC′), TK.similarstoragetype(B, TC′)) + ATT = AbstractTensorMap{scalartype(A), spacetype(A), numout(A), numin(A)} + BTT = AbstractTensorMap{scalartype(B), spacetype(B), numout(B), numin(B)} + # handle case with BraidingTensors, so that they assume the backing + # array type of the concrete element type + M = if eltype(A) == ATT + TK.similarstoragetype(B, TC′) + elseif eltype(B) == BTT + TK.similarstoragetype(A, TC′) + else + TK.similarstoragetype(TK.similarstoragetype(A, TC′), TK.similarstoragetype(B, TC′)) + end return if issparse(A) && issparse(B) sparseblocktensormaptype(S, N₁, N₂, M) else @@ -39,6 +48,12 @@ for TTA in (:AbstractTensorMap, :AbstractBlockTensorMap), TTB in (:AbstractTenso end end end +TO.tensorcontract_type( + TC, + A::AbstractTensorMap, pA::Index2Tuple, conjA::Bool, + B::AbstractBlockTensorMap, pB::Index2Tuple, conjB::Bool, + pAB::Index2Tuple{N₁, N₂}, + ) where {N₁, N₂} = TO.tensorcontract_type(TC, B, pB, conjB, A, pA, conjA, pAB) function similarblocktype(::Type{A}, ::Type{TT}) where {A, TT} return Core.Compiler.return_type(similar, Tuple{A, Type{TT}, NTuple{numind(TT), Int}}) From eb88c63217a07c35dd3ee9cec6a11826cf12ba77 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 2 Apr 2026 08:02:10 -0400 Subject: [PATCH 02/15] Fixes for BraidingTensor and svd --- src/linalg/factorizations.jl | 4 +++- src/tensors/indexmanipulations.jl | 20 ++++++++++++-------- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/src/linalg/factorizations.jl b/src/linalg/factorizations.jl index 7db8a33..cc0ea2b 100644 --- a/src/linalg/factorizations.jl +++ b/src/linalg/factorizations.jl @@ -171,7 +171,9 @@ end function MAK.initialize_output(::typeof(svd_compact!), t::AbstractBlockTensorMap, ::AbstractAlgorithm) V_cod = V_dom = infimum(fuse(codomain(t)), fuse(domain(t))) U = dense_similar(t, codomain(t) ← V_cod) - S = DiagonalTensorMap{real(scalartype(t))}(undef, V_cod) + Tr = real(scalartype(t)) + TAr = TK.similarstoragetype(t, Tr) + S = DiagonalTensorMap{Tr, spacetype(V_cod), TAr}(undef, V_cod) Vᴴ = dense_similar(t, V_dom ← domain(t)) return U, S, Vᴴ end diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 3dc82d8..902c58f 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -231,26 +231,26 @@ end Base.@constprop :aggressive function TK.insertleftunit( t::AbstractBlockTensorMap, i::Int = numind(t) + 1; - storage_type = TK.storagetype(t), kwargs... + kwargs... ) W = TK.insertleftunit(space(t), i; kwargs...) tdst = similar(t, W) for (I, v) in nonzero_pairs(t) I′ = CartesianIndex(TT.insertafter(I.I, i - 1, (1,))) - tdst[I′] = TK.insertleftunit(v, i; storage_type, kwargs...) + tdst[I′] = TK.insertleftunit(v, i; kwargs...) end return tdst end Base.@constprop :aggressive function TK.insertrightunit( t::AbstractBlockTensorMap, i::Int = numind(t) + 1; - storage_type = TK.storagetype(t), kwargs... + kwargs... ) W = TK.insertrightunit(space(t), i; kwargs...) tdst = similar(t, W) for (I, v) in nonzero_pairs(t) I′ = CartesianIndex(TT.insertafter(I.I, i, (1,))) - tdst[I′] = TK.insertrightunit(v, i; storage_type, kwargs...) + tdst[I′] = TK.insertrightunit(v, i; kwargs...) end return tdst end @@ -262,10 +262,14 @@ Base.@constprop :aggressive function TK.removeunit( tdst = similar(t, W) for (I, v) in nonzero_pairs(t) I′ = CartesianIndex(TT.deleteat(I.I, i)) - # pass the storagetype to account for the BraidingTensor - # case, to ensure the output TensorMap has the correct - # type of backing array - tdst[I′] = TK.removeunit(v, i; storage_type = TK.storagetype(t)) + if v isa TK.BraidingTensor + v′ = TK.removeunit(v, i) + tdst′ = similar(v′, TK.storagetype(t)) + copy!(tdst′, v′) + tdst[I′] = tdst′ + else + tdst[I′] = TK.removeunit(v, i) + end end return tdst end From 84c67b44722e5bd877ab30899494f189eb790c61 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 2 Apr 2026 08:03:04 -0400 Subject: [PATCH 03/15] Formatter --- src/tensors/tensoroperations.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tensors/tensoroperations.jl b/src/tensors/tensoroperations.jl index 0688936..981be3a 100644 --- a/src/tensors/tensoroperations.jl +++ b/src/tensors/tensoroperations.jl @@ -53,7 +53,7 @@ TO.tensorcontract_type( A::AbstractTensorMap, pA::Index2Tuple, conjA::Bool, B::AbstractBlockTensorMap, pB::Index2Tuple, conjB::Bool, pAB::Index2Tuple{N₁, N₂}, - ) where {N₁, N₂} = TO.tensorcontract_type(TC, B, pB, conjB, A, pA, conjA, pAB) +) where {N₁, N₂} = TO.tensorcontract_type(TC, B, pB, conjB, A, pA, conjA, pAB) function similarblocktype(::Type{A}, ::Type{TT}) where {A, TT} return Core.Compiler.return_type(similar, Tuple{A, Type{TT}, NTuple{numind(TT), Int}}) From 18ffe0b57a0b0a0a6c930a7c956ab0aacab0297b Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 23 Apr 2026 04:49:20 -0400 Subject: [PATCH 04/15] Cleanup removeunit --- src/tensors/indexmanipulations.jl | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 902c58f..c95bc8f 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -262,14 +262,7 @@ Base.@constprop :aggressive function TK.removeunit( tdst = similar(t, W) for (I, v) in nonzero_pairs(t) I′ = CartesianIndex(TT.deleteat(I.I, i)) - if v isa TK.BraidingTensor - v′ = TK.removeunit(v, i) - tdst′ = similar(v′, TK.storagetype(t)) - copy!(tdst′, v′) - tdst[I′] = tdst′ - else - tdst[I′] = TK.removeunit(v, i) - end + tdst[I′] = TK.removeunit(v, i) end return tdst end From 050cba41b76442fa95f364b966703bac1f8898b0 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 23 Apr 2026 04:51:43 -0400 Subject: [PATCH 05/15] Move over storagetype logic from TensorKit --- src/tensors/abstractblocktensor/abstracttensormap.jl | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/tensors/abstractblocktensor/abstracttensormap.jl b/src/tensors/abstractblocktensor/abstracttensormap.jl index 474a5a1..cf9b4e6 100644 --- a/src/tensors/abstractblocktensor/abstracttensormap.jl +++ b/src/tensors/abstractblocktensor/abstracttensormap.jl @@ -106,4 +106,11 @@ function Base.iterate(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, state... end Base.getindex(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, c::Sector) = block(iter.t, c) -TensorKit.storagetype(::Type{TT}) where {TT <: AbstractBlockTensorMap} = storagetype(eltype(TT)) +function TensorKit.storagetype(::Type{TT}) where {TT <: AbstractBlockTensorMap} + if eltype(TT) isa Union + TU = eltype(TT) + return promote_storagetype(TU.a, TU.b) + else + return storagetype(eltype(TT)) + end +end From 6c8328996003668afe9e087dd969971086d43a94 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 23 Apr 2026 05:46:53 -0400 Subject: [PATCH 06/15] Missing module qualifier --- src/tensors/abstractblocktensor/abstracttensormap.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/tensors/abstractblocktensor/abstracttensormap.jl b/src/tensors/abstractblocktensor/abstracttensormap.jl index cf9b4e6..2726d2f 100644 --- a/src/tensors/abstractblocktensor/abstracttensormap.jl +++ b/src/tensors/abstractblocktensor/abstracttensormap.jl @@ -109,8 +109,8 @@ Base.getindex(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, c::Sector) = blo function TensorKit.storagetype(::Type{TT}) where {TT <: AbstractBlockTensorMap} if eltype(TT) isa Union TU = eltype(TT) - return promote_storagetype(TU.a, TU.b) + return TensorKit.promote_storagetype(TU.a, TU.b) else - return storagetype(eltype(TT)) + return TensorKit.storagetype(eltype(TT)) end end From 3f46738f24628cb1c04faef3c81b9376a8c4a8ba Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 24 Apr 2026 03:41:56 -0400 Subject: [PATCH 07/15] Undo ugly formatting changes --- src/tensors/indexmanipulations.jl | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index c95bc8f..58e3cd7 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -102,6 +102,7 @@ for f! in (:add_permute!, :add_transpose!) end dstdata = parent(tdst) srcdata = permutedims(StridedView(parent(tsrc)), (p₁..., p₂...)) + @inbounds for I in eachindex(dstdata, srcdata) dstdata[I] = TK.$f!(dstdata[I], srcdata[I], (p₁, p₂), α, β, backend...) end @@ -230,8 +231,7 @@ function TK.add_braid!( end Base.@constprop :aggressive function TK.insertleftunit( - t::AbstractBlockTensorMap, i::Int = numind(t) + 1; - kwargs... + t::AbstractBlockTensorMap, i::Int = numind(t) + 1; kwargs... ) W = TK.insertleftunit(space(t), i; kwargs...) tdst = similar(t, W) @@ -243,8 +243,7 @@ Base.@constprop :aggressive function TK.insertleftunit( end Base.@constprop :aggressive function TK.insertrightunit( - t::AbstractBlockTensorMap, i::Int = numind(t) + 1; - kwargs... + t::AbstractBlockTensorMap, i::Int = numind(t) + 1; kwargs... ) W = TK.insertrightunit(space(t), i; kwargs...) tdst = similar(t, W) From c89afa3b0635d5a4ca5c1b57f2b0fe555d4bba82 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 24 Apr 2026 13:39:53 -0400 Subject: [PATCH 08/15] move auxiliary code --- src/BlockTensorKit.jl | 2 ++ src/auxiliary/blockarrays.jl | 30 ++++++++++++++++++++++++++++++ src/linalg/factorizations.jl | 26 ++------------------------ 3 files changed, 34 insertions(+), 24 deletions(-) create mode 100644 src/auxiliary/blockarrays.jl diff --git a/src/BlockTensorKit.jl b/src/BlockTensorKit.jl index e2bf9fb..4d354b8 100644 --- a/src/BlockTensorKit.jl +++ b/src/BlockTensorKit.jl @@ -37,6 +37,8 @@ import TensorOperations as TO import TupleTools as TT import MatrixAlgebraKit as MAK +include("auxiliary/blockarrays.jl") + # Spaces include("vectorspaces/sumspace.jl") include("vectorspaces/sumspaceindices.jl") diff --git a/src/auxiliary/blockarrays.jl b/src/auxiliary/blockarrays.jl new file mode 100644 index 0000000..c5fc625 --- /dev/null +++ b/src/auxiliary/blockarrays.jl @@ -0,0 +1,30 @@ +copy_dense(A) = copy_dense!(similar(first(A.blocks), size(A)), A) +function copy_dense!(Adense, A) + for bj in blockaxes(A, 2) + js = axes(A, 2)[bj] + for bi in blockaxes(A, 1) + a = view(A, bi, bj) + is = axes(A, 1)[bi] + Adense[is, js] = @view A[block_index...] + end + end + return Adense +end + +const BlockBlasMat{T <: MAK.BlasFloat} = BlockMatrix{T} + +function MAK.zero!(A::BlockBlasMat) + for bj in blockaxes(A, 2), bi in blockaxes(A, 1) + a = view(A, bi, bj) + MAK.zero!(a) + end + return A +end + +function MAK.one!(A::BlockBlasMat) + for bj in blockaxes(A, 2), bi in blockaxes(A, 1) + a = view(A, bi, bj) + bi == bj ? MAK.one!(a) : MAK.zero!(a) + end + return A +end diff --git a/src/linalg/factorizations.jl b/src/linalg/factorizations.jl index cc0ea2b..5a89a9d 100644 --- a/src/linalg/factorizations.jl +++ b/src/linalg/factorizations.jl @@ -4,27 +4,6 @@ import MatrixAlgebraKit as MAK # Type piracy for defining the MAK rules on BlockArrays! # ----------------------------------------------------- - -const BlockBlasMat{T <: MAK.BlasFloat} = BlockMatrix{T} - -function MatrixAlgebraKit.zero!(A::BlockBlasMat) - for bj in blockaxes(A, 2), bi in blockaxes(A, 1) - a = view(A, bi, bj) - MAK.zero!(a) - end - return A -end - -function MatrixAlgebraKit.one!(A::BlockBlasMat) - for bj in blockaxes(A, 2), bi in blockaxes(A, 1) - a = view(A, bi, bj) - bi == bj ? MAK.one!(a) : MAK.zero!(a) - end - return A -end - -_full(A) = Array(A) - for f in [ :svd_compact, :svd_full, :svd_vals, @@ -46,8 +25,7 @@ for f! in ( ) @eval function MAK.$f!(t::AbstractBlockTensorMap, F, alg::AbstractAlgorithm) TensorKit.foreachblock(t, F...) do _, (tblock, Fblocks...) - full_block = _full(tblock) - Fblocks′ = MAK.$f!(full_block, alg) + Fblocks′ = MAK.$f!(copy_dense(tblock), alg) # deal with the case where the output is not in-place for (b′, b) in zip(Fblocks′, Fblocks) b === b′ || copy!(b, b′) @@ -66,7 +44,7 @@ for f! in ( ) @eval function MAK.$f!(t::AbstractBlockTensorMap, N, alg::AbstractAlgorithm) TensorKit.foreachblock(t, N) do _, (tblock, Nblock) - Nblock′ = MAK.$f!(_full(tblock), alg) + Nblock′ = MAK.$f!(copy_dense(tblock), alg) # deal with the case where the output is not the same as the input Nblock === Nblock′ || copy!(Nblock, Nblock′) return nothing From 0c58d9ea5b4cfc613123b54b3c3e5316b693601c Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 24 Apr 2026 13:40:38 -0400 Subject: [PATCH 09/15] remove unnecessary auxiliary code --- ext/BlockTensorKitGPUArraysExt.jl | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/ext/BlockTensorKitGPUArraysExt.jl b/ext/BlockTensorKitGPUArraysExt.jl index 8de4451..ff5f217 100644 --- a/ext/BlockTensorKitGPUArraysExt.jl +++ b/ext/BlockTensorKitGPUArraysExt.jl @@ -3,33 +3,9 @@ module BlockTensorKitGPUArraysExt using BlockTensorKit, BlockArrays, GPUArrays, Strided using Strided: StridedViews using GPUArrays: KernelAbstractions -import BlockTensorKit: _full function KernelAbstractions.get_backend(BA::BlockArrays.BlockArray{T, N, A}) where {T, N, A <: AbstractArray{<:StridedView{T, N, <:AnyGPUArray}}} return KernelAbstractions.get_backend(first(BA.blocks)) end -function BlockTensorKit._full(A::BM) where {T <: Number, TA <: AnyGPUMatrix{T}, BM <: BlockMatrix{T, Matrix{TA}}} - arr = similar(first(A.blocks), size(A)) - # TODO -- should we use Threads here to parallelize these - # transfers in streams if possible? - for block_index in Iterators.product(blockaxes(A)...) - indices = getindex.(axes(A), block_index) - arr[indices...] = @view A[block_index...] - end - return arr -end - -# awful piracy but defined here as BlockArrays doesn't support this well -function Base.copyto!(dest::BM, src::TA) where {T <: Number, TA <: AnyGPUMatrix{T}, BM <: BlockMatrix{T, Matrix{TA}}} - # TODO -- should we use Threads here to parallelize these - # transfers in streams if possible? - for block_index in Iterators.product(blockaxes(dest)...) - indices = getindex.(axes(dest), block_index) - dest_view = @view dest[block_index...] - dest_view = src[indices...] - end - return dest -end - end From c9d5318d40ca18bad4f36a5441d68177c1e0e21f Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 24 Apr 2026 13:54:52 -0400 Subject: [PATCH 10/15] use similar_diagonal --- src/linalg/factorizations.jl | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/linalg/factorizations.jl b/src/linalg/factorizations.jl index 5a89a9d..667103b 100644 --- a/src/linalg/factorizations.jl +++ b/src/linalg/factorizations.jl @@ -125,7 +125,7 @@ end function MAK.initialize_output(::typeof(eigh_full!), t::AbstractBlockTensorMap, ::AbstractAlgorithm) V_D = ⊕(fuse(domain(t))) T = real(scalartype(t)) - D = DiagonalTensorMap{T}(undef, V_D) + D = TK.similar_diagonal(t, T, V_D) V = dense_similar(t, codomain(t) ← V_D) return D, V end @@ -133,7 +133,7 @@ end function MAK.initialize_output(::typeof(eig_full!), t::AbstractBlockTensorMap, ::AbstractAlgorithm) V_D = ⊕(fuse(domain(t))) Tc = complex(scalartype(t)) - D = DiagonalTensorMap{Tc}(undef, V_D) + D = TK.similar_diagonal(t, Tc, V_D) V = dense_similar(t, Tc, codomain(t) ← V_D) return D, V end @@ -149,9 +149,7 @@ end function MAK.initialize_output(::typeof(svd_compact!), t::AbstractBlockTensorMap, ::AbstractAlgorithm) V_cod = V_dom = infimum(fuse(codomain(t)), fuse(domain(t))) U = dense_similar(t, codomain(t) ← V_cod) - Tr = real(scalartype(t)) - TAr = TK.similarstoragetype(t, Tr) - S = DiagonalTensorMap{Tr, spacetype(V_cod), TAr}(undef, V_cod) + S = TK.similar_diagonal(t, real(scalartype(t)), V_cod) Vᴴ = dense_similar(t, V_dom ← domain(t)) return U, S, Vᴴ end From 2870ca31edbb73641e9be7cb4ccb0f58074fb78e Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 24 Apr 2026 13:55:03 -0400 Subject: [PATCH 11/15] revert tensorcontract_type changes --- src/tensors/tensoroperations.jl | 23 ++++------------------- 1 file changed, 4 insertions(+), 19 deletions(-) diff --git a/src/tensors/tensoroperations.jl b/src/tensors/tensoroperations.jl index 981be3a..80bb85a 100644 --- a/src/tensors/tensoroperations.jl +++ b/src/tensors/tensoroperations.jl @@ -21,26 +21,17 @@ end # tensoralloc_contract # -------------------- -for TTB in (:AbstractTensorMap, :AbstractBlockTensorMap) +for TTA in (:AbstractTensorMap, :AbstractBlockTensorMap), TTB in (:AbstractTensorMap, :AbstractBlockTensorMap) + TTA == TTB == :AbstractTensorMap && continue @eval function TO.tensorcontract_type( TC, - A::AbstractBlockTensorMap, ::Index2Tuple, ::Bool, + A::$TTA, ::Index2Tuple, ::Bool, B::$TTB, ::Index2Tuple, ::Bool, ::Index2Tuple{N₁, N₂}, ) where {N₁, N₂} S = TK.check_spacetype(A, B) TC′ = TK.promote_permute(TC, sectortype(S)) - ATT = AbstractTensorMap{scalartype(A), spacetype(A), numout(A), numin(A)} - BTT = AbstractTensorMap{scalartype(B), spacetype(B), numout(B), numin(B)} - # handle case with BraidingTensors, so that they assume the backing - # array type of the concrete element type - M = if eltype(A) == ATT - TK.similarstoragetype(B, TC′) - elseif eltype(B) == BTT - TK.similarstoragetype(A, TC′) - else - TK.similarstoragetype(TK.similarstoragetype(A, TC′), TK.similarstoragetype(B, TC′)) - end + M = TK.promote_storagetype(TK.similarstoragetype(A, TC′), TK.similarstoragetype(B, TC′)) return if issparse(A) && issparse(B) sparseblocktensormaptype(S, N₁, N₂, M) else @@ -48,12 +39,6 @@ for TTB in (:AbstractTensorMap, :AbstractBlockTensorMap) end end end -TO.tensorcontract_type( - TC, - A::AbstractTensorMap, pA::Index2Tuple, conjA::Bool, - B::AbstractBlockTensorMap, pB::Index2Tuple, conjB::Bool, - pAB::Index2Tuple{N₁, N₂}, -) where {N₁, N₂} = TO.tensorcontract_type(TC, B, pB, conjB, A, pA, conjA, pAB) function similarblocktype(::Type{A}, ::Type{TT}) where {A, TT} return Core.Compiler.return_type(similar, Tuple{A, Type{TT}, NTuple{numind(TT), Int}}) From 4ce3679dfab75754fe4545e67de8aeeec84bbf1e Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 24 Apr 2026 14:08:19 -0400 Subject: [PATCH 12/15] try and revert storagetype --- src/tensors/abstractblocktensor/abstracttensormap.jl | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/src/tensors/abstractblocktensor/abstracttensormap.jl b/src/tensors/abstractblocktensor/abstracttensormap.jl index 2726d2f..3a68f2e 100644 --- a/src/tensors/abstractblocktensor/abstracttensormap.jl +++ b/src/tensors/abstractblocktensor/abstracttensormap.jl @@ -106,11 +106,4 @@ function Base.iterate(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, state... end Base.getindex(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, c::Sector) = block(iter.t, c) -function TensorKit.storagetype(::Type{TT}) where {TT <: AbstractBlockTensorMap} - if eltype(TT) isa Union - TU = eltype(TT) - return TensorKit.promote_storagetype(TU.a, TU.b) - else - return TensorKit.storagetype(eltype(TT)) - end -end +TK.storagetype(::Type{TT}) where {TT <: AbstractBlockTensorMap} = storagetype(eltype(TT)) From 846c90d8a229daab2b39fe0519b5ef94d13200aa Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 24 Apr 2026 14:25:42 -0400 Subject: [PATCH 13/15] bump v0.3.10 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 30a4380..c0c0443 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "BlockTensorKit" uuid = "5f87ffc2-9cf1-4a46-8172-465d160bd8cd" -version = "0.3.9" +version = "0.3.10" authors = ["Lukas Devos and contributors"] [deps] From 61990ee10be1052684f24ae20773336b1d24f29b Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 24 Apr 2026 14:25:48 -0400 Subject: [PATCH 14/15] fix copy_dense! --- src/auxiliary/blockarrays.jl | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/auxiliary/blockarrays.jl b/src/auxiliary/blockarrays.jl index c5fc625..a2294ea 100644 --- a/src/auxiliary/blockarrays.jl +++ b/src/auxiliary/blockarrays.jl @@ -1,12 +1,9 @@ copy_dense(A) = copy_dense!(similar(first(A.blocks), size(A)), A) function copy_dense!(Adense, A) - for bj in blockaxes(A, 2) - js = axes(A, 2)[bj] - for bi in blockaxes(A, 1) - a = view(A, bi, bj) - is = axes(A, 1)[bi] - Adense[is, js] = @view A[block_index...] - end + for block_index in Iterators.product(blockaxes(A)...) + a = view(A, block_index...) + indices = getindex.(axes(A), block_index) + Adense[indices...] .= a end return Adense end From f48fa24997dff4cde0a7f1a546440ae80a277122 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 24 Apr 2026 15:24:29 -0400 Subject: [PATCH 15/15] fix with empty tensors --- src/auxiliary/blockarrays.jl | 1 - src/linalg/factorizations.jl | 4 ++-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/auxiliary/blockarrays.jl b/src/auxiliary/blockarrays.jl index a2294ea..94be967 100644 --- a/src/auxiliary/blockarrays.jl +++ b/src/auxiliary/blockarrays.jl @@ -1,4 +1,3 @@ -copy_dense(A) = copy_dense!(similar(first(A.blocks), size(A)), A) function copy_dense!(Adense, A) for block_index in Iterators.product(blockaxes(A)...) a = view(A, block_index...) diff --git a/src/linalg/factorizations.jl b/src/linalg/factorizations.jl index 667103b..1b06be9 100644 --- a/src/linalg/factorizations.jl +++ b/src/linalg/factorizations.jl @@ -25,7 +25,7 @@ for f! in ( ) @eval function MAK.$f!(t::AbstractBlockTensorMap, F, alg::AbstractAlgorithm) TensorKit.foreachblock(t, F...) do _, (tblock, Fblocks...) - Fblocks′ = MAK.$f!(copy_dense(tblock), alg) + Fblocks′ = MAK.$f!(copy_dense!(similar(tblock, size(tblock)), tblock), alg) # deal with the case where the output is not in-place for (b′, b) in zip(Fblocks′, Fblocks) b === b′ || copy!(b, b′) @@ -44,7 +44,7 @@ for f! in ( ) @eval function MAK.$f!(t::AbstractBlockTensorMap, N, alg::AbstractAlgorithm) TensorKit.foreachblock(t, N) do _, (tblock, Nblock) - Nblock′ = MAK.$f!(copy_dense(tblock), alg) + Nblock′ = MAK.$f!(copy_dense!(similar(tblock, size(tblock)), tblock), alg) # deal with the case where the output is not the same as the input Nblock === Nblock′ || copy!(Nblock, Nblock′) return nothing