[论文解读] Onesweep: A Faster Least Significant Digit Radix Sort for GPUs

Andy Adinets, Duane Merrill|arXiv (Cornell University)|Jun 3, 2022

Algorithms and Data Compression被引用 5

一句话总结

Onesweep 为 GPU 引入了一种新颖的最低有效数字（LSD）基数排序方法，通过使用单次遍历前缀和技术，减少了全局内存流量，每个数位分箱迭代仅需约 2n 次全局内存操作——相比之前方法的约 3n 次显著减少。在 NVIDIA A100 上，其对 256M 个随机 32 位键的排序达到 29.4 GKey/s 的性能，相比 CUB 提升了 1.5 倍，并在各种键分布下均优于 HRS。

ABSTRACT

We present Onesweep, a least-significant digit (LSD) radix sorting algorithm for large GPU sorting problems residing in global memory. Our parallel algorithm employs a method of single-pass prefix sum that only requires ~2n global read/write operations for each digit-binning iteration. This exhibits a significant reduction in last-level memory traffic versus contemporary GPU radix sorting implementations, where each iteration of digit binning requires two passes through the dataset totaling ~3n global memory operations. On the NVIDIA A100 GPU, our approach achieves 29.4 GKey/s when sorting 256M random 32-bit keys. Compared to CUB, the current state-of-the-art GPU LSD radix sort, our approach provides a speedup of ~1.5x. For 32-bit keys with varied distributions, our approach provides more consistent performance compared to HRS, the current state-of-the-art GPU MSD radix sort, and outperforms it in almost all cases.

研究动机与目标

• 通过最小化冗余内存操作，减少基于 GPU 的 LSD 基数排序中的全局内存流量。
• 设计一种高性能、可扩展的排序核函数，适用于全局内存中的大规模数据集。
• 在不同键分布下，性能超越现有最先进的 GPU 基数排序实现，兼具速度与一致性。
• 通过优化数位分箱阶段的内存访问模式，实现 GPU 上高效排序。

提出的方法

• 该算法使用单次遍历前缀和计算，消除了每个数位分箱迭代中两次独立内存遍历的需求。
• 它利用分段并行前缀和模式，以极低的全局内存流量高效计算分箱偏移量。
• 该方法采用两阶段核函数启动：第一阶段计算分箱计数，第二阶段使用计算出的偏移量执行独占扫描并散列到输出分箱。
• 它采用线程束级归约与基于 shuffle 的方法，以减少分支发散并最大化 GPU 占用率。
• 该设计将每个数位分箱步骤的全局内存操作减少至约 2n，显著降低最后一级缓存压力。
• 它支持 32 位键，并针对现实工作负载中常见的随机和多样化键分布进行了优化。

实验结果

研究问题

• RQ1 单次遍历前缀和技术能否将 GPU LSD 基数排序中的全局内存操作次数降至当前双次遍历标准以下？
• RQ2 减少内存流量对现代 GPU（如 A100）的性能有何影响？
• RQ3 所提出的方法能否在包括非均匀和偏斜数据在内的各种键分布下保持高性能？
• RQ4 与 CUB 和 HRS 相比，该算法在吞吐量和一致性方面表现如何？

主要发现

• Onesweep 在 NVIDIA A100 GPU 上对 256M 个随机 32 位键进行排序时达到 29.4 GKey/s 的性能，显著优于 CUB。
• 在相同条件下，该算法相比当前最先进的 GPU LSD 基数排序实现 CUB 提升了 1.5 倍。
• Onesweep 在各种键分布下的性能一致性优于 HRS（当前最先进的 GPU MSD 基数排序）。
• 该方法将每个数位分箱迭代的全局内存操作减少至约 2n，相比之前方法的约 3n 显著降低最后一级内存流量。