More tweaks

Author: kshyatt

ext/TensorKitCUDAExt/TensorKitCUDAExt.jl

@@ -10,13 +10,14 @@ using TensorKit.Factorizations
 using TensorKit.Strided
 using TensorKit.Factorizations: AbstractAlgorithm
 using TensorKit: SectorDict, tensormaptype, scalar, similarstoragetype, AdjointTensorMap, scalartype, project_symmetric_and_check
-import TensorKit: randisometry, rand, randn
+import TensorKit: randisometry, rand, randn, _copyto!, _add_general_kernel_nonthreaded!, blocktype
 
 using TensorKit: MatrixAlgebraKit
 
 using Random
 
 include("cutensormap.jl")
 include("truncation.jl")
+include("auxiliary.jl")
 
 end

ext/TensorKitCUDAExt/auxiliary.jl

@@ -0,0 +1,28 @@
+function TensorKit._copyto!(A::StridedView{TA, 1, <:CuArray{TA}}, B::StridedView{TB, 2, <:CuArray{TB}}) where {TA, TB}
+    length(A) == length(B) || throw(DimensionMismatch(lazy"length of A ($(length(A))) does not match length of B ($(length(B))"))
+
+    Adata = parent(A)
+    Astr = stride(A, 1)
+    IA = A.offset
+
+    Bdata = parent(B)
+    Bstr = strides(B)
+
+    IB_1 = B.offset
+    # build index arrays
+    IAs = Int[]
+    IBs = Int[]
+    @inbounds for _ in axes(B, 2)
+        IB = IB_1
+        for _ in axes(B, 1)
+            IA += Astr
+            append!(IAs, IA)
+            IB += Bstr[1]
+            append!(IBs, IB)
+        end
+        IB_1 += Bstr[2]
+    end
+    Adata[IAs] .= Bdata[IBs]
+
+    return A
+end

ext/TensorKitCUDAExt/cutensormap.jl

@@ -7,6 +7,17 @@ function CuTensorMap(t::TensorMap{T, S, N₁, N₂, A}) where {T, S, N₁, N₂,
     return CuTensorMap{T, S, N₁, N₂}(CuArray{T}(t.data), space(t))
 end
 
+#=function TensorKit.TensorMap{T, S₁, N₁, N₂, A}(
+    ::UndefInitializer, space::TensorMapSpace{S₂, N₁, N₂}
+) where {T, S₁, S₂ <: TensorKit.ElementarySpace, N₁, N₂, A <: CuVector{T}}
+        d = TensorKit.fusionblockstructure(space).totaldim
+        data = A(undef, d)
+    if !isbitstype(T)
+        zerovector!(data)
+    end
+    return TensorKit.TensorMap{T, S₂,  A}(data, space)
+end=#
+
 # project_symmetric! doesn't yet work for GPU types, so do this on the host, then copy
 function TensorKit.project_symmetric_and_check(::Type{T}, ::Type{A}, data::AbstractArray, V::TensorMapSpace; tol = sqrt(eps(real(float(eltype(data)))))) where {T, A <: CuVector{T}}
     h_t = TensorKit.TensorMapWithStorage{T, Vector{T}}(undef, V)
