diff --git a/src/linalg/linalg.jl b/src/linalg/linalg.jl index 0abd9a3..15ec1cf 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!!(getindex!(C, I), A, B, α, β) else C[I] = _mul!!(nothing, A, B, α, β) end diff --git a/src/tensors/abstractblocktensor/abstractarray.jl b/src/tensors/abstractblocktensor/abstractarray.jl index 46ce188..e2581f2 100644 --- a/src/tensors/abstractblocktensor/abstractarray.jl +++ b/src/tensors/abstractblocktensor/abstractarray.jl @@ -88,6 +88,8 @@ end getindex!(parent(t), I...) @propagate_inbounds getindex!(t::AbstractBlockTensorMap, I::CartesianIndex{N}) where {N} = getindex!(parent(t), I) +# missing blocks of a destination are obtained from the allocator of the operation +@propagate_inbounds getindex!(t::AbstractBlockTensorMap, I::CartesianIndex, allocator) = t[I] # slicing getindex needs to correctly allocate output blocktensor: @propagate_inbounds Base.getindex(t::AbstractBlockTensorMap, indices::Vararg{SliceIndex}) = diff --git a/src/tensors/blocktensor.jl b/src/tensors/blocktensor.jl index 3d2b6d9..244f162 100644 --- a/src/tensors/blocktensor.jl +++ b/src/tensors/blocktensor.jl @@ -219,3 +219,4 @@ end # ------- Base.haskey(t::BlockTensorMap, I::CartesianIndex) = checkbounds(Bool, t.data, I) Base.haskey(t::BlockTensorMap, i::Int) = checkbounds(Bool, t.data, i) +@propagate_inbounds getindex!(t::BlockTensorMap, I::CartesianIndex{N}) where {N} = t[I] diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 1194120..15f2243 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -24,7 +24,7 @@ end @inbounds for (I, v) in nonzero_pairs(tsrc) I′ = CartesianIndex(TT.getindices(I.I, p_lin)) tdst[I′] = TK.add_transform!( - tdst[I′], v, p_lin, transformer, α, One(), backend, allocator + getindex!(tdst, I′, allocator), v, p_lin, transformer, α, One(), backend, allocator ) end return tdst @@ -39,7 +39,7 @@ function TK.add_transform!( @inbounds for (I, v) in nonzero_pairs(tsrc) I′ = CartesianIndex(TT.getindices(I.I, p_lin)) tdst[I′] = TK.add_transform!( - tdst[I′], v, p, transformer, α, One(), backend, allocator + getindex!(tdst, I′, allocator), v, p, transformer, α, One(), backend, allocator ) end return tdst @@ -88,7 +88,9 @@ for f in (:permute, :transpose) p_lin = (p[1]..., p[2]...) @inbounds for (I, v) in nonzero_pairs(tsrc) I′ = CartesianIndex(TT.getindices(I.I, p_lin)) - tdst[I′] = TK.$f!(tdst[I′], v, p, α, One(), backend, allocator) + tdst[I′] = TK.$f!( + getindex!(tdst, I′, allocator), v, p, α, One(), backend, allocator + ) end return tdst end @@ -101,7 +103,9 @@ for f in (:permute, :transpose) p_lin = (p[1]..., p[2]...) @inbounds for (I, v) in nonzero_pairs(tsrc) I′ = CartesianIndex(TT.getindices(I.I, p)) - tdst[I′] = TK.$f!(tdst[I′], v, (p₁, p₂), α, One(), backend, allocator) + tdst[I′] = TK.$f!( + getindex!(tdst, I′, allocator), v, (p₁, p₂), α, One(), backend, allocator + ) end return tdst end @@ -147,7 +151,9 @@ end p_lin = (p[1]..., p[2]...) @inbounds for (I, v) in nonzero_pairs(tsrc) I′ = CartesianIndex(TT.getindices(I.I, p_lin)) - tdst[I′] = TK.braid!(tdst[I′], v, p, levels, α, One(), backend, allocator) + tdst[I′] = TK.braid!( + getindex!(tdst, I′, allocator), v, p, levels, α, One(), backend, allocator + ) end return tdst end @@ -160,7 +166,9 @@ function TK.braid!( p_lin = (p[1]..., p[2]...) @inbounds for (I, v) in nonzero_pairs(tsrc) I′ = CartesianIndex(TT.getindices(I.I, p_lin)) - tdst[I′] = TK.braid!(tdst[I′], v, p, levels, α, One(), backend, allocator) + tdst[I′] = TK.braid!( + getindex!(tdst, I′, allocator), v, p, levels, α, One(), backend, allocator + ) end return tdst end diff --git a/src/tensors/sparseblocktensor.jl b/src/tensors/sparseblocktensor.jl index d919caf..393c642 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,6 +193,30 @@ function Base.delete!(t::SparseBlockTensorMap{TT}, I::Vararg{Int, N}) where {TT, return delete!(t, CartesianIndex(I...)) end +# missing blocks of temporaries are obtained from 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) + +@propagate_inbounds function getindex!(t::SparseBlockTensorMap, I::CartesianIndex{N}) where {N} + istemp(t) || return getindex!(parent(t), I) + return haskey(t, I) ? t[I] : zerovector!(allocblock!(t, I, t.allocator)) +end +@propagate_inbounds getindex!( + t::SparseBlockTensorMap, I::CartesianIndex, ::TO.DefaultAllocator +) = getindex!(t, I) +@propagate_inbounds function getindex!(t::SparseBlockTensorMap, I::CartesianIndex, allocator) + return haskey(t, I) ? t[I] : zerovector!(allocblock!(t, I, allocator)) +end +function allocblock!(t::SparseBlockTensorMap, I::CartesianIndex, allocator) + V = eachspace(t)[I] + return t[I] = if istemp(t) + TO.tensoralloc(eltype(t), V, Val(true), t.allocator) + else + TO.tensoralloc(eltype(t), V, Val(false), allocator) + end +end + # Show # ---- function Base.showarg(io::IO, t::SparseBlockTensorMap, toplevel::Bool) diff --git a/src/tensors/tensoroperations.jl b/src/tensors/tensoroperations.jl index a0342c2..e6d14bd 100644 --- a/src/tensors/tensoroperations.jl +++ b/src/tensors/tensoroperations.jl @@ -25,7 +25,9 @@ function TO.tensoradd!( p_lin = (pA[1]..., pA[2]...) @inbounds for (I, v) in nonzero_pairs(A) I′ = CartesianIndex(TT.getindices(I.I, p_lin)) - C[I′] = TO.tensoradd!(C[I′], v, pA, conjA, α, One(), backend, allocator) + C[I′] = TO.tensoradd!( + getindex!(C, I′, allocator), v, pA, conjA, α, One(), backend, allocator + ) end return C end @@ -95,16 +97,119 @@ 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) || allocblock!(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)) + foreach(I -> getindex!(C, I, allocator), contract_keys(A, pA, B, pB, pAB)) + 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()) @@ -112,7 +217,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 @@ -150,7 +255,7 @@ function TK.trace_permute!( TT.getindices(Isrc.I, q₁) == TT.getindices(Isrc.I, q₂) || continue Idst = CartesianIndex(TT.getindices(Isrc.I, (p₁..., p₂...))) tdst[Idst] = TensorKit.trace_permute!( - tdst[Idst], vsrc, (p₁, p₂), (q₁, q₂), α, One(), backend + getindex!(tdst, Idst), vsrc, (p₁, p₂), (q₁, q₂), α, One(), backend ) end return tdst diff --git a/src/tensors/vectorinterface.jl b/src/tensors/vectorinterface.jl index f7fbe66..32bf017 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 # ----- 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