Allow cuDNN in non-CUDA non-system dir (apache#7608)

cuDNN is not a builtin library of the CUDA toolkit package.
The user can install it in the CUDA directory, the system
directory, or anywhere else. This patch relax the restriction
of locating cuDNN in the CUDA directory. This is helpfull
when trying out different versions of cuDNN.

cmake/config.cmake

@@ -174,8 +174,11 @@ set(USE_FLATBUFFERS_PATH none)
 # - /path/to/edgetpu: use specific path to edgetpu library
 set(USE_EDGETPU OFF)
 
-# Whether use CuDNN
-set(USE_CUDNN ON)
+# Possible values:
+# - ON: enable cuDNN with cmake's auto search in CUDA directory
+# - OFF: disable cuDNN
+# - /path/to/cudnn: use specific path to cuDNN path
+set(USE_CUDNN OFF)
 
 # Whether use cuBLAS
 set(USE_CUBLAS ON)

cmake/modules/CUDA.cmake

@@ -16,12 +16,12 @@
 # under the License.
 
 # CUDA Module
-find_cuda(${USE_CUDA})
+find_cuda(${USE_CUDA} ${USE_CUDNN})
 
 if(CUDA_FOUND)
   # always set the includedir when cuda is available
   # avoid global retrigger of cmake
-	include_directories(SYSTEM ${CUDA_INCLUDE_DIRS})
+  include_directories(SYSTEM ${CUDA_INCLUDE_DIRS})
 endif(CUDA_FOUND)
 
 if(USE_CUDA)
@@ -40,6 +40,7 @@ if(USE_CUDA)
 
   if(USE_CUDNN)
     message(STATUS "Build with cuDNN support")
+    include_directories(SYSTEM ${CUDA_CUDNN_INCLUDE_DIRS})
     file(GLOB CONTRIB_CUDNN_SRCS src/runtime/contrib/cudnn/*.cc)
     list(APPEND RUNTIME_SRCS ${CONTRIB_CUDNN_SRCS})
     list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUDNN_LIBRARY})

cmake/utils/FindCUDA.cmake

@@ -19,10 +19,12 @@
 # Enhanced version of find CUDA.
 #
 # Usage:
-#   find_cuda(${USE_CUDA})
+#   find_cuda(${USE_CUDA} ${USE_CUDNN})
 #
 # - When USE_CUDA=ON, use auto search
 # - When USE_CUDA=/path/to/cuda-path, use the cuda path
+# - When USE_CUDNN=ON, use auto search
+# - When USE_CUDNN=/path/to/cudnn-path, use the cudnn path
 #
 # Provide variables:
 #
@@ -32,10 +34,11 @@
 # - CUDA_CUDA_LIBRARY
 # - CUDA_CUDART_LIBRARY
 # - CUDA_NVRTC_LIBRARY
+# - CUDA_CUDNN_INCLUDE_DIRS
 # - CUDA_CUDNN_LIBRARY
 # - CUDA_CUBLAS_LIBRARY
 #
-macro(find_cuda use_cuda)
+macro(find_cuda use_cuda use_cudnn)
   set(__use_cuda ${use_cuda})
   if(${__use_cuda} MATCHES ${IS_TRUE_PATTERN})
     find_package(CUDA QUIET)
@@ -64,9 +67,6 @@ macro(find_cuda use_cuda)
       find_library(CUDA_NVRTC_LIBRARY nvrtc
         ${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
         ${CUDA_TOOLKIT_ROOT_DIR}/lib/Win32)
-      find_library(CUDA_CUDNN_LIBRARY cudnn
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib/Win32)
       find_library(CUDA_CUBLAS_LIBRARY cublas
         ${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
         ${CUDA_TOOLKIT_ROOT_DIR}/lib/Win32)
@@ -85,12 +85,6 @@ macro(find_cuda use_cuda)
         PATHS ${CUDA_TOOLKIT_ROOT_DIR}
         PATH_SUFFIXES lib lib64 targets/x86_64-linux/lib targets/x86_64-linux/lib/stubs lib64/stubs lib/x86_64-linux-gnu
         NO_DEFAULT_PATH)
-      find_library(CUDA_CUDNN_LIBRARY cudnn
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib64
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib
-        NO_DEFAULT_PATH)
-      # search default path if cannot find cudnn in non-default
-      find_library(CUDA_CUDNN_LIBRARY cudnn)
       find_library(CUDA_CUBLAS_LIBRARY cublas
         ${CUDA_TOOLKIT_ROOT_DIR}/lib64
         ${CUDA_TOOLKIT_ROOT_DIR}/lib
@@ -102,10 +96,38 @@ macro(find_cuda use_cuda)
         ${CUDA_TOOLKIT_ROOT_DIR}/lib
         NO_DEFAULT_PATH)
     endif(MSVC)
+
+    # find cuDNN
+    set(__use_cudnn ${use_cudnn})
+    if(${__use_cudnn} MATCHES ${IS_TRUE_PATTERN})
+      set(CUDA_CUDNN_INCLUDE_DIRS ${CUDA_INCLUDE_DIRS})
+      if(MSVC)
+        find_library(CUDA_CUDNN_LIBRARY cudnn
+          ${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
+          ${CUDA_TOOLKIT_ROOT_DIR}/lib/Win32)
+      else(MSVC)
+        find_library(CUDA_CUDNN_LIBRARY cudnn
+          ${CUDA_TOOLKIT_ROOT_DIR}/lib64
+          ${CUDA_TOOLKIT_ROOT_DIR}/lib
+          NO_DEFAULT_PATH)
+        # search default path if cannot find cudnn in non-default
+        find_library(CUDA_CUDNN_LIBRARY cudnn)
+      endif(MSVC)
+    elseif(IS_DIRECTORY ${__use_cudnn})
+      # cuDNN doesn't necessarily live in the CUDA dir
+      set(CUDA_CUDNN_ROOT_DIR ${__use_cudnn})
+      set(CUDA_CUDNN_INCLUDE_DIRS ${CUDA_CUDNN_ROOT_DIR}/include)
+      find_library(CUDA_CUDNN_LIBRARY cudnn
+        ${CUDA_CUDNN_ROOT_DIR}/lib64
+        ${CUDA_CUDNN_ROOT_DIR}/lib
+        NO_DEFAULT_PATH)
+    endif()
+
     message(STATUS "Found CUDA_TOOLKIT_ROOT_DIR=" ${CUDA_TOOLKIT_ROOT_DIR})
     message(STATUS "Found CUDA_CUDA_LIBRARY=" ${CUDA_CUDA_LIBRARY})
     message(STATUS "Found CUDA_CUDART_LIBRARY=" ${CUDA_CUDART_LIBRARY})
     message(STATUS "Found CUDA_NVRTC_LIBRARY=" ${CUDA_NVRTC_LIBRARY})
+    message(STATUS "Found CUDA_CUDNN_INCLUDE_DIRS=" ${CUDA_CUDNN_INCLUDE_DIRS})
     message(STATUS "Found CUDA_CUDNN_LIBRARY=" ${CUDA_CUDNN_LIBRARY})
     message(STATUS "Found CUDA_CUBLAS_LIBRARY=" ${CUDA_CUBLAS_LIBRARY})
     message(STATUS "Found CUDA_CUBLASLT_LIBRARY=" ${CUDA_CUBLASLT_LIBRARY})