From 3d93e15b601fdbb151df64e4d24a8e966719b71f Mon Sep 17 00:00:00 2001
From: Tom Lin <tom91136@gmail.com>
Date: Mon, 27 Nov 2023 08:15:54 +0000
Subject: [PATCH] Add Thrust changes from BabelStream Fix Thrust model
 compatibility with nvcc

---
 src/thrust/fasten.hpp  |  7 ++++++-
 src/thrust/model.cmake | 17 +++++++++++------
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/src/thrust/fasten.hpp b/src/thrust/fasten.hpp
index 521138e..80f1e9c 100644
--- a/src/thrust/fasten.hpp
+++ b/src/thrust/fasten.hpp
@@ -34,7 +34,12 @@ template <size_t PPWI> class IMPL_CLS final : public Bude<PPWI> {
          transforms_0 = transforms_0.data(), transforms_1 = transforms_1.data(), transforms_2 = transforms_2.data(), //
          transforms_3 = transforms_3.data(), transforms_4 = transforms_4.data(), transforms_5 = transforms_5.data(), //
          forcefield = forcefield.data(),                                                                             //
-         energies = energies.data()] __device__ __host__(const int group) {
+         energies = energies.data()] __device__
+#ifndef __NVCC__
+            // otherwise, we get "init-captures are not allowed for extended __host__ __device__"
+            __host__
+#endif
+        (const int group) {
           std::array<std::array<Vec4<float>, 3>, PPWI> transform = {};
           std::array<float, PPWI> etot = {};
 
diff --git a/src/thrust/model.cmake b/src/thrust/model.cmake
index 2a380f1..b014f44 100644
--- a/src/thrust/model.cmake
+++ b/src/thrust/model.cmake
@@ -39,11 +39,12 @@ macro(setup)
         # see CUDA.cmake, we're only adding a few Thrust related libraries here
 
         if (POLICY CMP0104)
-            cmake_policy(SET CMP0104 OLD)
+            cmake_policy(SET CMP0104 NEW)
         endif ()
 
+        set(CMAKE_CUDA_ARCHITECTURES  ${CUDA_ARCH})
         # add -forward-unknown-to-host-compiler for compatibility reasons
-        set(CMAKE_CUDA_FLAGS ${CMAKE_CUDA_FLAGS} "--expt-extended-lambda --expt-relaxed-constexpr -forward-unknown-to-host-compiler -arch=${CUDA_ARCH}" ${CUDA_EXTRA_FLAGS})
+        set(CMAKE_CUDA_FLAGS ${CMAKE_CUDA_FLAGS} "--expt-extended-lambda --expt-relaxed-constexpr" ${CUDA_EXTRA_FLAGS})
         enable_language(CUDA)
         # CMake defaults to -O2 for CUDA at Release, let's wipe that and use the global RELEASE_FLAG
         # appended later
@@ -52,7 +53,11 @@ macro(setup)
         message(STATUS "NVCC flags: ${CMAKE_CUDA_FLAGS} ${CMAKE_CUDA_FLAGS_${BUILD_TYPE}}")
 
 
+        # XXX NVHPC <= 21.9 has cub-config in `Linux_x86_64/21.9/cuda/11.4/include/cub/cmake`
+        # XXX NVHPC >= 22.3 has cub-config in `Linux_x86_64/22.3/cuda/11.6/lib64/cmake/cub/`
+        # same thing for thrust
         if (SDK_DIR)
+            list(APPEND CMAKE_PREFIX_PATH ${SDK_DIR})
             find_package(CUB REQUIRED CONFIG PATHS ${SDK_DIR}/cub)
             find_package(Thrust REQUIRED CONFIG PATHS ${SDK_DIR}/thrust)
         else ()
@@ -63,12 +68,12 @@ macro(setup)
         message(STATUS "Using Thrust backend: ${BACKEND}")
 
         # this creates the interface that we can link to
-        thrust_create_target(Thrust HOST CPP DEVICE ${BACKEND})
+        thrust_create_target(Thrust${BACKEND}
+                HOST CPP
+                DEVICE ${BACKEND})
         set_source_files_properties(${IMPL_SOURCES} PROPERTIES LANGUAGE CUDA)
-        message(WARNING ${IMPL_SOURCES})
+        register_link_library(Thrust${BACKEND})
 
-
-        register_link_library(Thrust)
     elseif (${THRUST_IMPL} STREQUAL "ROCM")
         if (SDK_DIR)
             find_package(rocprim REQUIRED CONFIG PATHS ${SDK_DIR}/rocprim)