Add CUDA NVRTC cases

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)

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)