add option to turn on/off of the NVTX3

MediosZ · Oct 24, 2024 · 6c752f9 · 6c752f9
1 parent be788c1
commit 6c752f9
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 5 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -12,6 +12,7 @@ option(MUDA_DEV "build muda example and unit test. if you're the developer, you
 option(MUDA_FORCE_CHECK "turn on muda runtime check for all mode (Debug/RelWithDebInfo/Release)" OFF)
 option(MUDA_WITH_CHECK "turn on muda runtime check when mode != Release" ON)
 option(MUDA_WITH_COMPUTE_GRAPH "turn on muda compute graph" OFF)
+option(MUDA_WITH_NVTX3 "turn on nividia tools extension library" OFF)
 
 if(MUDA_DEV)
   set(MUDA_BUILD_EXAMPLE ON)
@@ -64,6 +65,11 @@ if (MUDA_WITH_COMPUTE_GRAPH)
 else()
   target_compile_definitions(muda INTERFACE "-DMUDA_COMPUTE_GRAPH_ON=0")
 endif()
+if (MUDA_WITH_NVTX3)
+  target_compile_definitions(muda INTERFACE "-DMUDA_NVTX3_ON=1")
+else()
+  target_compile_definitions(muda INTERFACE "-DMUDA_NVTX3_ON=0")
+endif()
 
 
 if(MUDA_BUILD_EXAMPLE)

diff --git a/src/muda/launch/details/launch_base.inl b/src/muda/launch/details/launch_base.inl
@@ -1,4 +1,8 @@
 #include <muda/exception.h>
+#if MUDA_WITH_NVTX3
+#include <nvtx3/nvToolsExt.h>
+#include <nvtx3/nvToolsExtCuda.h>
+#endif
 #include <muda/compute_graph/compute_graph_accessor.h>
 #include <muda/compute_graph/compute_graph_var.h>
 #include <muda/graph/graph.h>
@@ -29,8 +33,10 @@ MUDA_INLINE MUDA_GENERIC LaunchCore::LaunchCore(cudaStream_t stream) MUDA_NOEXCE
 #endif
 }
 
+
 MUDA_INLINE void LaunchCore::push_range(const std::string& name)
 {
+#if MUDA_WITH_NVTX3
     MUDA_ASSERT(ComputeGraphBuilder::is_phase_none(),
                 "`push_range()` is meaningless in ComputeGraph");
 
@@ -42,15 +48,19 @@ MUDA_INLINE void LaunchCore::push_range(const std::string& name)
     eventAttrib.messageType           = NVTX_MESSAGE_TYPE_ASCII;
     eventAttrib.message.ascii         = name.c_str();
     nvtxRangePushEx(&eventAttrib);
+#endif
 }
 
 MUDA_INLINE void LaunchCore::pop_range()
 {
+#if MUDA_WITH_NVTX3
     MUDA_ASSERT(ComputeGraphBuilder::is_phase_none(),
                 "`pop_range()` is meaningless in ComputeGraph");
     nvtxRangePop();
+#endif
 }
 
+
 MUDA_INLINE void LaunchCore::record(cudaEvent_t e, int flag)
 {
     MUDA_ASSERT(ComputeGraphBuilder::is_phase_none(),

diff --git a/src/muda/launch/launch_base.h b/src/muda/launch/launch_base.h
@@ -10,8 +10,7 @@
 #include <cooperative_groups.h>
 
 #include <cuda_profiler_api.h>
-#include <nvtx3/nvToolsExt.h>
-#include <nvtx3/nvToolsExtCuda.h>
+
 #include <muda/type_traits/type_modifier.h>
 #include <muda/tools/launch_info_cache.h>
 
@@ -102,7 +101,6 @@ class LaunchBase : public LaunchCore
     T& push_range(const std::string& name);
     T& pop_range();
 
-
     // create a name for the following kernel launch
     // viewers will record this name for the sake of better recognization when debugging
     T& kernel_name(std::string_view name);

diff --git a/src/muda/profiler.h b/src/muda/profiler.h
@@ -8,8 +8,11 @@
 #include <device_launch_parameters.h>
 
 #include <cuda_profiler_api.h>
+
+#if MUDA_WITH_NVTX3
 #include <nvtx3/nvToolsExt.h>
 #include <nvtx3/nvToolsExtCuda.h>
+#endif
 
 #include <muda/check/check_cuda_errors.h>
 
@@ -38,6 +41,7 @@ class Profile
     }
     Profile(const std::string& name) MUDA_NOEXCEPT : need_pop(true)
     {
+#if MUDA_WITH_NVTX3
         nvtxEventAttributes_t eventAttrib = {0};
         eventAttrib.version               = NVTX_VERSION;
         eventAttrib.size                  = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
@@ -46,14 +50,16 @@ class Profile
         eventAttrib.messageType           = NVTX_MESSAGE_TYPE_ASCII;
         eventAttrib.message.ascii         = name.c_str();
         nvtxRangePushEx(&eventAttrib);
-
+#endif
         checkCudaErrors(cudaProfilerStart());
     }
     ~Profile()
     {
         checkCudaErrors(cudaProfilerStop());
+#if MUDA_WITH_NVTX3
         if(need_pop)
             nvtxRangePop();
+#endif
     }
 };
 
@@ -62,6 +68,7 @@ class RangeName
   public:
     RangeName(const std::string& name) MUDA_NOEXCEPT
     {
+#if MUDA_WITH_NVTX3
         nvtxEventAttributes_t eventAttrib = {0};
         eventAttrib.version               = NVTX_VERSION;
         eventAttrib.size                  = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
@@ -70,8 +77,13 @@ class RangeName
         eventAttrib.messageType           = NVTX_MESSAGE_TYPE_ASCII;
         eventAttrib.message.ascii         = name.c_str();
         nvtxRangePushEx(&eventAttrib);
+#endif
     }
 
-    ~RangeName() { nvtxRangePop(); }
+    ~RangeName() { 
+#if MUDA_WITH_NVTX3
+        nvtxRangePop();
+#endif
+    }
 };
 }  // namespace muda