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/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..2a9cf3d4ca39 --- /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 "" 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)