Remove memory resources from Statistics

cmake/cpm_override_packages.json

@@ -0,0 +1,3 @@
+{
+  "packages": {}
+}

cpp/CMakeLists.txt

@@ -122,6 +122,10 @@ include(../cmake/Modules/ConfigureCUDA.cmake) # set other CUDA compilation flags
 
 # ##################################################################################################
 # * dependencies ----------------------------------------------------------------------------------
+
+include("${rapids-cmake-dir}/cpm/package_override.cmake")
+rapids_cpm_package_override("${CMAKE_CURRENT_SOURCE_DIR}/../cmake/cpm_override_packages.json")
+
 rapids_cpm_init()
 
 rapids_find_package(

cpp/benchmarks/bench_comm.cpp

@@ -400,7 +400,9 @@ int main(int argc, char** argv) {
            << "/s | num_ops: " << args.num_ops << " | nranks: " << comm->nranks();
         log->print(ss.str());
     }
-    log->print(stats->report("Statistics (of the last run):"));
+    log->print(
+        stats->report(std::nullopt, std::nullopt, "Statistics (of the last run):")
+    );
 
 #ifdef RAPIDSMPF_HAVE_CUPTI
     // Save CUPTI monitoring results to CSV file

cpp/benchmarks/bench_shuffle.cpp

@@ -300,7 +300,9 @@ rapidsmpf::Duration do_run(
     auto const t0_elapsed = rapidsmpf::Clock::now();
     {
         RAPIDSMPF_NVTX_SCOPED_RANGE("Shuffling", total_num_partitions);
-        RAPIDSMPF_MEMORY_PROFILE(statistics, "shuffling");
+        if (args.enable_memory_profiler) {
+            RAPIDSMPF_MEMORY_PROFILE(statistics, br->device_mr(), "shuffling");
+        }
         rapidsmpf::shuffler::Shuffler shuffler(
             comm,
             0,  // op_id
@@ -538,7 +540,7 @@ int main(int argc, char** argv) {
     rapidsmpf::config::Options options{rapidsmpf::config::get_environment_variables()};
 
     set_current_rmm_resource(args.rmm_mr);
-    auto stat_enabled_mr = set_device_mem_resource_with_stats();
+    rapidsmpf::RmmResourceAdaptor stat_enabled_mr = set_device_mem_resource_with_stats();
 
     std::unordered_map<rapidsmpf::MemoryType, rapidsmpf::BufferResource::MemoryAvailable>
         memory_available{};
@@ -548,11 +550,7 @@ int main(int argc, char** argv) {
         };
     }
 
-    // Create statistics (enabled from the start) and pass to BufferResource so that
-    // all components (Shuffler, SpillManager, etc.) share the same statistics object.
-    auto stats = args.enable_memory_profiler
-                     ? std::make_shared<rapidsmpf::Statistics>(stat_enabled_mr)
-                     : std::make_shared<rapidsmpf::Statistics>(/* enable = */ true);
+    auto stats = std::make_shared<rapidsmpf::Statistics>(/* enable = */ true);
 
     // We're only going to measure the last run, so disable initially.
     stats->disable();
@@ -682,7 +680,16 @@ int main(int argc, char** argv) {
         }
         log->print(ss.str());
     }
-    log->print(stats->report("Statistics (of the last run):"));
+
+    if (args.enable_memory_profiler) {
+        log->print(
+            stats->report(stat_enabled_mr, std::nullopt, "Statistics (of the last run):")
+        );
+    } else {
+        log->print(
+            stats->report(std::nullopt, std::nullopt, "Statistics (of the last run):")
+        );
+    }
 
 #ifdef RAPIDSMPF_HAVE_CUPTI
     // Save CUPTI monitoring results to CSV file

cpp/benchmarks/streaming/bench_streaming_shuffle.cpp

@@ -363,13 +363,14 @@ int main(int argc, char** argv) {
         };
     }
 
-    auto stats = args.enable_memory_profiler
-                     ? std::make_shared<rapidsmpf::Statistics>(stat_enabled_mr)
-                     : std::make_shared<rapidsmpf::Statistics>(/* enable = */ true);
+    auto stats = std::make_shared<rapidsmpf::Statistics>(/* enable = */ true);
+
+    auto pinned_mr = args.pinned_mem_disable
+                         ? rapidsmpf::PinnedMemoryResource::Disabled
+                         : rapidsmpf::PinnedMemoryResource::make_if_available();
     auto br = std::make_shared<rapidsmpf::BufferResource>(
         stat_enabled_mr,
-        args.pinned_mem_disable ? rapidsmpf::PinnedMemoryResource::Disabled
-                                : rapidsmpf::PinnedMemoryResource::make_if_available(),
+        pinned_mr,
         std::move(memory_available),
         std::nullopt,
         std::make_shared<rmm::cuda_stream_pool>(
@@ -455,7 +456,17 @@ int main(int argc, char** argv) {
         }
         log.print(ss.str());
     }
-    log.print(ctx->statistics()->report("Statistics (of the last run):"));
+
+    if (args.enable_memory_profiler) {
+        log.print(ctx->statistics()->report(
+            stat_enabled_mr, pinned_mr, "Statistics (of the last run):"
+        ));
+    } else {
+        log.print(ctx->statistics()->report(
+            std::nullopt, std::nullopt, "Statistics (of the last run):"
+        ));
+    }
+
     if (!use_bootstrap) {
         RAPIDSMPF_MPI(MPI_Finalize());
     }

cpp/benchmarks/streaming/ndsh/bench_read.cpp

@@ -413,7 +413,9 @@ int main(int argc, char** argv) {
         std::chrono::duration<double> compute = end - start;
         timings.push_back(pipeline.count());
         timings.push_back(compute.count());
-        comm->logger()->print(ctx->statistics()->report());
+        comm->logger()->print(
+            ctx->statistics()->report(ctx->br()->device_mr(), ctx->br()->try_pinned_mr())
+        );
         ctx->statistics()->clear();
     }
 

cpp/benchmarks/streaming/ndsh/q01.cpp

@@ -435,7 +435,9 @@ int main(int argc, char** argv) {
         std::chrono::duration<double> compute = end - start;
         timings.push_back(pipeline.count());
         timings.push_back(compute.count());
-        comm->logger()->print(ctx->statistics()->report());
+        comm->logger()->print(
+            ctx->statistics()->report(ctx->br()->device_mr(), ctx->br()->try_pinned_mr())
+        );
         ctx->statistics()->clear();
     }
 

cpp/benchmarks/streaming/ndsh/q03.cpp

@@ -697,7 +697,9 @@ int main(int argc, char** argv) {
         std::chrono::duration<double> compute = end - start;
         timings.push_back(pipeline.count());
         timings.push_back(compute.count());
-        comm->logger()->print(ctx->statistics()->report());
+        comm->logger()->print(
+            ctx->statistics()->report(ctx->br()->device_mr(), ctx->br()->try_pinned_mr())
+        );
         ctx->statistics()->clear();
     }
     if (comm->rank() == 0) {

cpp/benchmarks/streaming/ndsh/q04.cpp

@@ -520,7 +520,9 @@ int main(int argc, char** argv) {
         std::chrono::duration<double> compute = end - start;
         timings.push_back(pipeline.count());
         timings.push_back(compute.count());
-        comm->logger()->print(ctx->statistics()->report());
+        comm->logger()->print(
+            ctx->statistics()->report(ctx->br()->device_mr(), ctx->br()->try_pinned_mr())
+        );
         ctx->statistics()->clear();
     }
 

cpp/benchmarks/streaming/ndsh/q09.cpp

@@ -655,7 +655,9 @@ int main(int argc, char** argv) {
         std::chrono::duration<double> compute = end - start;
         timings.push_back(pipeline.count());
         timings.push_back(compute.count());
-        comm->logger()->print(ctx->statistics()->report());
+        comm->logger()->print(
+            ctx->statistics()->report(ctx->br()->device_mr(), ctx->br()->try_pinned_mr())
+        );
         ctx->statistics()->clear();
     }
     if (comm->rank() == 0) {

cpp/benchmarks/streaming/ndsh/q21.cpp

@@ -997,7 +997,9 @@ int main(int argc, char** argv) {
         std::chrono::duration<double> compute = end - start;
         timings.push_back(pipeline.count());
         timings.push_back(compute.count());
-        comm->logger()->print(ctx->statistics()->report());
+        comm->logger()->print(
+            ctx->statistics()->report(ctx->br()->device_mr(), ctx->br()->try_pinned_mr())
+        );
         ctx->statistics()->clear();
     }
 

cpp/benchmarks/streaming/ndsh/utils.cpp

@@ -143,7 +143,7 @@ create_context(ProgramOptions& arguments, RmmResourceAdaptor&& mr) {
         memory_available[MemoryType::DEVICE] =
             LimitAvailableMemory{mr, static_cast<std::int64_t>(limit_size)};
     }
-    auto statistics = std::make_shared<Statistics>(mr);
+    auto statistics = std::make_shared<Statistics>(/* enable = */ true);
 
     RAPIDSMPF_EXPECTS(
         arguments.no_pinned_host_memory || is_pinned_memory_resources_supported(),

cpp/include/rapidsmpf/memory/buffer_resource.hpp

@@ -22,6 +22,7 @@
 #include <rapidsmpf/memory/host_memory_resource.hpp>
 #include <rapidsmpf/memory/memory_reservation.hpp>
 #include <rapidsmpf/memory/pinned_memory_resource.hpp>
+#include <rapidsmpf/memory/resource_types.hpp>
 #include <rapidsmpf/memory/spill_manager.hpp>
 #include <rapidsmpf/rmm_resource_adaptor.hpp>
 #include <rapidsmpf/statistics.hpp>
@@ -131,9 +132,18 @@ class BufferResource {
     /**
      * @brief Get the RMM pinned host memory resource.
      *
+     * @throws std::invalid_argument if no pinned memory resource is available.
      * @return Reference to the RMM resource used for pinned host allocations.
      */
-    [[nodiscard]] rmm::host_async_resource_ref pinned_mr();
+    [[nodiscard]] rmm::host_device_async_resource_ref pinned_mr();
+
+    /**
+     * @brief Get the pinned host memory resource if available.
+     *
+     * @return The pinned host memory resource as an `any_resource`, or `std::nullopt` if
+     * pinned host memory is not available.
+     */
+    [[nodiscard]] std::optional<any_host_device_resource> try_pinned_mr() const noexcept;
 
     /**
      * @brief Retrieves the memory availability function for a given memory type.

cpp/include/rapidsmpf/memory/resource_types.hpp

@@ -0,0 +1,19 @@
+/**
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <cuda/memory_resource>
+
+namespace rapidsmpf {
+
+/// @brief Owning type-erased device memory resource.
+using any_device_resource = cuda::mr::any_resource<cuda::mr::device_accessible>;
+
+/// @brief Owning type-erased host- and device-accessible memory resource.
+using any_host_device_resource =
+    cuda::mr::any_resource<cuda::mr::host_accessible, cuda::mr::device_accessible>;
+
+}  // namespace rapidsmpf