Make sparse block tensors consistent with TensorOperations allocators

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

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}) =

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]

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

src/tensors/sparseblocktensor.jl

@@ -2,27 +2,30 @@
     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
     function SparseBlockTensorMap{TT, 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)

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,24 +97,127 @@ 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())
     foreach(Base.Fix2(TO.tensorfree!, allocator), parent(t))
     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

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
 # -----