Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 5e86faa

Browse files
committedJun 26, 2024··
Enable using Thrust from CCCL
1 parent a018a6e commit 5e86faa

File tree

1 file changed

+35
-28
lines changed

1 file changed

+35
-28
lines changed
 

‎src/thrust/model.cmake

+35-28
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,24 @@
11

22
register_flag_optional(THRUST_IMPL
33
"Which Thrust implementation to use, supported options include:
4-
- CUDA (via https://github.com/NVIDIA/thrust)
4+
- CUDA (via https://github.com/NVIDIA/thrust or https://github.com/NVIDIA/CCCL)
55
- ROCM (via https://github.com/ROCmSoftwarePlatform/rocThrust)
66
"
77
"CUDA")
88

99
register_flag_optional(SDK_DIR
10-
"Path to the selected Thrust implementation (e.g `/opt/nvidia/hpc_sdk/Linux_x86_64/21.9/cuda/include` for NVHPC, `/opt/rocm` for ROCm)"
10+
"Path to the installation prefix for CCCL or Thrust (e.g `/opt/nvidia/hpc_sdk/Linux_x86_64/22.3/cuda/11.6` for NVHPC, or `/usr/local/cuda` for nvcc, or `/opt/rocm` for ROCm)"
1111
"")
1212

1313
register_flag_optional(BACKEND
1414
"[THRUST_IMPL==CUDA] CUDA's Thrust implementation supports the following backends:
15-
- CUDA
15+
- CUD
1616
- OMP
1717
- TBB
1818
"
1919
"CUDA")
2020

21-
register_flag_optional(MANAGED "Enabled managed memory mode."
21+
register_flag_optional(MANAGED "Enabled managed memory mode."
2222
"OFF")
2323

2424
register_flag_optional(CMAKE_CUDA_COMPILER
@@ -34,6 +34,9 @@ register_flag_optional(CUDA_EXTRA_FLAGS
3434
"[THRUST_IMPL==CUDA] Additional CUDA flags passed to nvcc, this is appended after `CUDA_ARCH`"
3535
"")
3636

37+
option(FETCH_CCCL "Fetch (download) the CCCL library. This uses CMake's FetchContent feature.
38+
Specify version by setting FETCH_CCCL_VERSION" OFF)
39+
set(FETCH_CCCL_VERSION "v2.4.0" CACHE STRING "Specify version of CCCL to use if FETCH_CCCL is ON")
3740

3841
macro(setup)
3942
set(CMAKE_CXX_STANDARD 14)
@@ -42,44 +45,48 @@ macro(setup)
4245
endif ()
4346

4447
if (${THRUST_IMPL} STREQUAL "CUDA")
45-
46-
# see CUDA.cmake, we're only adding a few Thrust related libraries here
47-
4848
if (POLICY CMP0104)
4949
cmake_policy(SET CMP0104 NEW)
5050
endif ()
51-
5251
set(CMAKE_CUDA_ARCHITECTURES ${CUDA_ARCH})
53-
# add -forward-unknown-to-host-compiler for compatibility reasons
5452
set(CMAKE_CUDA_FLAGS ${CMAKE_CUDA_FLAGS} "--expt-extended-lambda " ${CUDA_EXTRA_FLAGS})
5553
enable_language(CUDA)
56-
# CMake defaults to -O2 for CUDA at Release, let's wipe that and use the global RELEASE_FLAG
57-
# appended later
54+
# CMake defaults to -O2 for CUDA at Release, let's wipe that and use the global RELEASE_FLAG appended later
5855
wipe_gcc_style_optimisation_flags(CMAKE_CUDA_FLAGS_${BUILD_TYPE})
5956

6057
message(STATUS "NVCC flags: ${CMAKE_CUDA_FLAGS} ${CMAKE_CUDA_FLAGS_${BUILD_TYPE}}")
6158

62-
63-
# XXX NVHPC <= 21.9 has cub-config in `Linux_x86_64/21.9/cuda/11.4/include/cub/cmake`
64-
# XXX NVHPC >= 22.3 has cub-config in `Linux_x86_64/22.3/cuda/11.6/lib64/cmake/cub/`
65-
# same thing for thrust
6659
if (SDK_DIR)
60+
# CMake tries several subdirectories below SDK_DIR, see documentation:
61+
# https://cmake.org/cmake/help/latest/command/find_package.html#config-mode-search-procedure
6762
list(APPEND CMAKE_PREFIX_PATH ${SDK_DIR})
68-
find_package(CUB REQUIRED CONFIG PATHS ${SDK_DIR}/cub)
69-
find_package(Thrust REQUIRED CONFIG PATHS ${SDK_DIR}/thrust)
70-
else ()
71-
find_package(CUB REQUIRED CONFIG)
72-
find_package(Thrust REQUIRED CONFIG)
7363
endif ()
7464

75-
message(STATUS "Using Thrust backend: ${BACKEND}")
76-
77-
# this creates the interface that we can link to
78-
thrust_create_target(Thrust${BACKEND}
79-
HOST CPP
80-
DEVICE ${BACKEND})
81-
82-
register_link_library(Thrust${BACKEND})
65+
set(CCCL_THRUST_DEVICE_SYSTEM ${BACKEND} CACHE STRING "" FORCE)
66+
67+
# fetch CCCL if user wants to
68+
if (FETCH_CCCL)
69+
FetchContent_Declare(
70+
CCCL
71+
GIT_REPOSITORY https://github.com/nvidia/cccl.git
72+
GIT_TAG "${FETCH_CCCL_VERSION}"
73+
)
74+
FetchContent_MakeAvailable(CCCL)
75+
register_link_library(CCCL::CCCL)
76+
else()
77+
# try to find CCCL locally
78+
find_package(CCCL CONFIG)
79+
if (CCCL_FOUND)
80+
register_link_library(CCCL::CCCL)
81+
else()
82+
# backup: find legacy projects separately
83+
message(WARNING "No CCCL found on your system. Trying Thrust and CUB legacy targets.")
84+
find_package(CUB REQUIRED CONFIG)
85+
find_package(Thrust REQUIRED CONFIG)
86+
thrust_create_target(Thrust${BACKEND} HOST CPP DEVICE ${BACKEND})
87+
register_link_library(Thrust${BACKEND})
88+
endif()
89+
endif()
8390
elseif (${THRUST_IMPL} STREQUAL "ROCM")
8491
if (SDK_DIR)
8592
find_package(rocprim REQUIRED CONFIG PATHS ${SDK_DIR}/rocprim)

0 commit comments

Comments
 (0)
Please sign in to comment.