More updates for GPU friendliness

Author: kshyatt

src/tensors/abstractblocktensor/conversion.jl

@@ -1,6 +1,9 @@
 # Conversion
 # ----------
-function Base.convert(::Type{TensorMap}, t::AbstractBlockTensorMap)
+
+# use subtype of TensorMap here to support CuTensorMap
+# ROCTensorMap, etc.
+function Base.convert(::Type{<:TensorMap}, t::AbstractBlockTensorMap)
     S = spacetype(t)
     N₁, N₂ = numout(t), numin(t)
     cod = ProductSpace{S, N₁}(oplus.(codomain(t).spaces))
@@ -22,24 +25,15 @@ 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
-
 function Base.convert(::Type{TT}, t::AbstractTensorMap) where {TT <: AbstractBlockTensorMap}
     t isa TT && return t
     if t isa AbstractBlockTensorMap

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,18 @@ 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 -> getfield(b, :blocks), map(last, Bs)))
+    nz_B_ends = [map(b -> !iszero(b), B) for B in B_ends]
+    valid_Bs = filter(any, nz_B_ends)
+    isempty(valid_Bs) && return zero(TK.scalartype(t))
+    return only(last(first(valid_Bs)))
+end
+
 # tensoralloc_contract
 # --------------------
 for TTA in (:AbstractTensorMap, :AbstractBlockTensorMap), TTB in (:AbstractTensorMap, :AbstractBlockTensorMap)