Tracking memory resources (rapidsai#2973)

Detailed tracking of (almost) all allocations on device and host.

```C++
  // optionally pass an existing resource handle
  raft::resources res;

  // The tracking handle is a child of resource handle; it wraps all memory resources with statistics adaptors
  raft::memory_tracking_resources tracked(res, "allocations.csv", std::chrono::milliseconds(1));

  // All allocations are logged to a .csv as long as `tracked` is alive
  cuvs::neighbors::cagra::build(tracked, ...);
```
This produces a CSV file with sampled allocations with a timeline and NVTX correlation
```csv
timestamp_us,nvtx_depth,nvtx_range,host_current,host_total,pinned_current,pinned_total,managed_current,managed_total,device_current,device_total,workspace_current,workspace_total,large_workspace_current,large_workspace_total
198809,1,"hnsw::build<ACE>",20008,20008,0,0,0,0,148304,148304,0,0,0,0
199961,1,"hnsw::build<ACE>",20008,20008,0,0,0,0,15588304,15588304,0,0,0,0
201350,1,"hnsw::build<ACE>",0,20008,0,0,0,0,0,40385488,0,0,0,0
222216,3,"cagra::build_knn_graph<IVF-PQ>(5000000, 1536, 72)",1440000000,1440020008,0,0,0,0,0,40385488,0,0,0,0
273892,4,"ivf_pq::build(5000000, 1536)",1440020008,1440040016,0,0,0,0,40385488,80770976,0,0,0,0
304183,4,"ivf_pq::build(5000000, 1536)",1440020008,1440040016,0,0,0,0,40385488,80770976,0,0,4388567040,4388567040
309064,4,"ivf_pq::build(5000000, 1536)",1440020008,1440040016,0,0,0,0,53860384,94245872,0,0,4388567040,4388567040
334655,4,"ivf_pq::build(5000000, 1536)",1440020008,1440040016,0,0,0,0,67339295,107724783,0,0,4388567040,4388567040
385037,4,"ivf_pq::build(5000000, 1536)",1440020008,1440040016,0,0,0,0,74076743,114462231,0,0,4388567040,4388567040
386129,4,"ivf_pq::build(5000000, 1536)",1440020008,1440040016,0,0,0,0,80814199,121199687,0,0,4388567040,4388567040
402750,4,"ivf_pq::build(5000000, 1536)",1440020008,1440040016,0,0,0,0,46099768,126913967,0,0,4388567040,4388567040
...
```
This can later be visualized (the visualization script is not included in the PR):
<img width="2100" height="1350" alt="allocations" src="https://github.com/user-attachments/assets/3f0ab942-b49b-4e09-a0ea-9181725ae05e" />

#### Implementation overview

##### NVTX

Added thread-local tracking of NVTX range stack; the calling thread shares a handle to the sampling thread to correlate the NVTX range state with allocations.

##### Memory resource adaptors

- statistics adaptor: atomically counts allocations/deallocations for any `cuda::mr`-compatible resource
- notifying adaptor: sets a shared "notifier" state on each event

##### Resource monitor

A resource monitor registers a collection of resource statistics objects, a single NVTX range handle, and a single notifier state. It spawns a new thread to sample the resource statistics at a given rate (but only when the notifier is triggered). This thread writes to a CSV output stream.

##### Memory tracking resources

`raft::memory_tracking_resources` is a child of `raft::resources`, thus can be used as a drop-in replacement. It replaces all known memory resource for the duration of its lifetime and manages the output file or stream if necessary.


Depends on (and includes all changes of) rapidsai#2968

Authors:
  - Artem M. Chirkin (https://github.com/achirkin)

Approvers:
  - Tamas Bela Feher (https://github.com/tfeher)

URL: rapidsai#2973

Author: achirkin

cpp/bench/prims/CMakeLists.txt

@@ -1,6 +1,6 @@
 # =============================================================================
 # cmake-format: off
-# SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 # cmake-format: on
 # =============================================================================
@@ -67,7 +67,7 @@ function(ConfigureBench)
 endfunction()
 
 if(BUILD_PRIMS_BENCH)
-  ConfigureBench(NAME CORE_BENCH PATH core/bitset.cu core/copy.cu main.cpp)
+  ConfigureBench(NAME CORE_BENCH PATH core/bitset.cu core/copy.cu core/memory_tracking.cu main.cpp)
 
   ConfigureBench(NAME UTIL_BENCH PATH util/popc.cu main.cpp)
 

cpp/bench/prims/core/memory_tracking.cu

@@ -0,0 +1,129 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <common/benchmark.hpp>
+
+#include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
+#include <raft/core/resources.hpp>
+#include <raft/util/memory_tracking_resources.hpp>
+
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/resource_ref.hpp>
+
+#include <unistd.h>
+
+#include <chrono>
+#include <cstdlib>
+#include <filesystem>
+#include <memory>
+#include <vector>
+
+namespace raft::bench::core {
+
+struct tracking_inputs {
+  int num_allocs;
+  size_t alloc_size;
+  int64_t sample_rate_us;
+  bool batch;
+};
+
+struct tracking_overhead : public fixture {
+  tracking_overhead(const tracking_inputs& p) : fixture(true), params(p)
+  {
+    if (p.sample_rate_us >= 0) {
+      std::string tpl = (std::filesystem::temp_directory_path() / "raft_bench_XXXXXX").string();
+      int fd          = mkstemp(tpl.data());
+      if (fd != -1) close(fd);
+      tmp_path_ = std::move(tpl);
+      tracked_res_.emplace(handle, tmp_path_, std::chrono::microseconds{p.sample_rate_us});
+    }
+  }
+
+  ~tracking_overhead()
+  {
+    tracked_res_.reset();
+    if (!tmp_path_.empty()) { std::remove(tmp_path_.c_str()); }
+  }
+
+  void run_benchmark(::benchmark::State& state) override
+  {
+    state.counters["alloc_size"]     = params.alloc_size;
+    state.counters["sample_rate_us"] = params.sample_rate_us;
+    state.counters["batch"]          = params.batch;
+
+    run_allocs(state, tracked_res_ ? reinterpret_cast<raft::resources&>(*tracked_res_) : handle);
+
+    state.SetItemsProcessed(state.iterations() * params.num_allocs * 2);
+  }
+
+ private:
+  void run_allocs(::benchmark::State& state, raft::resources& res)
+  {
+    auto mr = raft::resource::get_workspace_resource_ref(res);
+    auto sv = raft::resource::get_cuda_stream(res);
+
+    if (params.batch) {
+      std::vector<void*> ptrs(params.num_allocs);
+      for (auto _ : state) {
+        auto t0 = std::chrono::high_resolution_clock::now();
+        for (int i = 0; i < params.num_allocs; i++)
+          ptrs[i] = mr.allocate(sv, params.alloc_size);
+        for (int i = params.num_allocs - 1; i >= 0; i--)
+          mr.deallocate(sv, ptrs[i], params.alloc_size);
+        state.SetIterationTime(
+          std::chrono::duration<double>(std::chrono::high_resolution_clock::now() - t0).count());
+      }
+    } else {
+      for (auto _ : state) {
+        auto t0 = std::chrono::high_resolution_clock::now();
+        for (int i = 0; i < params.num_allocs; i++) {
+          void* p = mr.allocate(sv, params.alloc_size);
+          mr.deallocate(sv, p, params.alloc_size);
+        }
+        state.SetIterationTime(
+          std::chrono::duration<double>(std::chrono::high_resolution_clock::now() - t0).count());
+      }
+    }
+  }
+
+  tracking_inputs params;
+  std::string tmp_path_;
+  std::optional<raft::memory_tracking_resources> tracked_res_ = std::nullopt;
+};
+
+const std::vector<tracking_inputs> inputs{
+  // ping-pong (isolates per-call overhead, pool recycles same block)
+  {10000, 256, -1, false},
+  {10000, 256, 0, false},
+  {10000, 256, 1, false},
+  {10000, 256, 10, false},
+  {10000, 256, 100, false},
+  {10000, 1 << 20, -1, false},
+  {10000, 1 << 20, 0, false},
+  {10000, 1 << 20, 1, false},
+  {10000, 1 << 20, 10, false},
+  {10000, 1 << 20, 100, false},
+  {1000, 1 << 26, -1, false},
+  {1000, 1 << 26, 0, false},
+  {1000, 1 << 26, 1, false},
+  {1000, 1 << 26, 10, false},
+  {1000, 1 << 26, 100, false},
+  // batch (allocate all, then deallocate all)
+  {10000, 256, -1, true},
+  {10000, 256, 0, true},
+  {10000, 256, 1, true},
+  {10000, 256, 10, true},
+  {10000, 256, 100, true},
+  {1000, 1 << 20, -1, true},
+  {1000, 1 << 20, 0, true},
+  {1000, 1 << 20, 1, true},
+  {1000, 1 << 20, 10, true},
+  {1000, 1 << 20, 100, true},
+};
+
+RAFT_BENCH_REGISTER(tracking_overhead, "", inputs);
+
+}  // namespace raft::bench::core

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

@@ -1,10 +1,12 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #pragma once
 
+#include <raft/core/detail/nvtx_range_stack.hpp>
+
 #include <rmm/cuda_stream_view.hpp>
 
 #ifdef NVTX_ENABLED
@@ -146,6 +148,7 @@ 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);
 }
 
 template <typename Domain, typename... Args>
@@ -168,12 +171,13 @@ inline void push_range(const char* format, Args... args)
 template <typename Domain>
 inline void pop_range()
 {
+  detail::range_name_stack_instance.pop();
   nvtxDomainRangePop(domain_store<Domain>::value());
 }
 
 }  // namespace raft::common::nvtx::detail
 