@@ -17,6 +28,10 @@ function TensorKit.project_symmetric_and_check(::Type{T}, ::Type{A}, data::Abstr
     return TensorKit.TensorMapWithStorage{T, A}(A(h_t.data), V)
 end
 
+function TensorKit.blocktype(::Type{<:CuTensorMap{T, S}}) where {T, S}
+    return SubArray{T, 1, CuVector{T, CUDA.DeviceMemory}, Tuple{UnitRange{Int}}, true}
+end
+
 for (fname, felt) in ((:zeros, :zero), (:ones, :one))
     @eval begin
         function CUDA.$fname(
@@ -102,9 +117,21 @@ function TensorKit.scalar(t::CuTensorMap{T, S, 0, 0}) where {T, S}
 end
 
 function Base.convert(
-        TT::Type{CuTensorMap{T, S, N₁, N₂}},
-        t::AbstractTensorMap{<:Any, S, N₁, N₂}
-    ) where {T, S, N₁, N₂}
+        TT::Type{TensorMap{T, S, N₁, N₂, A}},
+        t::TensorMap{T, S, N₁, N₂, AA}
+    ) where {T, S, N₁, N₂, A <: CuArray{T}, AA}
+    if typeof(t) === TT
+        return t
+    else
+        tnew = TT(undef, space(t))
+        return copy!(tnew, t)
+    end
+end
+
+function Base.convert(
+        TT::Type{TensorMap{T, S, N₁, N₂, A}},
+        t::AdjointTensorMap
+    ) where {T, S, N₁, N₂, A <: CuArray{T}}
     if typeof(t) === TT
         return t
     else
@@ -140,6 +167,8 @@ end
 
 TensorKit.promote_storage_rule(::Type{CuArray{T, N}}, ::Type{<:CuArray{T, N}}) where {T, N} =
     CuArray{T, N, CUDA.default_memory}
+TensorKit.promote_storage_rule(::Type{<:CuArray{T, N}}, ::Type{CuArray{T, N}}) where {T, N} =
+    CuArray{T, N, CUDA.default_memory}
 
 
 # CuTensorMap exponentation:
@@ -168,3 +197,21 @@ for f in (:sqrt, :log, :asin, :acos, :acosh, :atanh, :acoth)
         return tf
     end
 end
+
+function TensorKit._add_general_kernel_nonthreaded!(
+        tdst::CuTensorMap, tsrc::CuTensorMap, p, transformer::TensorKit.GenericTreeTransformer, α, β, backend...
+    )
+    # preallocate buffers
+    buffers = TensorKit.allocate_buffers(tdst, tsrc, transformer)
+
+    for subtransformer in transformer.data
+        # Special case without intermediate buffers whenever there is only a single block
+        if length(subtransformer[1]) == 1
+            TensorKit._add_transform_single!(tdst, tsrc, p, subtransformer, α, β, backend...)
+        else
+            cu_subtransformer = tuple(CUDA.adapt(CuArray, subtransformer[1]), subtransformer[2:end]...)
+            TensorKit._add_transform_multi!(tdst, tsrc, p, cu_subtransformer, buffers, α, β, backend...)
+        end
+    end
+    return nothing
+end

src/auxiliary/auxiliary.jl

@@ -60,7 +60,7 @@ end
 # Low-overhead implementation of `copyto!` for specific case of `stride(B, 1) < stride(B, 2)`
 # used in indexmanipulations: avoids the overhead of Strided.jl
 function _copyto!(A::StridedView{<:Any, 1}, B::StridedView{<:Any, 2})
-    length(A) == length(B) || throw(DimensionMismatch())
+    length(A) == length(B) || throw(DimensionMismatch(lazy"length of A ($(length(A))) does not match length of B ($(length(B))"))
 
     Adata = parent(A)
     Astr = stride(A, 1)

src/tensors/braidingtensor.jl

@@ -171,12 +171,15 @@ end
 has_shared_permute(t::BraidingTensor, ::Index2Tuple) = false
 function add_transform!(
         tdst::AbstractTensorMap,
-        tsrc::BraidingTensor, (p₁, p₂)::Index2Tuple,
+        tsrc::BraidingTensor{T, S},
+        (p₁, p₂)::Index2Tuple,
         fusiontreetransform,
         α::Number, β::Number, backend::AbstractBackend...
-    )
+    ) where {T, S}
+    tsrc_map = TensorMapWithStorage{scalartype(tdst), storagetype(tdst)}(undef, (tsrc.V2 ⊗ tsrc.V1) ← (tsrc.V1 ⊗ tsrc.V2))
+    copy!(tsrc_map, tsrc)
     return add_transform!(
-        tdst, TensorMap(tsrc), (p₁, p₂), fusiontreetransform, α, β,
+        tdst, tsrc_map, (p₁, p₂), fusiontreetransform, α, β,
         backend...
     )
 end

src/tensors/tensoroperations.jl

@@ -419,8 +419,10 @@ end
 # Scalar implementation
 #-----------------------
 function scalar(t::AbstractTensorMap{T, S, 0, 0}) where {T, S}
-    Bs = collect(blocks(t))
-    inds = findall(!iszero ∘ last, Bs)
-    isempty(inds) && return zero(scalartype(t))
-    return only(last(Bs[only(inds)]))
+    Bs = blocks(t)
+    B_ends = collect.(map(last, Bs))
+    nz_B_ends = [!iszero.(B) for B in B_ends]
+    valid_Bs = filter(any, B_ends)
+    isempty(valid_Bs) && return zero(scalartype(t))
+    return only(last(first(valid_Bs)))
 end

src/tensors/treetransformers.jl

@@ -46,7 +46,7 @@ function AbelianTreeTransformer(transform, p, Vdst, Vsrc)
 end
 
 const _GenericTransformerData{T, N} = Tuple{
-    Matrix{T},
+    DenseMatrix{T},
     Tuple{NTuple{N, Int}, Vector{Tuple{NTuple{N, Int}, Int}}},
     Tuple{NTuple{N, Int}, Vector{Tuple{NTuple{N, Int}, Int}}},
 }