From f3bb6bd9a2faed77a13654bb63be4e838c2ab535 Mon Sep 17 00:00:00 2001
From: Bernhard Manfred Gruber <bernhardmgruber@gmail.com>
Date: Wed, 24 Apr 2024 17:52:05 +0200
Subject: [PATCH] Add a standalone test for NVTX ranges

---
 cub/test/insert_nested_NVTX_range_guard.h |  2 +-
 cub/test/test_nvtx_standalone.cu          | 28 +++++++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)
 create mode 100644 cub/test/test_nvtx_standalone.cu

diff --git a/cub/test/insert_nested_NVTX_range_guard.h b/cub/test/insert_nested_NVTX_range_guard.h
index 1626aaff9b5..56d7aad6bc1 100644
--- a/cub/test/insert_nested_NVTX_range_guard.h
+++ b/cub/test/insert_nested_NVTX_range_guard.h
@@ -32,4 +32,4 @@ struct NestedNVTXRangeGuard
 #  define CUB_DETAIL_BEFORE_NVTX_RANGE_SCOPE(name)                              \
     ::cuda::std::optional<::NestedNVTXRangeGuard> __cub_nvtx3_reentrency_guard; \
     NV_IF_TARGET(NV_IS_HOST, __cub_nvtx3_reentrency_guard.emplace(name););
-#endif
+#endif // defined(__cpp_inline_variables)
diff --git a/cub/test/test_nvtx_standalone.cu b/cub/test/test_nvtx_standalone.cu
new file mode 100644
index 00000000000..96bed2bb120
--- /dev/null
+++ b/cub/test/test_nvtx_standalone.cu
@@ -0,0 +1,28 @@
+// The purpose of this test is to verify that CUB can use NVTX without any additional dependencies. It is built as part
+// of the unit tests, but can also be built standalone:
+
+// Compile (from current directory):
+//   nvcc test_nvtx_standalone.cu -I../../cub -I../../thrust -I../../libcudacxx/include -o nvtx_standalone
+// Profile & view:
+//   (nsys profile -o nvtx_standalone.nsys-rep -f true ./nvtx_standalone || true) && nsys-ui nvtx_standalone.nsys-rep
+
+#include <cub/device/device_for.cuh>
+
+#include <thrust/iterator/counting_iterator.h>
+
+struct Op
+{
+  _CCCL_HOST_DEVICE void operator()(int i) const
+  {
+    printf("%d\n", i);
+  }
+};
+
+int main()
+{
+  CUB_DETAIL_NVTX_RANGE_SCOPE("main");
+
+  thrust::counting_iterator<int> it{0};
+  cub::DeviceFor::ForEach(it, it + 16, Op{});
+  cudaDeviceSynchronize();
+}