Some more small changes for GPU support (#48)

* Updates for GPU compatibility for MPSKit

* Fixes for BraidingTensor and svd

* Formatter

* Cleanup removeunit

* Move over storagetype logic from TensorKit

* Missing module qualifier

* Undo ugly formatting changes

* move auxiliary code

* remove unnecessary auxiliary code

* use similar_diagonal

* revert tensorcontract_type changes

* try and revert storagetype

---------

Co-authored-by: lkdvos <ldevos98@gmail.com>

Author: kshyatt

Project.toml

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

src/BlockTensorKit.jl

@@ -37,6 +37,8 @@ import TensorOperations as TO
 import TupleTools as TT
 import MatrixAlgebraKit as MAK
 
+include("auxiliary/blockarrays.jl")
+
 # Spaces
 include("vectorspaces/sumspace.jl")
 include("vectorspaces/sumspaceindices.jl")

src/auxiliary/blockarrays.jl

@@ -0,0 +1,26 @@
+function copy_dense!(Adense, A)
+    for block_index in Iterators.product(blockaxes(A)...)
+        a = view(A, block_index...)
+        indices = getindex.(axes(A), block_index)
+        Adense[indices...] .= a
+    end
+    return Adense
+end
+
+const BlockBlasMat{T <: MAK.BlasFloat} = BlockMatrix{T}
+
+function MAK.zero!(A::BlockBlasMat)
+    for bj in blockaxes(A, 2), bi in blockaxes(A, 1)
+        a = view(A, bi, bj)
+        MAK.zero!(a)
+    end
+    return A
+end
+
+function MAK.one!(A::BlockBlasMat)
+    for bj in blockaxes(A, 2), bi in blockaxes(A, 1)
+        a = view(A, bi, bj)
+        bi == bj ? MAK.one!(a) : MAK.zero!(a)
+    end
+    return A
+end

src/linalg/factorizations.jl

@@ -4,25 +4,6 @@ import MatrixAlgebraKit as MAK
 
 # Type piracy for defining the MAK rules on BlockArrays!
 # -----------------------------------------------------
-
-const BlockBlasMat{T <: MAK.BlasFloat} = BlockMatrix{T}
-
-function MatrixAlgebraKit.zero!(A::BlockBlasMat)
-    for bj in blockaxes(A, 2), bi in blockaxes(A, 1)
-        a = view(A, bi, bj)
-        MAK.zero!(a)
-    end
-    return A
-end
-
-function MatrixAlgebraKit.one!(A::BlockBlasMat)
-    for bj in blockaxes(A, 2), bi in blockaxes(A, 1)
-        a = view(A, bi, bj)
-        bi == bj ? MAK.one!(a) : MAK.zero!(a)
-    end
-    return A
-end
-
 for f in
     [
         :svd_compact, :svd_full, :svd_vals,
@@ -44,7 +25,7 @@ for f! in (
     )
     @eval function MAK.$f!(t::AbstractBlockTensorMap, F, alg::AbstractAlgorithm)
         TensorKit.foreachblock(t, F...) do _, (tblock, Fblocks...)
-            Fblocks′ = MAK.$f!(Array(tblock), alg)
+            Fblocks′ = MAK.$f!(copy_dense!(similar(tblock, size(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′)
@@ -63,7 +44,7 @@ for f! in (
     )
     @eval function MAK.$f!(t::AbstractBlockTensorMap, N, alg::AbstractAlgorithm)
         TensorKit.foreachblock(t, N) do _, (tblock, Nblock)
-            Nblock′ = MAK.$f!(Array(tblock), alg)
+            Nblock′ = MAK.$f!(copy_dense!(similar(tblock, size(tblock)), tblock), alg)
             # deal with the case where the output is not the same as the input
             Nblock === Nblock′ || copy!(Nblock, Nblock′)
             return nothing
@@ -144,15 +125,15 @@ end
 function MAK.initialize_output(::typeof(eigh_full!), t::AbstractBlockTensorMap, ::AbstractAlgorithm)
     V_D = ⊕(fuse(domain(t)))
     T = real(scalartype(t))
-    D = DiagonalTensorMap{T}(undef, V_D)
+    D = TK.similar_diagonal(t, T, V_D)
     V = dense_similar(t, codomain(t) ← V_D)
     return D, V
 end
 
 function MAK.initialize_output(::typeof(eig_full!), t::AbstractBlockTensorMap, ::AbstractAlgorithm)
     V_D = ⊕(fuse(domain(t)))
     Tc = complex(scalartype(t))
-    D = DiagonalTensorMap{Tc}(undef, V_D)
+    D = TK.similar_diagonal(t, Tc, V_D)
     V = dense_similar(t, Tc, codomain(t) ← V_D)
     return D, V
 end
@@ -168,7 +149,7 @@ end
 function MAK.initialize_output(::typeof(svd_compact!), t::AbstractBlockTensorMap, ::AbstractAlgorithm)
     V_cod = V_dom = infimum(fuse(codomain(t)), fuse(domain(t)))
     U = dense_similar(t, codomain(t) ← V_cod)
-    S = DiagonalTensorMap{real(scalartype(t))}(undef, V_cod)
+    S = TK.similar_diagonal(t, real(scalartype(t)), V_cod)
     Vᴴ = dense_similar(t, V_dom ← domain(t))
     return U, S, Vᴴ
 end

src/tensors/abstractblocktensor/abstractarray.jl

@@ -328,7 +328,7 @@ Base.eltypeof(t::AbstractBlockTensorMap) = eltype(t)
     ) where {T <: AbstractTensorMap}
     catdims = Base.dims2cat(dims)
     V = space(Base._cat(dims, eachspace.(ts)...))
-    A = similar(ts[1], T, V)
+    A = similar(ts[1], TK.storagetype(ts[1]), V)
     shape = size(A)
     if count(!iszero, catdims)::Int > 1
         zerovector!(A)

src/tensors/abstractblocktensor/abstracttensormap.jl

@@ -106,4 +106,4 @@ function Base.iterate(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, state...
 end
 Base.getindex(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, c::Sector) = block(iter.t, c)
 
-TensorKit.storagetype(::Type{TT}) where {TT <: AbstractBlockTensorMap} = storagetype(eltype(TT))
+TK.storagetype(::Type{TT}) where {TT <: AbstractBlockTensorMap} = storagetype(eltype(TT))