Incremental

Author: kshyatt

ext/TensorKitCUDAExt/TensorKitCUDAExt.jl

@@ -1,6 +1,6 @@
 module TensorKitCUDAExt
 
-using CUDA, CUDA.CUBLAS, LinearAlgebra
+using CUDA, CUDA.CUBLAS, CUDA.CUSOLVER, LinearAlgebra
 using CUDA: @allowscalar
 using cuTENSOR: cuTENSOR
 import CUDA: rand as curand, rand! as curand!, randn as curandn, randn! as curandn!
@@ -18,6 +18,70 @@ using Random
 
 include("cutensormap.jl")
 
+# for ambiguity
+function Base.convert(A::Type{CuArray}, f::TensorKit.FusionTree{I, 0}) where {I}
+    return convert(A, TensorKit.fusiontensor(unit(I), unit(I), unit(I)))[1, 1, :]
+end
+function Base.convert(A::Type{CuArray}, f::TensorKit.FusionTree{I, 1}) where {I}
+    c = f.coupled
+    if f.isdual[1]
+        sqrtdc = TensorKit.sqrtdim(c)
+        Zcbartranspose = sqrtdc * convert(A, TensorKit.fusiontensor(dual(c), c, unit(c)))[:, :, 1, 1]
+        X = conj!(Zcbartranspose) # we want Zcbar^†
+    else
+        X = convert(A, TensorKit.fusiontensor(c, unit(c), c))[:, 1, :, 1, 1]
+    end
+    return X
+end
+# needed because the Int eltype isn't supported by CuTENSOR
+function Base.convert(A::Type{CuArray}, f::TensorKit.FusionTree{I, 2}) where {I}
+    a, b = f.uncoupled
+    isduala, isdualb = f.isdual
+    c = f.coupled
+    μ = (TensorKit.FusionStyle(I) isa TensorKit.GenericFusion) ? f.vertices[1] : 1
+    C = convert(A, TensorKit.fusiontensor(a, b, c))[:, :, :, μ]
+    X = C
+    fX = reinterpret(Float64, X)
+    if isduala
+        Za = convert(A, TensorKit.FusionTree((a,), a, (isduala,), ()))
+        # reinterpret all these as Float64 since cuTENSOR does not support Int64
+        fZa = reinterpret(Float64, Za)
+        @tensor fX[a′, b, c] := fZa[a′, a] * fX[a, b, c]
+    end
+    if isdualb
+        Zb = convert(A, TensorKit.FusionTree((b,), b, (isdualb,), ()))
+        fZb = reinterpret(Float64, Zb)
+        @tensor fX[a, b′, c] := fZb[b′, b] * fX[a, b, c]
+    end
+    return X
+end
+
+function Base.convert(A::Type{CuArray}, f::TensorKit.FusionTree{I, N}) where {I, N}
+    tailout = (f.innerlines[1], TensorKit.TupleTools.tail2(f.uncoupled)...)
+    isdualout = (false, TensorKit.TupleTools.tail2(f.isdual)...)
+    ftail = TensorKit.FusionTree(tailout, f.coupled, isdualout, Base.tail(f.innerlines), Base.tail(f.vertices))
+    Ctail = convert(A, ftail)
+    f₁ = TensorKit.FusionTree(
+        (f.uncoupled[1], f.uncoupled[2]), f.innerlines[1],
+        (f.isdual[1], f.isdual[2]), (), (f.vertices[1],)
+    )
+    C1 = convert(A, f₁)
+    dtail = size(Ctail)
+    d1 = size(C1)
+    X = similar(C1, (d1[1], d1[2], Base.tail(dtail)...))
+    trivialtuple = ntuple(identity, Val(N))
+    # reinterpret all these as Float64 since cuTENSOR does not support Int64
+    fX = reinterpret(Float64, X)
+    fC1 = reinterpret(Float64, C1)
+    fCtail = reinterpret(Float64, Ctail)
+    TensorKit.TensorOperations.tensorcontract!(
+        fX,
+        fC1, ((1, 2), (3,)), false,
+        fCtail, ((1,), Base.tail(trivialtuple)), false,
+        ((trivialtuple..., N + 1), ())
+    )
+    return X
+end
 # TODO
 # add VectorInterface extensions for proper CUDA promotion
 function TensorKit.VectorInterface.promote_add(TA::Type{<:CUDA.StridedCuMatrix{Tx}}, TB::Type{<:CUDA.StridedCuMatrix{Ty}}, α::Tα = TensorKit.VectorInterface.One(), β::Tβ = TensorKit.VectorInterface.One()) where {Tx, Ty, Tα, Tβ}

