Jutho / TensorOperations.jl

Julia package for tensor contractions and related operations
https://jutho.github.io/TensorOperations.jl/stable/
Other
443 stars 56 forks source link

GPU Documentation #104

Closed zjwegert closed 3 years ago

zjwegert commented 3 years ago

Hi, I'm having trouble getting the code working for tensor contractions on a GPU. See below:

A=randn(Float64, (3,20,5,3,4))
B=randn(Float64, (5,6,20,3))
@cutensor C1[a, g, e, d, f] := A[a, b, c, d, e] * B[c, f, b, g]

produces the error:

julia> @cutensor C1[a, g, e, d, f] := A[a, b, c, d, e] * B[c, f, b, g]
ERROR: MethodError: Cannot `convert` an object of type 
  Tuple{Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{}} to an object of type
  Tuple{Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{, 
Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64{},Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64,Int64}
Closest candidates are:
  convert(::Type{T}, ::Tuple{Vararg{Any, N}}) where {N, T<:Tuple} at essentials.jl:316
  convert(::Type{T}, ::T) where T<:Tuple at essentials.jl:315
  convert(::Type{var"#s4"} where var"#s4"<:Tuple, ::ChainRulesCore.Tangent{var"#s3", var"#s2"} where {var"#s3", var"#s2"<:Tuple}) at C:\Users\Zac\.julia\packages\ChainRulesCore\qbmEe\src\differentials\composite.jl:68     
  ...
Stacktrace:
 [1] _tuple_error(T::Type, x::NTuple{256, Int64})
   @ Base .\essentials.jl:313
 [2] convert(#unused#::Type{NTuple{288, Int64}}, x::NTuple{256, Int64})
   @ Base .\essentials.jl:320
 [3] CUDA.CUTENSOR.cutensorContractionDescriptor_t(fields::NTuple{256, Int64})
   @ CUDA.CUTENSOR C:\Users\Zac\.julia\packages\CUDA\3VnCC\lib\cutensor\libcutensor_common.jl:132
 [4] contract!(α::Bool, A::CuArray{Float64, 5}, CA::Symbol, B::CuArray{Float64, 4}, CB::Symbol, β::Bool, C::CuArray{Float64, 5}, oindA::Tuple{Int64, Int64, Int64}, cindA::Tuple{Int64, Int64}, oindB::Tuple{Int64, Int64}, cindB::Tuple{Int64, Int64}, indCinoAB::NTuple{5, Int64}, syms::Tuple{Symbol, Symbol, Symbol})
   @ TensorOperations C:\Users\Zac\.julia\packages\TensorOperations\nL308\src\implementation\cuarray.jl:202
 [5] contract!(α::Bool, A::CuArray{Float64, 5}, CA::Symbol, B::CuArray{Float64, 4}, CB::Symbol, β::Bool, C::CuArray{Float64, 5}, oindA::Tuple{Int64, Int64, Int64}, cindA::Tuple{Int64, Int64}, oindB::Tuple{Int64, Int64}, cindB::Tuple{Int64, Int64}, indleft::NTuple{5, Int64}, indright::Tuple{}, syms::Tuple{Symbol, Symbol, Symbol})
   @ TensorOperations C:\Users\Zac\.julia\packages\TensorOperations\nL308\src\implementation\stridedarray.jl:89
 [6] top-level scope
   @ none:1

Any suggestions?

Jutho commented 3 years ago

Which version of CUDA.jl is this, I've never seen this before, I guess something has changed in CUTENSOR or the way it is wrapped by CUDA.jl

Jutho commented 3 years ago

Maybe @maleadt has some time to chime in?

maleadt commented 3 years ago

The CUTENSOR handle has changed in size, so this is invalid now: https://github.com/Jutho/TensorOperations.jl/blob/6e9fa69244210565a46edb473f18ddc3174beb5a/src/implementation/cuarray.jl#L202

Turns out the handles didn't need to be zero-initialized anyway, so I just removed that: https://github.com/JuliaGPU/CUDA.jl/commit/bdab0d5363d31ee22a6532a6a806b32ad614110d

Jutho commented 3 years ago

That's quick! Thanks, this is very helpful. I will change my custom wrappers based on the way it is done in your recent implementation. Do I have to be worried about backwards compatibility, or will this also work on older versions of CUDA.jl?

maleadt commented 3 years ago

Nothing has changed on the side of the C wrappers themselves, only the size of the handle there has changed (in https://github.com/JuliaGPU/CUDA.jl/commit/ab741155990c98ff628bb184cc7af6b73a211089#diff-0cb7d3454209a70f1e260802e7e8c6a979f4b721cf5f9585cbaac8b92256a6e5, for CUTENSOR 1.3), so this should be backwards compatible with your code that's calling those wrappers directly.

Jutho commented 3 years ago

Ok, this should be fixed in TensorOperations.jl 3.2.1 (coincidentally the same as the current CUDA version), which is scheduled to be merged in the next round. Feel free to reopen if the problem is not solved completely or resurfaces.

zjwegert commented 3 years ago

Great! Thank you