Split out nvtx range impl according to owner's review

Pros: clarity, mutex free
cons: similar code duplication, double hooks from nvtx::push_range_name()

Author: huuanhhuyn

cpp/include/raft/core/detail/nvtx.hpp

@@ -1,11 +1,12 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #pragma once
 
 #include <raft/core/detail/macros.hpp>
+#include <raft/core/detail/nvtx_range_path_stack.hpp>
 #include <raft/core/detail/nvtx_range_stack.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
@@ -150,7 +151,8 @@ inline void push_range_name(const char* name)
   event_attrib.messageType           = NVTX_MESSAGE_TYPE_ASCII;
   event_attrib.message.ascii         = name;
   nvtxDomainRangePushEx(domain_store<Domain>::value(), &event_attrib);
-  detail::range_name_stack_instance.push(name);
+  detail::range_name_stack_instance.push(name);  // tracks inner range and depth, cross-thread
+  detail::full_range_stack_instance.push(name);  // tracks full range stack, thread-local
 }
 
 template <typename Domain, typename... Args>
@@ -174,6 +176,7 @@ template <typename Domain>
 inline void pop_range()
 {
   detail::range_name_stack_instance.pop();
+  detail::full_range_stack_instance.pop();
   nvtxDomainRangePop(domain_store<Domain>::value());
 }
 

cpp/include/raft/core/detail/nvtx_range_path_stack.hpp

@@ -0,0 +1,95 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#pragma once
+
+#include <raft/core/detail/macros.hpp>
+
+#include <atomic>
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace raft {
+namespace common::nvtx {
+
+namespace detail {
+
+/**
+ * Process-wide counter producing a unique id to every pushed range, so that
+ * two nvtx ranges sharing the same name can differentiate.
+ *
+ * Motivation: To sample different allocation stats over different runs of the
+ * same block of code. User can aggregate them into a distribution of allocations
+ * over multiple runs.
+ */
+RAFT_EXPORT inline std::atomic<std::uint64_t> range_instance_counter{0};
+
+/**
+ * @brief Per-thread NVTX range stack that records the full range path with
+ *        unique range instance ids.
+ */
+struct nvtx_full_range_stack {
+  void push(const char* name)
+  {
+    auto id = range_instance_counter.fetch_add(1, std::memory_order_relaxed) + 1;
+    stack_.emplace_back(id, name ? name : "");
+  }
+
+  void pop()
+  {
+    if (!stack_.empty()) { stack_.pop_back(); }
+  }
+
+  /** Innermost range name and stack depth (empty/0 when no range is active). */
+  [[nodiscard]] auto inner_range_and_depth() const noexcept -> std::pair<std::string, std::size_t>
+  {
+    if (stack_.empty()) { return {"", 0}; }
+    return {stack_.back().second, stack_.size()};
+  }
+
+  /** Full range path "name#id -> name#id -> ..." (empty when no range is active). */
+  [[nodiscard]] auto current_path() const -> std::string
+  {
+    std::string path;
+    for (auto const& [id, name] : stack_) {
+      if (!path.empty()) { path += " -> "; }
+      path += name + '#' + std::to_string(id);
+    }
+    return path;
+  }
+
+ private:
+  // (instance id, range name), outer -> inner (top).
+  std::vector<std::pair<std::uint64_t, std::string>> stack_{};
+};
+
+RAFT_EXPORT inline thread_local nvtx_full_range_stack full_range_stack_instance{};
+
+}  // namespace detail
+
+/**
+ * Mutex-free read of the current thread's innermost NVTX range name and stack depth.
+ *
+ * ONLY safe to call from the thread that owns this range stack (the current thread).
+ */
+RAFT_EXPORT inline auto thread_local_inner_range_and_depth() -> std::pair<std::string, std::size_t>
+{
+  return detail::full_range_stack_instance.inner_range_and_depth();
+}
+
+/**
+ * Mutex-free read of the current thread's full NVTX range path "name#id -> name#id -> ...".
+ *
+ * ONLY safe to call from the thread that owns this range stack (the current thread).
+ */
+RAFT_EXPORT inline auto thread_local_nvtx_full_path() -> std::string
+{
+  return detail::full_range_stack_instance.current_path();
+}
+
+}  // namespace common::nvtx
+}  // namespace raft

cpp/include/raft/core/detail/nvtx_range_stack.hpp

@@ -6,14 +6,12 @@
 
 #include <raft/core/detail/macros.hpp>
 
-#include <atomic>
 #include <cstddef>
-#include <cstdint>
 #include <memory>
 #include <mutex>
+#include <stack>
 #include <string>
 #include <utility>
