PAYATHON 2026

如何使用 CUB/CUDA 高效批量排序图像 tile，避免两次全局排序？

数据小许

可以使用 cub::DeviceSegmentedRadixSort::SortKeys 批量排序多个 tile。该接口接收各 segment 的起止偏移，排序结果只需在每个 segment 内部有序，不必先对整幅图像排序，再按 tile 重新分组。它不保证每个 segment 都由多个 block 协作处理，也不保证采用 decoupled lookback 或达到某个固定性能水平。

先区分容量限制和并行度限制

“一个 tile 无法完整放入 shared memory”和“必须由多个 block 同时处理一个 tile”是两个不同的条件。

tile 放不进 shared memory 并不妨碍分段排序：实现可以分批从 global memory 读取数据，无需同时放入整个 tile。至于是否有多个 block 协作，则需要核对具体实现。例如，CCCL v2.8.2 的 segmented radix sort 在每轮 radix pass 中用一个 block 处理一个 segment，通过遍历该 segment 完成统计和分发。它能够处理超过一次 block 内排序容量的数据，但不能据此推断单个 segment 内有多 block 并行。NVIDIA 对应版本源码

如果 tile 很多，segment 之间的并行可能已经足够；如果只有少量大 tile，这种分配方式可能限制 GPU 利用率。

将二维 tile 整理成连续 segment

分段排序使用的边界是半开区间：

[begin_offsets[s], end_offsets[s])

同一个 segment 内的数据必须在传入数组中连续。不同 segment 之间可以有空隙，但不能互相重叠。

图像通常按行存储，矩形 tile 的相邻两行之间往往夹着其他 tile 的像素。因此，不能直接把左上角和右下角的线性地址作为 segment 边界。

一种直接的做法是用 gather kernel 按 tile 顺序打包像素：

packed = [tile 0 的所有像素][tile 1 的所有像素]...
offsets = [0, tile 0 的结束位置, tile 1 的结束位置, ...]

若共有 S 个 tile，offsets 的长度就是 S + 1。边缘 tile 的实际像素数可能不同，应按实际大小计算前缀和。若上游计算能够直接写出这种布局，就可以省掉独立的 gather。

这里支持自定义迭代器的是边界参数。下例所用 SortKeys 重载的 key 输入和输出是指针，不能仅靠修改 offsets 表达跨行、带间隔的二维 tile。CUB 分段排序接口与参数定义

批量排序的核心调用

下面使用显式临时存储接口，调用形式对应 CCCL v2.8.2。假设像素已经打包为 unsigned int，相关设备内存均已分配，元素数量与偏移均在该重载的 int 范围内。CUDA_CHECK 表示项目中用于检查 CUDA 返回值的错误处理宏。

#include <cub/cub.cuh>
#include <cuda_runtime.h>

// 已初始化的变量：
// unsigned int* d_keys_in;
// unsigned int* d_keys_out;
// int* d_offsets;        // num_segments + 1 个设备端元素
// int num_items;         // d_offsets[num_segments] 的值
// int num_segments;
// cudaStream_t stream;

void* d_temp_storage = nullptr;
size_t temp_storage_bytes = 0;

CUDA_CHECK(cub::DeviceSegmentedRadixSort::SortKeys(
    nullptr, temp_storage_bytes,
    d_keys_in, d_keys_out,
    num_items, num_segments,
    d_offsets, d_offsets + 1,
    0, sizeof(unsigned int) * 8, stream));

CUDA_CHECK(cudaMalloc(&d_temp_storage, temp_storage_bytes));

CUDA_CHECK(cub::DeviceSegmentedRadixSort::SortKeys(
    d_temp_storage, temp_storage_bytes,
    d_keys_in, d_keys_out,
    num_items, num_segments,
    d_offsets, d_offsets + 1,
    0, sizeof(unsigned int) * 8, stream));

第一次调用只查询临时存储大小，不执行排序；第二次调用才执行排序。输出保留原有 segment 边界，每个 tile 的结果仍位于自己的区间。

输入、输出和 offsets 的内存范围必须满足接口的不重叠要求。排序及其后续操作完成后再释放临时存储；持续处理同类图像时，应尽量复用缓冲区。具体签名与数量限制以安装的 CUB 头文件为准。

从排序结果得到各 tile 的直方图

排序后还需要统计连续相等值的长度。仅做去重或 compact 会丢失频次，不能单独构成直方图。

一种处理流程是：

  1. 标记每个 segment 的首元素，以及 segment 内与前一个值不同的元素。
  2. 对标记执行 prefix scan，为每个 run 分配输出位置。
  3. 写出各 run 的像素值、起始位置和所属 tile。
  4. 用下一个 run 的起点减去当前起点得到频次；每个 tile 的最后一个 run 使用该 tile 的结束偏移计算。

判断 run 起点时必须包含 segment 边界条件：

is_run_start(i) =
    i 是所属 tile 的起点
    或者 keys[i] != keys[i - 1]

实现时应在 tile 起点直接判定为真，避免读取不存在的前驱。即使前一个 tile 的末值与后一个 tile 的首值相同，也不能合并这两个 run。

这样得到的是每个 tile 的稀疏直方图，即实际出现的值及其次数。如果需要完整的稠密直方图，还需分配目标 bin 数组、清零并写入计数；排序不会消除这部分存储需求。若 bin 代表值区间，应先明确像素到 bin ID 的映射，再按 bin ID 排序和计数。

如何选择接口并判断性能

也可以比较 cub::DeviceSegmentedSort::SortKeys。官方文档说明，DeviceSegmentedSort 会按 segment 大小分组并选择处理方式，适合考察中等长度或长度差异较大的 segments；DeviceSegmentedRadixSort 则针对较大的 segments 优化。前者不承诺固定使用哪一种底层排序算法。NVIDIA 接口说明

建议测量完整的 gather → segmented sort → run counting 流程。比较时保持 key 类型、有效位数、tile 数量和数据分布一致，并将临时存储分配移出稳定运行阶段。

若无符号整数实际只有部分有效位，可以在保持 bin 区分正确的前提下缩小 [begin_bit, end_bit)。不要直接把这一规则套用于有符号数或浮点数，因为其 radix 排序还涉及排序表示转换。

如果 segment 数量很少、每个 segment 很大，而且确实需要 segment 内多 block 并行，可以把“每个 tile 调用一次 cub::DeviceRadixSort，再通过多个 stream 调度”作为对照方案。并发调用需要互不冲突的数据和临时存储；是否有收益仍取决于资源占用和内存带宽。

分段接口可以避免额外的全局排序。是否满足性能目标，需要结合数据布局和实际 GPU 测量，不能仅凭是否使用 decoupled lookback 推导出 2 倍或 4 倍的差距。

备注：内容仅供参考。