Incremental fixes for MPSKit

ext/BlockTensorKitGPUArraysExt.jl

@@ -8,4 +8,15 @@ function KernelAbstractions.get_backend(BA::BlockArrays.BlockArray{T, N, A}) whe
     return KernelAbstractions.get_backend(first(BA.blocks))
 end
 
+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

src/linalg/factorizations.jl

@@ -25,11 +25,10 @@ for f! in (
     )
     @eval function MAK.$f!(t::AbstractBlockTensorMap, F, alg::AbstractAlgorithm)
         TensorKit.foreachblock(t, F...) do _, (tblock, Fblocks...)
-            dense_block = similar_dense(tblock)
-            Fblocks′ = MAK.$f!(copy_dense!(dense_block, tblock), alg)
+            Fblocks′ = MAK.$f!(copy_dense!(similar_dense(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′)
+                b === b′ || copyto!(b, b′)
             end
             return nothing
         end
@@ -45,10 +44,9 @@ for f! in (
     )
     @eval function MAK.$f!(t::AbstractBlockTensorMap, N, alg::AbstractAlgorithm)
         TensorKit.foreachblock(t, N) do _, (tblock, Nblock)
-            dense_block = similar_dense(tblock)
-            Nblock′ = MAK.$f!(copy_dense!(dense_block, tblock), alg)
+            Nblock′ = MAK.$f!(copy_dense!(similar_dense(tblock), tblock), alg)
             # deal with the case where the output is not the same as the input
-            Nblock === Nblock′ || copy!(Nblock, Nblock′)
+            Nblock === Nblock′ || copyto!(Nblock, Nblock′)
             return nothing
         end
         return N
@@ -190,3 +188,30 @@ for f! in (
     @eval MAK.$f!(::AbstractBlockTensorMap, x, ::DiagonalAlgorithm) =
         error("Blocktensors are incompatible with diagonal algorithm")
 end
+
+function TensorKit.Factorizations.truncate_domain!(tdst::AbstractBlockTensorMap, tsrc::AbstractBlockTensorMap, inds)
+    TensorKit.foreachblock(tdst, tsrc) do c, (dst_block, src_block)
+        I = get(inds, c, nothing)
+        dst_dense = copy_dense!(similar_dense(dst_block), dst_block)
+        src_dense = copy_dense!(similar_dense(src_block), src_block)
+        @assert !isnothing(I)
+        @views dst_dense .= src_dense[:, I]
+        # deal with the case where the output is not in-place
+        dst_dense === dst_block || copyto!(dst_block, dst_dense)
+        return nothing
+    end
+    return tdst
+end
+function TensorKit.Factorizations.truncate_codomain!(tdst::AbstractBlockTensorMap, tsrc::AbstractBlockTensorMap, inds)
+    TensorKit.foreachblock(tdst, tsrc) do c, (dst_block, src_block)
+        I = get(inds, c, nothing)
+        dst_dense = copy_dense!(similar_dense(dst_block), dst_block)
+        src_dense = copy_dense!(similar_dense(src_block), src_block)
+        @assert !isnothing(I)
+        @views dst_dense .= src_dense[I, :]
+        # deal with the case where the output is not in-place
+        dst_dense === dst_block || copyto!(dst_block, dst_dense)
+        return nothing
+    end
+    return tdst
+end

src/linalg/linalg.jl

@@ -239,3 +239,23 @@ function LinearAlgebra.isposdef!(t::AbstractBlockTensorMap)
     end
     return true
 end
+
+function LinearAlgebra.lmul!(D::DiagonalTensorMap, t::AbstractBlockTensorMap)
+    domain(D) == codomain(t) || throw(SpaceMismatch())
+    TensorKit.foreachblock(t, D) do c, (tblock, Dblock)
+        tblock′ = lmul!(Dblock, copy_dense!(similar_dense(tblock), tblock))
+        tblock === tblock′ || copyto!(tblock, tblock′)
+        return tblock
+    end
+    return t
+end
+
+function LinearAlgebra.rmul!(t::AbstractBlockTensorMap, D::DiagonalTensorMap)
+    codomain(D) == domain(t) || throw(SpaceMismatch())
+    TensorKit.foreachblock(t, D) do c, (tblock, Dblock)
+        tblock′ = rmul!(copy_dense!(similar_dense(tblock), tblock), Dblock)
+        tblock === tblock′ || copyto!(tblock, tblock′)
+        return tblock
+    end
+    return t
+end

src/tensors/abstractblocktensor/abstractarray.jl

@@ -288,7 +288,7 @@ function similar_tensormaptype(
     ) where {S}
     if eltype(t) === T && typeof(space(t)) === typeof(P)
         return T
-    elseif isconcretetype(T)
+    elseif isconcretetype(T) || T isa Union
         return tensormaptype(S, numout(P), numin(P), storagetype(T))
     else
         return AbstractTensorMap{scalartype(T), S, numout(P), numin(P)}

src/tensors/blocktensor.jl

@@ -161,6 +161,11 @@ end
 Base.eltype(::Type{<:BlockTensorMap{TT}}) where {TT} = TT
 Base.parent(t::BlockTensorMap) = t.data
 
+# handle this separately because the storagetype of `AbstractTensorMap` is
+# *always* Vector no matter the actual data storage type
+TK.storagetype(t::BlockTensorMap{AbstractTensorMap{E, S, N₁, N₂}}) where {E, S, N₁, N₂} =
+    foldl(TK.promote_storagetype, values(t.data))
+
 function Base.copyto!(
         dest::BlockTensorMap, Rdest::CartesianIndices,
         src::BlockTensorMap, Rsrc::CartesianIndices,

src/tensors/sparseblocktensor.jl

@@ -168,6 +168,11 @@ VI.scalartype(::Type{<:SparseBlockTensorMap{TT}}) where {TT} = scalartype(TT)
 Base.parent(t::SparseBlockTensorMap) = SparseTensorArray(t.data, space(t))
 Base.eltype(::Type{<:SparseBlockTensorMap{TT}}) where {TT} = TT
 
+# handle this separately because the storagetype of `AbstractTensorMap` is
+# *always* Vector no matter the actual data storage type
+TK.storagetype(t::SparseBlockTensorMap{AbstractTensorMap{E, S, N₁, N₂}}) where {E, S, N₁, N₂} =
+    foldl(TK.promote_storagetype, nonzero_values(t))
+
 issparse(::SparseBlockTensorMap) = true
 nonzero_keys(t::SparseBlockTensorMap) = keys(t.data)
 nonzero_values(t::SparseBlockTensorMap) = values(t.data)