Small fixes and missed tests for CUTENSORNET (#2713)

* Small fixes and missed tests for CUTENSORNET

* More fixes and tests

Author: kshyatt

lib/cutensornet/src/types.jl

@@ -51,7 +51,7 @@ mutable struct CuTensorDescriptor{T}
                                           T, desc_ref)
         obj = new{T}(desc_ref[])
         finalizer(cutensornetDestroyTensorDescriptor, obj)
-        obj
+        return obj
     end
 end
 CuTensorDescriptor(T::DataType, extents, strides, modes) = CuTensorDescriptor{T}(extents, strides, modes)
@@ -65,12 +65,14 @@ function Base.ndims(desc::CuTensorDescriptor)
 end
 
 function Base.size(desc::CuTensorDescriptor)
+    numModes = Ref{Int32}(C_NULL)
     extents  = Vector{Int64}(undef, ndims(desc))
     cutensornetGetTensorDetails(handle(), desc, numModes, C_NULL, C_NULL, extents, C_NULL)
     return tuple(extents...)
 end
 
 function Base.strides(desc::CuTensorDescriptor)
+    numModes = Ref{Int32}(C_NULL)
     strides  = Vector{Int64}(undef, ndims(desc))
     cutensornetGetTensorDetails(handle(), desc, numModes, C_NULL, C_NULL, C_NULL, strides)
     return tuple(strides...)
@@ -89,16 +91,16 @@ mutable struct CuTensorNetworkDescriptor
                                            extentsOut, stridesOut, modesOut, dataType, computeType, desc_ref)
         obj = new(desc_ref[])
         finalizer(cutensornetDestroyNetworkDescriptor, obj)
-        obj
+        return obj
     end
 end
 Base.unsafe_convert(::Type{cutensornetNetworkDescriptor_t}, desc::CuTensorNetworkDescriptor) = desc.handle
 
 function compute_type(T::DataType)
     if T == Float16
-        return Float32
-    elseif T == Float32
         return Float16
+    elseif T == Float32
+        return Float32
     elseif T == Float64
         return Float64
     end
@@ -133,7 +135,7 @@ mutable struct CuTensorSVDInfo
         cutensornetCreateTensorSVDInfo(handle(), info_ref)
         obj = new(info_ref[])
         finalizer(cutensornetDestroyTensorSVDInfo, obj)
-        obj
+        return obj
     end
 end
 Base.unsafe_convert(::Type{cutensornetTensorSVDInfo_t}, info::CuTensorSVDInfo) = info.handle
@@ -162,7 +164,7 @@ mutable struct CuTensorNetworkContractionOptimizerInfo
         cutensornetCreateContractionOptimizerInfo(handle(), net_desc, desc_ref)
         obj = new(desc_ref[])
         finalizer(cutensornetDestroyContractionOptimizerInfo, obj)
-        obj
+        return obj
     end
 end
 
@@ -175,7 +177,7 @@ mutable struct CuTensorNetworkWorkspaceDescriptor
         cutensornetCreateWorkspaceDescriptor(handle(), desc_ref)
         obj = new(desc_ref[])
         finalizer(cutensornetDestroyWorkspaceDescriptor, obj)
-        obj
+        return obj
     end
 end
 
@@ -188,7 +190,7 @@ mutable struct CuTensorNetworkContractionPlan
         cutensornetCreateContractionPlan(handle(), net_desc, info, ws_desc, desc_ref)
         obj = new(desc_ref[])
         finalizer(cutensornetDestroyContractionPlan, obj)
-        obj
+        return obj
     end
 end
 
@@ -250,7 +252,7 @@ mutable struct CuTensorNetworkContractionOptimizerConfig
             attr_buf = Ref(Base.getproperty(prefs, attr[1]))
             cutensornetContractionOptimizerConfigSetAttribute(handle(), desc_ref[], attr[2], attr_buf, sizeof(attr_buf))
         end
-        obj
+        return obj
     end
 end
 
@@ -287,7 +289,7 @@ mutable struct CuTensorSVDConfig
             attr_buf = Ref(Base.getproperty(prefs, attr[1]))
             cutensornetTensorSVDConfigSetAttribute(handle(), desc_ref[], attr[2], attr_buf, sizeof(attr_buf))
         end
-        obj
+        return obj
     end
 end
 function abs_cutoff(conf::CuTensorSVDConfig)
@@ -323,7 +325,7 @@ mutable struct CuTensorNetworkAutotunePreference
             attr_buf = Ref(Base.getproperty(prefs, attr[1]))
             cutensornetContractionAutotunePreferenceSetAttribute(handle(), pref_ref[], attr[2], attr_buf, sizeof(attr_buf))
         end
