mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
Enhancements
Add BLAS chooser Move cuDNN detection from Cuda -> FindCuDNN Refactor main C2 libs, should enable no-GPU build (untested)
This commit is contained in:
Simon Layton
@ -158,69 +158,6 @@ macro(caffe2_cuda_compile objlist_variable)
|
||||
set(${objlist_variable} ${cuda_objcs})
|
||||
endmacro()
|
||||
|
||||
################################################################################################
|
||||
# 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} ${CUDA_TOOLKIT_INCLUDE}
|
||||
DOC "Path to cuDNN include directory." )
|
||||
|
||||
# dynamic libs have different suffix in mac and linux
|
||||
if(APPLE)
|
||||
set(CUDNN_LIB_NAME "libcudnn.dylib")
|
||||
else()
|
||||
set(CUDNN_LIB_NAME "libcudnn.so")
|
||||
endif()
|
||||
|
||||
get_filename_component(__libpath_hist ${CUDA_CUDART_LIBRARY} PATH)
|
||||
find_library(CUDNN_LIBRARY NAMES ${CUDNN_LIB_NAME}
|
||||
PATHS ${CUDNN_ROOT} $ENV{CUDNN_ROOT} ${CUDNN_INCLUDE} ${__libpath_hist} ${__libpath_hist}/../lib
|
||||
DOC "Path to cuDNN library.")
|
||||
|
||||
if(CUDNN_INCLUDE AND CUDNN_LIBRARY)
|
||||
set(HAVE_CUDNN TRUE PARENT_SCOPE)
|
||||
set(CUDNN_FOUND TRUE PARENT_SCOPE)
|
||||
|
||||
file(READ ${CUDNN_INCLUDE}/cudnn.h CUDNN_VERSION_FILE_CONTENTS)
|
||||
|
||||
# cuDNN v3 and beyond
|
||||
string(REGEX MATCH "define CUDNN_MAJOR * +([0-9]+)"
|
||||
CUDNN_VERSION_MAJOR "${CUDNN_VERSION_FILE_CONTENTS}")
|
||||
string(REGEX REPLACE "define CUDNN_MAJOR * +([0-9]+)" "\\1"
|
||||
CUDNN_VERSION_MAJOR "${CUDNN_VERSION_MAJOR}")
|
||||
string(REGEX MATCH "define CUDNN_MINOR * +([0-9]+)"
|
||||
CUDNN_VERSION_MINOR "${CUDNN_VERSION_FILE_CONTENTS}")
|
||||
string(REGEX REPLACE "define CUDNN_MINOR * +([0-9]+)" "\\1"
|
||||
CUDNN_VERSION_MINOR "${CUDNN_VERSION_MINOR}")
|
||||
string(REGEX MATCH "define CUDNN_PATCHLEVEL * +([0-9]+)"
|
||||
CUDNN_VERSION_PATCH "${CUDNN_VERSION_FILE_CONTENTS}")
|
||||
string(REGEX REPLACE "define CUDNN_PATCHLEVEL * +([0-9]+)" "\\1"
|
||||
CUDNN_VERSION_PATCH "${CUDNN_VERSION_PATCH}")
|
||||
|
||||
if(NOT CUDNN_VERSION_MAJOR)
|
||||
set(CUDNN_VERSION "???")
|
||||
else()
|
||||
set(CUDNN_VERSION "${CUDNN_VERSION_MAJOR}.${CUDNN_VERSION_MINOR}.${CUDNN_VERSION_PATCH}")
|
||||
endif()
|
||||
|
||||
message(STATUS "Found cuDNN: ver. ${CUDNN_VERSION} found (include: ${CUDNN_INCLUDE}, library: ${CUDNN_LIBRARY})")
|
||||
|
||||
string(COMPARE LESS "${CUDNN_VERSION_MAJOR}" 3 cuDNNVersionIncompatible)
|
||||
if(cuDNNVersionIncompatible)
|
||||
message(FATAL_ERROR "cuDNN version >3 is required.")
|
||||
endif()
|
||||
|
||||
set(CUDNN_VERSION "${CUDNN_VERSION}" PARENT_SCOPE)
|
||||
mark_as_advanced(CUDNN_INCLUDE CUDNN_LIBRARY CUDNN_ROOT)
|
||||
|
||||
endif()
|
||||
endfunction()
|
||||
|
||||
################################################################################################
|
||||
### Non macro section
|
||||
################################################################################################
|
||||
@ -238,13 +175,6 @@ include_directories(SYSTEM ${CUDA_INCLUDE_DIRS})
|
||||
list(APPEND Caffe2_LINKER_LIBS ${CUDA_CUDART_LIBRARY}
|
||||
${CUDA_curand_LIBRARY} ${CUDA_CUBLAS_LIBRARIES})
|
||||
|
||||
# cudnn detection
|
||||
detect_cuDNN()
|
||||
if(HAVE_CUDNN)
|
||||
include_directories(SYSTEM ${CUDNN_INCLUDE})
|
||||
list(APPEND Caffe2_LINKER_LIBS ${CUDNN_LIBRARY})
|
||||
endif()
|
||||
|
||||
# setting nvcc arch flags
|
||||
caffe2_select_nvcc_arch_flags(NVCC_FLAGS_EXTRA)
|
||||
list(APPEND CUDA_NVCC_FLAGS ${NVCC_FLAGS_EXTRA})
|
||||
|
||||
@ -3,14 +3,28 @@ set(Caffe2_LINKER_LIBS "")
|
||||
|
||||
# ---[ Custom Protobuf
|
||||
include(cmake/ProtoBuf.cmake)
|
||||
|
||||
# ---[ Threads
|
||||
find_package(Threads REQUIRED)
|
||||
list(APPEND Caffe2_LINKER_LIBS ${CMAKE_THREAD_LIBS_INIT})
|
||||
|
||||
# ---[ ATLAS
|
||||
find_package(Atlas REQUIRED)
|
||||
include_directories(SYSTEM ${ATLAS_INCLUDE_DIRS})
|
||||
list(APPEND Caffe2_LINKER_LIBS ${ATLAS_LIBRARIES})
|
||||
# ---[ BLAS
|
||||
set(BLAS "Atlas" CACHE STRING "Selected BLAS library")
|
||||
set_property(CACHE BLAS PROPERTY STRINGS "Atlas;OpenBLAS;MKL")
|
||||
|
||||
if(BLAS STREQUAL "Atlas")
|
||||
find_package(Atlas REQUIRED)
|
||||
include_directories(SYSTEM ${ATLAS_INCLUDE_DIRS})
|
||||
list(APPEND Caffe2_LINKER_LIBS ${ATLAS_LIBRARIES})
|
||||
elseif(BLAS STREQUAL "OpenBLAS")
|
||||
find_package(OpenBLAS REQUIRED)
|
||||
include_directories(SYSTEM ${OpenBLAS_INCLUDE_DIR})
|
||||
list(APPEND Caffe2_LINKER_LIBS ${OpenBLAS_LIB})
|
||||
elseif(BLAS STREQUAL "MKL")
|
||||
find_package(MKL REQUIRED)
|
||||
include_directories(SYSTEM ${MKL_INCLUDE_DIR})
|
||||
list(APPEND Caffe2_LINKER_LIBS ${MKL_LIBRARIES})
|
||||
endif()
|
||||
|
||||
# ---[ Google-glog
|
||||
include("cmake/External/glog.cmake")
|
||||
@ -53,18 +67,6 @@ if(USE_LEVELDB)
|
||||
list(APPEND Caffe2_LINKER_LIBS ${Snappy_LIBRARIES})
|
||||
endif()
|
||||
|
||||
# ---[ CUDA
|
||||
include(cmake/Cuda.cmake)
|
||||
if(HAVE_CUDA)
|
||||
LIST(APPEND CUDA_NVCC_FLAGS -Xcompiler -std=c++11)
|
||||
LIST(APPEND CUDA_NVCC_FLAGS -gencode arch=compute_52,code=sm_52)
|
||||
endif()
|
||||
|
||||
# ---[ NCCL
|
||||
include("cmake/External/nccl.cmake")
|
||||
include_directories(SYSTEM ${NCCL_INCLUDE_DIRS})
|
||||
list(APPEND Caffe2_LINKER_LIBS ${NCCL_LIBRARIES})
|
||||
|
||||
# ---[ OpenCV
|
||||
if(USE_OPENCV)
|
||||
find_package(OpenCV QUIET COMPONENTS core highgui imgproc imgcodecs)
|
||||
@ -77,9 +79,6 @@ endif()
|
||||
# ---[ EIGEN
|
||||
include_directories(SYSTEM ${CMAKE_SOURCE_DIR}/third_party/eigen)
|
||||
|
||||
# ---[ pybind11
|
||||
include_directories(SYSTEM ${CMAKE_SOURCE_DIR}/third_party/pybind11/include)
|
||||
|
||||
# ---[ Python + Numpy
|
||||
find_package(PythonInterp 2.7)
|
||||
find_package(PythonLibs 2.7)
|
||||
@ -88,6 +87,9 @@ find_package(NumPy REQUIRED)
|
||||
include_directories(SYSTEM ${PYTHON_INCLUDE_DIRS} ${NUMPY_INCLUDE_DIRS})
|
||||
list(APPEND Caffe2_LINKER_LIBS ${PYTHON_LIBRARIES})
|
||||
|
||||
# ---[ pybind11
|
||||
include_directories(SYSTEM ${CMAKE_SOURCE_DIR}/third_party/pybind11/include)
|
||||
|
||||
# ---[ MPI
|
||||
if(USE_MPI)
|
||||
find_package(MPI)
|
||||
@ -105,10 +107,37 @@ if(OpenMP_FOUND)
|
||||
list(APPEND Caffe2_LINKER_LIBS ${OpenMP_CXX_FLAGS})
|
||||
endif()
|
||||
|
||||
# ---[ CUDA
|
||||
include(cmake/Cuda.cmake)
|
||||
if(HAVE_CUDA)
|
||||
LIST(APPEND CUDA_NVCC_FLAGS -Xcompiler -std=c++11)
|
||||
LIST(APPEND CUDA_NVCC_FLAGS -gencode arch=compute_52,code=sm_52)
|
||||
endif()
|
||||
|
||||
# ---[ CUDNN
|
||||
if(HAVE_CUDA)
|
||||
find_package(CuDNN REQUIRED)
|
||||
if(CUDNN_FOUND)
|
||||
include_directories(SYSTEM ${CUDNN_INCLUDE_DIRS})
|
||||
list(APPEND Caffe2_LINKER_LIBS ${CUDNN_LIBRARIES})
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# ---[ NCCL
|
||||
if(HAVE_CUDA)
|
||||
include("cmake/External/nccl.cmake")
|
||||
include_directories(SYSTEM ${NCCL_INCLUDE_DIRS})
|
||||
list(APPEND Caffe2_LINKER_LIBS ${NCCL_LIBRARIES})
|
||||
endif()
|
||||
|
||||
# ---[ CUB
|
||||
include_directories(SYSTEM ${CMAKE_SOURCE_DIR}/third_party/cub)
|
||||
if(HAVE_CUDA)
|
||||
include_directories(SYSTEM ${CMAKE_SOURCE_DIR}/third_party/cub)
|
||||
endif()
|
||||
|
||||
# ---[ CNMEM
|
||||
add_subdirectory(${CMAKE_SOURCE_DIR}/third_party/cnmem)
|
||||
include_directories(SYSTEM ${CMAKE_SOURCE_DIR}/third_party/cnmem/include)
|
||||
list(APPEND ${Caffe2_LINKER_LIBS} ${CMAKE_SOURCE_DIR}/third_party/cnmem/libcnmem.so)
|
||||
if(HAVE_CUDA)
|
||||
add_subdirectory(${CMAKE_SOURCE_DIR}/third_party/cnmem)
|
||||
include_directories(SYSTEM ${CMAKE_SOURCE_DIR}/third_party/cnmem/include)
|
||||
list(APPEND ${Caffe2_LINKER_LIBS} ${CMAKE_SOURCE_DIR}/third_party/cnmem/libcnmem.so)
|
||||
endif()
|
||||
|
||||
53
cmake/Modules/FindCuDNN.cmake
Normal file
53
cmake/Modules/FindCuDNN.cmake
Normal file
@ -0,0 +1,53 @@
|
||||
# - Try to find cuDNN
|
||||
#
|
||||
# The following variables are optionally searched for defaults
|
||||
# CUDNN_ROOT_DIR: Base directory where all cuDNN components are found
|
||||
#
|
||||
# The following are set after configuration is done:
|
||||
# CUDNN_FOUND
|
||||
# CUDNN_INCLUDE_DIRS
|
||||
# CUDNN_LIBRARIES
|
||||
# CUDNN_LIBRARY_DIRS
|
||||
|
||||
include(FindPackageHandleStandardArgs)
|
||||
|
||||
set(CUDNN_ROOT_DIR "" CACHE PATH "Folder contains NVIDIA cuDNN")
|
||||
|
||||
find_path(CUDNN_INCLUDE_DIR cudnn.h
|
||||
PATHS ${CUDNN_ROOT_DIR}
|
||||
PATH_SUFFIXES cuda/include include)
|
||||
|
||||
find_library(CUDNN_LIBRARY cudnn
|
||||
PATHS ${CUDNN_ROOT_DIR}
|
||||
PATH_SUFFIXES lib cuda/lib)
|
||||
|
||||
find_package_handle_standard_args(CUDNN DEFAULT_MSG CUDNN_INCLUDE_DIR CUDNN_LIBRARY)
|
||||
|
||||
if(CUDNN_FOUND)
|
||||
# get cuDNN version
|
||||
file(READ ${CUDNN_INCLUDE_DIR}/cudnn.h CUDNN_HEADER_CONTENTS)
|
||||
string(REGEX MATCH "define CUDNN_MAJOR * +([0-9]+)"
|
||||
CUDNN_VERSION_MAJOR "${CUDNN_HEADER_CONTENTS}")
|
||||
string(REGEX REPLACE "define CUDNN_MAJOR * +([0-9]+)" "\\1"
|
||||
CUDNN_VERSION_MAJOR "${CUDNN_VERSION_MAJOR}")
|
||||
string(REGEX MATCH "define CUDNN_MINOR * +([0-9]+)"
|
||||
CUDNN_VERSION_MINOR "${CUDNN_HEADER_CONTENTS}")
|
||||
string(REGEX REPLACE "define CUDNN_MINOR * +([0-9]+)" "\\1"
|
||||
CUDNN_VERSION_MINOR "${CUDNN_VERSION_MINOR}")
|
||||
string(REGEX MATCH "define CUDNN_PATCHLEVEL * +([0-9]+)"
|
||||
CUDNN_VERSION_PATCH "${CUDNN_HEADER_CONTENTS}")
|
||||
string(REGEX REPLACE "define CUDNN_PATCHLEVEL * +([0-9]+)" "\\1"
|
||||
CUDNN_VERSION_PATCH "${CUDNN_VERSION_PATCH}")
|
||||
# Assemble cuDNN version
|
||||
if(NOT CUDNN_VERSION_MAJOR)
|
||||
set(CUDNN_VERSION "?")
|
||||
else()
|
||||
set(CUDNN_VERSION "${CUDNN_VERSION_MAJOR}.${CUDNN_VERSION_MINOR}.${CUDNN_VERSION_PATCH}")
|
||||
endif()
|
||||
|
||||
set(CUDNN_INCLUDE_DIRS ${CUDNN_INCLUDE_DIR})
|
||||
set(CUDNN_LIBRARIES ${CUDNN_LIBRARY})
|
||||
message(STATUS "Found cuDNN: v${CUDNN_VERSION} (include: ${CUDNN_INCLUDE_DIR}, library: ${CUDNN_LIBRARY})")
|
||||
mark_as_advanced(CUDNN_ROOT_DIR CUDNN_LIBRARY CUDNN_INCLUDE_DIR)
|
||||
endif()
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
# - Try to find NCCL
|
||||
#
|
||||
# The following variables are optionally searched for defaults
|
||||
# NCCL_ROOT_DIR: Base directory where all GLOG components are found
|
||||
# NCCL_ROOT_DIR: Base directory where all NCCL components are found
|
||||
#
|
||||
# The following are set after configuration is done:
|
||||
# NCCL_FOUND
|
||||
|
||||
Reference in New Issue
Block a user