Extend memory tracking resources

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

@@ -1,17 +1,19 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * 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 <memory>
 #include <mutex>
-#include <stack>
 #include <string>
 #include <utility>
+#include <vector>
 
 namespace raft {
 namespace common::nvtx {
@@ -35,6 +37,16 @@ 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_);
@@ -45,38 +57,68 @@ 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)
+  void set(const char* name, std::size_t depth, std::string path)
   {
     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)
   {
-    stack_.emplace(name);
-    current_->set(name, stack_.size());
+    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());
   }
 
   void pop()
   {
-    if (!stack_.empty()) { stack_.pop(); }
-    current_->set(stack_.empty() ? nullptr : stack_.top().c_str(), stack_.size());
+    ensure_current();
+    if (!stack_.empty()) { stack_.pop_back(); }
+    current_->set(
+      stack_.empty() ? nullptr : stack_.back().second.c_str(), stack_.size(), build_path());
   }
 
-  auto current() const -> std::shared_ptr<const current_range> { return current_; }
+  [[nodiscard]] auto current() const -> std::shared_ptr<const current_range>
+  {
+    ensure_current();
+    return current_;
+  }
 
  private:
-  std::stack<std::string> stack_{};
-  std::shared_ptr<current_range> current_{std::make_shared<current_range>()};
+  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>()};
 };
 
-inline thread_local nvtx_range_name_stack range_name_stack_instance{};
+RAFT_EXPORT inline thread_local nvtx_range_name_stack range_name_stack_instance{};
 
 }  // namespace detail
 
@@ -85,7 +127,7 @@ inline thread_local nvtx_range_name_stack range_name_stack_instance{};
  * Pass the returned shared_ptr to another thread to read this thread's current NVTX range name at
  * any time.
  */
-inline auto thread_local_current_range() -> std::shared_ptr<const current_range>
+RAFT_EXPORT inline auto thread_local_current_range() -> std::shared_ptr<const current_range>
 {
   return detail::range_name_stack_instance.current();
 }