Fix CUDNN detection for CMake build

- Use FindCUDNN.cmake instead of the previous macro
- Previous macro did not detect CUDNN correctly on my system (Deep Learning AMI
  Ubuntu 18.04)
- We now explicitly fail if the user does not provide CUDNN and hasn't manually
  set -DUSE_CUDNN=0

CMakeLists.txt

@@ -502,14 +502,12 @@ add_subdirectory(${GTEST_ROOT})
 find_package(GTest REQUIRED)
 
 # cudnn detection
-if(USE_CUDNN AND USE_CUDA)
- detect_cuDNN()
- if(HAVE_CUDNN)
- add_definitions(-DUSE_CUDNN)
- include_directories(SYSTEM ${CUDNN_INCLUDE})
- list(APPEND mxnet_LINKER_LIBS ${CUDNN_LIBRARY})
- add_definitions(-DMSHADOW_USE_CUDNN=1)
- endif()
+if(USE_CUDNN)
+ find_package(CUDNN)
+ add_definitions(-DUSE_CUDNN)
+ include_directories(SYSTEM ${CUDNN_INCLUDE})
+ list(APPEND mxnet_LINKER_LIBS ${CUDNN_LIBRARY})
+ add_definitions(-DMSHADOW_USE_CUDNN=1)
 endif()
 
 if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/dmlc-core/cmake)

cmake/BuildTVM.cmake

@@ -98,18 +98,6 @@ set(USE_RANDOM OFF)
 # Whether use NNPack
 set(USE_NNPACK OFF)
 
-# Whether use CuDNN
-if(USE_CUDNN AND USE_CUDA)
- detect_cuDNN()
- if(HAVE_CUDNN)
- set(USE_CUDNN ON)
- else()
- set(USE_CUDNN OFF)
- endif()
-else()
- set(USE_CUDNN OFF)
-endif()
-
 # Whether use cuBLAS
 set(USE_CUBLAS OFF)
 

cmake/FirstClassLangCuda.cmake

@@ -23,34 +23,6 @@ if(USE_CXX14_IF_AVAILABLE)
  check_cxx_compiler_flag("-std=c++14" SUPPORT_CXX14)
 endif()
 
-################################################################################################
-# Short command for cuDNN detection. Believe it soon will be a part of CUDA toolkit distribution.
-# That's why not FindcuDNN.cmake file, but just the macro
-# Usage:
-# detect_cuDNN()
-function(detect_cuDNN)
- set(CUDNN_ROOT "" CACHE PATH "CUDNN root folder")
-
- find_path(CUDNN_INCLUDE cudnn.h
- PATHS ${CUDNN_ROOT} $ENV{CUDNN_ROOT}
- DOC "Path to cuDNN include directory." )
-
-
- find_library(CUDNN_LIBRARY NAMES libcudnn.so cudnn.lib # libcudnn_static.a
- PATHS ${CUDNN_ROOT} $ENV{CUDNN_ROOT} ${CUDNN_INCLUDE}
- PATH_SUFFIXES lib lib/x64
- DOC "Path to cuDNN library.")
-
- if(CUDNN_INCLUDE AND CUDNN_LIBRARY)
- set(HAVE_CUDNN TRUE PARENT_SCOPE)
- set(CUDNN_FOUND TRUE PARENT_SCOPE)
-
- mark_as_advanced(CUDNN_INCLUDE CUDNN_LIBRARY CUDNN_ROOT)
- message(STATUS "Found cuDNN (include: ${CUDNN_INCLUDE}, library: ${CUDNN_LIBRARY})")
- endif()
-endfunction()
-
-
 
 ################################################################################################
 # A function for automatic detection of GPUs installed (if autodetection is enabled)

cmake/Modules/FindCUDNN.cmake

@@ -0,0 +1,33 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+include(FindPackageHandleStandardArgs)
+
+set(CUDNN_ROOT "/usr/local/cuda/include" CACHE PATH "cuDNN root folder")
+
+find_path(CUDNN_INCLUDE cudnn.h
+ PATHS ${CUDNN_ROOT} $ENV{CUDNN_ROOT}
+ DOC "Path to cuDNN include directory." )
+
+find_library(CUDNN_LIBRARY NAMES libcudnn.so cudnn.lib # libcudnn_static.a
+ PATHS ${CUDNN_ROOT} $ENV{CUDNN_ROOT} ${CUDNN_INCLUDE}
+ PATH_SUFFIXES lib lib/x64 cuda/lib cuda/lib64 lib/x64
+ DOC "Path to cuDNN library.")
+
+find_package_handle_standard_args(CUDNN DEFAULT_MSG CUDNN_LIBRARY CUDNN_INCLUDE)
+
+mark_as_advanced(CUDNN_ROOT CUDNN_INCLUDE CUDNN_LIBRARY)