Comments and coverage

Author: kshyatt

ext/TensorKitCUDAExt/TensorKitCUDAExt.jl

@@ -10,7 +10,7 @@ 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
+import TensorKit: randisometry, rand, randn
 
 using TensorKit.MatrixAlgebraKit
 

ext/TensorKitCUDAExt/cutensormap.jl

@@ -3,12 +3,12 @@ const CuTensor{T, S, N} = CuTensorMap{T, S, N, 0}
 
 const AdjointCuTensorMap{T, S, N₁, N₂} = AdjointTensorMap{T, S, N₁, N₂, CuTensorMap{T, S, N₁, N₂}}
 
-function CuTensorMap{T, S, N₁, N₂}(t::TensorMap{T, S, N₁, N₂, A}) where {T, S, N₁, N₂, A}
+function CuTensorMap(t::TensorMap{T, S, N₁, N₂, A}) where {T, S, N₁, N₂, A}
     return CuTensorMap{T, S, N₁, N₂}(CuArray(t.data), t.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}}
+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)
     h_t = TensorKit.project_symmetric!(h_t, Array(data))
     # verify result
@@ -146,8 +146,9 @@ for randfun in (:rand, :randn, :randisometry)
     end
 end
 
-function Base.convert(::Type{CuTensorMap}, t::AbstractTensorMap)
-    return copy!(CuTensorMap{scalartype(t)}(undef, space(t)), t)
+function Base.convert(::Type{CuTensorMap}, t::AbstractTensorMap{T, S}) where {T, S}
+    d_data = CuArray(t.data)
+    return TensorMap{T}(d_data, t.space)
 end
 
 # Scalar implementation

src/tensors/linalg.jl

@@ -299,7 +299,7 @@ function LinearAlgebra.rank(
     r = 0 * dim(first(allunits(sectortype(t))))
     dim(t) == 0 && return r
     S = LinearAlgebra.svdvals(t)
-    tol = max(atol, rtol * mapreduce(maximum, max, values(S)))
+    tol = max(atol, rtol * maximum(S.data))
     for (c, b) in pairs(S)
         if !isempty(b)
             r += dim(c) * count(>(tol), b)
@@ -317,8 +317,8 @@ function LinearAlgebra.cond(t::AbstractTensorMap, p::Real = 2)
             return zero(real(float(scalartype(t))))
         end
         S = LinearAlgebra.svdvals(t)
-        maxS = mapreduce(maximum, max, values(S))
-        minS = mapreduce(minimum, min, values(S))
+        maxS = maximum(S.data)
+        minS = minimum(S.data)
         return iszero(maxS) ? oftype(maxS, Inf) : (maxS / minS)
     else
         throw(ArgumentError("cond currently only defined for p=2"))

src/tensors/tensor.jl

@@ -198,7 +198,7 @@ TensorMap(data::AbstractArray, codom::TensorSpace, dom::TensorSpace; kwargs...)
 Tensor(data::AbstractArray, codom::TensorSpace; kwargs...) =
     TensorMap(data, codom ← one(codom); kwargs...)
 
-function _project_symmetric_and_check(::Type{T}, ::Type{A}, data::AbstractArray, V::TensorMapSpace; tol = sqrt(eps(real(float(eltype(data)))))) where {T, A <: DenseVector{T}}
+function project_symmetric_and_check(::Type{T}, ::Type{A}, data::AbstractArray, V::TensorMapSpace; tol = sqrt(eps(real(float(eltype(data)))))) where {T, A <: DenseVector{T}}
     t = TensorMapWithStorage{T, A}(undef, V)
     t = project_symmetric!(t, data)
     # verify result
@@ -218,7 +218,7 @@ function TensorMapWithStorage{T, A}(
     sectortype(V) === Trivial &&
         return tensormaptype(spacetype(V), numout(V), numin(V), A)(reshape(data, length(data)), V)
 
-    return _project_symmetric_and_check(T, A, data, V; tol)
+    return project_symmetric_and_check(T, A, data, V; tol)
 end
 TensorMapWithStorage{T, A}(data::AbstractArray, codom::TensorSpace, dom::TensorSpace; kwargs...) where {T, A} =
     TensorMapWithStorage{T, A}(data, codom ← dom; kwargs...)

test/cuda/tensors.jl

@@ -4,6 +4,9 @@ using TensorKit, Combinatorics
 ad = adapt(Array)
 const CUDAExt = Base.get_extension(TensorKit, :TensorKitCUDAExt)
 @assert !isnothing(CUDAExt)
+const CuTensorMap = getglobal(CUDAExt, :CuTensorMap)
+const curand = getglobal(CUDAExt, :curand)
+const curandn = getglobal(CUDAExt, :curandn)
 
 @isdefined(TestSetup) || include("../setup.jl")
 using .TestSetup
@@ -41,6 +44,23 @@ for V in spacelist
         V1, V2, V3, V4, V5 = V
         @timedtestset "Basic tensor properties" begin
             W = V1 ⊗ V2 ⊗ V3 ⊗ V4 ⊗ V5
+            # test default pass-throughs
+            for f in (CUDA.zeros, CUDA.ones, curand, curandn)
+                t = @constinferred f(W)
+                @test scalartype(t) == Float64
+                @test codomain(t) == W
+                @test space(t) == (W ← one(W))
+                @test domain(t) == one(W)
+                @test typeof(t) == TensorMap{Float64, spacetype(t), 5, 0, CuVector{Float64, CUDA.DeviceMemory}}
+            end
+            for f in (rand, randn)
+                t = @constinferred f(CuVector{Float64}, W)
+                @test scalartype(t) == Float64
+                @test codomain(t) == W
+                @test space(t) == (W ← one(W))
+                @test domain(t) == one(W)
+                @test typeof(t) == TensorMap{Float64, spacetype(t), 5, 0, CuVector{Float64, CUDA.DeviceMemory}}
+            end
             for T in (Int, Float32, Float64, ComplexF32, ComplexF64)
                 t = @constinferred CUDA.zeros(T, W)
                 CUDA.@allowscalar begin
@@ -64,6 +84,26 @@ for V in spacelist
                 @test typeof(c) === sectortype(t)
             end
         end
+        @timedtestset "Conversion to/from host" begin
+            W = V1 ⊗ V2 ⊗ V3 ⊗ V4 ⊗ V5
+            for T in (Int, Float32, ComplexF64)
+                h_t = @constinferred rand(T, W)
+                t1 = convert(CuTensorMap, h_t)
+                @test collect(t1.data) == h_t.data
+                @test t1.space == h_t.space
+                @test scalartype(t1) == T
+                @test codomain(t1) == W
+                @test space(t1) == (W ← one(W))
+                @test domain(t1) == one(W)
+                t2 = CuTensorMap(h_t)
+                @test collect(t2.data) == h_t.data
+                @test t2.space == h_t.space
+                @test scalartype(t2) == T
+                @test codomain(t2) == W
+                @test space(t2) == (W ← one(W))
+                @test domain(t2) == one(W)
+            end
+        end
         @timedtestset "Tensor Dict conversion" begin
             W = V1 ⊗ V2 ⊗ V3 ← V4 ⊗ V5
             for T in (Int, Float32, ComplexF64)