向量化与性能可移植的快速排序 (2022)

# 向量化高性能可移植快速排序算法 来源：https://opensource.googleblog.com/2022/06/Vectorized%20and%20performance%20portable%20Quicksort.html 今天我们开源了一套排序代码，其对数字数组的排序速度比C++标准库的std::sort快约十倍，性能超越了当前最先进的架构特定算法，同时可在所有现代CPU架构上保持可移植性。下文将详细介绍我们的实现方法。 首先介绍背景知识。近年来列式数据库逐渐兴起，这类数据库将某一列的所有值连续存储，而非将单条记录或“行”的所有字段依次存储。这种存储方式能显著提升过滤和排序效率——而这正是SQL查询的关键基石。因此，我们的研究聚焦于这种数据布局。 既然排序算法已被广泛研究，我们如何还能实现十倍的性能提升？答案在于SIMD/向量指令。这类指令能在单次操作中处理多个独立元素——例如使用AVX-512指令集时可同时处理16个32位浮点数，或在Arm NEON架构上处理4个元素： [](https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEjqY5Zp8L4lAAwsahUx-H9A_6l6iyw5nSvuu6-HtJr1_Gnh0OtDeo_7OrlGfDkW5oVIHLVSTlRcvznaA3fBVi_GRD_iXbhryVoo8B6EbBIm3cZctlm3tyaJvGzNptPhvXbJ2TyUOTCoP3lXI3fH8fd4U4gsELVNP_KawFQNYnPbcbysAMRsP5BmaAnc/s1188/Screen%20Shot%202022-06-01%20at%201.08.53%20PM.png) 峰超级计算机.jpeg) (https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEinMQFHyn-5mc4if8A57r_6I3EdvqcjfzML2TFqHxBN0rO_Yiihd4RMdTbKNXZ79eiQk0hQ8SMHMc7kvQ_CTZVSvY2k2uoUgTHr9MnIaQY9KDGXzJRfmidRLKV3U5LHm-LAkU0XwPs_MvyCM9MQM5mk-LLGdjkLYLzeLmV7s3iwbydy7OyC_QMo46P6/s4216/CC%20BY%20Oak%20Ridge%20National%20Laboratory%20Summit_supercomputer_(44552259580).jpeg) 如果您已熟悉SIMD技术，可能知晓它常用于超级计算机、机器学习的线性代数运算、视频处理以及JPEG XL (https://en.wikipedia.org/wiki/JPEG_XL) 等图像编解码领域。但若SIMD操作仅针对独立元素，我们如何对需要重排相邻数组元素的排序任务进行向量化处理？ 假设我们拥有某种特殊的排序方法，例如能高效处理256个元素的数组。那么，对更大数组实施快速排序(https://en.wikipedia.org/wiki/Quicksort)算法的过程可分解为：先将数组划分为两个子数组——小于“枢轴值”（理想情况下为中位数）的元素为一组，其余元素为另一组；然后递归处理，直到子数组规模不超过256个元素，再使用我们的特殊方法进行排序。划分操作占据了大部分CPU时间，因此若能用SIMD技术加速该过程，就能实现高速排序。 令人欣喜的是，现代指令集（Arm SVE、RISC-V V、x86 AVX-512）都包含适用于划分操作的特殊指令。当输入一组“是/否”值（标记元素是否小于枢轴值）时，这种“压缩存储”指令能将对应“是”的元素连续存储到内存中。我们只需对“是/否”值进行逻辑取反，再次应用该指令即可将元素写入另一分区。这种策略已被应用于AVX-512专用快速排序(https://arxiv.org/pdf/1704.08579.pdf)。但对于不支持压缩存储指令的架构（如AVX2），先前研究(https://academic.oup.com/comjnl/article-abstract/59/1/83/2568602?login=false)已证明可通过排列指令模拟该操作。 我们基于这些技术，实现了首个支持三大架构六种指令集的向量化快速排序算法，且性能超越了先前的架构专用排序实现。我们的代码使用Highway(https://github.com/google/highway)库提供的可移植SIMD函数，无需为每个平台重写约3000行C++代码。Highway会在支持时使用压缩存储指令，否则自动采用等效的排列指令。与先前针对32位整数的state of the art(https://drops.dagstuhl.de/opus/volltexte/2021/13775/pdf/LIPIcs-SEA-2021-3.pdf)方案相比，我们支持16-128位全范围数据类型。 尽管采用单一可移植实现，我们在AVX2、AVX-512（Intel Skylake）和Arm NEON（Apple M1）架构上均创下速度纪录。对于百万量级的32/64/128位数字，在Apple M1上的处理速度可达499/471/466 MB/s；在配备AVX-512的3GHz Skylake处理器上，速度达到1123/1119/1120 MB/s。有趣的是，AVX-512比AVX2快1.4-1.6倍——这是零额外工作量获得的显著提升（Highway会自动检测CPU支持的指令集并选用最优指令）。在AVX2环境下，我们的测试速度为798 MB/s，而此前针对AVX2优化的最佳方案仅达到699 MB/s。作为对比，标准库在相同CPU上的处理速度仅为58/128/117 MB/s，因此我们根据数据类型实现了**9-19倍**的性能提升。 过去排序算法常被视为计算密集型任务。我们期待看到，当单核CPU能实现1 GB/s的排序速度时，将催生哪些新应用与新功能。采用Apache2许可证的源代码(https://github.com/google/highway/tree/master/hwy/contrib/sort)已发布于Github（如有疑问欢迎提交issue(https://github.com/google/highway/issues)），技术论文(https://arxiv.org/pdf/2205.05982.pdf)则提供了实现细节的详尽解析与评估（包括256元素的特殊处理案例）。 *作者：Jan Wassenberg – 脑机架构研究团队