Memory logging resources

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,103 @@
+/*
+ * 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 <raft/core/logger.hpp>  // RAFT_LOG_WARN
+
+#include <algorithm>
+#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;
+    std::string clean = name ? name : "";
+    if (clean.find(',') != std::string::npos) {
+      RAFT_LOG_WARN("NVTX range name '%s' contains ',' - removing it to keep CSV columns intact",
+                    clean.c_str());
+      clean.erase(std::remove(clean.begin(), clean.end(), ','), clean.end());
+    }
+    stack_.emplace_back(id, std::move(clean));
+  }
+
+  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 -> 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/memory_logging_resources.hpp

@@ -0,0 +1,195 @@
+/*
+ * 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 <raft/core/logger.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
+#include <raft/core/resource/managed_memory_resource.hpp>
+#include <raft/core/resource/pinned_memory_resource.hpp>
+#include <raft/core/resources.hpp>
+#include <raft/mr/host_device_resource.hpp>
+#include <raft/mr/host_memory_resource.hpp>
+#include <raft/mr/recording_adaptor.hpp>
+#include <raft/mr/recording_monitor.hpp>
+
+#include <rmm/mr/per_device_resource.hpp>
+#include <rmm/resource_ref.hpp>
+
+#include <fstream>
+#include <memory>
+#include <ostream>
+#include <string>
+
+namespace raft {
+
+/**
+ * @brief A resources handle that wraps all reachable memory resources with
+ *        recording_adaptor and logs every allocation/deallocation as a CSV row
+ *        from a background thread.
+ *
+ * Inherits from raft::resources, so it can be passed anywhere a
+ * raft::resources& is expected.
+ *
+ * Every allocation and deallocation is pushed as an event onto a thread-safe
+ * queue. The active NVTX range path is captured from the allocating/deallocating
+ * thread at the moment of the event — no mutex is taken on the NVTX stack because
+ * the read is always on the owning thread.  The per-pointer alloc_map that maps
+ * addresses to their allocation-time NVTX path still uses a mutex because
+ * deallocation can legitimately occur on a different thread than allocation.
+ *
+ * On construction the handle:
+ *   - Materializes all tracked resource types (host, device, pinned,
+ *     managed, workspace, large_workspace).
+ *   - Takes a snapshot of the original resources to keep them alive.
+ *   - Wraps each with a recording_adaptor.
+ *   - Replaces global host and device resources with tracked versions.
+ *   - Starts a background CSV writer (recording_monitor).
+ *
+ * On destruction the handle stops the monitor and restores the original
+ * global host and device resources.
+ */
+class memory_logging_resources : public resources {
+ public:
+  /**
+   * @brief Construct from an existing resources handle, logging to an ostream.
+   *        Every allocation and deallocation produces one CSV row.
+   */
+  memory_logging_resources(const resources& existing, std::ostream& out)
+    : memory_logging_resources(&existing, nullptr, &out)
+  {
+  }
+
+  /**
+   * @brief Construct from an existing resources handle, logging to a file.
+   *        Every allocation and deallocation produces one CSV row.
+   */
+  memory_logging_resources(const resources& existing, const std::string& file_path)
+    : memory_logging_resources(&existing, std::make_unique<std::ofstream>(file_path), nullptr)
+  {
+  }
+
+  ~memory_logging_resources() override
+  {
+    if (recorder_) recorder_->stop();
+    raft::mr::set_default_host_resource(old_host_);
+    rmm::mr::set_current_device_resource(old_device_);
+  }
+
+  memory_logging_resources(memory_logging_resources const&)            = delete;
+  memory_logging_resources(memory_logging_resources&&)                 = delete;
+  memory_logging_resources& operator=(memory_logging_resources const&) = delete;
+  memory_logging_resources& operator=(memory_logging_resources&&)      = delete;
+
+  /** @brief Access the recording monitor (always non-null after construction). */
+  [[nodiscard]] auto get_recorder() noexcept -> raft::mr::recording_monitor*
+  {
+    return recorder_.get();
+  }
+
+ private:
+  memory_logging_resources(const resources* existing,
+                           std::unique_ptr<std::ofstream> owned_stream,
+                           std::ostream* out_override)
+    : resources(existing ? *existing : resources{}),
+      owned_stream_(std::move(owned_stream)),
+      old_host_(raft::mr::get_default_host_resource()),
+      old_device_(rmm::mr::get_current_device_resource_ref())
+  {
+    std::ostream* outp = out_override;
+    if (!outp) { outp = static_cast<std::ostream*>(owned_stream_.get()); }
+    RAFT_LOG_INFO("memory_logging_resources: queue-based recording (every event captured)");
+    recorder_ = std::make_unique<raft::mr::recording_monitor>(*outp);
+    init_recording();
+  }
+
+  // Declaration order matters: snapshot_ is destroyed last (keeps original resource
+  // shared_ptrs alive); owned_stream_ outlives recorder_ (it writes to it);
+  // recorder_ is stopped in the destructor body before member destruction.
+  std::vector<std::shared_ptr<resource::resource_cell>> snapshot_;
+  std::unique_ptr<std::ofstream> owned_stream_;
+  std::unique_ptr<raft::mr::recording_monitor> recorder_;
+
+  raft::mr::host_resource old_host_;
+  raft::mr::device_resource old_device_;
+
+  using host_record_t = raft::mr::recording_adaptor<raft::mr::host_resource_ref>;
+  std::unique_ptr<host_record_t> host_adaptor_;
+
+  using device_record_t = raft::mr::recording_adaptor<rmm::device_async_resource_ref>;
+  std::unique_ptr<device_record_t> device_adaptor_;
+
+  void init_recording()
+  {
+    // Force-initialize lazily-created resources before we replace the global device MR,
+    // so their upstreams resolve against the original resource.
+    auto* ws          = raft::resource::get_workspace_resource(*this);
+    auto ws_free      = raft::resource::get_workspace_free_bytes(*this);
+    auto upstream_ref = ws->get_upstream_resource();
+    auto lws_ref      = raft::resource::get_large_workspace_resource_ref(*this);
+    auto pinned_ref   = raft::resource::get_pinned_memory_resource_ref(*this);
+    auto managed_ref  = raft::resource::get_managed_memory_resource_ref(*this);
+
+    snapshot_ = cells_;
+
+    auto queue = recorder_->get_queue();
+
+    // --- Host (global) ---
+    {
+      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 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, source_id});
+    }
+
+    // --- 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, 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 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 source_id = recorder_->register_source("workspace");
+      raft::resource::set_workspace_resource(
+        *this,
+        raft::mr::recording_adaptor<rmm::device_async_resource_ref>{upstream_ref, queue, source_id},
+        ws_free);
+    }
+
+    // --- 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, source_id});
+    }
+
+    recorder_->start();
+  }
+};
+
+}  // namespace raft