From 6b8d0ac5316bc64091a4ffd30f298d1d9ba3ef93 Mon Sep 17 00:00:00 2001
From: Leonard Lausen <lausen@amazon.com>
Date: Mon, 9 Dec 2019 06:45:13 +0000
Subject: [PATCH] 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                 | 14 ++++++--------
 cmake/BuildTVM.cmake           | 12 ------------
 cmake/FirstClassLangCuda.cmake | 28 ----------------------------
 cmake/Modules/FindCUDNN.cmake  | 33 +++++++++++++++++++++++++++++++++
 4 files changed, 39 insertions(+), 48 deletions(-)
 create mode 100644 cmake/Modules/FindCUDNN.cmake

diff --git a/CMakeLists.txt b/CMakeLists.txt
index aec5f21131d1..39dbd3dca2f0 100644
--- a/CMakeLists.txt
+++ b/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)
diff --git a/cmake/BuildTVM.cmake b/cmake/BuildTVM.cmake
index 7fcf706ffb7c..4bb749552f01 100644
--- a/cmake/BuildTVM.cmake
+++ b/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)
 
diff --git a/cmake/FirstClassLangCuda.cmake b/cmake/FirstClassLangCuda.cmake
index 8d79c2b63ad9..0eca1aff78d4 100644
--- a/cmake/FirstClassLangCuda.cmake
+++ b/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)
diff --git a/cmake/Modules/FindCUDNN.cmake b/cmake/Modules/FindCUDNN.cmake
new file mode 100644
index 000000000000..a8fda5c87d9a
--- /dev/null
+++ b/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)