More updates for GPU friendliness

Author: kshyatt

src/tensors/abstractblocktensor/conversion.jl

@@ -1,5 +1,6 @@
 # Conversion
 # ----------
+
 function Base.convert(::Type{TensorMap}, t::AbstractBlockTensorMap)
     S = spacetype(t)
     N₁, N₂ = numout(t), numin(t)
@@ -22,23 +23,17 @@ function Base.convert(::Type{TensorMap}, t::AbstractBlockTensorMap)
             indices = getindex.(blockax, Block.(Tuple(k)))
             arr_slice = arr[indices...]
             # need to check for empty since fusion tree pair might not be present
-            isempty(arr_slice) || copy!(arr_slice, v[f₁, f₂])
+            if !isempty(arr_slice)
+                arr[indices...] .= v[f₁, f₂]
+            end
         end
     end
 
     return tdst
 end
-
-function Base.convert(::Type{T}, t::AbstractBlockTensorMap) where {T <: TensorMap}
-    tdst = convert(TensorMap, t)
-    return convert(T, tdst)
-end
-
-# disambiguation
-function Base.convert(::Type{TensorMap{T, S, N₁, N₂, A}}, t::AB) where {T, S, N₁, N₂, A, AB <: AbstractBlockTensorMap}
-    tdst = convert(TensorMap, t)
-    return convert(T, tdst)
-end
+# use subtype of TensorMap here to support CuTensorMap
+# ROCTensorMap, etc.
+Base.convert(::Type{<:TensorMap}, t::AbstractBlockTensorMap) = convert(TensorMap, t)
 
 function Base.convert(::Type{TT}, t::AbstractTensorMap) where {TT <: AbstractBlockTensorMap}
     t isa TT && return t

src/tensors/blocktensor.jl

@@ -26,6 +26,10 @@ struct BlockTensorMap{TT <: AbstractTensorMap, E, S, N₁, N₂, N} <:
     end
 end
 
+# seems necessary to dispatch correctly onto the storage type of TT
+# AbstractBlockTensorMap doesn't have this TT field
+TensorKit.storagetype(::Type{<:BlockTensorMap{TT}}) where {TT <: AbstractTensorMap} = storagetype(TT)
+
 function BlockTensorMap{TT, E, S, N₁, N₂, N}(
         ::UndefInitializer, space::TensorMapSumSpace{S, N₁, N₂}
     ) where {TT, E, S, N₁, N₂, N}

src/tensors/sparseblocktensor.jl

@@ -63,6 +63,10 @@ function sparseblocktensormaptype(
     return SparseBlockTensorMap{TT}
 end
 
+# seems necessary to dispatch correctly onto the storage type of TT
+# AbstractBlockTensorMap doesn't have this TT field
+TensorKit.storagetype(::Type{<:SparseBlockTensorMap{TT}}) where {TT <: AbstractTensorMap} = storagetype(TT)
+
 # Constructors
 # ------------
 function SparseBlockTensorMap{TT}(

src/tensors/tensoroperations.jl

@@ -15,6 +15,17 @@ function TO.tensoradd_type(TC, A::AdjointBlockTensorMap, pA::Index2Tuple, conjA:
     return TO.tensoradd_type(TC, A', adjointtensorindices(A, pA), !conjA)
 end
 
+# copy blocks back to CPU/collect them into an array
+# seems necessary for GPU-backed BlockTensorMaps but
+# maybe not the most efficient approach?
+function TO.tensorscalar(t::AbstractBlockTensorMap{T, S, 0, 0}) where {T, S}
+    Bs = TK.blocks(t)
+    B_ends = collect.(map(b -> collect.(getfield(b, :blocks)), map(last, Bs)))
+    inds = findall(!iszero ∘ last, B_ends)
+    isempty(inds) && return zero(TKscalartype(t))
+    return only(last(B_ends[only(inds)]))
+end
+
 # tensoralloc_contract
 # --------------------
 for TTA in (:AbstractTensorMap, :AbstractBlockTensorMap), TTB in (:AbstractTensorMap, :AbstractBlockTensorMap)