CUDA tensors updates

Author: kshyatt

Project.toml

@@ -28,7 +28,7 @@ cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 GPUArrays = {rev = "master", url = "https://github.com/JuliaGPU/GPUArrays.jl"}
 MatrixAlgebraKit = {rev = "ksh/tk", url = "https://github.com/QuantumKitHub/MatrixAlgebraKit.jl"}
 AMDGPU = {rev = "master", url = "https://github.com/JuliaGPU/AMDGPU.jl"}
-cuTENSOR = {subdir = "lib/cutensor", url = "https://github.com/JuliaGPU/CUDA.jl"}
+cuTENSOR = {subdir = "lib/cutensor", url = "https://github.com/JuliaGPU/CUDA.jl", rev="master"}
 
 [extensions]
 TensorKitAMDGPUExt = "AMDGPU"

ext/TensorKitCUDAExt/TensorKitCUDAExt.jl

@@ -8,8 +8,10 @@ import CUDA: rand as curand, rand! as curand!, randn as curandn, randn! as curan
 using TensorKit
 import TensorKit.VectorInterface: scalartype as vi_scalartype
 using TensorKit.Factorizations
+using TensorKit.Strided
 using TensorKit.Factorizations: AbstractAlgorithm
 using TensorKit: SectorDict, tensormaptype, scalar, similarstoragetype, AdjointTensorMap
+import TensorKit: randisometry
 
 using TensorKit.MatrixAlgebraKit
 

ext/TensorKitCUDAExt/cutensormap.jl

@@ -75,6 +75,12 @@ function CuTensorMap(
     ) where {S}
     return CuTensorMap(data, codom ← dom)
 end
+function CuTensorMap(data::DenseVector{T}, V::TensorMapSpace{S, N₁, N₂}) where {T,S,N₁,N₂}
+    return CuTensorMap{T, S, N₁, N₂}(data, V)
+end
+function CuTensorMap(data::CuArray{T}, V::TensorMapSpace{S, N₁, N₂}) where {T,S,N₁,N₂}
+    return CuTensorMap{T, S, N₁, N₂}(vec(data), V)
+end
 
 for (fname, felt) in ((:zeros, :zero), (:ones, :one))
     @eval begin
@@ -158,6 +164,55 @@ for randfun in (:curand, :curandn)
     end
 end
 
+for randfun in (:rand, :randn, :randisometry)
+    randfun! = Symbol(randfun, :!)
+    @eval begin
+        # converting `codomain` and `domain` into `HomSpace`
+        function $randfun(::Type{A}, codomain::TensorSpace{S},
+                               domain::TensorSpace{S}) where {A <: CuArray, S<:IndexSpace}
+            return $randfun(A, codomain ← domain)
+        end
+        function $randfun(::Type{T}, ::Type{A}, codomain::TensorSpace{S},
+                          domain::TensorSpace{S}) where {T,S<:IndexSpace, A<:CuArray{T}}
+            return $randfun(T, A, codomain ← domain)
+        end
+        function $randfun(rng::Random.AbstractRNG, ::Type{T}, ::Type{A},
+                          codomain::TensorSpace{S},
+                          domain::TensorSpace{S}) where {T,S<:IndexSpace, A<:CuArray{T}}
+            return $randfun(rng, T, A, codomain ← domain)
+        end
+
+        # accepting single `TensorSpace`
+        $randfun(::Type{A}, codomain::TensorSpace) where {A<:CuArray} = $randfun(A, codomain ← one(codomain))
+        function $randfun(::Type{T}, ::Type{A}, codomain::TensorSpace) where {T, A<:CuArray{T}}
+            return $randfun(T, A, codomain ← one(codomain))
+        end
+        function $randfun(rng::Random.AbstractRNG, ::Type{T},
+                          ::Type{A}, codomain::TensorSpace) where {T, A<:CuArray{T}}
+            return $randfun(rng, T, A, codomain ← one(domain))
+        end
+
+        # filling in default eltype
+        $randfun(::Type{A}, V::TensorMapSpace) where {A<:CuArray} = $randfun(eltype(A), A, V)
+        function $randfun(rng::Random.AbstractRNG, ::Type{A}, V::TensorMapSpace) where {A<:CuArray}
+            return $randfun(rng, eltype(A), A, V)
+        end
+
+        # filling in default rng
+        function $randfun(::Type{T}, ::Type{A}, V::TensorMapSpace) where {T, A<:CuArray{T}}
+            return $randfun(Random.default_rng(), T, A, V)
+        end
+
+        # implementation
+        function $randfun(rng::Random.AbstractRNG, ::Type{T},
+                          ::Type{A}, V::TensorMapSpace) where {T, A<:CuArray{T}}
+            t = CuTensorMap{T}(undef, V)
+            $randfun!(rng, t)
+            return t
+        end
+    end
+end
+
 # converters
 # ----------
 function Base.convert(::Type{CuTensorMap}, d::Dict{Symbol, Any})