-#include <vector>
 
 namespace raft {
 namespace common::nvtx {
@@ -37,16 +35,6 @@ class current_range {
     return {value_, depth_};
   }
 
-  /**
-   * Read the full nvtx range path with instance ids, formatted as
-   * "name#id > name#id > ..." (empty when no range is active).
-   */
-  auto get_path() const -> std::string
-  {
-    std::lock_guard lock(mu_);
-    return path_;
-  }
-
   operator std::string() const
   {
     std::lock_guard lock(mu_);
@@ -57,81 +45,35 @@ class current_range {
   mutable std::mutex mu_;
   std::string value_;
   std::size_t depth_{0};
-  std::string path_;
 
-  void set(const char* name, std::size_t depth, std::string path)
+  void set(const char* name, std::size_t depth)
   {
     std::lock_guard lock(mu_);
     value_ = name ? name : "";
     depth_ = depth;
-    path_  = std::move(path);
   }
 };
 
 namespace detail {
 
-RAFT_EXPORT inline std::atomic<std::uint64_t> range_instance_counter{0};
-
 struct nvtx_range_name_stack {
   void push(const char* name)
   {
-    ensure_current();
-    auto id = range_instance_counter.fetch_add(1, std::memory_order_relaxed) + 1;
-    stack_.emplace_back(id, name ? name : "");
-    current_->set(stack_.back().second.c_str(), stack_.size(), build_path());
+    stack_.emplace(name);
+    current_->set(name, stack_.size());
   }
 
   void pop()
   {
-    ensure_current();
-    if (!stack_.empty()) { stack_.pop_back(); }
-    current_->set(
-      stack_.empty() ? nullptr : stack_.back().second.c_str(), stack_.size(), build_path());
-  }
-
-  [[nodiscard]] auto current() const -> std::shared_ptr<const current_range>
-  {
-    ensure_current();
-    return current_;
-  }
-
-  /**
-   * Innermost NVTX range name and stack depth.
-   * Mutex-free — ONLY safe to call from the thread that owns this stack.
-   */
-  [[nodiscard]] auto current_name_and_depth() const noexcept -> std::pair<std::string, std::size_t>
-  {
-    if (stack_.empty()) { return {"", 0}; }
-    return {stack_.back().second, stack_.size()};
+    if (!stack_.empty()) { stack_.pop(); }
+    current_->set(stack_.empty() ? nullptr : stack_.top().c_str(), stack_.size());
   }
 
-  /**
-   * Full NVTX range path "name#id > name#id > ...".
-   * Mutex-free — ONLY safe to call from the thread that owns this stack.
-   */
-  [[nodiscard]] auto current_path() const -> std::string { return build_path(); }
+  auto current() const -> std::shared_ptr<const current_range> { return current_; }
 
  private:
-  void ensure_current() const
-  {
-    if (!current_) { current_ = std::make_shared<current_range>(); }
-  }
-
-  // Serialize the active stack as "name#id > name#id > ..." (outer -> inner).
-  [[nodiscard]] auto build_path() const -> std::string
-  {
-    std::string path;
-    for (auto const& [id, name] : stack_) {
-      if (!path.empty()) { path += " > "; }
-      path += name;
-      path += '#';
-      path += std::to_string(id);
-    }
-    return path;
-  }
-
-  std::vector<std::pair<std::uint64_t, std::string>> stack_{};
-  mutable std::shared_ptr<current_range> current_{std::make_shared<current_range>()};
+  std::stack<std::string> stack_{};
+  std::shared_ptr<current_range> current_{std::make_shared<current_range>()};
 };
 
 RAFT_EXPORT inline thread_local nvtx_range_name_stack range_name_stack_instance{};
@@ -148,25 +90,5 @@ RAFT_EXPORT inline auto thread_local_current_range() -> std::shared_ptr<const cu
   return detail::range_name_stack_instance.current();
 }
 
-/**
- * Mutex-free read of the current thread's innermost NVTX range name and stack depth.
- * ONLY safe to call from the thread that owns this range stack (the current thread).
- * Use instead of thread_local_current_range()->get() when no cross-thread sharing is needed.
- */
-RAFT_EXPORT inline auto thread_local_current_name_and_depth() -> std::pair<std::string, std::size_t>
-{
-  return detail::range_name_stack_instance.current_name_and_depth();
-}
-
-/**
- * Mutex-free read of the current thread's full NVTX range path "name#id > name#id > ...".
- * ONLY safe to call from the thread that owns this range stack (the current thread).
- * Use instead of thread_local_current_range()->get_path() when no cross-thread sharing is needed.
- */
-RAFT_EXPORT inline auto thread_local_current_path() -> std::string
-{
-  return detail::range_name_stack_instance.current_path();
-}
-
 }  // namespace common::nvtx
 }  // namespace raft

cpp/include/raft/core/memory_logging_resources.hpp

@@ -138,51 +138,54 @@ class memory_logging_resources : public resources {
 
     // --- Host (global) ---
     {
-      int id        = recorder_->register_source("host");
-      host_adaptor_ = std::make_unique<host_record_t>(old_host_, queue, id);
+      int source_id = recorder_->register_source("host");
+      host_adaptor_ = std::make_unique<host_record_t>(old_host_, queue, source_id);
       raft::mr::set_default_host_resource(*host_adaptor_);
     }
 
     // --- Pinned ---
     {
-      int id = recorder_->register_source("pinned");
+      int source_id = recorder_->register_source("pinned");
       raft::resource::set_pinned_memory_resource(
         *this,
-        raft::mr::recording_adaptor<raft::mr::host_device_resource_ref>{pinned_ref, queue, id});
+        raft::mr::recording_adaptor<raft::mr::host_device_resource_ref>{
+          pinned_ref, queue, source_id});
     }
 
     // --- Managed ---
     {
-      int id = recorder_->register_source("managed");
+      int source_id = recorder_->register_source("managed");
       raft::resource::set_managed_memory_resource(
         *this,
-        raft::mr::recording_adaptor<raft::mr::host_device_resource_ref>{managed_ref, queue, id});
+        raft::mr::recording_adaptor<raft::mr::host_device_resource_ref>{
+          managed_ref, queue, source_id});
     }
 
     // --- Device (global) ---
     {
       // Invalidate the cached thrust policy — its resource_ref will be stale
       // once we replace the global device resource.
       cells_[resource::resource_type::THRUST_POLICY] = std::make_shared<resource::resource_cell>();
-      int id                                         = recorder_->register_source("device");
-      device_adaptor_ = std::make_unique<device_record_t>(old_device_, queue, id);
+      int source_id                                  = recorder_->register_source("device");
+      device_adaptor_ = std::make_unique<device_record_t>(old_device_, queue, source_id);
       rmm::mr::set_current_device_resource(*device_adaptor_);
     }
 
     // --- Workspace (track upstream to preserve limiting_resource_adaptor) ---
     {
-      int id = recorder_->register_source("workspace");
+      int source_id = recorder_->register_source("workspace");
       raft::resource::set_workspace_resource(
         *this,
-        raft::mr::recording_adaptor<rmm::device_async_resource_ref>{upstream_ref, queue, id},
+        raft::mr::recording_adaptor<rmm::device_async_resource_ref>{upstream_ref, queue, source_id},
         ws_free);
     }
 
     // --- Large workspace ---
     {
-      int id = recorder_->register_source("large_workspace");
+      int source_id = recorder_->register_source("large_workspace");
       raft::resource::set_large_workspace_resource(
-        *this, raft::mr::recording_adaptor<rmm::device_async_resource_ref>{lws_ref, queue, id});
+        *this,
+        raft::mr::recording_adaptor<rmm::device_async_resource_ref>{lws_ref, queue, source_id});
     }
 
     recorder_->start();

