make CuDNN finders respect library major version (pytorch#5399)

Author: soumith

aten/cmake/FindCuDNN.cmake

@@ -15,13 +15,21 @@ include(FindPackageHandleStandardArgs)
 
 set(CUDNN_ROOT_DIR "" CACHE PATH "Folder contains NVIDIA cuDNN")
 
-find_path(CUDNN_INCLUDE_DIR cudnn.h
+if($ENV{CUDNN_INCLUDE_DIR})
+  SET(CUDNN_INCLUDE_DIR $ENV{CUDNN_INCLUDE_DIR})
+else($ENV{CUDNN_INCLUDE_DIR})
+  find_path(CUDNN_INCLUDE_DIR cudnn.h
     HINTS ${CUDNN_ROOT_DIR} ${CUDA_TOOLKIT_ROOT_DIR}
     PATH_SUFFIXES cuda/include include)
+endif($ENV{CUDNN_INCLUDE_DIR})
 
-find_library(CUDNN_LIBRARY cudnn
+if($ENV{CUDNN_LIBRARY})
+  SET(CUDNN_LIBRARY $ENV{CUDNN_LIBRARY})
+else($ENV{CUDNN_LIBRARY})
+  find_library(CUDNN_LIBRARY cudnn
     HINTS ${CUDNN_LIB_DIR} ${CUDNN_ROOT_DIR} ${CUDA_TOOLKIT_ROOT_DIR}
     PATH_SUFFIXES lib lib64 cuda/lib cuda/lib64 lib/x64)
+endif($ENV{CUDNN_LIBRARY})
 
 find_package_handle_standard_args(
     CUDNN DEFAULT_MSG CUDNN_INCLUDE_DIR CUDNN_LIBRARY)

setup.py

@@ -709,8 +709,7 @@ def run(self):
         "torch/csrc/cuda/python_nccl.cpp",
     ]
 if WITH_CUDNN:
-    main_libraries += ['cudnn']
-    library_dirs.insert(0, CUDNN_LIB_DIR)
+    main_libraries += [CUDNN_LIBRARY]
     # NOTE: these are at the front, in case there's another cuDNN in CUDA path
     include_dirs.insert(0, CUDNN_INCLUDE_DIR)
     if not IS_WINDOWS:

tools/setup_helpers/cudnn.py

@@ -35,6 +35,22 @@
     if IS_CONDA:
         lib_paths.append(os.path.join(CONDA_DIR, 'lib'))
         include_paths.append(os.path.join(CONDA_DIR, 'include'))
+    for path in include_paths:
+        if path is None or not os.path.exists(path):
+            continue
+        include_file_path = os.path.join(path, 'cudnn.h')
+        if os.path.exists(include_file_path):
+            CUDNN_INCLUDE_DIR = path
+            CUDNN_INCLUDE_VERSION = -1
+            with open(include_file_path) as f:
+                for line in f:
+                    if "#define CUDNN_MAJOR" in line:
+                        CUDNN_INCLUDE_VERSION = int(line.split()[-1])
+                        break
+            if CUDNN_INCLUDE_VERSION == -1:
+                raise AssertionError("Could not find #define CUDNN_MAJOR in " + include_file_path)
+            break
+
     for path in lib_paths:
         if path is None or not os.path.exists(path):
             continue
@@ -45,18 +61,11 @@
                 CUDNN_LIB_DIR = path
                 break
         else:
-            libraries = sorted(glob.glob(os.path.join(path, 'libcudnn*')))
+            libraries = sorted(glob.glob(os.path.join(path, 'libcudnn*' + str(CUDNN_INCLUDE_VERSION) + "*")))
             if libraries:
                 CUDNN_LIBRARY = libraries[0]
                 CUDNN_LIB_DIR = path
                 break
-    for path in include_paths:
-        if path is None or not os.path.exists(path):
-            continue
-        if os.path.exists((os.path.join(path, 'cudnn.h'))):
-            CUDNN_INCLUDE_DIR = path
-            break
-
     # Specifying the library directly will overwrite the lib directory
     library = os.getenv('CUDNN_LIBRARY')
     if library is not None and os.path.exists(library):