frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 12:54:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactored CUDA detection a bit

Browse Source 
Refactoring, minor fixes
This commit is contained in:
Simon Layton
2016-12-12 12:29:00 -05:00
parent 9f35f47411
commit 681267b66a
7 changed files with 39 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
cmake/Cuda.cmake
View File

@ -1,6 +1,6 @@
# Known NVIDIA GPU achitectures Caffe2 can be compiled for.
# This list will be used for CUDA_ARCH_NAME = All option
set(Caffe2_known_gpu_archs "20 21(20) 30 35 50")
set(Caffe2_known_gpu_archs "20 21(20) 30 35 50 52 60 61")
################################################################################################
# A function for automatic detection of GPUs installed (if autodetection is enabled)
@ -52,7 +52,7 @@ endfunction()
# caffe_select_nvcc_arch_flags(out_variable)
function(caffe2_select_nvcc_arch_flags out_variable)
# List of arch names
set(__archs_names "Fermi" "Kepler" "Maxwell" "All" "Manual")
set(__archs_names "Kepler" "Maxwell" "Pascal" "All" "Manual")
set(__archs_name_default "All")
if(NOT CMAKE_CROSSCOMPILING)
list(APPEND __archs_names "Auto")
@ -79,9 +79,7 @@ function(caffe2_select_nvcc_arch_flags out_variable)
unset(CUDA_ARCH_PTX CACHE)
endif()
if(${CUDA_ARCH_NAME} STREQUAL "Fermi")
set(__cuda_arch_bin "20 21(20)")
elseif(${CUDA_ARCH_NAME} STREQUAL "Kepler")
if(${CUDA_ARCH_NAME} STREQUAL "Kepler")
set(__cuda_arch_bin "30 35")
elseif(${CUDA_ARCH_NAME} STREQUAL "Maxwell")
set(__cuda_arch_bin "50")
@ -201,10 +199,11 @@ foreach(diag cc_clobber_ignored integer_sign_change useless_using_declaration se
list(APPEND CUDA_NVCC_FLAGS -Xcudafe --diag_suppress=${diag})
endforeach()
# setting default testing device
if(NOT CUDA_TEST_DEVICE)
set(CUDA_TEST_DEVICE -1)
endif()
# Set C++11 support
list(APPEND CUDA_NVCC_FLAGS "-Xcompiler -fPIC")
# Set :expt-relaxed-constexpr to suppress Eigen warnings
list(APPEND CUDA_NVCC_FLAGS "--expt-relaxed-constexpr")
mark_as_advanced(CUDA_BUILD_CUBIN CUDA_BUILD_EMULATION CUDA_VERBOSE_BUILD)
mark_as_advanced(CUDA_SDK_ROOT_DIR CUDA_SEPARABLE_COMPILATION)

10
cmake/Dependencies.cmake
View File

@ -11,6 +11,7 @@ list(APPEND Caffe2_LINKER_LIBS ${CMAKE_THREAD_LIBS_INIT})
# ---[ BLAS
set(BLAS "Atlas" CACHE STRING "Selected BLAS library")
set_property(CACHE BLAS PROPERTY STRINGS "Atlas;OpenBLAS;MKL")
list(APPEND Caffe2_LINKER_LIBS cblas)
if(BLAS STREQUAL "Atlas")
find_package(Atlas REQUIRED)
@ -109,10 +110,6 @@ 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)
@ -127,6 +124,7 @@ endif()
if(HAVE_CUDA)
include("cmake/External/nccl.cmake")
include_directories(SYSTEM ${NCCL_INCLUDE_DIRS})
message(STATUS "NCCL: ${NCCL_LIBRARIES}")
list(APPEND Caffe2_LINKER_LIBS ${NCCL_LIBRARIES})
endif()
@ -139,5 +137,7 @@ endif()
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)
# message(STATUS "cnmem: ${CMAKE_SOURCE_DIR}/third_party/cnmem/libcnmem.so")
# message(STATUS "${CMAKE_CURRENT_BINARY_DIR}")
list(APPEND Caffe2_LINKER_LIBS "${CMAKE_CURRENT_BINARY_DIR}/third_party/cnmem/libcnmem.so")
endif()