Update to CUDNN 8.8.1 for CUDA 12 compatibility. (#1792)

Also re-instate init error messages on subpackages when CUDA is functional.

Author: maleadt

.buildkite/pipeline.yml

@@ -57,7 +57,7 @@ steps:
                 - examples
         agents:
           queue: "juliagpu"
-          cuda: "11.0"
+          cuda: "*"
         if: build.message !~ /\[skip tests\]/ && !build.pull_request.draft
         timeout_in_minutes: 120
         matrix:
@@ -95,7 +95,7 @@ steps:
                 - examples
         agents:
           queue: "juliagpu"
-          cuda: "11.0"
+          cuda: "*"
         if: build.message !~ /\[skip tests\]/ && !build.pull_request.draft
         timeout_in_minutes: 120
         commands: |
@@ -117,10 +117,46 @@ steps:
             println("+++ :julia: Running tests")
             Pkg.test()'
 
+      - label: "{{matrix}} on CUDA 12"
+        matrix:
+            - "cuDNN"
+        plugins:
+          - JuliaCI/julia#v1:
+              version: 1.6
+          - JuliaCI/julia-coverage#v1:
+              codecov: true
+              dirs:
+                - src
+                - lib
+                - examples
+        agents:
+          queue: "juliagpu"
+          cuda: "*"
+        if: build.message !~ /\[skip tests\]/ && !build.pull_request.draft
+        timeout_in_minutes: 120
+        commands: |
+          julia -e '
+            using Pkg
+
+            println("--- :julia: Instantiating project")
+            Pkg.activate(joinpath(pwd(), "lib", lowercase("{{matrix}}")))
+            if "{{matrix}}" == "cuTensorNet"
+              # HACK: cuTensorNet depends on a development version of cuTENSOR
+              Pkg.develop(path=joinpath(pwd(), "lib", "cutensor"))
+            end
+            Pkg.develop(path=pwd())
+            Pkg.instantiate()
+
+            using CUDA
+            CUDA.set_runtime_version!(v"12.0")
+
+            println("+++ :julia: Running tests")
+            Pkg.test()'
+
   - group: ":telescope: Downstream"
     depends_on: "cuda"
     steps:
-      - label: "NNlibCUDA.jl"
+      - label: "NNlibCUDA.jl on CUDA 12"
         plugins:
           - JuliaCI/julia#v1:
               version: 1.6
@@ -136,18 +172,20 @@ steps:
 
             println("--- :julia: Instantiating project")
             Pkg.develop(path=pwd())
+            Pkg.develop(path=joinpath(pwd(), "lib", "cudnn"))
             Pkg.add(name="NNlibCUDA", rev="master")
             Pkg.instantiate()
 
+            using CUDA
+            CUDA.set_runtime_version!(v"12.0")
+
             println("+++ :julia: Running tests")
             Pkg.test("NNlibCUDA"; coverage=true)'
         agents:
           queue: "juliagpu"
           cuda: "*"
         if: build.message !~ /\[skip tests\]/ && !build.pull_request.draft
         timeout_in_minutes: 60
-        soft_fail:
-          - exit_status: 1
 
   - group: ":eyes: Special"
     depends_on: "cuda"
@@ -216,7 +254,7 @@ steps:
                 - examples
         agents:
           queue: "juliagpu"
-          cuda: "11.0"
+          cuda: "*"
         env:
           JULIA_CUDA_USE_COMPAT: 'false'  # NVIDIA bug #3418723: injection tools prevent probing libcuda
         commands: |

lib/cudnn/Project.toml

@@ -9,7 +9,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CUDNN_jll = "62b44479-cb7b-5706-934f-f13b2eb2e645"
 
 [compat]
-CUDNN_jll = "~8.6"
+CUDNN_jll = "~8.8"
 CUDA = "~4.0"
 CEnum = "0.2, 0.3, 0.4"
 julia = "1.6"

lib/cudnn/src/cuDNN.jl

@@ -24,7 +24,7 @@ function has_cudnn(show_reason::Bool=false)
     precompiling && return
 
     if !CUDNN_jll.is_available()
-        show_reason && error("cuDNN library not found")
+        show_reason && error("cuDNN JLL not available")
         return false
     end
     return true
@@ -161,8 +161,10 @@ function __init__()
     precompiling = ccall(:jl_generating_output, Cint, ()) != 0
     precompiling && return
 
+    CUDA.functional() || return
+
     if !CUDNN_jll.is_available()
-        #@error "cuDNN is not available for your platform ($(Base.BinaryPlatforms.triplet(CUDNN_jll.host_platform)))"
+        @error "cuDNN is not available for your platform ($(Base.BinaryPlatforms.triplet(CUDNN_jll.host_platform)))"
         return
     end
 

lib/cudnn/src/libcudnn.jl

@@ -164,6 +164,7 @@ const cudnnTensorTransformDescriptor_t = Ptr{cudnnTensorTransformStruct}
     CUDNN_DATA_BOOLEAN = 11
     CUDNN_DATA_FP8_E4M3 = 12
     CUDNN_DATA_FP8_E5M2 = 13
+    CUDNN_DATA_FAST_FLOAT_FOR_FP8 = 14
 end
 
 @cenum cudnnMathType_t::UInt32 begin
@@ -3554,6 +3555,38 @@ end
     @ccall libcudnn.cudnnCnnTrainVersionCheck()::cudnnStatus_t
 end
 
+const CUDNN_MAX_SM_MAJOR_NUMBER = 9
+
+const CUDNN_MAX_SM_MINOR_NUMBER = 0
+
+const CUDNN_SM_50 = 500
+
+const CUDNN_SM_52 = 520
+
+const CUDNN_SM_53 = 530
+
+const CUDNN_SM_60 = 600
+
+const CUDNN_SM_61 = 610
+
+const CUDNN_SM_62 = 620
+
+const CUDNN_SM_70 = 700
+
+const CUDNN_SM_72 = 720
+
+const CUDNN_SM_75 = 750
+
+const CUDNN_SM_80 = 800
+
+const CUDNN_SM_86 = 860
+
+const CUDNN_SM_87 = 870
+
+const CUDNN_SM_89 = 890
+
+const CUDNN_SM_90 = 900
+
 const CUDNN_DIM_MAX = 8
 
 const CUDNN_LRN_MIN_N = 1

lib/custatevec/src/cuStateVec.jl

@@ -13,7 +13,7 @@ export has_custatevec
 
 function has_custatevec(show_reason::Bool=false)
     if !cuQuantum_jll.is_available()
-        show_reason && error("cuStateVec library not found")
+        show_reason && error("cuStateVec JLL not available")
         return false
     end
     return true
@@ -109,8 +109,10 @@ function __init__()
     precompiling = ccall(:jl_generating_output, Cint, ()) != 0
     precompiling && return
 
+    CUDA.functional() || return
+
     if !cuQuantum_jll.is_available()
-        #@error "cuQuantum is not available for your platform ($(Base.BinaryPlatforms.triplet(cuQuantum_jll.host_platform)))"
+        @error "cuQuantum is not available for your platform ($(Base.BinaryPlatforms.triplet(cuQuantum_jll.host_platform)))"
         return
     end
 

lib/cutensor/src/cuTENSOR.jl

@@ -14,7 +14,7 @@ export has_cutensor
 
 function has_cutensor(show_reason::Bool=false)
     if !CUTENSOR_jll.is_available()
-        show_reason && error("cuTENSOR library not found")
+        show_reason && error("cuTENSOR JLL not available")
         return false
     end
     return true
@@ -93,8 +93,10 @@ function __init__()
     precompiling = ccall(:jl_generating_output, Cint, ()) != 0
     precompiling && return
 
+    CUDA.functional() || return
+
     if !CUTENSOR_jll.is_available()
-        #@error "cuTENSOR is not available for your platform ($(Base.BinaryPlatforms.triplet(CUTENSOR_jll.host_platform)))"
+        @error "cuTENSOR is not available for your platform ($(Base.BinaryPlatforms.triplet(CUTENSOR_jll.host_platform)))"
         return
     end
 

lib/cutensornet/src/cuTensorNet.jl

@@ -17,7 +17,7 @@ export has_cutensornet
 
 function has_cutensornet(show_reason::Bool=false)
     if !cuQuantum_jll.is_available()
-        show_reason && error("cuTensorNet library not found")
+        show_reason && error("cuTensorNet JLL not available")
         return false
     end
     return true
@@ -110,8 +110,10 @@ function __init__()
     precompiling = ccall(:jl_generating_output, Cint, ()) != 0
     precompiling && return
 
+    CUDA.functional() || return
+
     if !cuQuantum_jll.is_available()
-        #@error "cuQuantum is not available for your platform ($(Base.BinaryPlatforms.triplet(cuQuantum_jll.host_platform)))"
+        @error "cuQuantum is not available for your platform ($(Base.BinaryPlatforms.triplet(cuQuantum_jll.host_platform)))"
         return
     end
 

res/wrap/Manifest.toml

@@ -40,10 +40,10 @@ uuid = "4f82f1eb-248c-5f56-a42e-99106d144614"
 version = "12.1.0+0"
 
 [[CUDNN_jll]]
-deps = ["Artifacts", "CUDA_Runtime_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "Pkg", "TOML"]
-git-tree-sha1 = "57011df4fce448828165e566af9befa2ea94350a"
+deps = ["Artifacts", "CUDA_Runtime_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"]
+git-tree-sha1 = "2918fbffb50e3b7a0b9127617587afa76d4276e8"
 uuid = "62b44479-cb7b-5706-934f-f13b2eb2e645"
-version = "8.6.0+3"
+version = "8.8.1+0"
 
 [[CUTENSOR_jll]]
 deps = ["Artifacts", "CUDA_Runtime_jll", "CompilerSupportLibraries_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "Pkg", "TOML"]

res/wrap/libcudnn_prologue.jl

@@ -1,4 +1,4 @@
-# CUDNN uses CUDA runtime objects, which are compatible with our driver usage
+# cuDNN uses CUDA runtime objects, which are compatible with our driver usage
 const cudaStream_t = CUstream
 
 # outlined functionality to avoid GC frame allocation