此操作旨在在 GPU 内存不足时以流式方式使用,在此情况下,我们希望最大程度地减少小的 cudaMemcpy 调用次数以及与 cudf 表相关联的所有元数据的跟踪。由于内存限制,所有 thrust 和 scratch 内存分配都仅使用传入的内存资源,而不使用每设备内存资源。
// 创建一个 table_view
cudf::table_view tv = ...;
// 选择一个内存资源(可选)。此内存资源用于 scratch/thrust 临时
// 数据。在内存受限的情况下,此资源可用于在程序开始时为 chunked_pack 预留 scratch 内存。
// for `chunked_pack` at the beginning of a program.
auto mr = cudf::get_current_device_resource_ref();
// 定义每个分块的缓冲区大小:缓冲区越大,此算法可以占用的 SM 越多。
// occupied by this algorithm.
//
// 在内部,GPU 的工作单位是 1MB 的批次。当我们实例化 cudf::chunked_pack 时,
// source table_view 的所有 1MB 批次都会预先计算。此外,
// chunked_pack 会根据给定的 user_buffer_size 缓冲区计算遍历所有这些批次(分块)所需的迭代次数。
// chunked_pack calculates the number of iterations that are required to go through all those
// batches given a `user_buffer_size` buffer. The number of 1MB batches in each iteration (chunk)
//
std::size_t user_buffer_size = 128*1024*1024;
auto chunked_packer = cudf::chunked_pack::create(tv, user_buffer_size, mr);
std::size_t host_offset = 0;
auto host_buffer = ...; // 获取您要复制到的 host 缓冲区
while (chunked_packer->has_next()) {
// 获取大小为 user_buffer_size 的用户缓冲区
cudf::device_span<uint8_t> user_buffer = ...;
std::size_t bytes_copied = chunked_packer->next(user_buffer);
// 缓冲区将最多容纳 user_buffer_size 字节的连续打包输入 table_view 的内容。
// of the contiguously packed input `table_view`. You are now free to copy
// 您现在可以自由地将此内存复制到其他地方,例如 host。
cudaMemcpyAsync(
host_buffer.data() + host_offset,
user_buffer.data(),
bytes_copied,
cudaMemcpyDefault,
stream);
host_offset += bytes_copied;
}