Revert D16914345: [pytorch][PR] Move the detection of cuDNN to FindCUDNN.cmake

Differential Revision:
D16914345

Original commit changeset: fd261478c01d

fbshipit-source-id: b933ad7ed49028ab9ac6976c3ae768132dc9bacb

CMakeLists.txt

@@ -109,9 +109,6 @@ option(CAFFE2_STATIC_LINK_CUDA "Statically link CUDA libraries" OFF)
 cmake_dependent_option(
     USE_CUDNN "Use cuDNN" ON
     "USE_CUDA" OFF)
-cmake_dependent_option(
-    USE_STATIC_CUDNN "Use cuDNN static libraries" OFF
-    "USE_CUDNN" OFF)
 option(USE_FBGEMM "Use FBGEMM (quantized 8-bit server operators)" ON)
 option(USE_FFMPEG "Use ffmpeg" OFF)
 option(USE_GFLAGS "Use GFLAGS" OFF)

cmake/Dependencies.cmake

@@ -1270,7 +1270,7 @@ if (NOT INTERN_BUILD_MOBILE)
     MESSAGE(STATUS "CuDNN not found. Compiling without CuDNN support")
     set(AT_CUDNN_ENABLED 0)
   ELSE()
-    include_directories(SYSTEM ${CUDNN_INCLUDE_PATH})
+    include_directories(SYSTEM ${CUDNN_INCLUDE_DIRS})
     set(AT_CUDNN_ENABLED 1)
   ENDIF()
 

cmake/Modules_CUDA_fix/FindCUDNN.cmake

@@ -1,50 +0,0 @@
-# Find the CUDNN libraries
-#
-# The following variables are optionally searched for defaults
-#  CUDNN_ROOT: Base directory where CUDNN is found
-#  CUDNN_INCLUDE_DIR: Directory where CUDNN header is searched for
-#  CUDNN_LIBRARY: Directory where CUDNN library is searched for
-#  CUDNN_STATIC: Are we looking for a static library? (default: no)
-#
-# The following are set after configuration is done:
-#  CUDNN_FOUND
-#  CUDNN_INCLUDE_PATH
-#  CUDNN_LIBRARY_PATH
-#
-
-include(FindPackageHandleStandardArgs)
-
-set(CUDNN_ROOT $ENV{CUDNN_ROOT_DIR} CACHE PATH "Folder containing NVIDIA cuDNN")
-if (DEFINED $ENV{CUDNN_ROOT_DIR})
-  message(WARNING "CUDNN_ROOT_DIR is deprecated. Please set CUDNN_ROOT instead.")
-endif()
-list(APPEND CUDNN_ROOT $ENV{CUDNN_ROOT_DIR} ${CUDA_TOOLKIT_ROOT_DIR})
-
-# Compatible layer for CMake <3.12. CUDNN_ROOT will be accounted in for searching paths and libraries for CMake >=3.12.
-list(APPEND CMAKE_PREFIX_PATH ${CUDNN_ROOT})
-
-set(CUDNN_INCLUDE_DIR $ENV{CUDNN_INCLUDE_DIR} CACHE PATH "Folder containing NVIDIA cuDNN header files")
-
-find_path(CUDNN_INCLUDE_PATH cudnn.h
-  HINTS ${CUDNN_INCLUDE_DIR}
-  PATH_SUFFIXES cuda/include include)
-
-option(CUDNN_STATIC "Look for static CUDNN" OFF)
-if (CUDNN_STATIC)
-  set(CUDNN_LIBNAME "libcudnn_static.a")
-else()
-  set(CUDNN_LIBNAME "cudnn")
-endif()
-
-set(CUDNN_LIBRARY $ENV{CUDNN_LIBRARY} CACHE PATH "Path to the cudnn library file (e.g., libcudnn.so)")
-if (CUDNN_LIBRARY MATCHES ".*cudnn_static.a" AND NOT CUDNN_STATIC)
-  message(WARNING "CUDNN_LIBRARY points to a static library (${CUDNN_LIBRARY}) but CUDNN_STATIC is OFF.")
-endif()
-
-find_library(CUDNN_LIBRARY_PATH ${CUDNN_LIBNAME}
-  PATHS ${CUDNN_LIBRARY}
-  PATH_SUFFIXES lib lib64 cuda/lib cuda/lib64 lib/x64)
-
-find_package_handle_standard_args(CUDNN DEFAULT_MSG CUDNN_LIBRARY_PATH CUDNN_INCLUDE_PATH)
-
-mark_as_advanced(CUDNN_ROOT CUDNN_INCLUDE_DIR CUDNN_LIBRARY)

cmake/public/cuda.cmake

@@ -93,14 +93,40 @@ if(CUDA_FOUND)
 endif()
 
 # Find cuDNN.
-if(CAFFE2_STATIC_LINK_CUDA OR USE_STATIC_CUDNN)
-  set(CUDNN_STATIC ON)
+if(CAFFE2_STATIC_LINK_CUDA)
+  SET(CUDNN_LIBNAME "libcudnn_static.a")
 else()
