adding pack_to_host

Signed-off-by: niranda perera <niranda.perera@gmail.com>

Author: nirandaperera

cpp/include/rapidsmpf/integrations/cudf/partition.hpp

@@ -242,4 +242,36 @@ PackedData chunked_pack(
     cudf::table_view const& table, Buffer& bounce_buf, MemoryReservation& data_res
 );
 
+/// @brief The minimum buffer size for `cudf::chunked_pack`.
+constexpr size_t cudf_chunked_pack_min_buffer_size = 1 << 20;  ///< 1 MiB
+
+/**
+ * @brief Pack a table to host memory using  `cudf::pack` or `cudf::chunked_pack`.
+ *
+ * If device memory reservation can be made for the estimated table size, `cudf::pack`
+ * is used. Otherwise, `cudf::chunked_pack` is used with a bounce buffer size of the
+ * estimated table size * @p cpack_buf_size_factor (with at least
+ * `cudf_chunked_pack_min_buffer_size`).
+ *
+ * @param table The table to pack.
+ * @param stream CUDA stream used for device memory operations and kernel launches.
+ * @param host_data_res Memory reservation for the host data buffer.
+ * @param cpack_buf_size_factor The factor to use for the chunked pack buffer size.
+ * Default is 0.1, i.e. 10% of the estimated table size.
+ * @param cpack_buf_mem_types The memory types to use for the bounce buffer. Default is
+ * `DEVICE_ACCESSIBLE_MEMORY_TYPES`.
+ *
+ * @return A `PackedData` containing the packed table.
+ *
+ * @throws std::invalid_argument If the memory reservation is not host accessible.
+ * @throws std::runtime_error If the memory reservation fails.
+ */
+std::unique_ptr<PackedData> pack_to_host(
+    cudf::table_view const& table,
+    rmm::cuda_stream_view stream,
+    MemoryReservation& host_data_res,
+    float cpack_buf_size_factor = 0.1,
+    std::span<MemoryType const> cpack_buf_mem_types = DEVICE_ACCESSIBLE_MEMORY_TYPES
+);
+
 }  // namespace rapidsmpf

cpp/include/rapidsmpf/memory/buffer_resource.hpp

@@ -211,7 +211,9 @@ class BufferResource {
                 return std::move(res);
             }
         }
