Use cudf::pack with pinned mr in TableChunk::copy

cpp/include/rapidsmpf/memory/buffer_resource.hpp

@@ -134,7 +134,7 @@ class BufferResource {
      *
      * @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 Retrieves the memory availability function for a given memory type.
@@ -298,7 +298,7 @@ class BufferResource {
     );
 
     /**
-     * @brief Move device buffer data into a Buffer.
+     * @brief Move device or pinned host buffer data into a Buffer.
      *
      * This operation is cheap; no copy is performed. The resulting Buffer resides in
      * device memory.
@@ -310,10 +310,18 @@ class BufferResource {
      * @param data Unique pointer to the device buffer.
      * @param stream CUDA stream associated with the new Buffer. Use or synchronize with
      * this stream when operating on the Buffer.
+     * @param mem_type The memory type of the underlying @p data. A device accessible
+     * memory type is required (ie. MemoryType::DEVICE or MemoryType::PINNED_HOST).
      * @return Unique pointer to the resulting Buffer.
+     *
+     * @throws std::invalid_argument If the memory type is invalid.
+     * @throws std::invalid_argument If @p mem_type is MemoryType::PINNED_HOST and the
+     * pinned memory resource is not available.
      */
     std::unique_ptr<Buffer> move(
-        std::unique_ptr<rmm::device_buffer> data, rmm::cuda_stream_view stream
+        std::unique_ptr<rmm::device_buffer> data,
+        rmm::cuda_stream_view stream,
+        MemoryType mem_type = MemoryType::DEVICE
     );
 
     /**

cpp/include/rapidsmpf/memory/host_buffer.hpp

@@ -213,7 +213,7 @@ class HostBuffer {
     static HostBuffer from_rmm_device_buffer(
         std::unique_ptr<rmm::device_buffer> pinned_host_buffer,
         rmm::cuda_stream_view stream,
-        PinnedMemoryResource& mr
+        rmm::host_device_async_resource_ref mr
     );
 
   private:

cpp/src/memory/buffer_resource.cpp

@@ -82,7 +82,7 @@ rmm::host_async_resource_ref BufferResource::host_mr() noexcept {
     return host_mr_;
 }
 
-rmm::host_async_resource_ref BufferResource::pinned_mr() {
+rmm::host_device_async_resource_ref BufferResource::pinned_mr() {
     RAPIDSMPF_EXPECTS(
         pinned_mr_, "no pinned memory resource is available", std::invalid_argument
     );
@@ -196,14 +196,37 @@ std::unique_ptr<Buffer> BufferResource::allocate(
 }
 
 std::unique_ptr<Buffer> BufferResource::move(
-    std::unique_ptr<rmm::device_buffer> data, rmm::cuda_stream_view stream
+    std::unique_ptr<rmm::device_buffer> data,
+    rmm::cuda_stream_view stream,
+    MemoryType mem_type
 ) {
     auto upstream = data->stream();
     if (upstream.value() != stream.value()) {
         cuda_stream_join(stream, upstream);
         data->set_stream(stream);
     }
-    return std::unique_ptr<Buffer>(new Buffer(std::move(data), MemoryType::DEVICE));
+
+    if (mem_type == MemoryType::DEVICE) {
+        return std::unique_ptr<Buffer>(new Buffer(std::move(data), MemoryType::DEVICE));
+    } else if (mem_type == MemoryType::PINNED_HOST) {
+        RAPIDSMPF_EXPECTS(
+            pinned_mr_ != PinnedMemoryResource::Disabled,
+            "pinned memory resource is not available",
+            std::invalid_argument
+        );
+
+        auto pinned_host_buffer = std::make_unique<HostBuffer>(
+            HostBuffer::from_rmm_device_buffer(std::move(data), stream, pinned_mr())
+        );
+        return std::unique_ptr<Buffer>(
+            new Buffer(std::move(pinned_host_buffer), stream, MemoryType::PINNED_HOST)
+        );
+    } else {
+        RAPIDSMPF_FAIL(
+            "Invalid memory type: " + std::string(to_string(mem_type)),
+            std::invalid_argument
+        );
+    }
 }
 
 std::unique_ptr<Buffer> BufferResource::move(

cpp/src/memory/host_buffer.cpp

@@ -141,7 +141,7 @@ HostBuffer HostBuffer::from_owned_vector(
 HostBuffer HostBuffer::from_rmm_device_buffer(
     std::unique_ptr<rmm::device_buffer> pinned_host_buffer,
     rmm::cuda_stream_view stream,
-    PinnedMemoryResource& mr
+    rmm::host_device_async_resource_ref mr
 ) {
     RAPIDSMPF_EXPECTS(
         pinned_host_buffer != nullptr,

cpp/src/streaming/cudf/table_chunk.cpp

@@ -161,6 +161,10 @@ TableChunk TableChunk::copy(MemoryReservation& reservation) const {
     // 2. The chunk is available and the data is a generic cudf table that is
     //    not already packed. In this case, the table data must first be packed
     //    before copying it to host or pinned memory.
+    //    a. reservation in pinned memory - Use cudf::pack to directly copy the table
+    //    data to pinned memory.
+    //    b. reservation in host memory - Use cudf::pack to copy the table data to
+    //    intermediate device memory and then copy to host memory.
     //
     // 3. The chunk data is already packed (packed_data_ != nullptr).
     //    In this case, we simply use buffer_copy() to copy the packed data
@@ -184,9 +188,34 @@ TableChunk TableChunk::copy(MemoryReservation& reservation) const {
                 br->release(reservation, nbytes);
                 return TableChunk(std::move(table), stream());
             }
-        case MemoryType::HOST:
-        case MemoryType::PINNED_HOST:
-            // Case 2.
+        case MemoryType::PINNED_HOST:  // Case 2a.
+            if (packed_data_ == nullptr) {  // data is in device memory as a table
+                auto stream = this->stream();
+
+                StreamOrderedTiming timing{stream, br->statistics()};
+
+                // use cudf pack with pinned mr, because it is more device memory
+                // efficient. When the pinned mr is warmed up, the performance is in par
+                // with packing to device memory and copying.
+                auto packed_pinned = cudf::pack(table_view(), stream, br->pinned_mr());
+                auto nbytes = packed_pinned.gpu_data->size();
+
+                br->statistics()->record_copy(
+                    MemoryType::DEVICE, MemoryType::PINNED_HOST, nbytes, std::move(timing)
+                );
+                // update the provided `reservation`
+                br->release(reservation, nbytes);
+                auto host_buffer = br->move(
+                    std::move(packed_pinned.gpu_data), stream, MemoryType::PINNED_HOST
+                );
+                return TableChunk(
+                    std::make_unique<PackedData>(
+                        std::move(packed_pinned.metadata), std::move(host_buffer)
+                    )
+                );
+            }  // else fall through to the default Case 3.
+            break;
+        case MemoryType::HOST:  // Case 2b.
             if (packed_data_ == nullptr) {
                 // We use libcudf's pack() to serialize `table_view()` into a
                 // packed_columns and then we move the packed_columns' gpu_data to a
@@ -218,6 +247,7 @@ TableChunk TableChunk::copy(MemoryReservation& reservation) const {
                 packed_data->data = br->move(std::move(packed_data->data), reservation);
                 return TableChunk(std::move(packed_data));
             }
+            // else fall through to the Case 3.
             break;
         default:
             RAPIDSMPF_FAIL("MemoryType: unknown");
@@ -226,7 +256,8 @@ TableChunk TableChunk::copy(MemoryReservation& reservation) const {
     // Note, `is_available() == false` implies `packed_data_ != nullptr`.
     RAPIDSMPF_EXPECTS(packed_data_ != nullptr, "something went wrong");
 
-    // Case 3.
+    // Case 3. The chunk data is already packed (packed_data_ != nullptr). We need
+    // to copy the packed data into the reservation-specified memory type.
     auto const nbytes = packed_data_->data->size;
     auto metadata = std::make_unique<std::vector<std::uint8_t>>(*packed_data_->metadata);
     auto data = br->allocate(nbytes, packed_data_->stream(), reservation);