From befd889626583003d4c9c210d2253eaa2789d6ca Mon Sep 17 00:00:00 2001
From: Peter Bell <peterbell10@live.co.uk>
Date: Wed, 10 Aug 2022 21:50:12 +0100
Subject: [PATCH] Prefer CUDAToolkit over CUDA if available

In newer CMake versions `find_package(CUDA)` is deprecated in favor of
`find_package(CUDATookit)` and `enable_language(CUDA)`.

This first tries to find `CUDAToolkit` and falls back to `CUDA` if
it's not available.
---
 cmake/Cuda.cmake               | 24 ++++++++++++++++++++++++
 tensorpipe/CMakeLists.txt      | 10 +---------
 tensorpipe/test/CMakeLists.txt | 16 ++++++++++++----
 3 files changed, 37 insertions(+), 13 deletions(-)
 create mode 100644 cmake/Cuda.cmake

diff --git a/cmake/Cuda.cmake b/cmake/Cuda.cmake
new file mode 100644
index 000000000..4ead37492
--- /dev/null
+++ b/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()
diff --git a/tensorpipe/CMakeLists.txt b/tensorpipe/CMakeLists.txt
index efcffc257..7efb62119 100644
--- a/tensorpipe/CMakeLists.txt
+++ b/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)
diff --git a/tensorpipe/test/CMakeLists.txt b/tensorpipe/test/CMakeLists.txt
index 53e8643fd..ef76c441a 100644
--- a/tensorpipe/test/CMakeLists.txt
+++ b/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)