Add CCCL_DISABLE_NVTX macro (#2173)

Fixes: #2172
NVIDIA · Aug 16, 2024 · ba9e9bb · ba9e9bb
1 parent 4a5dcc4
commit ba9e9bb
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 4 deletions.
diff --git 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
@@ -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();
+}