CMake: Prefer CUDAToolkit over CUDA if available

cmake/Cuda.cmake

@@ -0,0 +1,24 @@
+# Result variables:
+# TP_CUDA_LINK_LIBRARIES is list of dependent libraries to be linked
+# TP_CUDA_INCLUDE_DIRS is list of include paths to be used
+
+# find_package will respect <Package>_ROOT variables
+if(CMAKE_VERSION VERSION_GREATER_EQUAL 3.12.0)
+  cmake_policy(SET CMP0074 NEW)
+endif()
+
+# Try to find the newer CUDAToolkit, otherwise use the deprecated find_package(CUDA)
+find_package(CUDAToolkit)
+if(CUDAToolkit_FOUND)
+  if(CUDA_USE_STATIC_CUDA_RUNTIME)
+    set(TP_CUDA_LINK_LIBRARIES CUDA::cudart_static)
+  else()
+    set(TP_CUDA_LINK_LIBRARIES CUDA::cudart)
+  endif()
+  set(TP_CUDA_INCLUDE_DIRS "${CUDAToolkit_INCLUDE_DIR}")
+else()
+  find_package(CUDA REQUIRED)
+
+  set(TP_CUDA_LINK_LIBRARIES ${CUDA_LIBRARIES})
+  set(TP_CUDA_INCLUDE_DIRS ${CUDA_INCLUDE_DIRS})
+endif()

tensorpipe/CMakeLists.txt

@@ -228,15 +228,7 @@ if(TP_USE_CUDA)
   # TP_PUBLIC_HDRS is the list of public header files that we need to install.
   set(TP_CUDA_PUBLIC_HDRS)
 
-  # TP_LINK_LIBRARIES is list of dependent libraries to be linked
-  set(TP_CUDA_LINK_LIBRARIES)
-
-  # TP_INCLUDE_DIRS is list of include path to be used
-  set(TP_CUDA_INCLUDE_DIRS)
-
-  find_package(CUDA REQUIRED)
-  list(APPEND TP_CUDA_LINK_LIBRARIES ${CUDA_LIBRARIES})
-  list(APPEND TP_CUDA_INCLUDE_DIRS ${CUDA_INCLUDE_DIRS})
+  include(Cuda)
 
   list(APPEND TP_CUDA_SRCS
     common/cuda_buffer.cc)

tensorpipe/test/CMakeLists.txt

@@ -70,9 +70,9 @@ if(TP_ENABLE_CMA)
 endif()
 
 if(TP_USE_CUDA)
-  find_package(CUDA REQUIRED)
-  list(APPEND TP_TEST_LINK_LIBRARIES ${CUDA_LIBRARIES})
-  list(APPEND TP_TEST_INCLUDE_DIRS ${CUDA_INCLUDE_DIRS})
+  include(Cuda)
+  list(APPEND TP_TEST_LINK_LIBRARIES ${TP_CUDA_LINK_LIBRARIES})
+  list(APPEND TP_TEST_INCLUDE_DIRS ${TP_CUDA_INCLUDE_DIRS})
   list(APPEND TP_TEST_COMPILE_DEFS TP_USE_CUDA)
 
   list(APPEND TP_TEST_SRCS
@@ -98,7 +98,15 @@ if(TP_USE_CUDA)
     channel/cuda_gdr/cuda_gdr_test.cc
     )
 
-  cuda_add_library(tensorpipe_cuda_kernel channel/kernel.cu)
+  if(CMAKE_VERSION VERSION_GREATER_EQUAL 3.10)
+    enable_language(CUDA)
+    add_library(tensorpipe_cuda_kernel channel/kernel.cu)
+  else()
+    cuda_add_library(tensorpipe_cuda_kernel channel/kernel.cu)
+  endif()
+  target_link_libraries(tensorpipe_cuda_kernel PRIVATE ${TP_CUDA_LINK_LIBRARIES})
+  target_include_directories(tensorpipe_cuda_kernel PRIVATE ${TP_CUDA_INCLUDE_DIRS})
+
   list(APPEND TP_TEST_LINK_LIBRARIES tensorpipe_cuda_kernel)
 
   list(APPEND TP_TEST_LINK_LIBRARIES tensorpipe_cuda)