(tvm) branch main updated: [Cmake] Allow using custom CCCL path for thrust (#16816)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new cd60f6d4fe [Cmake] Allow using custom CCCL path for thrust (#16816)
cd60f6d4fe is described below

commit cd60f6d4feb56abbc61abf050a398a041639ded4
Author: Wuwei Lin <wu...@apache.org>
AuthorDate: Fri Mar 29 10:29:56 2024 -0700

    [Cmake] Allow using custom CCCL path for thrust (#16816)
    
    * [CMake] Allow using custom CCCL path for thrust
---
 cmake/config.cmake       | 4 ++++
 cmake/modules/CUDA.cmake | 9 +++++++--
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/cmake/config.cmake b/cmake/config.cmake
index 2666185fce..9207204997 100644
--- a/cmake/config.cmake
+++ b/cmake/config.cmake
@@ -333,6 +333,10 @@ set(USE_VTA_TSIM OFF)
 set(USE_VTA_FPGA OFF)
 
 # Whether use Thrust
+# Possible values:
+# - ON: enable Thrust with cmake's auto search
+# - OFF: disable Thrust
+# - /path/to/cccl: use specific path to CCCL
 set(USE_THRUST OFF)
 
 # Whether use cuRAND
diff --git a/cmake/modules/CUDA.cmake b/cmake/modules/CUDA.cmake
index 84f466f591..7d7283641e 100644
--- a/cmake/modules/CUDA.cmake
+++ b/cmake/modules/CUDA.cmake
@@ -71,9 +71,14 @@ if(USE_CUDA)
 
   if(USE_THRUST)
     message(STATUS "Build with Thrust support")
-    set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-extended-lambda")
     tvm_file_glob(GLOB CONTRIB_THRUST_SRC src/runtime/contrib/thrust/*.cu)
-    list(APPEND RUNTIME_SRCS ${CONTRIB_THRUST_SRC})
+    add_library(tvm_thrust_objs OBJECT ${CONTRIB_THRUST_SRC})
+    target_compile_options(tvm_thrust_objs PRIVATE 
$<$<COMPILE_LANGUAGE:CUDA>:--expt-extended-lambda>)
+    if (NOT USE_THRUST MATCHES ${IS_TRUE_PATTERN})
+      find_package(CCCL REQUIRED COMPONENTS Thrust)
+      target_link_libraries(tvm_thrust_objs PRIVATE CCCL::Thrust)
+    endif()
+    list(APPEND TVM_RUNTIME_EXT_OBJS $<TARGET_OBJECTS:tvm_thrust_objs>)
   endif(USE_THRUST)
 
   if(USE_CURAND)