-  set(CUDNN_STATIC OFF)
+  SET(CUDNN_LIBNAME "cudnn")
+endif()
+include(FindPackageHandleStandardArgs)
+
+if(DEFINED ENV{CUDNN_ROOT_DIR})
+  set(CUDNN_ROOT_DIR $ENV{CUDNN_ROOT_DIR} CACHE PATH "Folder contains NVIDIA cuDNN")
+else()
+  set(CUDNN_ROOT_DIR "" CACHE PATH "Folder contains NVIDIA cuDNN")
 endif()
 
-find_package(CUDNN)
+if(DEFINED ENV{CUDNN_INCLUDE_DIR})
+  set(CUDNN_INCLUDE_DIR $ENV{CUDNN_INCLUDE_DIR})
+else()
+  find_path(CUDNN_INCLUDE_DIR cudnn.h
+    HINTS ${CUDNN_ROOT_DIR} ${CUDA_TOOLKIT_ROOT_DIR}
+    PATH_SUFFIXES cuda/include include)
+endif()
 
+if(DEFINED ENV{CUDNN_LIBRARY})
+  set(CUDNN_LIBRARY $ENV{CUDNN_LIBRARY})
+  if (CUDNN_LIBRARY MATCHES ".*cudnn_static.a")
+    SET(CUDNN_STATIC_LINKAGE ON)
+  endif()
+else()
+  find_library(CUDNN_LIBRARY ${CUDNN_LIBNAME}
+    HINTS ${CUDNN_ROOT_DIR} ${CUDA_TOOLKIT_ROOT_DIR}
+    PATH_SUFFIXES lib lib64 cuda/lib cuda/lib64 lib/x64)
+endif()
+
+find_package_handle_standard_args(
+    CUDNN DEFAULT_MSG CUDNN_INCLUDE_DIR CUDNN_LIBRARY)
 if(NOT CUDNN_FOUND)
   message(WARNING
     "Caffe2: Cannot find cuDNN library. Turning the option off")
@@ -129,7 +155,7 @@ endif()
 # ---[ Extract versions
 if(CAFFE2_USE_CUDNN)
   # Get cuDNN version
-  file(READ ${CUDNN_INCLUDE_PATH}/cudnn.h CUDNN_HEADER_CONTENTS)
+  file(READ ${CUDNN_INCLUDE_DIR}/cudnn.h CUDNN_HEADER_CONTENTS)
   string(REGEX MATCH "define CUDNN_MAJOR * +([0-9]+)"
                CUDNN_VERSION_MAJOR "${CUDNN_HEADER_CONTENTS}")
   string(REGEX REPLACE "define CUDNN_MAJOR * +([0-9]+)" "\\1"
@@ -149,7 +175,7 @@ if(CAFFE2_USE_CUDNN)
     set(CUDNN_VERSION
         "${CUDNN_VERSION_MAJOR}.${CUDNN_VERSION_MINOR}.${CUDNN_VERSION_PATCH}")
   endif()
-  message(STATUS "Found cuDNN: v${CUDNN_VERSION}  (include: ${CUDNN_INCLUDE_PATH}, library: ${CUDNN_LIBRARY_PATH})")
+  message(STATUS "Found cuDNN: v${CUDNN_VERSION}  (include: ${CUDNN_INCLUDE_DIR}, library: ${CUDNN_LIBRARY})")
   if(CUDNN_VERSION VERSION_LESS "7.0.0")
     message(FATAL_ERROR "PyTorch requires cuDNN 7 and above.")
   endif()
@@ -206,10 +232,10 @@ if(CAFFE2_USE_CUDNN)
   add_library(caffe2::cudnn UNKNOWN IMPORTED)
   set_property(
       TARGET caffe2::cudnn PROPERTY IMPORTED_LOCATION
-      ${CUDNN_LIBRARY_PATH})
+      ${CUDNN_LIBRARY})
   set_property(
       TARGET caffe2::cudnn PROPERTY INTERFACE_INCLUDE_DIRECTORIES
-      ${CUDNN_INCLUDE_PATH})
+      ${CUDNN_INCLUDE_DIR})
 endif()
 
 # curand

torch/CMakeLists.txt

@@ -193,8 +193,8 @@ if (USE_CUDNN)
     # Basically, this is the case where $CUDA_HOME/lib64 has an old or
     # incompatible libcudnn.so, which we can inadvertently link to if
     # we're not careful.
-    list(INSERT 0 TORCH_PYTHON_LINK_LIBRARIES ${CUDNN_LIBRARY_PATH})
-    list(INSERT 0 TORCH_PYTHON_INCLUDE_DIRECTORIES ${CUDNN_INCLUDE_PATH})
+    list(INSERT 0 TORCH_PYTHON_LINK_LIBRARIES ${CUDNN_LIBRARY})
+    list(INSERT 0 TORCH_PYTHON_INCLUDE_DIRECTORIES ${CUDNN_INCLUDE_DIR})
 endif()
 
 if (USE_NUMPY)