Support cuTENSOR contractors for 1D views (#2650)

Author: kshyatt

lib/cutensor/src/types.jl

@@ -194,13 +194,15 @@ function CUDA.unsafe_free!(plan::CuTensorPlan)
 end
 
 
+const CUTENSOR_ALIGNMENT = UInt32(128)
+
 ## descriptor
 
 mutable struct CuTensorDescriptor
     handle::cutensorTensorDescriptor_t
     # inner constructor handles creation and finalizer of the descriptor
     function CuTensorDescriptor(sz::Vector{Int64}, st::Vector{Int64}, eltype::DataType,
-                                alignmentRequirement::UInt32=UInt32(128))
+                                alignmentRequirement::UInt32=CUTENSOR_ALIGNMENT)
         desc = Ref{cutensorTensorDescriptor_t}()
         length(st) == (N = length(sz)) || throw(ArgumentError("size and stride vectors must have the same length"))
         cutensorCreateTensorDescriptor(handle(), desc, N, sz, st, eltype, alignmentRequirement)
@@ -236,7 +238,12 @@ mutable struct CuTensor{T, N}
     inds::Vector{Int32}
 
     function CuTensor{T, N}(data::CuArray{T,N}, inds::Vector) where {T<:Number, N}
-        new(data, inds)
+        if !iszero(UInt(pointer(data)) % CUTENSOR_ALIGNMENT)
+            @warn "The data for this CuTensor does not obey the CUTENSOR alignment requirement of $CUTENSOR_ALIGNMENT. An explicit copy will be made to ensure the requirement is met."
+            return new(copy(data), inds)
+        else
+            return new(data, inds)
+        end
     end
 end
 

lib/cutensor/test/contractions.jl

@@ -164,4 +164,33 @@ eltypes = [(Float32, Float32, Float32, Float32),
     end
 end
 
+# https://github.com/JuliaGPU/CUDA.jl/issues/2407
+@testset "contractions of views" begin
+    @testset for (eltyA, eltyB, eltyC, eltyCompute) in eltypes
+        dimsA = (16,)
+        dimsB = (4,)
+        dimsC = (8,)
+        A = rand(eltyA, dimsA)
+        B = rand(eltyB, dimsB)
+        C = rand(eltyC, dimsC)
+        dA = CuArray(A)
+        dB = CuArray(B)
+        dC = CuArray(C)
+        dD = CuArray(C)
+        vA = @view dA[1:4]
+        vB = @view dB[4:4]
+        vC = @view dC[3:6]
+        vD = @view dD[3:6]
+        tA = CuTensor(reshape(vA, (4, 1)), [1, 2])
+        tB = CuTensor(reshape(vB, (1, 1)), [3, 2])
+        tC = CuTensor(reshape(vC, (1, 4)), [3, 1])
+        mul!(tC, tA, tB)
+        tA2 = CuTensor(copy(vA), [1, 2])
+        tB2 = CuTensor(copy(vB), [3, 2])
+        tD = CuTensor(copy(vD), [3, 1])
+        mul!(tD, tA2, tB2)
+        @test tD.data ≈ tD.data
+    end
+end
+
 end