From ba9e9bbc20dca1ac49b5da6b4d1716d85b4f495e Mon Sep 17 00:00:00 2001
From: Bernhard Manfred Gruber <bernhardmgruber@gmail.com>
Date: Fri, 16 Aug 2024 19:33:47 +0200
Subject: [PATCH] Add `CCCL_DISABLE_NVTX` macro (#2173)

Fixes: #2172
---
 cub/cub/detail/nvtx.cuh        | 15 +++++++++++----
 cub/test/test_nvtx_disabled.cu | 19 +++++++++++++++++++
 2 files changed, 30 insertions(+), 4 deletions(-)
 create mode 100644 cub/test/test_nvtx_disabled.cu

diff --git a/cub/cub/detail/nvtx.cuh b/cub/cub/detail/nvtx.cuh
index d570df3adc..a8422263fa 100644
--- a/cub/cub/detail/nvtx.cuh
+++ b/cub/cub/detail/nvtx.cuh
@@ -37,11 +37,16 @@
 #  pragma system_header
 #endif // no system header
 
+#ifdef DOXYGEN_SHOULD_SKIP_THIS // Only parse this during doxygen passes:
+//! When this macro is defined, no NVTX ranges are emitted by CCCL
+#  define CCCL_DISABLE_NVTX
+#endif // DOXYGEN_SHOULD_SKIP_THIS
+
 // Enable the functionality of this header if:
 // * The NVTX3 C API is available in CTK
-// * NVTX is not explicitly disabled
+// * NVTX is not explicitly disabled (via CCCL_DISABLE_NVTX or NVTX_DISABLE)
 // * C++14 is availabl for cuda::std::optional
-#if __has_include(<nvtx3/nvToolsExt.h> ) && !defined(NVTX_DISABLE) && _CCCL_STD_VER >= 2014
+#if __has_include(<nvtx3/nvToolsExt.h> ) && !defined(CCCL_DISABLE_NVTX) && !defined(NVTX_DISABLE) && _CCCL_STD_VER >= 2014
 // Include our NVTX3 C++ wrapper if not available from the CTK
 #  if __has_include(<nvtx3/nvtx3.hpp>) // TODO(bgruber): replace by a check for the first CTK version shipping the header
 #    include <nvtx3/nvtx3.hpp>
@@ -96,7 +101,9 @@ CUB_NAMESPACE_END
 #    define CUB_DETAIL_NVTX_RANGE_SCOPE_IF(condition, name)
 #    define CUB_DETAIL_NVTX_RANGE_SCOPE(name)
 #  endif // NVTX3_CPP_DEFINITIONS_V1_0
-#else // __has_include(<nvtx3/nvToolsExt.h> ) && !defined(NVTX_DISABLE) && _CCCL_STD_VER >= 2014
+#else // __has_include(<nvtx3/nvToolsExt.h> ) && !defined(CCCL_DISABLE_NVTX) && !defined(NVTX_DISABLE) && _CCCL_STD_VER
+      // >= 2014
 #  define CUB_DETAIL_NVTX_RANGE_SCOPE_IF(condition, name)
 #  define CUB_DETAIL_NVTX_RANGE_SCOPE(name)
-#endif // __has_include(<nvtx3/nvToolsExt.h> ) && !defined(NVTX_DISABLE) && _CCCL_STD_VER >= 2014
+#endif // __has_include(<nvtx3/nvToolsExt.h> ) && !defined(CCCL_DISABLE_NVTX) && !defined(NVTX_DISABLE) && _CCCL_STD_VER
+       // >= 2014
diff --git a/cub/test/test_nvtx_disabled.cu b/cub/test/test_nvtx_disabled.cu
new file mode 100644
index 0000000000..c6eba196b1
--- /dev/null
+++ b/cub/test/test_nvtx_disabled.cu
@@ -0,0 +1,19 @@
+#define CUB_DETAIL_BEFORE_NVTX_RANGE_SCOPE(name) static_assert(false, "");
+#define CCCL_DISABLE_NVTX
+
+#include <cub/device/device_for.cuh>
+
+#include <thrust/iterator/counting_iterator.h>
+
+#include <cuda/std/functional>
+
+#if defined(CCCL_DISABLE_NVTX) && defined(NVTX_VERSION)
+#  error "NVTX was included somewhere even though it is turned off via CCCL_DISABLE_NVTX"
+#endif // defined(CCCL_DISABLE_NVTX) && defined(NVTX_VERSION)
+
+int main()
+{
+  thrust::counting_iterator<int> it{0};
+  cub::DeviceFor::ForEach(it, it + 16, ::cuda::std::negate<int>{});
+  cudaDeviceSynchronize();
+}