如何用CUB实现类似thrust::reduce_by_key的按键归约功能？

CUB确实没有提供和thrust::reduce_by_key完全等价的单接口函数，但通过组合它的几个核心原语，完全可以实现你的需求——将键值对归约为唯一键+对应归约结果的数组，同时满足你对内存、流和自定义内核交互的精细控制。

具体实现分为四个关键步骤：

• 排序键值对
  CUB的分段归约依赖连续的同键段，所以首先用cub::DeviceRadixSort::SortPairs对输入的键值对按键排序。这一步会把相同键的值聚合到连续区间，同时保证值和键的对应关系不被打乱。

• 生成段信息与唯一键
  排序完成后，用cub::DeviceRunLengthEncode::Encode处理排序后的键数组。这个原语会输出：

  • 去重后的唯一键数组（即最终需要的键集合）
  • 每个同键段的长度（可选）
  • 各段的起始偏移数组（给后续分段归约使用）

• 执行分段归约
  基于排序后的值数组和上一步得到的段起始偏移，调用cub::DeviceSegmentedReduce系列接口（比如Sum，或自定义归约算子），对每个同键段的值执行归约，得到对应唯一键的结果。

• 自定义扩展（可选）
  如果需要非默认的归约逻辑，你可以自定义归约 functor 传入DeviceSegmentedReduce；同时所有CUB设备接口都支持指定CUDA流，让操作在你指定的流中异步执行，方便和自定义内核协同调度。

以下是简化的流程伪代码：

// 输入：d_keys（原始键数组）、d_values（原始值数组）、num_items（元素总数）
// 输出：d_unique_keys（唯一键数组）、d_reduced_values（归约结果数组）、num_unique（唯一键数量）

// 1. 分配排序用临时内存（先调用一次获取所需大小，再实际分配）
size_t temp_sort_bytes = 0;
cub::DeviceRadixSort::SortPairs(nullptr, temp_sort_bytes, d_keys, d_sorted_keys, d_values, d_sorted_values, num_items);
cudaMalloc(&d_temp_sort, temp_sort_bytes);

// 执行键值对排序（可指定CUDA流）
cub::DeviceRadixSort::SortPairs(d_temp_sort, temp_sort_bytes, d_keys, d_sorted_keys, d_values, d_sorted_values, num_items, 0, stream);

// 2. 分配运行长度编码用临时内存
size_t temp_rle_bytes = 0;
cub::DeviceRunLengthEncode::Encode(nullptr, temp_rle_bytes, d_sorted_keys, d_unique_keys, d_segment_lengths, d_segment_offsets, num_items, num_unique, stream);
cudaMalloc(&d_temp_rle, temp_rle_bytes);

// 生成段信息与唯一键
cub::DeviceRunLengthEncode::Encode(d_temp_rle, temp_rle_bytes, d_sorted_keys, d_unique_keys, d_segment_lengths, d_segment_offsets, num_items, num_unique, stream);

// 3. 分配分段归约用临时内存
size_t temp_reduce_bytes = 0;
cub::DeviceSegmentedReduce::Sum(nullptr, temp_reduce_bytes, d_sorted_values, d_reduced_values, num_unique, d_segment_offsets, d_segment_offsets + 1, stream);
cudaMalloc(&d_temp_reduce, temp_reduce_bytes);

// 执行分段归约
cub::DeviceSegmentedReduce::Sum(d_temp_reduce, temp_reduce_bytes, d_sorted_values, d_reduced_values, num_unique, d_segment_offsets, d_segment_offsets + 1, stream);

注意事项：

• CUB的所有设备API都需要先调用一次获取临时内存大小，再分配内存执行实际操作，这让你可以自主控制临时内存的分配位置（比如使用统一内存或特定设备内存）。
• 如果原始键数组已经是有序的，可以跳过排序步骤，直接进行后续操作，减少计算开销。

这种组合方式完全覆盖thrust::reduce_by_key的功能，同时能满足你对内存、流调度和自定义内核交互的精细控制需求。

内容的提问来源于stack exchange，提问作者Treeman