-        RAPIDSMPF_FAIL("failed to reserve memory", std::runtime_error);
+        RAPIDSMPF_FAIL(
+            "failed to reserve memory " + std::to_string(size), std::runtime_error
+        );
     }
 
     /**

cpp/include/rapidsmpf/memory/memory_type.hpp

@@ -40,6 +40,20 @@ constexpr std::array<MemoryType, 2> SPILL_TARGET_MEMORY_TYPES{
     {MemoryType::PINNED_HOST, MemoryType::HOST}
 };
 
+/**
+ * @brief Memory types that are device accessible in the order of preference.
+ */
+constexpr std::array<MemoryType, 2> DEVICE_ACCESSIBLE_MEMORY_TYPES{
+    {MemoryType::DEVICE, MemoryType::PINNED_HOST}
+};
+
+/**
+ * @brief Memory types that are host accessible in the order of preference.
+ */
+constexpr std::array<MemoryType, 2> HOST_ACCESSIBLE_MEMORY_TYPES{
+    {MemoryType::PINNED_HOST, MemoryType::HOST}
+};
+
 /**
  * @brief Get the memory types with preference lower than or equal to @p mem_type.
  *
@@ -65,6 +79,28 @@ static_assert(std::ranges::equal(
     leq_memory_types(static_cast<MemoryType>(-1)), std::ranges::empty_view<MemoryType>{}
 ));
 
+/**
+ * @brief Check if the memory type is host accessible.
+ *
+ * @param mem_type The memory type.
+ * @return True if the memory type is host accessible, false otherwise.
+ */
+constexpr bool is_host_accessible(MemoryType mem_type) noexcept {
+    return std::ranges::find(HOST_ACCESSIBLE_MEMORY_TYPES, mem_type)
+           != std::ranges::end(HOST_ACCESSIBLE_MEMORY_TYPES);
+};
+
+/**
+ * @brief Check if the memory type is device accessible.
+ *
+ * @param mem_type The memory type.
+ * @return True if the memory type is device accessible, false otherwise.
+ */
+constexpr bool is_device_accessible(MemoryType mem_type) noexcept {
+    return std::ranges::find(DEVICE_ACCESSIBLE_MEMORY_TYPES, mem_type)
+           != std::ranges::end(DEVICE_ACCESSIBLE_MEMORY_TYPES);
+};
+
 /**
  * @brief Get the name of a MemoryType.
  *

cpp/src/integrations/cudf/partition.cpp

@@ -272,10 +272,11 @@ PackedData chunked_pack(
     cudf::table_view const& table, Buffer& bounce_buf, MemoryReservation& data_res
 ) {
     RAPIDSMPF_EXPECTS(
-        bounce_buf.mem_type() == MemoryType::DEVICE,
-        "bounce buffer must be in device memory",
+        is_device_accessible(bounce_buf.mem_type()),
+        "bounce buffer is not device accessible",
         std::invalid_argument
     );
+
     // all copies will be done on the bounce buffer's stream
     auto stream = bounce_buf.stream();
     auto* br = data_res.br();
@@ -316,4 +317,70 @@ PackedData chunked_pack(
     return {packer.build_metadata(), std::move(data_buf)};
 }
 
+std::unique_ptr<PackedData> pack_to_host(
+    cudf::table_view const& table,
+    rmm::cuda_stream_view stream,
+    MemoryReservation& host_data_res,
+    float chunked_pack_buffer_size_factor,
+    std::span<MemoryType const> cpack_buf_mem_types
+) {
+    RAPIDSMPF_EXPECTS(
+        is_host_accessible(host_data_res.mem_type()),
+        "memory reservation is not host accessible",
+        std::invalid_argument
+    );
+
+    auto* br = host_data_res.br();
+
+    size_t est_table_size = estimated_memory_usage(table, stream);
+    {
+        // make a device reservation for packing
+        auto [pack_res, overbooking] =
+            br->reserve(MemoryType::DEVICE, est_table_size, true);
+
+        if (overbooking == 0) {
+            // if there is enough memory to pack the table, use `cudf::pack`
+            auto packed_columns = cudf::pack(table, stream, br->device_mr());
+            // clear the reservation as we are done with it.
+            pack_res.clear();
+
+            // note that this is a device buffer, so we need to move it to host memory
+            auto packed_data = std::make_unique<PackedData>(
+                std::move(packed_columns.metadata),
+                br->move(std::move(packed_columns.gpu_data), stream)
+            );
+
+            // Handle the case where `cudf::pack` allocates slightly more than
+            // the input size. This can occur because cudf uses aligned
+            // allocations, which may exceed the requested size. To
+            // accommodate this, we allow some wiggle room.
+            if (packed_data->data->size > host_data_res.size()) {
+                if (packed_data->data->size
+                    <= host_data_res.size() + total_packing_wiggle_room(table))
+                {
+                    host_data_res =
+                        br->reserve(
+                              host_data_res.mem_type(), packed_data->data->size, true
+                        )
+                            .first;
+                }
+            }
+
+            // finally copy the packed data device buffer to HOST memory
+            packed_data->data = br->move(std::move(packed_data->data), host_data_res);
+            return packed_data;
+        }
+    }
+
+    // there is not enough memory to use cudf::pack. Use chunked_pack.
+    auto chunk_size = std::max(
+        static_cast<size_t>(est_table_size * chunked_pack_buffer_size_factor),
+        cudf_chunked_pack_min_buffer_size
+    );
+    auto bounce_res = br->reserve_or_fail(chunk_size, cpack_buf_mem_types);
+    auto bounce_buf = br->allocate(chunk_size, stream, bounce_res);
+
+    return std::make_unique<PackedData>(chunked_pack(table, *bounce_buf, host_data_res));
+}
+
 }  // namespace rapidsmpf

cpp/src/streaming/cudf/table_chunk.cpp

@@ -171,57 +171,7 @@ TableChunk TableChunk::copy(MemoryReservation& reservation) const {
                     // serialize `table_view()` into a packed_columns and then we move
                     // the packed_columns' gpu_data to a new host buffer.
 
-                    // make a reservation for packing
-                    auto [pack_res, overbooking] = br->reserve(
-                        MemoryType::DEVICE,
-                        estimated_memory_usage(table_view(), stream()),
-                        true
-                    );
-
-                    if (overbooking > 0) {
-                        // there is not enough memory to pack the table.
-                        size_t avail_dev_mem = pack_res.size() - overbooking;
-                        RAPIDSMPF_EXPECTS(
-                            avail_dev_mem > 1 << 20,
-                            "not enough device memory for the bounce buffer",
-                            std::runtime_error
-                        );
-                        auto bounce_buf = br->allocate(avail_dev_mem, stream(), pack_res);
-
-                        packed_data = std::make_unique<PackedData>(
-                            chunked_pack(table_view(), *bounce_buf, reservation)
-                        );
-                    } else {
-                        // if there is enough memory to pack the table, use `cudf::pack`
-                        auto packed_columns =
-                            cudf::pack(table_view(), stream(), br->device_mr());
-                        // clear the reservation as we are done with it.
-                        pack_res.clear();
-                        packed_data = std::make_unique<PackedData>(
-                            std::move(packed_columns.metadata),
-                            br->move(std::move(packed_columns.gpu_data), stream())
-                        );
-
-                        // Handle the case where `cudf::pack` allocates slightly more than
-                        // the input size. This can occur because cudf uses aligned
-                        // allocations, which may exceed the requested size. To
-                        // accommodate this, we allow some wiggle room.
-                        if (packed_data->data->size > reservation.size()) {
-                            if (packed_data->data->size
-                                <= reservation.size()
-                                       + total_packing_wiggle_room(table_view()))
-                            {
-                                reservation =
-                                    br->reserve(
-                                          MemoryType::HOST, packed_data->data->size, true
-                                    )
-                                        .first;
-                            }
-                        }
-                        // finally copy the packed data device buffer to HOST memory
-                        packed_data->data =
-                            br->move(std::move(packed_data->data), reservation);
-                    }
+                    packed_data = pack_to_host(table_view(), stream(), reservation);
                 }
                 return TableChunk(std::move(packed_data));
             }