[BE] Don't copy CuDNN libs twice (#115872)

malfet · pytorchmergebot · commit 4ea7430ffb0b · 2023-12-15T09:47:14.000Z
- It was installed twice : once in `/usr/local/cuda/lib64` folder and 2nd time in `/usr/lib64` - And don't install CuDNN headers thrice, only in `/usr/local/cuda/includa` - Error on unknown CUDA version - Modify bazel builds to look for cudnn in `/usr/local/cuda` folder Pull Request resolved: #115872 Approved by: https://github.com/huydhn
diff --git a/.ci/docker/common/install_cudnn.sh b/.ci/docker/common/install_cudnn.sh
@@ -2,26 +2,23 @@
 
 if [[ ${CUDNN_VERSION} == 8 ]]; then
     # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
-    mkdir tmp_cudnn && cd tmp_cudnn
-    CUDNN_NAME="cudnn-linux-x86_64-8.3.2.44_cuda11.5-archive"
+    mkdir tmp_cudnn
+    pushd tmp_cudnn
     if [[ ${CUDA_VERSION:0:4} == "12.1" ]]; then
         CUDNN_NAME="cudnn-linux-x86_64-8.9.2.26_cuda12-archive"
         curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/${CUDNN_NAME}.tar.xz
     elif [[ ${CUDA_VERSION:0:4} == "11.8" ]]; then
         CUDNN_NAME="cudnn-linux-x86_64-8.7.0.84_cuda11-archive"
         curl --retry 3 -OLs https://developer.download.nvidia.com/compute/redist/cudnn/v8.7.0/local_installers/11.8/${CUDNN_NAME}.tar.xz
     else
-        curl --retry 3 -OLs https://developer.download.nvidia.com/compute/redist/cudnn/v8.3.2/local_installers/11.5/${CUDNN_NAME}.tar.xz
+        print "Unsupported CUDA version ${CUDA_VERSION}"
+        exit 1
     fi
 
     tar xf ${CUDNN_NAME}.tar.xz
-    cp -a ${CUDNN_NAME}/include/* /usr/include/
     cp -a ${CUDNN_NAME}/include/* /usr/local/cuda/include/
-    cp -a ${CUDNN_NAME}/include/* /usr/include/x86_64-linux-gnu/
-
     cp -a ${CUDNN_NAME}/lib/* /usr/local/cuda/lib64/
-    cp -a ${CUDNN_NAME}/lib/* /usr/lib/x86_64-linux-gnu/
-    cd ..
+    popd
     rm -rf tmp_cudnn
     ldconfig
 fi
diff --git a/WORKSPACE b/WORKSPACE
@@ -243,7 +243,7 @@ new_local_repository(
 new_local_repository(
     name = "cudnn",
     build_file = "@//third_party:cudnn.BUILD",
-    path = "/usr/",
+    path = "/usr/local/cuda",
 )
 
 local_repository(
diff --git a/third_party/cudnn.BUILD b/third_party/cudnn.BUILD
@@ -12,7 +12,7 @@ cc_library(
 
 cc_import(
     name = "cudnn_lib",
-    shared_library = "lib/x86_64-linux-gnu/libcudnn.so",
+    shared_library = "lib64/libcudnn.so",
     visibility = ["//visibility:private"],
 )
 
