feat: add multiple_blocks_allocation RAII handle for fixed_size_memory_resource

cpp/include/rmm/mr/detail/fixed_size_memory_resource_impl.hpp

@@ -12,12 +12,18 @@
 #include <cuda/memory_resource>
 
 #include <cstddef>
+#include <memory>
 #include <mutex>
 #include <utility>
 #include <vector>
 
 namespace RMM_NAMESPACE {
 namespace mr {
+
+// Forward declarations for friend access from detail::fixed_size_memory_resource_impl
+class fixed_size_memory_resource;
+class multiple_blocks_allocation;
+
 namespace detail {
 
 /**
@@ -79,6 +85,11 @@ class fixed_size_memory_resource_impl final
   std::pair<std::size_t, std::size_t> free_list_summary(free_list const& blocks);
 
  private:
+  friend class RMM_NAMESPACE::mr::multiple_blocks_allocation;
+
+  // Caller must hold get_mutex().
+  void deallocate_blocks_async_unsafe(std::vector<std::byte*>&& blocks, cuda_stream_view stream);
+
   free_list blocks_from_upstream(cuda_stream_view stream);
 
   void release();

cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp

@@ -290,7 +290,7 @@ class stream_ordered_memory_resource : public crtp<PoolResource> {
     bool operator<(stream_event_pair const& rhs) const { return event < rhs.event; }
   };
 
- private:
+ protected:
   /**
    * @brief get a unique CUDA event (possibly new) associated with `stream`
    *

cpp/include/rmm/mr/fixed_size_memory_resource.hpp

@@ -4,13 +4,17 @@
  */
 #pragma once
 
+#include <rmm/cuda_stream_view.hpp>
 #include <rmm/detail/export.hpp>
 #include <rmm/mr/detail/fixed_size_memory_resource_impl.hpp>
 #include <rmm/resource_ref.hpp>
 
 #include <cuda/memory_resource>
+#include <cuda/std/span>
 
 #include <cstddef>
+#include <memory>
+#include <vector>
 
 namespace RMM_NAMESPACE {
 namespace mr {
@@ -84,6 +88,119 @@ class RMM_EXPORT fixed_size_memory_resource
 static_assert(cuda::mr::resource_with<fixed_size_memory_resource, cuda::mr::device_accessible>,
               "fixed_size_memory_resource does not satisfy the cuda::mr::resource concept");
 
+/**
+ * @brief RAII handle for an allocation that may span multiple fixed-size blocks from a
+ *        `fixed_size_memory_resource`.
+ *
+ * Returned by `allocate_blocks_async`. When destroyed, all blocks are returned to the
+ * memory resource on the same stream used for allocation. Move and copy are disabled to
+ * prevent double deallocation. Holds a `fixed_size_memory_resource` (which has shared,
+ * refcounted ownership of the underlying pool) so the pool outlives the handle.
+ */
+class RMM_EXPORT multiple_blocks_allocation {
+ public:
+  /**
+   * @brief Allocate device memory spanning one or more fixed-size blocks, stream-ordered.
+   *
+   * Use this for allocations larger than a single block. The allocation is ordered on
+   * `stream`; deallocation (when the returned handle is destroyed) is also ordered on
+   * the same stream. A single event is recorded for the whole allocation, so there is no
+   * per-block event overhead.
+   *
+   * @param mr The `fixed_size_memory_resource` that supplies blocks. Copied by value since
+   *        `fixed_size_memory_resource` has refcounted shared ownership.
+   * @param size Minimum number of bytes to allocate. Will be rounded up to a multiple of
+   *        block size (see `get_block_size()` on `*mr`).
+   * @param stream CUDA stream on which the allocation is ordered.
+   * @return Unique handle to the allocation; destroys to deallocate. Empty (zero-size)
+   *         allocation returns a valid handle with size 0 and no blocks.
+   * @throw Any exception from allocating blocks. Blocks successfully taken from the pool
+   *        before the failure are returned to the pool on `stream` (same ordering as normal
+   *        deallocation).
+   */
+  [[nodiscard]] static std::unique_ptr<multiple_blocks_allocation> make_async(
+    fixed_size_memory_resource mr, std::size_t size, cuda_stream_view stream);
+
+  ~multiple_blocks_allocation();
+
+  multiple_blocks_allocation(multiple_blocks_allocation const&)            = delete;
+  multiple_blocks_allocation& operator=(multiple_blocks_allocation const&) = delete;
+  multiple_blocks_allocation(multiple_blocks_allocation&&)                 = delete;
+  multiple_blocks_allocation& operator=(multiple_blocks_allocation&&)      = delete;
+
+  /**
+   * @brief Number of bytes requested for this allocation.
+   *
+   * @return Requested size in bytes.
+   */
+  constexpr std::size_t size() const noexcept { return size_; }
+
+  /**
+   * @brief Total capacity in bytes (number of blocks × block size).
+   *
+   * @return Capacity in bytes; always >= size().
+   */
+  std::size_t capacity() const noexcept { return block_size() * blocks_.size(); }
+
+  /**
+   * @brief Size in bytes of each block in this allocation.
+   *
+   * @return Block size (same as the memory resource's get_block_size()).
+   */
+  std::size_t block_size() const noexcept { return mr_->get_block_size(); }
+
+  /**
+   * @brief Non-owning view of the underlying block pointers.
+   *
+   * @return Span of device pointers, one per block; each block has size block_size().
+   */
+  [[nodiscard]] cuda::std::span<std::byte* const> get_blocks() const noexcept
+  {
+    return {blocks_.data(), blocks_.size()};
+  }
+
+  /**
+   * @brief Span over the i-th block's bytes.
+   *
+   * @param i Block index in [0, get_blocks().size()).
+   * @return Span of std::byte over the i-th block.
+   */
+  cuda::std::span<std::byte> operator[](std::size_t i) const
+  {
+    return {blocks_[i], mr_->get_block_size()};
+  }
+
+  /**
+   * @brief Span over the i-th block's bytes with bounds checking.
+   *
+   * @param i Block index.
+   * @return Span of std::byte over the i-th block.
+   * @throws std::out_of_range if i >= number of blocks.
+   */
+  cuda::std::span<std::byte> at(std::size_t i) const
+  {
+    return {blocks_.at(i), mr_->get_block_size()};
+  }
+
+  /**
+   * @brief Stream on which this allocation is ordered.
+   *
+   * @return The stream passed to make_async.
+   */
+  constexpr cuda_stream_view stream() const noexcept { return stream_; }
+
+ private:
+  multiple_blocks_allocation(std::size_t size,
+                             std::vector<std::byte*> buffers,
+                             cuda_stream_view stream,
+                             fixed_size_memory_resource mr);
+
+  std::vector<std::byte*> blocks_;
+  std::size_t const size_;
+  cuda_stream_view stream_;
+  fixed_size_memory_resource mr_;
+};
+
 /** @} */  // end of group
 }  // namespace mr
 }  // namespace RMM_NAMESPACE

cpp/src/mr/detail/fixed_size_memory_resource_impl.cpp

@@ -5,10 +5,14 @@
 
 #include <rmm/aligned.hpp>
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/detail/error.hpp>
 #include <rmm/detail/logging_assert.hpp>
 #include <rmm/mr/detail/fixed_size_memory_resource_impl.hpp>
+#include <rmm/mr/detail/stream_ordered_memory_resource.hpp>
 
 #include <cuda/iterator>
+#include <cuda/std/algorithm>
+#include <cuda_runtime_api.h>
 
 #include <cstddef>
 #include <mutex>
@@ -107,6 +111,21 @@ std::pair<std::size_t, std::size_t> fixed_size_memory_resource_impl::free_list_s
                            : std::make_pair(block_size_, blocks.size() * block_size_);
 }
 
+void fixed_size_memory_resource_impl::deallocate_blocks_async_unsafe(
+  std::vector<std::byte*>&& blocks, cuda_stream_view stream)
+{
+  if (blocks.empty()) { return; }
+
+  free_list blocks_free_list;
+  cuda::std::ranges::for_each(blocks, [this, &blocks_free_list](std::byte* ptr) {
+    blocks_free_list.insert(this->free_block(ptr, get_block_size()));
+  });
+
+  auto stream_event = get_event(stream);
+  RMM_ASSERT_CUDA_SUCCESS(cudaEventRecord(stream_event.event, stream.value()));
+  this->insert_blocks(std::move(blocks_free_list), stream);
+}
+
 #ifdef RMM_DEBUG_PRINT
 void fixed_size_memory_resource_impl::print()
 {

cpp/src/mr/fixed_size_memory_resource.cpp

@@ -3,9 +3,14 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+#include <rmm/detail/logging_assert.hpp>
 #include <rmm/mr/fixed_size_memory_resource.hpp>
 
+#include <cuda/cmath>
+
 #include <cstddef>
+#include <mutex>
+#include <vector>
 
 namespace RMM_NAMESPACE {
 namespace mr {
@@ -29,5 +34,55 @@ std::size_t fixed_size_memory_resource::get_block_size() const noexcept
   return get().get_block_size();
 }
 
+// multiple_blocks_allocation
+
+multiple_blocks_allocation::multiple_blocks_allocation(std::size_t size,
+                                                       std::vector<std::byte*> buffers,
+                                                       cuda_stream_view stream,
+                                                       fixed_size_memory_resource mr)
+  : blocks_(std::move(buffers)), size_(size), stream_(stream), mr_(std::move(mr))
+{
+  RMM_LOGGING_ASSERT(size_ <= mr_->get_block_size() * blocks_.size());
+  RMM_LOGGING_ASSERT(blocks_.empty() ||
+                     blocks_.size() == cuda::ceil_div(size_, mr_->get_block_size()));
+}
+
+multiple_blocks_allocation::~multiple_blocks_allocation()
+{
+  if (!blocks_.empty()) {
+    std::lock_guard<std::mutex> lock(mr_->get_mutex());
+    mr_->deallocate_blocks_async_unsafe(std::move(blocks_), stream_);
+  }
+}
+
+std::unique_ptr<multiple_blocks_allocation> multiple_blocks_allocation::make_async(
+  fixed_size_memory_resource mr, std::size_t size, cuda_stream_view stream)
+{
+  if (size == 0) {
+    return std::unique_ptr<multiple_blocks_allocation>(
+      new multiple_blocks_allocation(0, {}, stream, std::move(mr)));
+  }
+
+  auto& self = *mr;
+  std::lock_guard<std::mutex> lock(self.get_mutex());
+
+  auto stream_event            = self.get_event(stream);
+  std::size_t const num_blocks = cuda::ceil_div(size, self.get_block_size());
+  std::vector<std::byte*> blocks;
+  blocks.reserve(num_blocks);
+  try {
+    for (std::size_t i = 0; i < num_blocks; ++i) {
+      blocks.push_back(
+        static_cast<std::byte*>(self.get_block(self.get_block_size(), stream_event).pointer()));
+    }
+  } catch (...) {
+    self.deallocate_blocks_async_unsafe(std::move(blocks), stream);
+    throw;
+  }
+
+  return std::unique_ptr<multiple_blocks_allocation>(
+    new multiple_blocks_allocation(size, std::move(blocks), stream, std::move(mr)));
+}
+
 }  // namespace mr
 }  // namespace RMM_NAMESPACE

cpp/tests/CMakeLists.txt

@@ -145,6 +145,7 @@ ConfigureTest(SYSTEM_MR_REF_TEST mr/mr_ref_system_tests.cpp)
 ConfigureTest(PINNED_MR_REF_TEST mr/mr_ref_pinned_tests.cpp)
 ConfigureTest(LOGGING_MR_REF_TEST mr/mr_ref_logging_tests.cpp)
 ConfigureTest(FIXED_SIZE_MR_REF_TEST mr/mr_ref_fixed_size_tests.cpp)
+ConfigureTest(FIXED_SIZE_MR_TEST mr/fixed_size_mr_test.cpp)
 ConfigureTest(DEFAULT_MR_REF_TEST mr/mr_ref_default_tests.cpp)
 
 # general adaptor tests