Make sparse block tensors consistent with TensorOperations allocators

Project.toml

@@ -1,6 +1,6 @@
 name = "BlockTensorKit"
 uuid = "5f87ffc2-9cf1-4a46-8172-465d160bd8cd"
-version = "0.3.20"
+version = "0.3.21"
 authors = ["Lukas Devos <ldevos98@gmail.com> and contributors"]
 
 [deps]

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

src/tensors/abstractblocktensor/abstracttensormap.jl

@@ -122,10 +122,10 @@ function _blockdims(P::ProductSumSpace{S, N}, c::Sector) where {S, N}
 end
 
 zeroblock(::Type{TT}, d₁::Int, d₂::Int) where {TT <: AbstractTensorMap} =
-    zerovector!(similar(storagetype(TT), d₁, d₂))
+    zerovector!(similar(storagetype(TT), d₁, d₂))::TK.blocktype(TT)
 function zeroblock(::Type{TT}, d₁::Int, d₂::Int) where {TT <: TensorMap}
     data = zerovector!(similar(storagetype(TT), d₁ * d₂))
-    return reshape(view(data, 1:(d₁ * d₂)), (d₁, d₂))
+    return reshape(view(data, 1:(d₁ * d₂)), (d₁, d₂))::TK.blocktype(TT)
 end
 zeroblock(::Type{<:AdjointTensorMap{<:Any, <:Any, <:Any, <:Any, TT}}, d₁::Int, d₂::Int) where {TT} =
     zeroblock(TT, d₂, d₁)'

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 are marked by `istemp`, as their blocks are obtained from and
+released to the allocator of the operation.
 """
 struct SparseBlockTensorMap{TT <: AbstractTensorMap, E, S, N₁, N₂, N} <:
     AbstractBlockTensorMap{E, S, N₁, N₂}
     data::Dict{CartesianIndex{N}, TT}
     space::TensorMapSumSpace{S, N₁, N₂}
+    istemp::Bool
 
     # uninitialized constructor
     function SparseBlockTensorMap{TT, E, S, N₁, N₂, N}(
-            ::UndefBlocksInitializer, space::TensorMapSumSpace{S, N₁, N₂}
+            ::UndefBlocksInitializer, space::TensorMapSumSpace{S, N₁, N₂}; istemp::Bool = false
         ) 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, istemp)
     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, false)
     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
@@ -213,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
@@ -222,6 +226,19 @@ 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) the allocator of the operation,
+# so kernels without allocator cannot add them
+istemp(t::AbstractBlockTensorMap) = false
+istemp(t::SparseBlockTensorMap) = t.istemp
+
+function newblock(t::AbstractBlockTensorMap, I::CartesianIndex, allocator = nothing)
+    V = eachspace(t)[I]
+    istemp(t) || return TO.tensoralloc(eltype(t), V, Val(false), something(allocator, TO.DefaultAllocator()))
+    isnothing(allocator) &&
+        throw(ArgumentError("cannot add block $I to a temporary without its allocator"))
+    return TO.tensoralloc(eltype(t), V, Val(true), allocator)
+end
+
 # Show
 # ----
 function Base.showarg(io::IO, t::SparseBlockTensorMap, toplevel::Bool)

src/tensors/tensoroperations.jl

@@ -97,24 +97,129 @@ 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
+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; istemp = istemp === Val(true)
+    )
+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())
     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
 

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
 # -----
@@ -49,12 +53,12 @@ end
 # blockwise updates
 # -----------------
 @inline function _addblock!(
-        f, t::AbstractBlockTensorMap, I::CartesianIndex, β::Number = One(), allocator = TO.DefaultAllocator()
+        f, t::AbstractBlockTensorMap, I::CartesianIndex, β::Number = One(), allocator = nothing
     )
     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,11 @@ 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(β)
+    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