From 4638d7125d034f24c0b7d1606e5812f462b911c1 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 13 Oct 2026 15:13:53 -0400 Subject: [PATCH 1/2] Allocate sparse block tensors consistently through the allocator Sparse temporaries keep their allocator: blocks they gain later are obtained from it, and only temporaries release their blocks in `tensorfree!`. `tensoralloc_add`/`tensoralloc_contract` allocate the blocks a sparse result will receive up front, blocks are stored as the type the allocator hands out, and a sparse `tensorcontract!` destination receives the product blocks before contracting, or goes through a temporary when it is not a BLAS destination. Co-Authored-By: Claude Opus 5.5 --- src/linalg/linalg.jl | 2 + src/tensors/sparseblocktensor.jl | 32 +++++++-- src/tensors/tensoroperations.jl | 115 +++++++++++++++++++++++++++++-- src/tensors/vectorinterface.jl | 13 ++-- test/linalg/tensoroperations.jl | 91 +++++++++++++++++++++++- 5 files changed, 237 insertions(+), 16 deletions(-) diff --git a/src/linalg/linalg.jl b/src/linalg/linalg.jl index f0296e8..9f6cbb7 100644 --- a/src/linalg/linalg.jl +++ b/src/linalg/linalg.jl @@ -118,6 +118,8 @@ end ) return if haskey(C, I) C[I] = _mul!!(C[I], A, B, α, β) + elseif istemp(C) + C[I] = _mul!!(newblock(C, I), A, B, α, Zero()) else C[I] = _mul!!(nothing, A, B, α, β) end diff --git a/src/tensors/sparseblocktensor.jl b/src/tensors/sparseblocktensor.jl index e429159..89804ad 100644 --- a/src/tensors/sparseblocktensor.jl +++ b/src/tensors/sparseblocktensor.jl @@ -2,19 +2,22 @@ struct SparseBlockTensorMap{TT<:AbstractTensorMap{E,S,N₁,N₂}} <: AbstractBlockTensorMap{E,S,N₁,N₂} Sparse `SparseBlockTensorMap` type that stores tensors of type `TT` in a sparse dictionary. +Temporaries of TensorOperations keep their `allocator` (otherwise `nothing`), from which their +blocks are obtained and to which they are released. """ struct SparseBlockTensorMap{TT <: AbstractTensorMap, E, S, N₁, N₂, N} <: AbstractBlockTensorMap{E, S, N₁, N₂} data::Dict{CartesianIndex{N}, TT} space::TensorMapSumSpace{S, N₁, N₂} + allocator::Any # uninitialized constructor function SparseBlockTensorMap{TT, E, S, N₁, N₂, N}( - ::UndefBlocksInitializer, space::TensorMapSumSpace{S, N₁, N₂} + ::UndefBlocksInitializer, space::TensorMapSumSpace{S, N₁, N₂}; allocator = nothing ) where {E, S, N₁, N₂, N, TT <: AbstractTensorMap{E, S, N₁, N₂}} @assert N₁ + N₂ == N "SparseBlockTensorMap: data has wrong number of dimensions" data = Dict{CartesianIndex{N}, TT}() - return new{TT, E, S, N₁, N₂, N}(data, space) + return new{TT, E, S, N₁, N₂, N}(data, space, allocator) end # constructor from data @@ -22,7 +25,7 @@ struct SparseBlockTensorMap{TT <: AbstractTensorMap, E, S, N₁, N₂, N} <: data::Dict{CartesianIndex{N}, TT}, space::TensorMapSumSpace{S, N₁, N₂} ) where {E, S, N₁, N₂, N, TT <: AbstractTensorMap{E, S, N₁, N₂}} @assert N₁ + N₂ == N "SparseBlockTensorMap: data has wrong number of dimensions" - return new{TT, E, S, N₁, N₂, N}(data, space) + return new{TT, E, S, N₁, N₂, N}(data, space, nothing) end end @@ -34,10 +37,11 @@ end # uninitialized constructor function SparseBlockTensorMap{TT}( - ::Union{UndefBlocksInitializer, UndefInitializer}, space::TensorMapSumSpace{S, N₁, N₂} + ::Union{UndefBlocksInitializer, UndefInitializer}, space::TensorMapSumSpace{S, N₁, N₂}; + kwargs... ) where {E, S, N₁, N₂, TT <: AbstractTensorMap{E, S, N₁, N₂}} N = N₁ + N₂ - return SparseBlockTensorMap{TT, E, S, N₁, N₂, N}(undef_blocks, space) + return SparseBlockTensorMap{TT, E, S, N₁, N₂, N}(undef_blocks, space; kwargs...) end # constructor from data @@ -189,7 +193,7 @@ end # `blocks` first store the (nonempty) entries they write to. function storeentries!(t::SparseBlockTensorMap, keep = Returns(true)) for (I, V) in pairs(eachspace(t)) - !haskey(t, I) && dim(V) > 0 && keep(I) && (t[I] = zerovector!(similar(eltype(t), V))) + !haskey(t, I) && dim(V) > 0 && keep(I) && (t[I] = zerovector!(newblock(t, I))) end return t end @@ -222,6 +226,22 @@ for f in (:(Random.rand!), :(Random.randn!), :(Random.randexp!), :(TK.randisomet invoke($f, Tuple{Random.AbstractRNG, AbstractTensorMap}, rng, storeentries!(t)) end +# blocks of temporaries are obtained from (and released to) their own allocator, those of +# other destinations from the allocator of the operation, if any +istemp(t::AbstractBlockTensorMap) = false +istemp(t::SparseBlockTensorMap) = !isnothing(t.allocator) + +function newblock(t::AbstractBlockTensorMap, I::CartesianIndex, allocator = TO.DefaultAllocator()) + V = eachspace(t)[I] + istemp(t) && return TO.tensoralloc(eltype(t), V, Val(true), t.allocator) + return TO.tensoralloc(eltype(t), V, Val(false), allocator) +end + +@propagate_inbounds function getindex!(t::SparseBlockTensorMap, I::CartesianIndex{N}) where {N} + istemp(t) || return getindex!(parent(t), I) + return haskey(t, I) ? t[I] : (t[I] = zerovector!(newblock(t, I))) +end + # Show # ---- function Base.showarg(io::IO, t::SparseBlockTensorMap, toplevel::Bool) diff --git a/src/tensors/tensoroperations.jl b/src/tensors/tensoroperations.jl index ad9c357..4ff483f 100644 --- a/src/tensors/tensoroperations.jl +++ b/src/tensors/tensoroperations.jl @@ -97,16 +97,121 @@ function similarblocktype(::Type{A}, ::Type{TT}) where {A, TT} return Core.Compiler.return_type(similar, Tuple{A, Type{TT}, NTuple{numind(TT), Int}}) end +# blocks are stored as the type the allocator hands out, e.g. `PtrArray`-backed temporaries +function allocated_blocktype(::Type{TT}, structure, istemp::Val, allocator) where {TT} + V = eltype(SumSpaceIndices(structure)) + TT′ = Base.promote_op(TO.tensoralloc, Type{TT}, V, typeof(istemp), typeof(allocator)) + TA = AbstractTensorMap{scalartype(TT), spacetype(TT), numout(TT), numin(TT)} + return typeintersect(TT′, TA) +end + function TO.tensoralloc( ::Type{BT}, structure::TensorMapSumSpace, istemp::Val, allocator = TO.DefaultAllocator() ) where {BT <: AbstractBlockTensorMap} - C = BT(undef_blocks, structure) - issparse(C) && return C # don't fill up sparse blocks - blockallocator(V) = TO.tensoralloc(eltype(C), V, istemp, allocator) - map!(blockallocator, parent(C), eachspace(C)) + TT = eltype(BT) + C = BlockTensorMap{allocated_blocktype(TT, structure, istemp, allocator)}( + undef_blocks, structure + ) + map!(V -> TO.tensoralloc(TT, V, istemp, allocator), parent(C), eachspace(C)) return C end +# sparse tensors start out empty, and temporaries remember their allocator +function TO.tensoralloc( + ::Type{BT}, structure::TensorMapSumSpace, istemp::Val, allocator = TO.DefaultAllocator() + ) where {BT <: SparseBlockTensorMap} + TT = allocated_blocktype(eltype(BT), structure, istemp, allocator) + return SparseBlockTensorMap{TT}( + undef_blocks, structure; allocator = istemp === Val(true) ? allocator : nothing + ) +end + +# sparse results allocate the blocks that will be written up front +const BlockOrAdjoint = Union{AbstractBlockTensorMap, AdjointBlockTensorMap} + +_tensoralloc(ttype, structure, keys, istemp::Val, allocator) = + TO.tensoralloc(ttype, structure, istemp, allocator) +function _tensoralloc( + ttype::Type{<:SparseBlockTensorMap}, structure, keys, istemp::Val, allocator + ) + C = TO.tensoralloc(ttype, structure, istemp, allocator) + foreach(I -> haskey(C, I) || (C[I] = newblock(C, I, allocator)), keys) + return C +end + +function contract_keys(A, pA::Index2Tuple, B, pB::Index2Tuple, pAB::Index2Tuple) + OB = NTuple{length(pB[2]), Int} + openB = Dict{NTuple{length(pB[1]), Int}, Vector{OB}}() + for IB in nonzero_keys(B) + push!(get!(Vector{OB}, openB, TT.getindices(IB.I, pB[1])), TT.getindices(IB.I, pB[2])) + end + p = TO.linearize(pAB) + keys = Set{CartesianIndex{length(p)}}() + for IA in nonzero_keys(A) + oA = TT.getindices(IA.I, pA[1]) + for oB in get(openB, TT.getindices(IA.I, pA[2]), ()) + push!(keys, CartesianIndex(TT.getindices((oA..., oB...), p))) + end + end + return keys +end + +function TO.tensoralloc_add( + TC, A::BlockOrAdjoint, pA::Index2Tuple, conjA::Bool, + istemp::Val = Val(false), allocator = TO.DefaultAllocator() + ) + ttype = TO.tensoradd_type(TC, A, pA, conjA) + structure = TO.tensoradd_structure(A, pA, conjA) + p = TO.linearize(pA) + keys = (CartesianIndex(TT.getindices(I.I, p)) for I in nonzero_keys(A)) + return _tensoralloc(ttype, structure, keys, istemp, allocator) +end + +function TO.tensoralloc_contract( + TC, A::BlockOrAdjoint, pA::Index2Tuple, conjA::Bool, + B::BlockOrAdjoint, pB::Index2Tuple, conjB::Bool, + pAB::Index2Tuple, istemp::Val = Val(false), allocator = TO.DefaultAllocator() + ) + ttype = TO.tensorcontract_type(TC, A, pA, conjA, B, pB, conjB, pAB) + structure = TO.tensorcontract_structure(A, pA, conjA, B, pB, conjB, pAB) + keys = contract_keys(A, pA, B, pB, pAB) + return _tensoralloc(ttype, structure, keys, istemp, allocator) +end + +# contract directly into `C` when it is a valid BLAS destination, otherwise go through a +# temporary holding only the product blocks instead of a copy of all blocks of `C` +function TO.tensorcontract!( + C::SparseBlockTensorMap, + A::BlockOrAdjoint, pA::Index2Tuple, conjA::Bool, + B::BlockOrAdjoint, pB::Index2Tuple, conjB::Bool, + pAB::Index2Tuple, α::Number, β::Number, backend, allocator + ) + ipAB = TO.oindABinC(pAB, pA, pB) + if TO.isblasdestination(C, ipAB) || TO.isblasdestination(C, reverse(ipAB)) + for I in contract_keys(A, pA, B, pB, pAB) + haskey(C, I) || (C[I] = zerovector!(newblock(C, I, allocator))) + end + return _tensorcontract!(C, A, pA, conjA, B, pB, conjB, pAB, α, β, backend, allocator) + end + N₁ = length(pA[1]) + pAB′ = (ntuple(identity, N₁), ntuple(i -> N₁ + i, length(pB[2]))) + AB = TO.tensoralloc_contract( + scalartype(C), A, pA, conjA, B, pB, conjB, pAB′, Val(true), allocator + ) + _tensorcontract!(AB, A, pA, conjA, B, pB, conjB, pAB′, One(), Zero(), backend, allocator) + TO.tensoradd!(C, AB, pAB, false, α, β, backend, allocator) + TO.tensorfree!(AB, allocator) + return C +end +function _tensorcontract!(C, A, pA, conjA, B, pB, conjB, pAB, α, β, backend, allocator) + return @invoke TO.tensorcontract!( + C::AbstractTensorMap, + A::AbstractTensorMap, pA::Index2Tuple, conjA::Bool, + B::AbstractTensorMap, pB::Index2Tuple, conjB::Bool, + pAB::Index2Tuple, α::Number, β::Number, backend::Any, allocator::Any + ) +end + # tensorfree! # ----------- function TO.tensorfree!(t::BlockTensorMap, allocator = TO.DefaultAllocator()) @@ -114,7 +219,7 @@ function TO.tensorfree!(t::BlockTensorMap, allocator = TO.DefaultAllocator()) return nothing end function TO.tensorfree!(t::SparseBlockTensorMap, allocator = TO.DefaultAllocator()) - foreach(Base.Fix2(TO.tensorfree!, allocator), nonzero_values(t)) + istemp(t) && foreach(Base.Fix2(TO.tensorfree!, allocator), nonzero_values(t)) return nothing end diff --git a/src/tensors/vectorinterface.jl b/src/tensors/vectorinterface.jl index e7d0b6e..409ff36 100644 --- a/src/tensors/vectorinterface.jl +++ b/src/tensors/vectorinterface.jl @@ -1,7 +1,11 @@ # zerovector # ---------- VI.zerovector!(t::AbstractBlockTensorMap) = (zerovector!(parent(t)); t) -VI.zerovector!(t::SparseBlockTensorMap) = (empty!(t.data); t) +# blocks of temporaries belong to an allocator, and are kept until `tensorfree!` +function VI.zerovector!(t::SparseBlockTensorMap) + istemp(t) ? foreach(zerovector!, nonzero_values(t)) : empty!(t.data) + return t +end # scale # ----- @@ -54,7 +58,7 @@ end t[I] = if haskey(t, I) f(t[I], isone(β) ? One() : β) else - f(TO.tensoralloc(eltype(t), eachspace(t)[I], Val(false), allocator), Zero()) + f(newblock(t, I, allocator), Zero()) end return t end @@ -65,11 +69,12 @@ function _scale_untouched!( (issparse(tsrc) && !isone(β)) || return t q = TT.invperm(p) istouched(I) = haskey(tsrc, CartesianIndex(TT.getindices(I.I, q))) - if issparse(t) && iszero(β) + # blocks of temporaries are zeroed rather than dropped, as they belong to an allocator + if issparse(t) && iszero(β) && !istemp(t) filter!(istouched ∘ first, parent(t)) else for (I, v) in nonzero_pairs(t) - istouched(I) || scale!(v, β) + istouched(I) || (iszero(β) ? zerovector!(v) : scale!(v, β)) end end return t diff --git a/test/linalg/tensoroperations.jl b/test/linalg/tensoroperations.jl index 67ad59a..4e972ca 100644 --- a/test/linalg/tensoroperations.jl +++ b/test/linalg/tensoroperations.jl @@ -3,7 +3,7 @@ using BlockTensorKit using TensorKit using TensorOperations using Random -using LinearAlgebra: norm +using LinearAlgebra: norm, mul! ## Vtr = ( @@ -150,3 +150,92 @@ end @test convert(TensorMap, F1) ≈ F2 end end + +## + +# temporaries are `PtrArray`s, such that converting them into other storage is detected +struct TrackingAllocator + live::Set{UInt} # temporaries that were handed out and not freed yet + foreign::Base.RefValue{Int} # freed arrays that were not handed out as temporaries +end +TrackingAllocator() = TrackingAllocator(Set{UInt}(), Ref(0)) +function TensorOperations.tensoralloc( + ::Type{A}, structure, ::Val{istemp}, allocator::TrackingAllocator + ) where {A <: AbstractArray, istemp} + istemp || return TensorOperations.tensoralloc(A, structure, Val(false)) + x = fill!( + TensorOperations.tensoralloc(A, structure, Val(true), TensorOperations.ManualAllocator()), + NaN + ) + push!(allocator.live, objectid(x)) + return x +end +function TensorOperations.tensorfree!(x::AbstractArray, allocator::TrackingAllocator) + if objectid(x) in allocator.live + delete!(allocator.live, objectid(x)) + TensorOperations.tensorfree!(x, TensorOperations.ManualAllocator()) + else + allocator.foreign[] += 1 + end + return nothing +end +function check_allocator(f) + allocator = TrackingAllocator() + @test f(allocator) ≈ f(TensorOperations.DefaultAllocator()) + @test allocator.foreign[] == 0 + return @test isempty(allocator.live) +end + +@testset "allocator only frees what it allocated ($(sectortype(a)))" for (a, b, c) in ( + (SumSpace(ℂ^2, ℂ^1), SumSpace(ℂ^1, ℂ^2, ℂ^1), SumSpace(ℂ^2, ℂ^3)), + ( + SumSpace(Vect[FermionParity](0 => 1, 1 => 1), Vect[FermionParity](0 => 2)), + SumSpace(Vect[FermionParity](1 => 2), Vect[FermionParity](0 => 1, 1 => 1)), + SumSpace(Vect[FermionParity](0 => 1), Vect[FermionParity](1 => 1)), + ), + ) + A = sprand(Float64, a ⊗ b ← c, 0.4) + B = sprand(Float64, c ⊗ b' ← a, 0.4) + D = sprand(Float64, a ← b, 0.4) + E = sprand(Float64, a ← b, 0.4) + X = sprand(Float64, a ⊗ b ← a ⊗ b, 0.5) + + check_allocator() do al + return @tensor allocator = al R[x; q] := A[x z; w] * B[w z; y] * D[y; q] + end + check_allocator() do al # temporary receiving blocks outside of its initial pattern + return @tensor allocator = al R[x; q] := A[x z; w] * B[w z; y] * (D[y; q] + 2 * E[y; q]) + end + check_allocator() do al # several blocks of `X` are traced into the same output block + return @tensor allocator = al R[x; q] := X[x z; y z] * D[y; q] + end + check_allocator() do al + return @planar allocator = al R[x; q] := A[x z; w] * B[w z; y] * D[y; q] + end + check_allocator() do al # kernels without allocator argument fill an empty temporary + T = TensorOperations.tensoralloc(typeof(D), a ← a, Val(true), al) + mul!(T, D, D') + TensorOperations.tensortrace!( + T, X, ((1,), (3,)), ((2,), (4,)), false, 1, 1, TensorOperations.DefaultBackend(), al + ) + n = norm(T) + TensorOperations.tensorfree!(T, al) + return n + end + + # permuted and BLAS-compatible destinations, without and with blocks outside of the product + @tensor C1[x y; p q] := A[x q; w] * B[w p; y] + @tensor C2[x q; p y] := A[x q; w] * B[w p; y] + for C in (spzeros(Float64, space(C1)), sprand(Float64, space(C1), 1.0)) + check_allocator() do al + C′ = copy(C) + return @tensor allocator = al C′[x y; p q] += A[x q; w] * B[w p; y] + end + end + for C in (spzeros(Float64, space(C2)), sprand(Float64, space(C2), 1.0)) + check_allocator() do al + C′ = copy(C) + return @tensor allocator = al C′[x q; p y] += A[x q; w] * B[w p; y] + end + end +end From 8ba311fe1857d0fd7e06a06f6e7e36bb5246ab1b Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 13 Oct 2026 15:37:05 -0400 Subject: [PATCH 2/2] Keep the blocks of temporaries in `one!`/`isomorphism!`/`isometry!` Off-diagonal blocks of a temporary belong to its allocator, so they are kept (and zeroed) instead of dropped. Zero blocks for missing entries are allocated as `similarstoragetype`, since storage types such as `PtrArray` cannot be constructed from a size. Co-Authored-By: Claude Opus 5.5 --- src/tensors/abstractblocktensor/abstracttensormap.jl | 10 ++++++---- src/tensors/sparseblocktensor.jl | 2 +- test/linalg/tensoroperations.jl | 12 ++++++++++++ 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/src/tensors/abstractblocktensor/abstracttensormap.jl b/src/tensors/abstractblocktensor/abstracttensormap.jl index 49622ce..970f63c 100644 --- a/src/tensors/abstractblocktensor/abstracttensormap.jl +++ b/src/tensors/abstractblocktensor/abstracttensormap.jl @@ -105,7 +105,8 @@ function _subblock(t::AbstractBlockTensorMap, entries, (f₁, f₂)::Tuple{Fusio V = W[I] sz = (dims(codomain(V), f₁.uncoupled)..., dims(domain(V), f₂.uncoupled)...) if prod(sz) == 0 || (issparse(t) && !haskey(entries, I)) - return sreshape(StridedView(zerovector!(similar(storagetype(t), prod(sz)))), sz) + data = similar(TK.similarstoragetype(storagetype(t), scalartype(t)), prod(sz)) + return sreshape(StridedView(zerovector!(data)), sz) else return getsubblock(entries[I], (f₁, f₂)) end @@ -121,10 +122,11 @@ function _blockdims(P::ProductSumSpace{S, N}, c::Sector) where {S, N} ) end +# allocated as `similarstoragetype`, as the storage type itself need not be constructible zeroblock(::Type{TT}, d₁::Int, d₂::Int) where {TT <: AbstractTensorMap} = - zerovector!(similar(storagetype(TT), d₁, d₂)) + zerovector!(similar(TK.similarstoragetype(TT, scalartype(TT)), d₁, d₂)) function zeroblock(::Type{TT}, d₁::Int, d₂::Int) where {TT <: TensorMap} - data = zerovector!(similar(storagetype(TT), d₁ * d₂)) + data = zerovector!(similar(TK.similarstoragetype(TT, scalartype(TT)), d₁ * d₂)) return reshape(view(data, 1:(d₁ * d₂)), (d₁, d₂)) end zeroblock(::Type{<:AdjointTensorMap{<:Any, <:Any, <:Any, <:Any, TT}}, d₁::Int, d₂::Int) where {TT} = @@ -136,7 +138,7 @@ TensorKit.hasblock(t::AbstractBlockTensorMap, c::Sector) = c in blocksectors(t) Base.@assume_effects :foldable function _blockeltype(::Type{TT}) where {TT <: AbstractBlockTensorMap} T = scalartype(TT) - B = TK.blocktype(eltype(TT)) + B = Union{TK.blocktype(eltype(TT)), Base.promote_op(zeroblock, Type{eltype(TT)}, Int, Int)} return B <: AbstractMatrix{T} ? B : AbstractMatrix{T} # safeguard against type-instability end function TensorKit.blocktype(::Type{TT}) where {TT <: AbstractBlockTensorMap} diff --git a/src/tensors/sparseblocktensor.jl b/src/tensors/sparseblocktensor.jl index 89804ad..56c0151 100644 --- a/src/tensors/sparseblocktensor.jl +++ b/src/tensors/sparseblocktensor.jl @@ -217,7 +217,7 @@ for f in (:(TK.one!), :(TK.isomorphism!), :(TK.isometry!)) @eval function $f(t::SparseBlockTensorMap) diagonal = diagonalentries(t) invoke($f, Tuple{AbstractTensorMap}, storeentries!(t, in(diagonal))) - filter!(in(diagonal) ∘ first, t.data) + istemp(t) || filter!(in(diagonal) ∘ first, t.data) return t end end diff --git a/test/linalg/tensoroperations.jl b/test/linalg/tensoroperations.jl index 4e972ca..7b8ce23 100644 --- a/test/linalg/tensoroperations.jl +++ b/test/linalg/tensoroperations.jl @@ -223,6 +223,18 @@ end return n end + for f! in (TensorKit.one!, TensorKit.isomorphism!, TensorKit.isometry!) + check_allocator() do al # blocks of a temporary off the diagonal are kept + T = TensorOperations.tensoralloc_add(Float64, X, ((1, 2), (3, 4)), false, Val(true), al) + TensorOperations.tensoradd!( + T, X, ((1, 2), (3, 4)), false, 1, 0, TensorOperations.DefaultBackend(), al + ) + n = norm(f!(T)) + TensorOperations.tensorfree!(T, al) + return n + end + end + # permuted and BLAS-compatible destinations, without and with blocks outside of the product @tensor C1[x y; p q] := A[x q; w] * B[w p; y] @tensor C2[x q; p y] := A[x q; w] * B[w p; y]