ext/TensorKitCUDAExt/cutensormap.jl

@@ -11,8 +11,6 @@ function TensorKit.tensormaptype(S::Type{<:IndexSpace}, N₁, N₂, TorA::Type{<
     end
 end
 
-TensorKit.matrixtype(::Type{<:TensorMap{T, S, N₁, N₂, A}}) where {T, S, N₁, N₂, A <: CuVector{T}} = CuMatrix{T}
-
 function CuTensorMap{T}(::UndefInitializer, V::TensorMapSpace{S, N₁, N₂}) where {T, S, N₁, N₂}
     return CuTensorMap{T, S, N₁, N₂}(undef, V)
 end
@@ -213,6 +211,10 @@ end
 TensorKit.scalartype(A::StridedCuArray{T}) where {T} = T
 TensorKit.scalartype(::Type{<:CuTensorMap{T}}) where {T} = T
 TensorKit.scalartype(::Type{<:CuArray{T}}) where {T} = T
+TensorKit.densevectortype(::Type{<:TensorMap{T, S, N₁, N₂, A}}) where {T, S, N₁, N₂, A <: CuVector{T}} = A
+TensorKit.densevectortype(::Type{<:CuArray{T}}) where {T} = CuVector{T}
+TensorKit.matrixtype(::Type{<:TensorMap{T, S, N₁, N₂, A}}) where {T, S, N₁, N₂, A <: CuVector{T}} = CuMatrix{T}
+TensorKit.matrixtype(::Type{CuArray{T}}) where {T} = CuMatrix{T}
 
 function TensorKit.similarstoragetype(TT::Type{<:CuTensorMap{TTT, S, N₁, N₂}}, ::Type{T}) where {TTT, T, S, N₁, N₂}
     return CuVector{T, CUDA.DeviceMemory}
@@ -261,7 +263,7 @@ end
 function Base.convert(::Type{CuArray}, t::AbstractTensorMap)
     I = sectortype(t)
     if I === Trivial
-        convert(CuArray, t[])
+        CUDA.@allowscalar convert(CuArray, t[])
     else
         cod = codomain(t)
         dom = domain(t)
@@ -271,8 +273,33 @@ function Base.convert(::Type{CuArray}, t::AbstractTensorMap)
         for (f₁, f₂) in fusiontrees(t)
             F = convert(CuArray, (f₁, f₂))
             Aslice = StridedView(A)[axes(cod, f₁.uncoupled)..., axes(dom, f₂.uncoupled)...]
-            add!(Aslice, StridedView(TensorKit._kron(convert(CuArray, t[f₁, f₂]), F)))
+            CUDA.@allowscalar add!(Aslice, StridedView(TensorKit._kron(convert(CuArray, t[f₁, f₂]), F)))
         end
         return A
     end
 end
+
+# CuTensorMap exponentation:
+function TensorKit.exp!(t::CuTensorMap)
+    domain(t) == codomain(t) ||
+        error("Exponential of a tensor only exist when domain == codomain.")
+    for (c, b) in blocks(t)
+        copy!(b, parent(Base.exp(Hermitian(b))))
+    end
+    return t
+end
+
+# functions that don't map ℝ to (a subset of) ℝ
+for f in (:sqrt, :log, :asin, :acos, :acosh, :atanh, :acoth)
+    sf = string(f)
+    @eval function Base.$f(t::CuTensorMap)
+        domain(t) == codomain(t) ||
+            throw(SpaceMismatch("`$($sf)` of a tensor only exist when domain == codomain"))
+        T = complex(float(scalartype(t)))
+        tf = similar(t, T)
+        for (c, b) in blocks(t)
+            copy!(block(tf, c), parent($f(Hermitian(b))))
+        end
+        return tf
+    end
+end

src/tensors/abstracttensor.jl

@@ -47,9 +47,10 @@ Return the type of vector that stores the data of a tensor.
 
 @doc """
     matrixtype(t::AbstractTensorMap) -> Type{A<:AbstractVector}
-    matrixtype(T::Type{<:AbstractTensorMap}) -> Type{A<:AbstractVector}
+    matrixtrype(T::Type{<:AbstractTensorMap}) -> Type{A<:AbstractVector}
 
-Return the type of **matrix** that stores the data of a tensor.
+Return the type of matrix that stores the data of a tensor, for conversion
+to/from dictionaries.
 """ matrixtype
 
 similarstoragetype(TT::Type{<:AbstractTensorMap}) = similarstoragetype(TT, scalartype(TT))
@@ -181,8 +182,8 @@ end
 #------------------------------------------------------------
 InnerProductStyle(t::AbstractTensorMap) = InnerProductStyle(typeof(t))
 storagetype(t::AbstractTensorMap) = storagetype(typeof(t))
-matrixtype(t::AbstractTensorMap) = matrixtype(typeof(t))
 blocktype(t::AbstractTensorMap) = blocktype(typeof(t))
+matrixtype(t::AbstractTensorMap) = matrixtype(typeof(t))
 similarstoragetype(t::AbstractTensorMap, T = scalartype(t)) = similarstoragetype(typeof(t), T)
 
 numout(t::AbstractTensorMap) = numout(typeof(t))
@@ -633,7 +634,8 @@ function Base.convert(::Type{Array}, t::AbstractTensorMap)
         for (f₁, f₂) in fusiontrees(t)
             F = convert(Array, (f₁, f₂))
             Aslice = StridedView(A)[axes(cod, f₁.uncoupled)..., axes(dom, f₂.uncoupled)...]
-            add!(Aslice, StridedView(_kron(convert(Array, t[f₁, f₂]), F)))
+            tf₁f₂ = convert(Array, t[f₁, f₂])
+            add!(Aslice, StridedView(_kron(tf₁f₂, F)))
         end
         return A
     end

src/tensors/tensor.jl

@@ -67,10 +67,17 @@ Return the type of the storage `A` of the tensor map.
 """
 storagetype(::Type{<:TensorMap{T, S, N₁, N₂, A}}) where {T, S, N₁, N₂, A <: DenseVector{T}} = A
 """
-    matrixtype(::Union{T,Type{T}}) where {T<:TensorMap} -> Type{A<:Vector}
+    densevectortype(::Union{T,Type{T}}) where {T<:TensorMap} -> Type{A<:Vector}
 
 Return the type of the storage `A` of the tensor map.
 """
+densevectortype(::Type{<:TensorMap{T, S, N₁, N₂, A}}) where {T, S, N₁, N₂, A <: Vector{T}} = A
+densevectortype(::Type{<:Array{T}}) where {T} = Vector{T}
+
+"""
+    matrixtype(::Union{T,Type{T}}) where {T<:TensorMap} -> Type{A<:Vector}
+Return the matrix analogue type of the storage `A` of the tensor map.
+"""
 matrixtype(::Type{<:TensorMap{T, S, N₁, N₂, A}}) where {T, S, N₁, N₂, A <: Vector{T}} = Matrix{T}
 
 dim(t::TensorMap) = length(t.data)
@@ -332,6 +339,8 @@ function TensorMap(
     if ndims(data) == 1 && length(data) == dim(V)
         if data isa DenseVector # refer to specific data-capturing constructor
             return TensorMap{T}(data, V)
+        elseif data isa DenseArray # can be reshaped to a dense vector
+            return TensorMap{T}(vec(data), V)
         else
             return TensorMap{T}(collect(data), V)
         end
@@ -351,11 +360,26 @@ function TensorMap(
         return TensorMap(reshape(data, length(data)), V)
     end
 
-    t = TensorMap{T}(undef, codom, dom)
-    arraydata = reshape(collect(data), arraysize)
-    t = project_symmetric!(t, arraydata)
-    if !isapprox(arraydata, convert(Array, t); atol = tol)
-        throw(ArgumentError("Data has non-zero elements at incompatible positions"))
+    arraydata = reshape(data, arraysize)
+    if data isa DenseVector
+        t = TensorMap{T, S, N₁, N₂, typeof(data)}(undef, codom ← dom) # to ensure the correct storage type
+        t = project_symmetric!(t, arraydata)
+        if !isapprox(Array(arraydata), convert(Array, t); atol = tol)
+            throw(ArgumentError("Data has non-zero elements at incompatible positions"))
+        end
+    elseif data isa DenseArray # can be reshaped into a DenseVector
+        A = densevectortype(typeof(data))
+        t = TensorMap{T, S, N₁, N₂, A}(undef, codom ← dom) # to ensure the correct storage type
+        t = project_symmetric!(t, arraydata)
+        if !isapprox(Array(arraydata), convert(Array, t); atol = tol)
+            throw(ArgumentError("Data has non-zero elements at incompatible positions"))
+        end
+    else
+        t = TensorMap{T}(undef, codom ← dom) # to ensure the correct storage type
+        t = project_symmetric!(t, arraydata)
+        if !isapprox(arraydata, convert(Array{scalartype(t)}, t); atol = tol)
+            throw(ArgumentError("Data has non-zero elements at incompatible positions"))
+        end
     end
     return t
 end