cpp/include/raft/mr/recording_adaptor.hpp

@@ -5,7 +5,7 @@
 #pragma once
 
 #include <raft/core/detail/macros.hpp>
-#include <raft/core/detail/nvtx_range_stack.hpp>  // thread_local_current_path, thread_local_current_name_and_depth
+#include <raft/core/detail/nvtx_range_path_stack.hpp>  // thread_local_nvtx_full_path, thread_local_inner_range_and_depth
 #include <raft/mr/recording_monitor.hpp>   // allocation_event, allocation_event_queue
 #include <raft/mr/statistics_adaptor.hpp>  // resource_stats (atomic counters, reused)
 
@@ -53,7 +53,7 @@ class recording_adaptor : public cuda::forward_property<recording_adaptor<Upstre
   {
     std::string path = "";
     if (ptr != nullptr) {
-      path = raft::common::nvtx::thread_local_current_path();
+      path = raft::common::nvtx::thread_local_nvtx_full_path();
       if (!path.empty()) {
         std::lock_guard<std::mutex> lock(alloc_map_->mtx);
         alloc_map_->paths[ptr] = path;
@@ -78,7 +78,7 @@ class recording_adaptor : public cuda::forward_property<recording_adaptor<Upstre
   // Enqueue an event.  Called on the allocating/deallocating thread — mutex-free NVTX read.
   void emit(std::string nvtx_full_range, std::int64_t signed_bytes) noexcept
   {
-    auto [name, depth] = raft::common::nvtx::thread_local_current_name_and_depth();
+    auto [name, depth] = raft::common::nvtx::thread_local_inner_range_and_depth();
     allocation_event event;
     event.source_id       = source_id_;
     event.current         = stats_->bytes_current.load(std::memory_order_relaxed);