diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index d957176f6f9..8a55f0e35a6 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -6,7 +6,6 @@ set(Caffe2_DIR_EXCLUDES "/mpi" # Exclude for now "/experiments" # Exclude for now "/python" # handled below - "/cuda_rtc" # TODO: Handle this somehow ) set(Caffe2_FILE_EXCLUDES "/operators/fully_connected_op_sparse.cc" @@ -115,7 +114,7 @@ add_library(Caffe2_CPU ${Caffe2_CPU_SRCS}) target_link_libraries(Caffe2_CPU ${Caffe2_LINK_LIBS}) list(APPEND Caffe2_LINK_LIBS Caffe2_CPU) -# CUDA library +# ---[ CUDA library if(HAVE_CUDA) CUDA_ADD_LIBRARY(Caffe2_GPU ${Caffe2_GPU_SRCS}) target_link_libraries(Caffe2_GPU ${Caffe2_LINK_LIBS} glog gflags atlas cblas) diff --git a/cmake/Cuda.cmake b/cmake/Cuda.cmake index 1cd3f20486c..a45e9d0dbd2 100644 --- a/cmake/Cuda.cmake +++ b/cmake/Cuda.cmake @@ -175,10 +175,26 @@ include_directories(SYSTEM ${CUDA_INCLUDE_DIRS}) list(APPEND Caffe2_LINKER_LIBS ${CUDA_CUDART_LIBRARY} ${CUDA_curand_LIBRARY} ${CUDA_CUBLAS_LIBRARIES}) +# find libcuda.so and lbnvrtc.so +find_library(CUDA_CUDA_LIB cuda + PATHS ${CUDA_TOOLKIT_ROOT_DIR} + PATH_SUFFIXES lib lib64) +find_library(CUDA_NVRTC_LIB nvrtc + PATHS ${CUDA_TOOLKIT_ROOT_DIR} + PATH_SUFFIXES lib lib64) + # setting nvcc arch flags caffe2_select_nvcc_arch_flags(NVCC_FLAGS_EXTRA) list(APPEND CUDA_NVCC_FLAGS ${NVCC_FLAGS_EXTRA}) message(STATUS "Added CUDA NVCC flags for: ${NVCC_FLAGS_EXTRA_readable}") +if(CUDA_CUDA_LIB) + message(STATUS "Found libcuda: ${CUDA_CUDA_LIB}") + list(APPEND Caffe2_LINKER_LIBS ${CUDA_CUDA_LIB}) +endif() +if(CUDA_NVRTC_LIB) + message(STATUS "Found libnvrtc: ${CUDA_NVRTC_LIB}") + list(APPEND Caffe2_LINKER_LIBS ${CUDA_NVRTC_LIB}) +endif() # disable some nvcc diagnostic that apears in boost, glog, glags, opencv, etc. foreach(diag cc_clobber_ignored integer_sign_change useless_using_declaration set_but_not_used)