Test combined internal/user-side use of NVTX

cub/test/CMakeLists.txt

@@ -26,6 +26,9 @@ option(METAL_BUILD_EXAMPLES OFF)
 option(METAL_BUILD_TESTS OFF)
 CPMAddPackage("gh:brunocodutra/[email protected]")
 
+CPMAddPackage(NAME NVTX GITHUB_REPOSITORY NVIDIA/NVTX GIT_TAG release-v3 DOWNLOAD_ONLY True SYSTEM)
+add_subdirectory(${NVTX_SOURCE_DIR}/c ${NVTX_BINARY_DIR}/c)
+
 find_package(CUDAToolkit)
 
 set(curand_default OFF)
@@ -276,6 +279,9 @@ function(cub_add_test target_name_var test_name test_src cub_target launcher_id)
 
     add_executable(${test_target} "${test_src}")
     target_link_libraries(${test_target} ${cub_target})
+    if (${test_src} MATCHES "test_nvtx_in_usercode.cu")
+      target_link_libraries(${test_target} nvtx3-cpp)
+    endif()
     cub_clone_target_properties(${test_target} ${cub_target})
     target_include_directories(${test_target} PRIVATE "${CUB_SOURCE_DIR}/test")
     target_compile_definitions(${test_target} PRIVATE CUB_DETAIL_DEBUG_ENABLE_SYNC)

cub/test/test_nvtx_in_usercode.cu

@@ -0,0 +1,20 @@
+#include <cub/device/device_for.cuh> // internal include of NVTX
+#include <thrust/iterator/counting_iterator.h>
+#include <nvtx3/nvtx3.hpp> // user-side include of NVTX, retrieved elsewhere
+
+struct Op
+{
+  _CCCL_HOST_DEVICE void operator()(int i) const
+  {
+    printf("%d\n", i);
+  }
+};
+
+int main()
+{
+  nvtx3::scoped_range range("user-range"); // user-side use of NVTX
+
+  thrust::counting_iterator<int> it{0};
+  cub::DeviceFor::ForEach(it, it + 16, Op{}); // internal use of NVTX
+  cudaDeviceSynchronize();
+}