Enable CUB NVTX ranges only if the NVTX3 C++ API V1 is available

NVIDIA · May 16, 2024 · 3e5ea18 · 3e5ea18
1 parent 06de31f
commit 3e5ea18
Showing 1 changed file with 16 additions and 5 deletions.
diff --git a/cub/cub/detail/nvtx.cuh b/cub/cub/detail/nvtx.cuh
@@ -37,18 +37,22 @@
 #  pragma system_header
 #endif // no system header
 
-// Enable the functionality of this header if
+// Include the NVTX3 C++ API if:
 // * The NVTX3 C API is available in CTK
 // * NVTX is not explicitly disabled
-// * 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(NVTX_DISABLE)
 // 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>
 #  else // __has_include(<nvtx3/nvtx3.hpp>)
 #    include "nvtx3.hpp"
 #  endif // __has_include(<nvtx3/nvtx3.hpp>)
+#endif // __has_include(<nvtx3/nvToolsExt.h>) && !defined(NVTX_DISABLE) && _CCCL_STD_VER > 2011
 
+// Enable the functionality of this header if:
+// * The NVTX3 C++ API V1 is available
+// * C++14 is availabl for cuda::std::optional
+#if defined(NVTX3_CPP_DEFINITIONS_V1_0) && _CCCL_STD_VER >= 2014
 #  include <cuda/std/optional>
 
 CUB_NAMESPACE_BEGIN
@@ -82,7 +86,14 @@ CUB_NAMESPACE_END
       (void) __cub_nvtx3_range;)
 
 #  define CUB_DETAIL_NVTX_RANGE_SCOPE(name) CUB_DETAIL_NVTX_RANGE_SCOPE_IF(true, name)
-#else // __has_include(<nvtx3/nvToolsExt.h>) && !defined(NVTX_DISABLE) && _CCCL_STD_VER > 2011
+#else // defined(NVTX3_CPP_DEFINITIONS_V1_0) && _CCCL_STD_VER >= 2014
+// Tell the user we don't support their NVTX3 version.
+#  if defined(_CCCL_COMPILER_MSVC)
+#    pragma message("warning: CUB only supports NVTX3 V1. Please open an issue on GitHub.")
+#  else
+#    warning CUB only supports NVTX3 V1. Please open an issue on GitHub.
+#  endif
+
 #  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 > 2011
+#endif // defined(NVTX3_CPP_DEFINITIONS_V1_0) && _CCCL_STD_VER >= 2014