Add ChunkBatch interface to aggregate small messages before sending

cpp/include/rapidsmpf/buffer/buffer.hpp

@@ -5,6 +5,7 @@
 #pragma once
 
 #include <array>
+#include <cstddef>
 #include <memory>
 #include <variant>
 #include <vector>
@@ -17,6 +18,7 @@ namespace rapidsmpf {
 
 class BufferResource;
 
+
 /// @brief Enum representing the type of memory.
 enum class MemoryType : int {
     DEVICE = 0,  ///< Device memory
@@ -64,7 +66,7 @@ class Buffer {
      *
      * @throws std::logic_error if the buffer does not manage host memory.
      */
-    [[nodiscard]] constexpr std::unique_ptr<std::vector<uint8_t>> const& host() const {
+    [[nodiscard]] constexpr HostStorageT const& host() const {
         if (const auto* ref = std::get_if<HostStorageT>(&storage_)) {
             return *ref;
         } else {
@@ -79,7 +81,7 @@ class Buffer {
      *
      * @throws std::logic_error if the buffer does not manage device memory.
      */
-    [[nodiscard]] constexpr std::unique_ptr<rmm::device_buffer> const& device() const {
+    [[nodiscard]] constexpr DeviceStorageT const& device() const {
         if (const auto* ref = std::get_if<DeviceStorageT>(&storage_)) {
             return *ref;
         } else {
@@ -122,6 +124,48 @@ class Buffer {
         );
     }
 
+    /**
+     * @brief Copy a slice of the buffer to a new buffer.
+     *
+     * @param offset Offset in bytes from the start of the buffer.
+     * @param length Length in bytes of the slice.
+     * @param stream CUDA stream to use for the copy.
+     * @returns A new buffer containing the copied slice.
+     */
+    [[nodiscard]] std::unique_ptr<Buffer> copy_slice(
+        std::ptrdiff_t offset, std::ptrdiff_t length, rmm::cuda_stream_view stream
+    ) const;
+
+    /**
+     * @brief Copy a slice of the buffer to a new buffer.
+     *
+     * @param target Memory type of the new buffer.
+     * @param offset Offset in bytes from the start of the buffer.
+     * @param length Length in bytes of the slice.
+     * @param stream CUDA stream to use for the copy.
+     * @returns A new buffer containing the copied slice.
+     */
+    [[nodiscard]] std::unique_ptr<Buffer> copy_slice(
+        MemoryType target,
+        std::ptrdiff_t offset,
+        std::ptrdiff_t length,
+        rmm::cuda_stream_view stream
+    ) const;
+
+    /**
+     * @brief Copy the buffer to a destination buffer with a given offset.
+     *
+     * @param dest Destination buffer.
+     * @param offset Offset of the destination buffer.
+     * @param stream CUDA stream to use for the copy.
+     * @returns Number of bytes written to the destination buffer.
+     *
+     * @throws std::logic_error if copy violates the bounds of the destination buffer.
+     */
+    [[nodiscard]] std::ptrdiff_t copy_to(
+        Buffer& dest, std::ptrdiff_t offset, rmm::cuda_stream_view stream
+    ) const;
+
     /// @brief Buffer has a move ctor but no copy or assign operator.
     Buffer(Buffer&&) = default;
     Buffer(Buffer const&) = delete;

cpp/include/rapidsmpf/shuffler/chunk.hpp

@@ -65,7 +65,10 @@ class Chunk {
     );
 
     /**
-     * @brief Construct a new chunk of a partition.
+     * @brief Construct a new chunk with no data.
+     *
+     * This is used to indicate the number of chunks that were sent for a particular
+     * partition.
      *
      * @param pid The ID of the partition this chunk is part of.
      * @param cid The ID of the chunk.
@@ -83,6 +86,8 @@ class Chunk {
         /// If not zero, the number of chunks of the partition expected to get from the
         /// sending rank. Ignored when it is zero.
         std::size_t expected_num_chunks;
+        /// If known, the size of the metadata buffer (in bytes).
+        std::size_t metadata_size;
         /// If known, the size of the gpu data buffer (in bytes).
         std::size_t gpu_data_size;
     };
@@ -94,6 +99,17 @@ class Chunk {
      */
     [[nodiscard]] std::unique_ptr<std::vector<uint8_t>> to_metadata_message() const;
 
+    /**
+     * @brief Serializes this chunk into a given metadata message buffer.
+     *
+     * @param msg The metadata message as a serialized byte vector.
+     * @param offset The offset in the message buffer to start writing.
+     * @returns The number of bytes written to the message buffer.
+     */
+    [[nodiscard]] std::ptrdiff_t to_metadata_message(
+        std::vector<uint8_t>& msg, std::ptrdiff_t offset
+    ) const;
+
     /**
      * @brief Construct a new chunk from a metadata message.
      *