@@ -250,3 +305,26 @@ function LinearAlgebra.isposdef(t::CuTensorMap)
     end
     return true
 end
+
+
+# Conversion to CuArray:
+#----------------------
+# probably not optimized for speed, only for checking purposes
+function Base.convert(::Type{CuArray}, t::AbstractTensorMap)
+    I = sectortype(t)
+    if I === Trivial
+        convert(CuArray, t[])
+    else
+        cod = codomain(t)
+        dom = domain(t)
+        T = sectorscalartype(I) <: Complex ? complex(scalartype(t)) :
+            sectorscalartype(I) <: Integer ? scalartype(t) : float(scalartype(t))
+        A = CUDA.zeros(T, dims(cod)..., dims(dom)...)
+        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)))
+        end
+        return A
+    end
+end

src/auxiliary/random.jl

@@ -1,17 +1,20 @@
 """
-    randisometry([::Type{T}=Float64], dims::Dims{2}) -> Array{T,2}
-    randhaar([::Type{T}=Float64], dims::Dims{2}) -> Array{T,2}
+    randisometry([::Type{T}=Float64], [::Type{A}=Matrix{T}], dims::Dims{2}) -> A
+    randhaar([::Type{T}=Float64], [::Type{A}=Matrix{T}], dims::Dims{2}) -> A
 
 Create a random isometry of size `dims`, uniformly distributed according to the Haar measure.
 
 See also [`randuniform`](@ref) and [`randnormal`](@ref).
 """
-randisometry(dims::Base.Dims{2}) = randisometry(Float64, dims)
+randisometry(dims::Base.Dims{2}) = randisometry(Float64, Matrix{Float64}, dims)
 function randisometry(::Type{T}, dims::Base.Dims{2}) where {T <: Number}
     return randisometry(Random.default_rng(), T, dims)
 end
-function randisometry(rng::Random.AbstractRNG, ::Type{T}, dims::Base.Dims{2}) where {T <: Number}
-    return randisometry!(rng, Matrix{T}(undef, dims))
+function randisometry(::Type{T}, ::Type{A}, dims::Base.Dims{2}) where {T <: Number, A<:AbstractArray{T}}
+    return randisometry(Random.default_rng(), T, A, dims)
+end
+function randisometry(rng::Random.AbstractRNG, ::Type{T}, ::Type{A}, dims::Base.Dims{2}) where {T <: Number, A<:AbstractArray{T}}
+    return randisometry!(rng, A(undef, dims))
 end
 
 randisometry!(A::AbstractMatrix) = randisometry!(Random.default_rng(), A)

src/tensors/linalg.jl

@@ -281,7 +281,7 @@ function _norm(blockiter, p::Real, init::Real)
         end
     elseif p > 0
         nᵖ = mapreduce(+, blockiter; init = init) do (c, b)
-            return isempty(b) ? init : oftype(init, dim(c) * LinearAlgebra.normp(b, p)^p)
+            return isempty(b) ? init : oftype(init, dim(c) * LinearAlgebra.norm(b, p)^p)
         end
         return (nᵖ)^inv(oftype(nᵖ, p))
     else
@@ -431,7 +431,7 @@ function exp!(t::TensorMap)
     domain(t) == codomain(t) ||
         error("Exponential of a tensor only exist when domain == codomain.")
     for (c, b) in blocks(t)
-        copy!(b, LinearAlgebra.exp!(b))
+        copy!(b, exp!(b))
     end
     return t
 end

src/tensors/tensor.jl

@@ -357,7 +357,7 @@ function TensorMap(
     ) where {S}
     return TensorMap(data, codom ← dom; kwargs...)
 end
-function Tensor(data::AbstractArray, codom::TensorSpace, ; kwargs...)
+function Tensor(data::AbstractArray, codom::TensorSpace; kwargs...)
     return TensorMap(data, codom ← one(codom); kwargs...)
 end