From 4ea7430ffb0b37e84108c70f90f402f5fd37262e Mon Sep 17 00:00:00 2001
From: Nikita Shulga <2453524+malfet@users.noreply.github.com>
Date: Fri, 15 Dec 2023 09:47:10 +0000
Subject: [PATCH] [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 | 13 +++++--------
 WORKSPACE                          |  2 +-
 third_party/cudnn.BUILD            |  2 +-
 3 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/.ci/docker/common/install_cudnn.sh b/.ci/docker/common/install_cudnn.sh
index a4dfa26f0ef..f654c9fee24 100644
--- a/.ci/docker/common/install_cudnn.sh
+++ b/.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
diff --git a/WORKSPACE b/WORKSPACE
index edd67aeb1fd..412ba2fbf04 100644
--- 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
index 03736508a68..82b4526c21e 100644
--- 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"],
 )