-#else  // NVTX_ENABLED
+#else   // NVTX_ENABLED
 
 namespace raft::common::nvtx::detail {
 
@@ -188,5 +192,4 @@ inline void pop_range()
 }
 
 }  // namespace raft::common::nvtx::detail
-
 #endif  // NVTX_ENABLED

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

@@ -0,0 +1,90 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#pragma once
+
+#include <cstddef>
+#include <memory>
+#include <mutex>
+#include <stack>
+#include <string>
+#include <utility>
+
+namespace raft::common::nvtx {
+
+namespace detail {
+struct nvtx_range_name_stack;
+}  // namespace detail
+
+/**
+ * Shared, read-only handle to the current NVTX range name of a specific thread
+ * (set internally by one thread, read publicly by zero or more threads).
+ */
+class current_range {
+  friend detail::nvtx_range_name_stack;
+
+ public:
+  /** Read the current range name and stack depth (safe to call from any thread). */
+  auto get() const -> std::pair<std::string, std::size_t>
+  {
+    std::lock_guard lock(mu_);
+    return {value_, depth_};
+  }
+
+  operator std::string() const
+  {
+    std::lock_guard lock(mu_);
+    return value_;
+  }
+
+ private:
+  mutable std::mutex mu_;
+  std::string value_;
+  std::size_t depth_{0};
+
+  void set(const char* name, std::size_t depth)
+  {
+    std::lock_guard lock(mu_);
+    value_ = name ? name : "";
+    depth_ = depth;
+  }
+};
+
+namespace detail {
+
+struct nvtx_range_name_stack {
+  void push(const char* name)
+  {
+    stack_.emplace(name);
+    current_->set(name, stack_.size());
+  }
+
+  void pop()
+  {
+    if (!stack_.empty()) { stack_.pop(); }
+    current_->set(stack_.empty() ? nullptr : stack_.top().c_str(), stack_.size());
+  }
+
+  auto current() const -> std::shared_ptr<const current_range> { return current_; }
+
+ private:
+  std::stack<std::string> stack_{};
+  std::shared_ptr<current_range> current_{std::make_shared<current_range>()};
+};
+
+inline thread_local nvtx_range_name_stack range_name_stack_instance{};
+
+}  // namespace detail
+
+/**
+ * Get a read-only handle to this thread's current NVTX range name.
+ * 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>
+{
+  return detail::range_name_stack_instance.current();
+}
+
+}  // namespace raft::common::nvtx