Merge RmmResourceAdaptor into BufferResource as a CCCL resource

cpp/CMakeLists.txt

@@ -206,7 +206,6 @@ add_library(
   src/memory/spill_manager.cpp
   src/pausable_thread_loop.cpp
   src/progress_thread.cpp
-  src/rmm_resource_adaptor.cpp
   src/rrun/rrun.cpp
   src/shuffler/chunk.cpp
   src/shuffler/finish_counter.cpp

cpp/benchmarks/bench_shuffle.cpp

@@ -536,7 +536,7 @@ int main(int argc, char** argv) {
     rapidsmpf::config::Options options{rapidsmpf::config::get_environment_variables()};
 
     set_current_rmm_resource(args.rmm_mr);
-    rapidsmpf::RmmResourceAdaptor stat_enabled_mr = set_device_mem_resource_with_stats();
+    auto stat_enabled_mr = set_device_mem_resource_with_stats();
 
     std::unordered_map<rapidsmpf::MemoryType, std::int64_t> memory_limits{};
     if (args.device_mem_limit_mb >= 0) {
@@ -548,7 +548,7 @@ int main(int argc, char** argv) {
     // We're only going to measure the last run, so disable initially.
     stats->disable();
     rapidsmpf::BufferResource br{
-        stat_enabled_mr,
+        std::move(stat_enabled_mr),
         args.pinned_mem_disable ? rapidsmpf::PinnedMemoryResource::Disabled
                                 : rapidsmpf::PinnedMemoryResource::make_if_available(),
         std::move(memory_limits),
@@ -663,7 +663,7 @@ int main(int argc, char** argv) {
            << " | out_parts: " << args.num_output_partitions
            << " | nranks: " << comm->nranks();
         if (args.enable_memory_profiler) {
-            auto record = stat_enabled_mr.get_main_record();
+            auto record = br.get_main_record();
             ss << " | device memory peak: " << rapidsmpf::format_nbytes(record.peak())
                << " | device memory total: "
                << rapidsmpf::format_nbytes(
@@ -675,9 +675,7 @@ int main(int argc, char** argv) {
     }
 
     if (args.enable_memory_profiler) {
-        log->print(stats->report(
-            {.mr = stat_enabled_mr, .header = "Statistics (of the last run):"}
-        ));
+        log->print(stats->report({.mr = br, .header = "Statistics (of the last run):"}));
     } else {
         log->print(stats->report({.header = "Statistics (of the last run):"}));
     }

cpp/benchmarks/streaming/bench_streaming_shuffle.cpp

@@ -370,7 +370,7 @@ int main(int argc, char** argv) {
                          ? rapidsmpf::PinnedMemoryResource::Disabled
                          : rapidsmpf::PinnedMemoryResource::make_if_available();
     auto br = std::make_shared<rapidsmpf::BufferResource>(
-        stat_enabled_mr,
+        std::move(stat_enabled_mr),
         pinned_mr,
         std::move(memory_limits),
         std::nullopt,
@@ -447,7 +447,7 @@ int main(int argc, char** argv) {
            << " | out_parts: " << args.num_output_partitions
            << " | nranks: " << comm->nranks();
         if (args.enable_memory_profiler) {
-            auto record = stat_enabled_mr.get_main_record();
+            auto record = br->get_main_record();
             ss << " | device memory peak: " << rapidsmpf::format_nbytes(record.peak())
                << " | device memory total: "
                << rapidsmpf::format_nbytes(
@@ -461,7 +461,7 @@ int main(int argc, char** argv) {
     auto statistics = ctx->statistics();
     if (args.enable_memory_profiler) {
         log.print(statistics->report({
-            .mr = stat_enabled_mr,
+            .mr = *br,
             .pinned_mr = pinned_mr,
             .header = "Statistics (of the last run):",
         }));

cpp/benchmarks/streaming/ndsh/q21.cpp

@@ -30,7 +30,6 @@
 #include <rapidsmpf/communicator/communicator.hpp>
 #include <rapidsmpf/integrations/cudf/bloom_filter.hpp>
 #include <rapidsmpf/nvtx.hpp>
-#include <rapidsmpf/rmm_resource_adaptor.hpp>
 #include <rapidsmpf/streaming/coll/allgather.hpp>
 #include <rapidsmpf/streaming/core/actor.hpp>
 #include <rapidsmpf/streaming/core/channel.hpp>

cpp/benchmarks/streaming/ndsh/utils.cpp

@@ -33,7 +33,6 @@
 #include <rapidsmpf/error.hpp>
 #include <rapidsmpf/memory/buffer_resource.hpp>
 #include <rapidsmpf/memory/pinned_memory_resource.hpp>
-#include <rapidsmpf/rmm_resource_adaptor.hpp>
 #include <rapidsmpf/streaming/core/context.hpp>
 #include <rapidsmpf/streaming/cudf/table_chunk.hpp>
 

cpp/benchmarks/utils/rmm_utils.hpp

@@ -6,14 +6,15 @@
 
 #include <string>
 
+#include <cuda/memory_resource>
+
 #include <rmm/mr/cuda_async_memory_resource.hpp>
 #include <rmm/mr/cuda_memory_resource.hpp>
 #include <rmm/mr/managed_memory_resource.hpp>
 #include <rmm/mr/per_device_resource.hpp>
 #include <rmm/mr/pool_memory_resource.hpp>
 
 #include <rapidsmpf/error.hpp>
-#include <rapidsmpf/rmm_resource_adaptor.hpp>
 
 /**
  * @brief Create and set a RMM memory resource as the current device resource.
@@ -45,13 +46,18 @@ inline void set_current_rmm_resource(std::string const& name) {
 }
 
 /**
- * @brief Create a statistics-enabled device memory resource wrapping the current
- * device resource, and set it as the current device resource.
+ * @brief Return the current device resource as a CCCL `any_resource`.
+ *
+ * Compatibility shim for benchmarks that previously wrapped the current device
+ * resource in `RmmResourceAdaptor` to gain statistics. Tracking is now part of
+ * `BufferResource` itself, so callers pass the returned `any_resource` directly
+ * to a `BufferResource` constructor.
  *
- * @return A RmmResourceAdaptor (shared ownership) for accessing statistics.
+ * @return The current device resource as a type-erased CCCL resource.
  */
-[[nodiscard]] inline rapidsmpf::RmmResourceAdaptor set_device_mem_resource_with_stats() {
-    rapidsmpf::RmmResourceAdaptor adaptor{rmm::mr::get_current_device_resource_ref()};
-    rmm::mr::set_current_device_resource(adaptor);
-    return adaptor;
+[[nodiscard]] inline cuda::mr::any_resource<cuda::mr::device_accessible>
+set_device_mem_resource_with_stats() {
+    return cuda::mr::any_resource<cuda::mr::device_accessible>{
+        rmm::mr::get_current_device_resource_ref()
+    };
 }

cpp/include/rapidsmpf/detail/rmm_resource_adaptor_impl.hpp

@@ -29,11 +29,12 @@
 namespace rapidsmpf::detail {
 
 /**
- * @brief Implementation class for RmmResourceAdaptor.
+ * @brief Implementation class for instrumented RMM memory resources.
  *
  * Holds all mutable state for memory tracking. This class satisfies the CCCL
- * `cuda::mr::resource` concept and is held by `RmmResourceAdaptor` via
- * `cuda::mr::shared_resource` for reference-counted ownership.
+ * `cuda::mr::resource` concept and is the building block used internally by
+ * both `BufferResource` (for device memory tracking) and `PinnedMemoryResource`
+ * (for pinned-host tracking with in-place storage of `cuda::pinned_memory_pool`).
  *
  * @tparam PrimaryMR The type of the primary memory resource. Use a concrete
  * resource type (e.g. `cuda::pinned_memory_pool`) to store the resource
@@ -96,25 +97,37 @@ class RmmResourceAdaptorImpl {
         return primary_mr_;
     }
 
-    /// @copydoc RmmResourceAdaptor::get_main_record
+    /**
+     * @brief Returns a copy of the main memory record (lifetime-of-resource stats).
+     *
+     * @return A copy of the main `ScopedMemoryRecord`.
+     */
     [[nodiscard]] ScopedMemoryRecord get_main_record() const {
         std::lock_guard<std::mutex> lock(mutex_);
         return main_record_;
     }
 
-    /// @copydoc RmmResourceAdaptor::current_allocated
+    /**
+     * @brief Total number of currently allocated bytes.
+     *
+     * @return Currently outstanding allocated bytes.
+     */
     [[nodiscard]] std::int64_t current_allocated() const noexcept {
         std::lock_guard<std::mutex> lock(mutex_);
         return main_record_.current();
     }
 
-    /// @copydoc RmmResourceAdaptor::begin_scoped_memory_record
+    /// @brief Push a new scoped memory record onto the current thread's stack.
     void begin_scoped_memory_record() {
         std::lock_guard<std::mutex> lock(mutex_);
         record_stacks_[std::this_thread::get_id()].emplace();
     }
 
-    /// @copydoc RmmResourceAdaptor::end_scoped_memory_record
+    /**
+     * @brief Pop and return the topmost scoped memory record on the current thread.
+     *
+     * @return The popped `ScopedMemoryRecord`.
+     */
     ScopedMemoryRecord end_scoped_memory_record() {
         std::lock_guard lock(mutex_);
         auto& stack = record_stacks_.at(std::this_thread::get_id());

cpp/include/rapidsmpf/memory/buffer.hpp

@@ -24,6 +24,10 @@
 
 namespace rapidsmpf {
 
+namespace detail {
+class BufferResourceImpl;
+}  // namespace detail
+
 /**
  * @brief Buffer representing device or host memory.
  *
@@ -46,6 +50,7 @@ namespace rapidsmpf {
  */
 class Buffer {
     friend class BufferResource;
+    friend class detail::BufferResourceImpl;
 
   public:
     /// @brief Storage type for a device buffer.