pytorch/cmake/Modules_CUDA_fix/FindCUDNN.cmake
Hong Xu 92750acb88 Move the detection of cuDNN to FindCUDNN.cmake (#24938)
Summary:
Currently they sit together with other code in cuda.cmake. This commit is the first step toward cleaning up cuDNN detection in our build system.

Another attempt to https://github.com/pytorch/pytorch/issues/24293,  which breaks manywheels build because it does not handle `USE_STATIC_CUDNN` properly.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/24938

Differential Revision: D17070920

Pulled By: ezyang

fbshipit-source-id: a4d017a3505c102d9c435a73ae62332e4336c52e
2019-08-27 06:51:52 -07:00

50 lines
1.8 KiB
CMake

# 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)