diff --git a/cpp/include/raft/core/detail/nvtx.hpp b/cpp/include/raft/core/detail/nvtx.hpp index 334e18d665..2efa4bf6d1 100644 --- a/cpp/include/raft/core/detail/nvtx.hpp +++ b/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 +#include #include #include @@ -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::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 @@ -174,6 +176,7 @@ template inline void pop_range() { detail::range_name_stack_instance.pop(); + detail::full_range_stack_instance.pop(); nvtxDomainRangePop(domain_store::value()); } diff --git a/cpp/include/raft/core/detail/nvtx_range_path_stack.hpp b/cpp/include/raft/core/detail/nvtx_range_path_stack.hpp new file mode 100644 index 0000000000..5fb5901f37 --- /dev/null +++ b/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 +#include // RAFT_LOG_WARN + +#include +#include +#include +#include +#include +#include +#include + +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 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 + { + 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> 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 +{ + 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 diff --git a/cpp/include/raft/core/memory_logging_resources.hpp b/cpp/include/raft/core/memory_logging_resources.hpp new file mode 100644 index 0000000000..4785aacf47 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include + +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(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 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(owned_stream_.get()); } + RAFT_LOG_INFO("memory_logging_resources: queue-based recording (every event captured)"); + recorder_ = std::make_unique(*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> snapshot_; + std::unique_ptr owned_stream_; + std::unique_ptr recorder_; + + raft::mr::host_resource old_host_; + raft::mr::device_resource old_device_; + + using host_record_t = raft::mr::recording_adaptor; + std::unique_ptr host_adaptor_; + + using device_record_t = raft::mr::recording_adaptor; + std::unique_ptr 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(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{ + pinned_ref, queue, source_id}); + } + + // --- Managed --- + { + int source_id = recorder_->register_source("managed"); + raft::resource::set_managed_memory_resource( + *this, + raft::mr::recording_adaptor{ + 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(); + int source_id = recorder_->register_source("device"); + device_adaptor_ = std::make_unique(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{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{lws_ref, queue, source_id}); + } + + recorder_->start(); + } +}; + +} // namespace raft diff --git a/cpp/include/raft/mr/recording_adaptor.hpp b/cpp/include/raft/mr/recording_adaptor.hpp new file mode 100644 index 0000000000..2239273a81 --- /dev/null +++ b/cpp/include/raft/mr/recording_adaptor.hpp @@ -0,0 +1,158 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include // thread_local_nvtx_full_path, thread_local_inner_range_and_depth +#include +#include // allocation_event, allocation_event_queue + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace raft { +namespace mr { + +/** + * @brief Resource adaptor that records each allocation/deallocation as an event, + * and associates it with the active NVTX range AT THE TIME OF THE EVENT. + */ +template +class recording_adaptor : public cuda::forward_property, Upstream> { + // Map an allocated address to the nvtx stack range responsible for the allocation. + // It allows the deallocation event to be tagged with the same range, even if the responsible + // range has ended by the time of deallocation. + // Map is shared across threads (e.g. one thread allocates and another deallocates) + struct address_range_map { + std::mutex mtx; + std::unordered_map paths; + }; + + Upstream upstream_; + std::shared_ptr queue_; + std::shared_ptr alloc_map_; + int source_id_; + std::shared_ptr> current_bytes_; + + // Record the alloc-time NVTX path for this pointer. + // Called on the allocating thread — mutex-free NVTX read is safe. + // The mutex protects the shared map from concurrent alloc/dealloc threads. + auto record_allocation(void* ptr) -> std::string + { + std::string path = ""; + if (ptr != nullptr) { + path = raft::common::nvtx::thread_local_nvtx_full_path(); + if (!path.empty()) { + std::lock_guard lock(alloc_map_->mtx); + alloc_map_->paths[ptr] = path; + } + } + return path; + } + + // Returns the NVTX path recorded at alloc time for this pointer, then removes it. + auto forget_allocation(void* ptr) -> std::string + { + std::string path = ""; + std::lock_guard lock(alloc_map_->mtx); + auto it = alloc_map_->paths.find(ptr); + if (it != alloc_map_->paths.end()) { + path = std::move(it->second); + alloc_map_->paths.erase(it); + } + return path; + } + + // Enqueue an event. This is called on the allocating/deallocating thread + void emit(std::string nvtx_full_range, + std::int64_t signed_bytes, + std::chrono::steady_clock::time_point const& timestamp) + { + auto [name, depth] = raft::common::nvtx::thread_local_inner_range_and_depth(); + allocation_event event{ + .timestamp = timestamp, + .source_id = source_id_, + .current_bytes = + current_bytes_->fetch_add(signed_bytes, std::memory_order_relaxed) + signed_bytes, + .delta_bytes = signed_bytes, + .nvtx_depth = depth, + .nvtx_inner_range = std::move(name), + .nvtx_full_range = std::move(nvtx_full_range), + }; + queue_->push(std::move(event)); + } + + public: + recording_adaptor(Upstream upstream, std::shared_ptr queue, int source_id) + : upstream_(std::move(upstream)), + queue_(std::move(queue)), + alloc_map_(std::make_shared()), + source_id_(source_id), + current_bytes_{std::make_shared>(0)} + { + RAFT_EXPECTS(queue_ != nullptr, "event queue must be initialized"); + } + + void* allocate_sync(std::size_t bytes, std::size_t alignment = alignof(std::max_align_t)) + { + std::chrono::steady_clock::time_point timestamp = std::chrono::steady_clock::now(); + void* ptr = upstream_.allocate_sync(bytes, alignment); + emit(record_allocation(ptr), static_cast(bytes), timestamp); + return ptr; + } + + void deallocate_sync(void* ptr, + std::size_t bytes, + std::size_t alignment = alignof(std::max_align_t)) + { + std::chrono::steady_clock::time_point timestamp = std::chrono::steady_clock::now(); + upstream_.deallocate_sync(ptr, bytes, alignment); + emit(forget_allocation(ptr), -static_cast(bytes), timestamp); + } + + template , int> = 0> + void* allocate(cuda::stream_ref stream, + std::size_t bytes, + std::size_t alignment = alignof(std::max_align_t)) + { + std::chrono::steady_clock::time_point timestamp = std::chrono::steady_clock::now(); + void* ptr = upstream_.allocate(stream, bytes, alignment); + emit(record_allocation(ptr), static_cast(bytes), timestamp); + return ptr; + } + + template , int> = 0> + void deallocate(cuda::stream_ref stream, + void* ptr, + std::size_t bytes, + std::size_t alignment = alignof(std::max_align_t)) + { + std::chrono::steady_clock::time_point timestamp = std::chrono::steady_clock::now(); + upstream_.deallocate(stream, ptr, bytes, alignment); + emit(forget_allocation(ptr), -static_cast(bytes), timestamp); + } + + [[nodiscard]] bool operator==(recording_adaptor const& other) const noexcept + { + return upstream_ == other.upstream_; + } + + [[nodiscard]] auto upstream_resource() noexcept -> Upstream& { return upstream_; } + [[nodiscard]] auto upstream_resource() const noexcept -> Upstream const& { return upstream_; } +}; + +} // namespace mr +} // namespace raft diff --git a/cpp/include/raft/mr/recording_monitor.hpp b/cpp/include/raft/mr/recording_monitor.hpp new file mode 100644 index 0000000000..6c93dda98a --- /dev/null +++ b/cpp/include/raft/mr/recording_monitor.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 +#include // RAFT_EXPECTS +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace raft { +namespace mr { + +/** + * @brief A single allocation or deallocation event, captured on the allocating thread. + */ +struct allocation_event { + std::chrono::steady_clock::time_point timestamp{}; //< (de)allocation timestamp + int source_id{0}; //< resource (host, pinned, workspace, etc.) which event belongs to + std::int64_t current_bytes{0}; //< live bytes after this event (i.e. incl. delta bytes) + std::int64_t delta_bytes{0}; //< signed delta bytes by this event (+alloc / -free) + std::size_t nvtx_depth{0}; //< NVTX stack depth at event time + std::string nvtx_inner_range{}; //< NVTX range name active at event time + std::string + nvtx_full_range{}; //< NVTX full path "name#id -> ..." captured at the ALLOCATION time +}; + +/** + * @brief Thread-safe multi-producer / single-consumer queue of allocation_events. + */ +class allocation_event_queue { + public: + /** @brief Append an event (any thread). */ + void push(allocation_event event) + { + { + std::lock_guard lock(mtx_); + events_.push_back(std::move(event)); + } + cv_.notify_one(); + } + + /** + * @brief Block until events are available or the queue is stopped, then move + * all pending events into `out`. + * + * @return true once the queue is stopped AND drained (consumer should exit) + */ + bool wait_and_take(std::vector& out) + { + std::unique_lock lock(mtx_); + cv_.wait(lock, [this] { return stopped_ || !events_.empty(); }); + out.clear(); + out.swap(events_); + return stopped_ && out.empty(); + } + + /** @brief Signal the consumer to drain and exit. */ + void stop() + { + { + std::lock_guard lock(mtx_); + stopped_ = true; + } + cv_.notify_all(); + } + + private: + std::mutex mtx_; + std::condition_variable cv_; + std::vector events_; + bool stopped_{false}; +}; + +/** + * @brief Consumes allocation_events from a queue and writes one CSV row per + * event from a background thread. + */ +class recording_monitor { + public: + explicit recording_monitor(std::ostream& out) : out_(out) {} + + ~recording_monitor() { stop(); } + + recording_monitor(recording_monitor const&) = delete; + recording_monitor& operator=(recording_monitor const&) = delete; + + [[nodiscard]] auto get_queue() const noexcept -> std::shared_ptr + { + return queue_; + } + + /** + * @brief Register a named source and return its id (column-group index). + * Must be called before start(). + */ + auto register_source(std::string name) -> int + { + RAFT_EXPECTS(name.find(',') == std::string::npos, + "source name must not contain ',' (delimiter). This would break CSV columns: '%s'", + name.c_str()); + int source_id = static_cast(sources_.size()); + sources_.push_back(std::move(name)); + source_current_.push_back(0); // last-known live bytes for this source (carried forward) + return source_id; + } + + void start() + { + if (worker_.joinable()) { return; } + write_header(); + // Start the background thread that consumes events from the queue + // and writes one CSV row per event. + worker_ = std::thread([this] { run(); }); + } + + void stop() + { + if (!worker_.joinable()) { return; } + queue_->stop(); // drains the queue and causes the worker to exit its loop + worker_.join(); + } + + private: + void write_header() + { + out_ << "timestamp_us,source"; + for (auto const& name : sources_) { + out_ << ',' << name << "_current_bytes"; + } + out_ << ",delta_bytes"; + out_ << ",nvtx_depth,nvtx_inner_range,nvtx_full_range\n"; + out_.flush(); + } + + void run() + { + std::vector batch; + while (true) { + bool finished = queue_->wait_and_take(batch); + for (auto const& event : batch) { + write_row(event); + } + out_.flush(); + if (finished) { break; } + } + } + + void write_row(allocation_event const& event) + { + if (event.source_id < 0 || static_cast(event.source_id) >= sources_.size()) { + RAFT_LOG_WARN("Event source id %d is out-of-bound (number of sources = %zu)", + event.source_id, + sources_.size()); + return; + } + + // timestamp since start [us] + out_ << std::chrono::duration_cast(event.timestamp - start_).count(); + out_ << "," << sources_[event.source_id]; + // live bytes per source (last-known value for each) + source_current_[event.source_id] = event.current_bytes; + for (auto const& current_bytes : source_current_) { + out_ << "," << current_bytes; + } + // delta bytes + out_ << "," << event.delta_bytes; + // nvtx + out_ << "," << event.nvtx_depth; + out_ << ",\"" << event.nvtx_inner_range << "\""; + out_ << ",\"" << event.nvtx_full_range << "\"\n"; + } + + std::ostream& out_; + std::shared_ptr queue_{std::make_shared()}; + std::vector sources_; // pinned, workspace, host, etc. + std::vector source_current_; // last-known live bytes per source (carried forward) + std::chrono::steady_clock::time_point start_{std::chrono::steady_clock::now()}; + std::thread worker_; +}; + +} // namespace mr +} // namespace raft diff --git a/cpp/tests/core/monitor_resources.cu b/cpp/tests/core/monitor_resources.cu index c7be4bd285..043206ff6b 100644 --- a/cpp/tests/core/monitor_resources.cu +++ b/cpp/tests/core/monitor_resources.cu @@ -1,30 +1,40 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include +#include +#include #include +#include +#include +#include #include #include #include #include +#include +#include +#include #include #include #include namespace { -TEST(MemoryTrackingResources, TracksDeviceAllocations) -{ - using namespace std::chrono_literals; +namespace nvtx = raft::common::nvtx; +using namespace std::chrono_literals; +constexpr std::size_t MiB = std::size_t{1024} * 1024; +TEST(MemoryTrackingResources, SamplingSingleThread) +{ std::ostringstream oss; { raft::resources res; - raft::resource::set_workspace_to_pool_resource(res, 1024 * 1024); + raft::resource::set_workspace_to_pool_resource(res, 1 * MiB); raft::memory_tracking_resources tracked(res, oss, 1ms); @@ -49,4 +59,129 @@ TEST(MemoryTrackingResources, TracksDeviceAllocations) << output; } +TEST(MemoryLoggingResources, RecordingSingleThread) +{ + std::ostringstream oss; + { + raft::resources res; + raft::memory_logging_resources tracked(res, oss); + { + nvtx::range r{"1. expect 10 KB"}; + auto matrix = raft::make_host_vector(tracked, 10 * 1024); + } + { + nvtx::range r{"2. expect 100 MiB"}; + auto vector = raft::make_host_vector(tracked, 100 * MiB); + } + { + nvtx::range r{"3. expect 4 MiB"}; + auto matrix = raft::make_host_vector(tracked, 4 * MiB); + } + } // tracked destroyed: stops the recorder and flushes the CSV + + auto output = oss.str(); + EXPECT_NE(output.find("timestamp_us"), std::string::npos); + EXPECT_NE(output.find("source"), std::string::npos); + EXPECT_NE(output.find("host_current_bytes"), std::string::npos); + EXPECT_NE(output.find("device_current_bytes"), std::string::npos); + EXPECT_NE(output.find("workspace_current_bytes"), std::string::npos); + EXPECT_NE(output.find("delta_bytes"), std::string::npos); + EXPECT_NE(output.find("nvtx_depth"), std::string::npos); + EXPECT_NE(output.find("nvtx_inner_range"), std::string::npos); + EXPECT_NE(output.find("nvtx_full_range"), std::string::npos); + + auto num_lines = std::count(output.begin(), output.end(), '\n'); + constexpr size_t NUM_ALLOCS = 3; + constexpr size_t NUM_DEALLOCS = NUM_ALLOCS; + constexpr size_t NUM_HEADER_LINES = 1; + constexpr size_t NUM_LINES_EXPECTED = NUM_ALLOCS + NUM_DEALLOCS + NUM_HEADER_LINES; + EXPECT_GE(num_lines, NUM_LINES_EXPECTED) + << "Expected at least " << NUM_LINES_EXPECTED + << " data records (allocation + deallocation + header); got " << num_lines << " lines" + << std::endl + << "content: " << std::endl + << output; +} + +// --------------------------------------------------------------------------- +// Parallel-thread stress tests +// --------------------------------------------------------------------------- + +TEST(MemoryLoggingResources, RecordingParallelThreads) +{ + constexpr int kNumThreads = 64; + constexpr int kNumIters = 200; + constexpr std::size_t kAllocSize = 256 * 1024; // 256 KiB + + std::ostringstream oss; + { + raft::resources res; + raft::memory_logging_resources tracked(res, oss); + + // Each thread tags its allocations with a distinct NVTX range. + // NVTX range read is mutex-free: each allocating thread reads its own stack. + auto run = [&](int t) { + for (int i = 0; i < kNumIters; ++i) { + std::string label = std::string("thread-") + std::to_string(t); + nvtx::range r{label.c_str()}; + auto vec = raft::make_host_vector(tracked, kAllocSize); + } + }; + + std::vector workers; + for (int t = 0; t < kNumThreads; ++t) { + workers.emplace_back(run, t); + } + for (auto& w : workers) + w.join(); + } // tracked destroyed: drains queue, flushes CSV + + std::string output = oss.str(); + auto num_lines = std::count(output.begin(), output.end(), '\n'); + + // Each iteration: one alloc row + one dealloc row; plus one header line. + const int expected_rows = kNumThreads * kNumIters * 2 + 1; + EXPECT_GE(num_lines, expected_rows) + << "Expected at least " << expected_rows << " lines; got " << num_lines << "\noutput:\n" + << output; +} + +TEST(MemoryTrackingResources, SamplingParallelThreads) +{ + constexpr int kNumThreads = 64; + constexpr int kNumIters = 200; + constexpr std::size_t kAllocSize = 256 * 1024; // 256 KiB + + std::ostringstream oss; + { + raft::resources res; + raft::memory_tracking_resources tracked(res, oss, 1us); + + auto run = [&](int t) { + for (int i = 0; i < kNumIters; ++i) { + std::string label = std::string("thread-") + std::to_string(t); + nvtx::range r{label.c_str()}; + auto vec = raft::make_host_vector(tracked, kAllocSize); + } + }; + + std::vector workers; + for (int t = 0; t < kNumThreads; ++t) { + workers.emplace_back(run, t); + } + for (auto& w : workers) + w.join(); + } + + std::string output = oss.str(); + auto num_lines = std::count(output.begin(), output.end(), '\n'); + + // Sampling approach drops many rows + const int max_num_rows = kNumThreads * kNumIters * 2 + 1; + EXPECT_TRUE(3 < num_lines && num_lines < max_num_rows) + << "Expected at least several rows. It is expected when many rows are dropped; got " + << num_lines << "\noutput:\n" + << output; +} + } // namespace