-        obj
+        return obj
     end
 end
 Base.unsafe_convert(::Type{cutensornetContractionAutotunePreference_t}, prefs::CuTensorNetworkAutotunePreference)   = prefs.handle
@@ -336,14 +338,14 @@ mutable struct CuTensorNetworkSliceGroup
         cutensornetCreateSliceGroupFromIDRange(handle(), sliceStart, sliceStop, sliceStep, group_ref)
         obj = new(group_ref[])
         finalizer(cutensornetDestroySliceGroup, obj)
-        obj
+        return obj
     end
     function CuTensorNetworkSliceGroup(slices::Vector{Int64})
         group_ref = Ref{cutensornetSliceGroup_t}()
         cutensornetCreateSliceGroupFromIDs(handle(), pointer(slices), pointer(slices, length(slices)), group_ref)
         obj = new(group_ref[])
         finalizer(cutensornetDestroySliceGroup, obj)
-        obj
+        return obj
     end
 end
 Base.unsafe_convert(::Type{cutensornetSliceGroup_t}, prefs::CuTensorNetworkSliceGroup)   = prefs.handle

lib/cutensornet/test/runtests.jl

@@ -15,6 +15,37 @@ import cuTensorNet: CuTensorNetwork, rehearse_contraction, perform_contraction!,
 using TensorOperations
 
 @testset "cuTensorNet" begin
+    @testset "Helpers and types" begin
+        @test convert(cuTensorNet.cutensornetComputeType_t, Int8)  == cuTensorNet.CUTENSORNET_COMPUTE_8I
+        @test convert(cuTensorNet.cutensornetComputeType_t, UInt8) == cuTensorNet.CUTENSORNET_COMPUTE_8U
+        @test convert(cuTensorNet.cutensornetComputeType_t, Float16) == cuTensorNet.CUTENSORNET_COMPUTE_16F
+        @test convert(cuTensorNet.cutensornetComputeType_t, Int32)  == cuTensorNet.CUTENSORNET_COMPUTE_32I
+        @test convert(cuTensorNet.cutensornetComputeType_t, UInt32) == cuTensorNet.CUTENSORNET_COMPUTE_32U
+        @test_throws ArgumentError("cuTensorNet type equivalent for compute type ComplexF64 does not exist!") convert(cuTensorNet.cutensornetComputeType_t, ComplexF64)
+        @test convert(Type, cuTensorNet.CUTENSORNET_COMPUTE_8I) == Int8
+        @test convert(Type, cuTensorNet.CUTENSORNET_COMPUTE_8U) == UInt8
+        @test convert(Type, cuTensorNet.CUTENSORNET_COMPUTE_16F) == Float16
+        @test convert(Type, cuTensorNet.CUTENSORNET_COMPUTE_32F) == Float32
+        @test convert(Type, cuTensorNet.CUTENSORNET_COMPUTE_32U) == UInt32
+        @test convert(Type, cuTensorNet.CUTENSORNET_COMPUTE_32I) == Int32
+        @test convert(Type, cuTensorNet.CUTENSORNET_COMPUTE_64F) == Float64
+
+        
+        modesA = ['m', 'h', 'k', 'n']
+        extent = Dict{Char, Int}()
+        extent['m'] = 96;
+        extent['n'] = 96;
+        extent['h'] = 64;
+        extent['k'] = 64;
+        extentsA = [extent[mode] for mode in modesA]
+        @testset for elty in [Float32, Float64, ComplexF32, ComplexF64]
+            A = CUDA.rand(elty, extentsA...)
+            descA = cuTensorNet.CuTensorDescriptor(A, modesA)
+            @test ndims(descA) == ndims(A)
+            @test size(descA)  == size(A)
+            @test strides(descA) == strides(A)
+        end
+    end
     n = 8
     m = 16
     k = 32
@@ -84,6 +115,10 @@ using TensorOperations
             @test cuTensorNet.reduced_extent(info)   == n
             @test cuTensorNet.discarded_weight(info) ≈ 0.0
             @test collect(U)*diagm(collect(S))*collect(V) ≈ collect(A)
+            config = cuTensorNet.CuTensorSVDConfig()
+            @test cuTensorNet.abs_cutoff(config) == 0.0
+            @test cuTensorNet.rel_cutoff(config) == 0.0
+            @test cuTensorNet.normalization(config) == cuTensorNet.CUTENSORNET_TENSOR_SVD_NORMALIZATION_NONE
         end
         @testset "GateSplit" begin
             a = 16