字符串数据的高效处理对于许多数据科学应用至关重要。为了从字符串数据中提取有价值的信息, RAPIDS libcudf 提供了加速字符串数据转换的强大工具。libcudf 是一个基于 C++ 的 GPU DataFrame 库,专用于加载、连接、聚合和过滤数据。本文将带你从零开始,掌握使用 libcudf 进行字符串转换的两种核心方法:通用 API 和自定义内核,助你解锁 GPU 上的极致性能。
字符串数据的重要性与 Arrow 格式简介
在数据科学中,字符串数据涵盖语音、文本、基因序列、日志记录等多种信息类型。在机器学习和特征工程中处理字符串数据时,经常需要对其进行规范化和转换,才能应用于特定场景。libcudf 提供了通用 API 和设备端实用工具,支持广泛的自定义字符串操作。
要在 GPU 上高效处理字符串,首先需要了解 libcudf 的存储方式。libcudf 采用 Arrow 格式在设备内存中存储字符串数据,它将字符串列表示为两个子列:chars 和 offsets(图 1)。
图 1. 箭头格式使用 chars 和 offsets 子列表示字符串列的示意图
- chars 列:以 UTF-8 编码的字符字节形式保存字符串数据,并在内存中连续存储。
- offsets 列:包含递增的整数序列,这些整数是字节位置,用于标识字符数据数组中每个字符串的起始位置。最后一个 offset 元素是 chars 列的总字节数。这意味着第 i 行的单个字符串大小定义为 (offsets[i+1] - offsets[i])。
小提示:理解 Arrow 格式是编写高效自定义内核的关键,因为你需要直接操作 offset 和 char 两个子列,而非像处理普通表格那样对待字符串。
字符串编校函数示例:一个具体的转换场景
为了演示字符串转换的示例,考虑一个函数,该函数接收两个输入字符串列,并生成一个经过修订的输出字符串列。
输入数据的格式如下:一个“名称”列包含由空格分隔的名字和姓氏,另一个“可见性”列包含“公共”或“私有”状态。
我们建议使用“redact”函数对输入数据进行操作,以生成由姓氏首字母 + 空格 + 完整名字组成的输出数据。但如果对应的可见性列为“私有”,则输出字符串应完全修订为“X X”。
表 1. “编校”字符串转换示例:输入名称和可见性字符串列,输出部分或完全编校的数据
方法一:使用 libcudf API 转换字符串(通用 API 方法)
首先,可以使用 libcudf strings API 完成字符串转换。通用 API 是比较性能的良好起点和基线。
工作原理:API 函数操作整个字符串列,每个函数至少启动一个内核,每个字符串分配一个线程。每个线程在 GPU 上并行处理一行数据,并输出一行作为新输出列的一部分。
步骤详解:
要使用通用 API 完成修订示例函数,请执行以下步骤:
- 使用 contains 将“可见性”字符串列转换为布尔列。
- 每当布尔列中对应的行条目为 false 时,通过复制“X X”从名称列创建新的字符串列。
- 将“修订”列拆分为名字和姓氏列。
- 将姓氏的第一个字符切成姓氏首字母。
- 通过用空格分隔符连接姓氏缩写列和名字列来构建输出列。
// convert the visibility label into a boolean
auto const visible = cudf::string_scalar(std::string("public"));
auto const allowed = cudf::strings::contains(visibilities, visible);
// redact names
auto const redaction = cudf::string_scalar(std::string("X X"));
auto const redacted = cudf::copy_if_else(names, redaction, allowed->view());
// split the first name and last initial into two columns
auto const sv = cudf::strings_column_view(redacted->view())
auto const first_last = cudf::strings::split(sv);
auto const first = first_last->view().column(0);
auto const last = first_last->view().column(1);
auto const last_initial = cudf::strings::slice_strings(last, 0, 1);
// assemble a result column
auto const tv = cudf::table_view({last_initial->view(), first});
auto result = cudf::strings::concatenate(tv, std::string(" "));
在包含 600K 行数据的 A6000 上,此方法大约需要 3.5 毫秒。本例使用 contains、copy_if_else、split、slice_strings 和 concatenate 完成自定义字符串转换。
性能分析:使用 Nsight Systems 进行的分析表明,split 函数耗时最长,其次是
和 concatenate。
图 2. 修订示例中 Nsight Systems 的分析数据,显示端到端字符串处理每秒可达 6 亿个元素。这些区域对应每个功能关联的 NVTX 范围,浅蓝色范围对应 CUDA 内核运行周期。
常见问题 1:为什么通用 API 方法比较慢?
因为每个 API 调用(如 split、slice)都会启动一个新的 CUDA 内核,并创建至少一个新的字符串列,这在全局设备内存中产生了大量的读写和分配开销。如果你的目标是极致性能,就需要考虑自定义内核。
方法二:使用自定义内核转换字符串
libcudf strings API 是一个快速高效的字符串转换工具包,但有时性能关键型函数需要运行得更快。libcudf strings API 中额外工作的一个关键来源是在全局设备内存中为每个 API 调用创建至少一个新的字符串列,这为将多个 API 调用合并到自定义内核提供了机会。
核心原则:libcudf strings 列是不可变的。不能就地更改字符串列,因为字符字节是连续存储的,对字符串长度的任何更改都会使偏移量数据无效。因此,redact_kernel 自定义内核通过使用 libcudf 列工厂来构建 offsets 和 chars 子列,从而生成一个新的字符串列。
第一步:编写内核并分配内存(包含 malloc 的版本 —— 这是一个错误示范)
在第一种方法中,使用内核内的 malloc 调用在动态设备内存中创建每行的输出字符串。自定义内核输出是一个指向每行输出的设备指针向量,该向量用作字符串列工厂的输入。
// note the column_device_view inputs to the kernel
__global__ void redact_kernel(cudf::column_device_view const d_names,
cudf::column_device_view const d_visibilities,
cudf::string_view redaction,
cudf::string_view* d_output)
{
// get index for this thread
auto index = threadIdx.x + blockIdx.x * blockDim.x;
if (index >= d_names.size()) return;
auto const visible = cudf::string_view("public", 6);
auto const name = d_names.element(index);
auto const vis = d_visibilities.element(index);
if (vis == visible) {
auto const space_idx = name.find(' ');
auto const first = name.substr(0, space_idx);
auto const last_initial = name.substr(space_idx + 1, 1);
auto const output_size = first.size_bytes() + last_initial.size_bytes() + 1;
char* output_ptr = static_cast(malloc(output_size));
// build output string
d_output[index] = cudf::string_view{output_ptr, output_size};
memcpy(output_ptr, last_initial.data(), last_initial.size_bytes());
output_ptr += last_initial.size_bytes();
*output_ptr++ = ' ';
memcpy(output_ptr, first.data(), first.size_bytes());
} else {
d_output[index] = cudf::string_view{redaction.data(), redaction.size_bytes()};
}
}
__global__ void free_kernel(cudf::string_view redaction, cudf::string_view* d_output, int count)
{
auto index = threadIdx.x + blockIdx.x * blockDim.x;
if (index >= count) return;
auto ptr = const_cast(d_output[index].data());
if (ptr != redaction.data()) free(ptr); // free everything that does match the redaction string
}
糟糕的性能表现:在包含 600K 行数据的 A6000 上,这种方法大约需要 108 毫秒,比通用 API 方法慢 30 倍以上。
redact_kernel 60.3ms
free_kernel 45.5ms
make_strings_column 0.5ms
问题根源:这里的主要瓶颈是两个内核中的 malloc/free 调用。CUDA 动态设备内存需要同步内核中的 malloc/free 调用,导致并行执行退化为顺序执行。
第二步:预先分配工作内存以消除瓶颈(预分配内核)
优化策略:在启动内核之前,用预先分配的工作内存替换内核中的 malloc/free 调用,从而消除 malloc/free 瓶颈。
对于编校示例,每个字符串的输出大小不应大于输入字符串本身,因为逻辑只删除字符。因此,可以使用与输入缓冲区大小相同的单个设备内存缓冲区,并利用输入偏移定位每行位置。
访问字符串列的偏移量需要先用 cudf::strings_column_view 包装 cudf::column_view,再调用其 offsets_begin 方法。也可以使用 chars_size 方法获取 chars 子列的大小。然后在调用内核之前,预先分配 rmm::device_uvector 来存储字符输出数据。
auto const scv = cudf::strings_column_view(names);
auto const offsets = scv.offsets_begin();
auto working_memory = rmm::device_uvector(scv.chars_size(), stream);
预分配内核代码:
__global__ void redact_kernel(cudf::column_device_view const d_names,
cudf::column_device_view const d_visibilities,
cudf::string_view redaction,
char* working_memory,
cudf::offset_type const* d_offsets,
cudf::string_view* d_output)
{
auto index = threadIdx.x + blockIdx.x * blockDim.x;
if (index >= d_names.size()) return;
auto const visible = cudf::string_view("public", 6);
auto const name = d_names.element(index);
auto const vis = d_visibilities.element(index);
if (vis == visible) {
auto const space_idx = name.find(' ');
auto const first = name.substr(0, space_idx);
auto const last_initial = name.substr(space_idx + 1, 1);
auto const output_size = first.size_bytes() + last_initial.size_bytes() + 1;
// resolve output string location
char* output_ptr = working_memory + d_offsets[index];
d_output[index] = cudf::string_view{output_ptr, output_size};
// build output string into output_ptr
memcpy(output_ptr, last_initial.data(), last_initial.size_bytes());
output_ptr += last_initial.size_bytes();
*output_ptr++ = ' ';
memcpy(output_ptr, first.data(), first.size_bytes());
} else {
d_output[index] = cudf::string_view{redaction.data(), redaction.size_bytes()};
}
}
内核输出 cudf::string_view 对象的矢量,该矢量被传递给 cudf::make_strings_column 工厂函数。此函数的第二个参数用于标识输出列中的空条目。本文中的示例没有空条目,因此使用了 nullptr 占位符 cudf::string_view{nullptr,0}。
auto str_ptrs = rmm::device_uvector(names.size(), stream);
redact_kernel<<>>(*d_names,
*d_visibilities,
d_redaction.value(),
working_memory.data(),
offsets,
str_ptrs.data());
auto result = cudf::make_strings_column(str_ptrs, cudf::string_view{nullptr,0}, stream);
在包含 600K 行数据的 A6000 上,此方法大约需要 1.1 毫秒,比基线高出 2 倍以上。大致细分如下:
redact_kernel 66us
make_strings_column 400us
剩余时间花在 cudaMalloc、cudaFree、cudaMemcpy 上,这是管理 rmm::device_uvector 临时实例的典型开销。如果保证所有输出字符串的大小都与输入字符串相同或更小,则此方法很有效。
小提示:使用 RAPIDS RMM 切换到批量工作内存分配是一个显著的改进,也是自定义字符串函数的一个很好的解决方案。
第三步:优化列创建以加快计算时间(最优化内核)
有没有办法进一步改进?现在的瓶颈是 cudf::make_strings_column 工厂函数,它从 cudf::string_view 对象的矢量构建两个字符串列组件 offsets 和 chars。
在 libcudf 中,包含了许多工厂函数来构建字符串列。前面示例中使用的工厂函数获取 cudf::string_view 对象的 cudf::device_span,然后通过对底层字符数据执行 gather 来构造列,以构建偏移量和字符子列。rmm::device_uvector 可自动转换为 cudf::device_span,无需复制任何数据。
核心优化:如果直接构建字符矢量和偏移矢量,则可以使用不同的工厂函数,它只需创建字符串列,而不需要 gather 来复制数据。
Step A:sizes_kernel 传递输入数据,以计算每个输出行的确切输出大小
__global__ void sizes_kernel(cudf::column_device_view const d_names,
cudf::column_device_view const d_visibilities,
cudf::size_type* d_sizes)
{
auto index = threadIdx.x + blockIdx.x * blockDim.x;
if (index >= d_names.size()) return;
auto const visible = cudf::string_view("public", 6);
auto const redaction = cudf::string_view("X X", 3);
auto const name = d_names.element(index);
auto const vis = d_visibilities.element(index);
cudf::size_type result = redaction.size_bytes(); // init to redaction size
if (vis == visible) {
auto const space_idx = name.find(' ');
auto const first = name.substr(0, space_idx);
auto const last_initial = name.substr(space_idx + 1, 1);
result = first.size_bytes() + last_initial.size_bytes() + 1;
}
d_sizes[index] = result;
}
Step B:通过执行就地 exclusive_scan 将输出大小转换为偏移量
请注意,offsets 矢量是用 names.size()+1 元素创建的。最后一项是字节总数(所有大小加在一起),而第一项是 0。这两项都由 exclusive_scan 调用处理。从 offsets 列的最后一个条目检索 chars 列的大小,以构建字符矢量。
// create offsets vector
auto offsets = rmm::device_uvector(names.size() + 1, stream);
// compute output sizes
sizes_kernel<<>>(
*d_names, *d_visibilities, offsets.data());
thrust::exclusive_scan(rmm::exec_policy(stream), offsets.begin(), offsets.end(), offsets.begin());
Step C:优化后的 redact_kernel 逻辑,接受输出 d_offsets 矢量来解析每行的输出位置
__global__ void redact_kernel(cudf::column_device_view const d_names,
cudf::column_device_view const d_visibilities,
cudf::size_type const* d_offsets,
char* d_chars)
{
auto index = threadIdx.x + blockIdx.x * blockDim.x;
if (index >= d_names.size()) return;
auto const visible = cudf::string_view("public", 6);
auto const redaction = cudf::string_view("X X", 3);
// resolve output_ptr using the offsets vector
char* output_ptr = d_chars + d_offsets[index];
auto const name = d_names.element(index);
auto const vis = d_visibilities.element(index);
if (vis == visible) {
auto const space_idx = name.find(' ');
auto const first = name.substr(0, space_idx);
auto const last_initial = name.substr(space_idx + 1, 1);
auto const output_size = first.size_bytes() + last_initial.size_bytes() + 1;
// build output string
memcpy(output_ptr, last_initial.data(), last_initial.size_bytes());
output_ptr += last_initial.size_bytes();
*output_ptr++ = ' ';
memcpy(output_ptr, first.data(), first.size_bytes());
} else {
memcpy(output_ptr, redaction.data(), redaction.size_bytes());
}
}
Step D:构建最终列(零复制)
从 d_offsets 列的最后一个条目检索输出 d_chars 列的大小,以分配字符矢量。内核使用预先计算的偏移量向量启动,并返回填充的字符向量。最后,libcudf strings 列工厂创建输出字符串列。
cudf::size_type output_size = offsets.back_element(stream);
auto chars = rmm::device_uvector(output_size, stream);
redact_kernel<<>>(
*d_names, *d_visibilities, offsets.data(), chars.data());
// from pre-assembled offsets and character buffers
auto result = cudf::make_strings_column(names.size(), std::move(offsets), std::move(chars));
惊人性能:在包含 600K 行数据的 A6000 上,此方法需要大约 300 us(0.3 ms)的时间,比以前的方法提高了 2 倍多。
为什么两次计算反而更快?你可能会注意到 sizes_kernel 和 redact_kernel 共享很多相同的逻辑:一次测量输出的大小,然后再次填充输出。从代码质量的角度来看,将转换重构为由大小和编校内核调用的设备函数是有益的。从性能的角度来看,你可能会惊讶地看到转换的计算成本是原来的两倍。但 内存管理和更高效的列创建的好处往往超过两次执行转换的计算成本。
常见问题 2:为什么使用两遍法计算大小和填充数据反而比一遍法更快?
因为一遍法(使用malloc或预分配工作内存)需要从string_view向量构建列,这会触发额外的gather操作来重新排列数据,产生了大量的内存读写。而两遍法直接构建了精确的offsets和chars数组,调用make_strings_column工厂时不再需要任何数据复制,可以直接“吞下”这两个数组。这种内存布局上的优化,避免了多次数据搬移,从而大幅提升了吞吐量。
性能对比与峰值分析
表 2 显示了本文讨论的四种解决方案的计算时间、内核计数和处理字节数。“内核启动总数”反映了启动的内核总数,包括计算内核和辅助内核。“处理的总字节数”是累积的 DRAM 读写吞吐量,“处理的最小字节数”为测试输入和输出的平均每行 37.9 字节。理想的“内存带宽受限”情况假设 768 GB/s 带宽,即 A6000 的理论峰值吞吐量。
表 2. 本文讨论的四种解决方案的计算时间、内核计数和处理字节数
由于内核启动次数减少,处理的总字节数减少,“优化内核”提供了最高的吞吐量。有了高效的自定义内核,内核的总启动次数从 31 次减少到 4 次,处理的总字节数从输入加输出大小的 12.6 倍减少到 1.75 倍。因此,定制内核的吞吐量比用于编校转换的通用字符串 API 高出 10 倍以上。
RMM 池内存资源的作用:RAPIDS Memory Manager (RMM) 中的池内存资源是另一个可用于提高性能的工具。上述示例使用默认的“CUDA 内存资源”分配和释放全局设备内存。然而,分配工作内存所需的时间增加了字符串转换步骤之间的延迟。RMM 中的“池内存资源”通过预先分配一个大的内存池,并在处理过程中根据需要分配子分配来减少延迟。
对于 CUDA 内存资源,“Optimized Kernel”显示了 10-15 倍的加速,由于分配大小的增加,在较高的行数时开始下降(图 3)。使用池内存资源可以缓解这种影响,并比 libcudf strings API 方法保持 15x-25x 的速度提升。
图 3. 使用默认 CUDA 内存资源(实线)和池内存资源(虚线)从自定义内核“预分配内核”和“优化内核”加速,而使用默认 CUDA 存储资源的 libcudf 字符串 API
利用池内存资源,两遍算法的端到端内存吞吐量接近理论极限。“优化内核”的吞吐量达到 320-340 GB/s ,使用输入大小加上输出大小和计算时间进行测量(图 4)。
两遍法首先测量输出元素的大小,分配内存,然后用输出设置内存。给定两遍处理算法,“优化内核”中的实现性能接近内存带宽限制。“端到端内存吞吐量”定义为输入加输出大小(GB)除以计算时间 * RTX A6000 内存带宽(768 GB/s)。
常见问题 3:我何时应该使用通用 API,何时应该编写自定义内核?
当你的数据量较小(如几万行)或对性能要求不苛刻时,通用 API 是绝佳选择:代码简洁、开发快速、易于维护。当你的数据量巨大(如百万行以上)、对延迟敏感(如实时推理管道),或者你的转换逻辑非常特殊(不是简单的拼接、切片等标准操作)时,就应该考虑编写自定义内核。自定义内核能带来 10 倍以上的性能提升,但在代码复杂度和开发周期上需要更多投入。
关键要点
本文演示了在 libcudf 中编写高效字符串数据转换的两种方法。libcudf 通用 API 对于开发人员来说快速而直接,并且提供了良好的性能。libcudf 还提供了设计用于自定义内核的设备端实用程序,在本例中,解锁性能提高了 10 倍以上。
- 通用 API 方法:入门快,适合快速原型开发和较小数据集。但过多内核启动和列创建会带来额外开销。
- 自定义内核(malloc版):并行性能因动态内存分配而大幅退化,应尽量避免。
- 自定义内核(预分配版):消除了malloc瓶颈,性能提升2倍。
- 自定义内核(优化版 - 两遍法):通过直接构建offsets和chars数组,避免了数据gather,性能再提升2倍以上,端到端吞吐量接近内存带宽极限。
- 结合 RMM池内存资源,可以进一步减少内存分配延迟,在不同数据规模下都能保持稳定的高性能。
