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

- 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: https://github.com/pytorch/pytorch/pull/115872
Approved by: https://github.com/huydhn

.ci/docker/common/install_cudnn.sh

@@ -2,8 +2,8 @@
 
 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
@@ -11,17 +11,14 @@ if [[ ${CUDNN_VERSION} == 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

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(

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"],
 )