Add a strict Buffer Resource back-reference to Host and Pinned memory resources

cpp/benchmarks/bench_comm.cpp

@@ -326,7 +326,7 @@ int main(int argc, char** argv) {
     rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource_ref();
     auto br = BufferResource::create(
         mr,
-        PinnedMemoryResource::Disabled,
+        PinnedMemoryDisabled,
         {},
         std::chrono::milliseconds{1},
         std::make_shared<rmm::cuda_stream_pool>(

cpp/benchmarks/bench_memory_resources.cpp

@@ -5,14 +5,18 @@
 
 #include <cstdlib>
 #include <cstring>
+#include <optional>
 
 #include <benchmark/benchmark.h>
 
+#include <rmm/cuda_stream_pool.hpp>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_buffer.hpp>
 #include <rmm/mr/cuda_memory_resource.hpp>
+#include <rmm/mr/per_device_resource.hpp>
 
 #include <rapidsmpf/error.hpp>
+#include <rapidsmpf/memory/buffer_resource.hpp>
 #include <rapidsmpf/memory/cuda_memcpy_async.hpp>
 #include <rapidsmpf/memory/host_memory_resource.hpp>
 #include <rapidsmpf/memory/pinned_memory_resource.hpp>
@@ -96,6 +100,19 @@ class NewDelete {
     friend void get_property(NewDelete const&, cuda::mr::host_accessible) noexcept {}
 };
 
+// Build a buffer resource with the given (optional) pinned pool properties.
+std::shared_ptr<rapidsmpf::BufferResource> make_pinned_buffer_resource(
+    std::optional<rapidsmpf::PinnedPoolProperties> props
+) {
+    return rapidsmpf::BufferResource::create(
+        rmm::mr::get_current_device_resource_ref(),
+        std::move(props),
+        {},
+        std::nullopt,  // disable the periodic spill-check thread
+        std::make_shared<rmm::cuda_stream_pool>(1, rmm::cuda_stream::flags::non_blocking)
+    );
+}
+
 // Helper function to create a type-erased host memory resource.
 cuda::mr::any_resource<cuda::mr::host_accessible> create_host_memory_resource(
     ResourceType const& resource_type
@@ -104,16 +121,20 @@ cuda::mr::any_resource<cuda::mr::host_accessible> create_host_memory_resource(
     case ResourceType::NEW_DELETE:
         return NewDelete{};
     case ResourceType::HOST_MEMORY_RESOURCE:
-        return rapidsmpf::HostMemoryResource{};
+        {
+            auto br = make_pinned_buffer_resource(rapidsmpf::PinnedMemoryDisabled);
+            return br->host_mr();  // br is kept alive by the back-reference
+        }
     case ResourceType::PINNED_MEMORY_RESOURCE:
         {
-            auto mr = rapidsmpf::PinnedMemoryResource::make_if_available();
+            auto br = make_pinned_buffer_resource(rapidsmpf::PinnedPoolProperties{});
+            auto mr = br->try_pinned_mr();
             RAPIDSMPF_EXPECTS(
                 mr.has_value(),
                 "pinned memory is not supported on this system",
                 std::runtime_error
             );
-            return *mr;
+            return *mr;  // br is kept alive by the back-reference
         }
     default:
         RAPIDSMPF_FAIL("Unknown memory resource type");
@@ -442,9 +463,8 @@ void BM_PinnedFirstAlloc_InitialPoolSize(benchmark::State& state) {
 
     for (auto _ : state) {
         state.PauseTiming();
-        auto mr = rapidsmpf::PinnedMemoryResource::make_if_available(
-            rapidsmpf::get_current_numa_node(), props
-        );
+        auto br = make_pinned_buffer_resource(props);
+        auto mr = br->try_pinned_mr();
         state.ResumeTiming();
         void* ptr = mr->allocate(stream, allocation_size);
         stream.synchronize();
@@ -499,13 +519,11 @@ void BM_PinnedPoolInit_InitialPoolSize(benchmark::State& state) {
     };
 
     for (auto _ : state) {
-        auto mr = rapidsmpf::PinnedMemoryResource::make_if_available(
-            rapidsmpf::get_current_numa_node(), props
-        );
-        benchmark::DoNotOptimize(mr);
-        // Destroy mr at end of iteration (pool teardown excluded from timing).
+        auto br = make_pinned_buffer_resource(props);
+        benchmark::DoNotOptimize(br);
+        // Destroy br at end of iteration (pool teardown excluded from timing).
         state.PauseTiming();
-        mr.reset();
+        br.reset();
         state.ResumeTiming();
     }
 

cpp/include/rapidsmpf/memory/buffer_resource.hpp

@@ -91,9 +91,12 @@ class BufferResource : public std::enable_shared_from_this<BufferResource> {
      * allocations are tracked for memory-limit accounting and statistics, use
      * `BufferResource::device_mr()` instead of the original memory resource after
      * construction.
-     * @param pinned_mr Pinned host memory resource used for
-     * `MemoryType::PINNED_HOST` allocations, or `PinnedMemoryResource::Disabled` to
-     * disable pinned allocations.
+     * @param pinned_pool_properties Configuration for the pinned host memory pool
+     * used for `MemoryType::PINNED_HOST` allocations, or `PinnedMemoryDisabled` to
+     * disable pinned allocations. The pinned resource is constructed internally and
+     * owned by the `BufferResource`. When a value is provided, pinned host memory
+     * must be supported on the system (see `is_pinned_memory_resources_supported()`);
+     * otherwise a `std::runtime_error` is thrown.
      * @param memory_limits Maximum allocation limits in bytes per `MemoryType`. Missing
      * entries are treated as unlimited.
      * @param periodic_spill_check Interval between periodic spill checks. `std::nullopt`
@@ -102,10 +105,12 @@ class BufferResource : public std::enable_shared_from_this<BufferResource> {
      * explicit CUDA stream.
      * @param statistics Statistics instance used for runtime metrics.
      * @return A newly constructed `BufferResource` owned by `std::shared_ptr`.
+     * @throws std::runtime_error if `pinned_pool_properties` has a value but pinned
+     * host memory is not supported on this system.
      */
     [[nodiscard]] static std::shared_ptr<BufferResource> create(
         cuda::mr::any_resource<cuda::mr::device_accessible> device_mr,
-        std::optional<PinnedMemoryResource> pinned_mr = PinnedMemoryResource::Disabled,
+        std::optional<PinnedPoolProperties> pinned_pool_properties = PinnedMemoryDisabled,
         std::unordered_map<MemoryType, std::int64_t> memory_limits = {},
         std::optional<Duration> periodic_spill_check = std::chrono::milliseconds{1},
         std::shared_ptr<rmm::cuda_stream_pool> stream_pool = std::make_shared<
@@ -217,6 +222,11 @@ class BufferResource : public std::enable_shared_from_this<BufferResource> {
      * @brief Get the RMM host memory resource.
      *
      * @return Reference to the RMM resource used for host allocations.
+     *
+     * @note Lifetime semantics are identical to `device_mr()`. See its
+     * `@par CCCL lifetime semantics` section for details. In brief, the returned
+     * `resource_ref` is non-owning. Promote it to a `any_host_resource` to extend the
+     * `BufferResource` lifetime.
      */
     [[nodiscard]] rmm::host_async_resource_ref host_mr() noexcept;
 
@@ -225,16 +235,22 @@ class BufferResource : public std::enable_shared_from_this<BufferResource> {
      *
      * @throws std::invalid_argument if no pinned memory resource is available.
      * @return Reference to the RMM resource used for pinned host allocations.
+     *
+     * @note Lifetime semantics are identical to `device_mr()`. See its
+     * `@par CCCL lifetime semantics` section for details. In brief, the returned
+     * `resource_ref` is non-owning. Promote it to a `any_host_device_resource` to extend
+     * the `BufferResource` lifetime.
      */
     [[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.
+     * @return The `PinnedMemoryResource` is available, or `std::nullopt` if pinned host
+     * memory is not available. The returned handle keeps this `BufferResource` alive as
+     * long as the handle (or any copy) exists.
      */
-    [[nodiscard]] std::optional<any_host_device_resource> try_pinned_mr() const noexcept;
+    [[nodiscard]] std::optional<PinnedMemoryResource> try_pinned_mr() const;
 
     /**
      * @brief Returns the currently available memory for a given memory type, in bytes.
@@ -340,9 +356,7 @@ class BufferResource : public std::enable_shared_from_this<BufferResource> {
     [[nodiscard]] MemoryReservation reserve_or_fail(std::size_t size, Range mem_types) {
         // try to reserve memory from the given order
         for (auto const& mem_type : mem_types) {
-            if (mem_type == MemoryType::PINNED_HOST
-                && pinned_mr_ == PinnedMemoryResource::Disabled)
-            {
+            if (mem_type == MemoryType::PINNED_HOST && !pinned_mr_.has_value()) {
                 // Pinned host memory is only available if the memory resource is
                 // available.
                 continue;

cpp/include/rapidsmpf/memory/host_memory_resource.hpp

@@ -13,9 +13,12 @@
 #include <rmm/resource_ref.hpp>
 
 #include <rapidsmpf/error.hpp>
+#include <rapidsmpf/memory/back_ref_mixin.hpp>
 
 namespace rapidsmpf {
 
+class BufferResource;
+
 /**
  * @brief Host memory resource using standard CPU allocation.
  *
@@ -29,9 +32,8 @@ namespace rapidsmpf {
  * buffers. The hint is applied via `madvise(MADV_HUGEPAGE)` and may be ignored
  * by the kernel depending on system configuration or resource availability.
  */
-class HostMemoryResource {
+class HostMemoryResource : public BackRefMixin<BufferResource> {
   public:
-    HostMemoryResource() = default;
     ~HostMemoryResource() = default;
 
     HostMemoryResource(HostMemoryResource const&) = default;  ///< Copyable.
@@ -138,6 +140,12 @@ class HostMemoryResource {
     friend void get_property(
         HostMemoryResource const&, cuda::mr::host_accessible
     ) noexcept {}
+
+  private:
+    /// @brief Default construct. Private: only `BufferResource` creates instances.
+    HostMemoryResource() = default;
+
+    friend class BufferResource;
 };
 
 static_assert(cuda::mr::resource<HostMemoryResource>);

cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp

@@ -6,6 +6,7 @@
 
 #include <cstddef>
 #include <memory>
+#include <optional>
 
 #include <cuda.h>
 #include <cuda_runtime_api.h>
@@ -20,6 +21,7 @@
 #include <rapidsmpf/config.hpp>
 #include <rapidsmpf/detail/rmm_resource_adaptor_impl.hpp>
 #include <rapidsmpf/error.hpp>
+#include <rapidsmpf/memory/back_ref_mixin.hpp>
 #include <rapidsmpf/system_info.hpp>
 #include <rapidsmpf/utils/misc.hpp>
 
@@ -32,6 +34,8 @@
 
 namespace rapidsmpf {
 
+class BufferResource;
+
 /**
  * @brief Checks if the PinnedMemoryResource is supported for the current CUDA version.
  *
@@ -75,8 +79,40 @@ struct PinnedPoolProperties {
 
     /// @brief Maximum size of the pool. `std::nullopt` means no limit.
     std::optional<std::size_t> max_pool_size = std::nullopt;
+
+    /// @brief NUMA node from which pinned memory should be allocated. Defaults to
+    /// the NUMA node of the calling thread.
+    int numa_id = get_current_numa_node();
 };
 
+/**
+ * @brief Sentinel used to disable pinned host memory.
+ *
+ * Pass this in place of a `PinnedPoolProperties` (e.g. to `BufferResource::create()`)
+ * to disable pinned host memory allocations.
+ */
+inline constexpr std::optional<PinnedPoolProperties> PinnedMemoryDisabled{};
+
+/**
+ * @brief Parse pinned memory pool properties from configuration options.
+ *
+ * Recognized options:
+ * - "pinned_memory": enable pinned memory.
+ * - "pinned_initial_pool_size" (bytes or percentage): initial pool size.
+ *   - Byte values (e.g. "1 MiB") are applied literally.
+ *   - Percentages (e.g. "10%") are relative to `get_host_memory_per_gpu()`.
+ * - "pinned_max_pool_size" (bytes, percentage, or disabled): maximum pool size.
+ *   - Byte and percentages uses the same parsing rules as "pinned_initial_pool_size".
+ *   - A disabled value (e.g. "off") leaves the pool unbounded.
+ *
+ * @param options Configuration options.
+ * @return The parsed `PinnedPoolProperties` when "pinned_memory" is enabled,
+ * otherwise `std::nullopt` (pinned host memory disabled).
+ */
+std::optional<PinnedPoolProperties> pinned_pool_properties_from_options(
+    config::Options options
+);
+
 /**
  * @brief Memory resource that provides pinned (page-locked) host memory using a pool.
  *
@@ -91,47 +127,12 @@ struct PinnedPoolProperties {
  */
 class PinnedMemoryResource final
     : public cuda::mr::shared_resource<
-          detail::RmmResourceAdaptorImpl<cuda::pinned_memory_pool>> {
+          detail::RmmResourceAdaptorImpl<cuda::pinned_memory_pool>>,
+      public BackRefMixin<BufferResource> {
     using shared_base = cuda::mr::shared_resource<
         detail::RmmResourceAdaptorImpl<cuda::pinned_memory_pool>>;
 
   public:
-    /// @brief Sentinel value indicating that pinned host memory is disabled.
-    static constexpr std::nullopt_t Disabled = std::nullopt;
-
-    /**
-     * @brief Create a pinned memory resource if the system supports pinned memory.
-     *
-     * @param numa_id The NUMA node to associate with the resource. Defaults to the
-     * current NUMA node.
-     * @param pool_properties Properties for configuring the pinned memory pool.
-     *
-     * @return A `PinnedMemoryResource` when supported, otherwise `std::nullopt`.
-     *
-     * @see PinnedMemoryResource::PinnedMemoryResource
-     */
-    static std::optional<PinnedMemoryResource> make_if_available(
-        int numa_id = get_current_numa_node(), PinnedPoolProperties pool_properties = {}
-    );
-
-    /**
-     * @brief Construct from configuration options.
-     *
-     * Recognized options:
-     * - "pinned_memory": enable pinned memory.
-     * - "pinned_initial_pool_size" (bytes or percentage): initial pool size.
-     *   - Byte values (e.g. "1 MiB") are applied literally.
-     *   - Percentages (e.g. "10%") are relative to `get_host_memory_per_gpu()`.
-     * - "pinned_max_pool_size" (bytes, percentage, or disabled): maximum pool size.
-     *   - Byte and percentages uses the same parsing rules as "pinned_initial_pool_size".
-     *   - A disabled value (e.g. "off") leaves the pool unbounded.
-     *
-     * @param options Configuration options.
-     * @return A `PinnedMemoryResource` if pinned memory is enabled and supported;
-     * otherwise `std::nullopt`.
-     */
-    static std::optional<PinnedMemoryResource> from_options(config::Options options);
-
     /**
      * @brief Allocates pinned host memory associated with a CUDA stream.
      *
@@ -217,16 +218,16 @@ class PinnedMemoryResource final
     /**
      * @brief Construct a pinned (page-locked) host memory resource.
      *
-     * Private — use `make_if_available` or `from_options` to obtain an instance.
+     * Private: use `BufferResource` to construct instances.
      *
-     * @param numa_id NUMA node from which memory should be allocated.
-     * @param pool_properties Properties for configuring the pinned memory pool.
+     * @param pool_properties Properties for configuring the pinned memory pool,
+     * including the NUMA node from which memory should be allocated.
      *
      * @throws std::invalid_argument If pinned host memory pools are not supported.
      */
-    PinnedMemoryResource(
-        int numa_id = get_current_numa_node(), PinnedPoolProperties pool_properties = {}
-    );
+    explicit PinnedMemoryResource(PinnedPoolProperties pool_properties);
+
+    friend class BufferResource;
 
     PinnedPoolProperties pool_properties_;  ///< properties used to configure the pool
 };

cpp/include/rapidsmpf/memory/resource_types.hpp

@@ -16,6 +16,9 @@ using any_device_resource = cuda::mr::any_resource<cuda::mr::device_accessible>;
 using any_host_device_resource =
     cuda::mr::any_resource<cuda::mr::host_accessible, cuda::mr::device_accessible>;
 
+/// @brief Owning type-erased host memory resource.
+using any_host_resource = cuda::mr::any_resource<cuda::mr::host_accessible>;
+
 /**
  * @brief Check whether a type-erased memory resource is host-accessible.
  *