人人都应该了解 SIMD
SIMD(单指令多数据)一直被认为很复杂。我遇到过很多非常优秀的软件工程师,他们认为它过于复杂难以学习,或是一种只适用于极致性能软件的小众优化,在日常编程中派不上用场。
我认为这种看法是错的。SIMD 可以简单易懂,常见的用来加速朴素 for 循环的“一次处理 N 个值”的 SIMD 代码几乎总是遵循同样的基本形态。一旦掌握了基础,编写 SIMD 就和写 for 循环一样简单。而当它并不简单时,通常就意味着暂时可以先跳过它。
每一位开发者至少都应该掌握这些 SIMD 知识。
本文使用 Zig 作为示例,但其中的内容是通用的,适用于任何编程语言。对 SIMD 指令的支持因编程语言而异,我希望未来有更多编程语言能够暴露这些通用概念!
我很讨厌现在每篇文章都得声明这一点,但我还是想说明,本文完全由人工手写,未使用任何 AI 辅助。
背景:什么是 SIMD?
如果你已经知道 SIMD 是什么,可以跳过本节。
SIMD允许 CPU 并行处理多个值。例如,CPU 不必一次只比较一个字节,而是可以用一条指令同时比较 4 个、8 个甚至更多字节。
如果你在代码中见过这样的循环:
for (byte in bytes) { /* ... */ }
for (character in string) { /* ... */ }
for (value in array) { /* ... */ }这里就有使用 SIMD 的机会。SIMD 会把它变成这样:
for (8 byte chunk in bytes) { /* ... */ }这会带来与并行度直接对应的局部加速:处理数据的速度提升 4 倍、8 倍甚至更快。
要让这种优化真正产生收益,唯一的硬性要求是你需要定期处理足够大量的字节。如果这些 for 循环处理的数据始终只有几个或几十个字节,那就没必要这么做。但如果要遍历的是数百、数千、数百万字节,收益就会非常巨大。
以上就是基础知识。像 simdutf 和 simdjson 这样的项目则把这一点发挥到了极致,使用了难以理解的 SIMD 技巧。但你不需要编写那样的算法也能从 SIMD 中获益。常见情况要简单得多。
常见形态
常见的“一次处理 N 个值”的 SIMD 代码遵循相同的五个步骤:
- 广播你需要的任何常量,并初始化向量累加器(如果有的话)。
- 每次循环处理一个向量宽度的输入块。
- 在所有通道上并行执行比较或算术运算。
- 根据需要归约或存储向量结果。
- 用标量尾部处理剩余元素。标量尾部就是你向量化之前的普通循环,只不过它只处理无法填满整个向量的剩余部分。
随着你做得越来越多,你会开始自然地将每个 for 循环分解为这五个步骤,编写 SIMD 也会变得几乎像编写标量循环一样自然。
真实示例
我们来看一个来自 Ghostty 的真实示例。我们将分别查看标量实现、SIMD 实现,然后再将其映射回上面提到的常见形态。
我有一个已解码的码位切片,想要一直消费,直到遇到小于等于 0xF 的值(即 C0 控制字符)。2 终端大多是待打印的普通字符,所以我们尝试将它们批量处理。因此,这个循环会尽可能快地找到下一段可打印序列的末尾。
标量循环只有一行:
while (end < cps.len and cps[end] > 0xF) end += 1;它一次只处理一个码位,非常容易理解。
下面是不使用特定 CPU 内在函数的通用向量版本,没有任何注释。我稍后会详细解释它。
if (simd.lanes(u32)) |lanes| {
const V = @Vector(lanes, u32);
const threshold: V = @splat(0xF);
while (end + lanes <= cps.len) : (end += lanes) {
const values: V = cps[end..][0..lanes].*;
const greater_than_threshold = values > threshold;
if (@reduce(.And, greater_than_threshold)) continue;
const mask: std.meta.Int(.unsigned, lanes) = @bitCast(greater_than_threshold);
end += @ctz(~mask);
break;
}
}
while (end < cps.len and cps[end] > 0xF) end += 1;多了 12 行代码。
这可以将循环的吞吐量提升至多 4 倍(使用 ARM NEON,包括 Apple Silicon)、8 倍(使用 AVX2,即大多数现代 x86 CPU)和 16 倍(使用 AVX-512,即部分 Intel CPU 以及 AMD Zen 4 及更新型号)。
在一台搭载 AVX2 的 Intel 台式机上,从终端程序到最终终端状态的真实端到端吞吐量中,这带来的加速更接近 5 倍。由于 SIMD 代码周围的其他开销,你总会损失一部分理想加速比,但是……这仍然是 5 倍!
好吧,我知道对于不熟悉这些概念的人来说,那 12 行代码看起来会非常陌生。所以现在让我们回过头来,一步步地解释,并将其直接映射到前面提到的形态上。
步骤 1:广播常量
我们先从前三行开始:
if (simd.lanes(u32)) |lanes| {
const V = @Vector(lanes, u32);
const threshold: V = @splat(0xF);simd.lanes(u32) 是 Ghostty 中的一个辅助函数,它返回目标 CPU 一次能处理的 u32 值的数量。这些单独的值被称为通道。在 ARM 上它返回 4,AVX2 返回 8,AVX-512 返回 16。如果目标没有我们想使用的向量大小,它会返回 null,我们就会跳过所有这些代码,完全不做 SIMD 工作。
@Vector(lanes, u32) 创建向量类型。如果 lanes 是 8,那么 V 就是一个包含八个可由 CPU 并行操作的 u32 值的单个值。依此类推。
最后,我们需要将每个值与 0xF 进行比较。向量比较要求两边都是向量,所以 @splat(0xF) 会将 0xF 复制(即广播)到每个通道中。结果是一个如下所示的向量:
{ 0xF, 0xF, 0xF, 0xF, 0xF, 0xF, 0xF, 0xF }这就是步骤 1:准备向量类型并广播所有常量。有些算法还会在这里初始化向量累加器,但这个算法不需要。
步骤 2:一次循环处理一个向量
接下来,我们一次循环处理一个完整的向量:
while (end + lanes <= cps.len) : (end += lanes) {
const values: V = cps[end..][0..lanes].*;如果 lanes 是 8,我们只在至少还剩八个值时才进入循环。在循环内部,我们将这八个值加载到向量 values 中。每次循环结束时,end += lanes 会向前移动八个值,而不是一个。
对完整向量的要求很重要。如果只剩下五个值,我们就无法加载一个八通道的向量。有各种技巧可以处理这种情况,但我们采用最简单的做法,通过标量尾部来处理它们,这一点我会在步骤 5 中解释。
这就是步骤 2:一次加载并循环处理一个向量宽度的输入块。你可以在这里看到通道数量带来的加速!
步骤 3:执行 SIMD 操作
现在我们执行比较:
const greater_than_threshold = values > threshold;values 和 threshold 都是向量,所以这会映射为一个向量操作(一条真正的向量 CPU 指令)。这一个 > 会将 values 中的每个通道与 threshold 中对应的通道进行比较。如果有八个通道,这相当于执行八次标量比较 cps[end] > 0xF,但它只用一条 CPU 指令就完成了。4
结果是另一个每个通道包含一个布尔值的向量。从概念上看,它看起来像这样:
values: { 0x41, 0x42, 0x43, 0x0A, 0x44, 0x45, 0x46, 0x47 }
threshold: { 0xF, 0xF, 0xF, 0xF, 0xF, 0xF, 0xF, 0xF }
greater_than_threshold: { true, true, true, false, true, true, true, true }这就是真正的 SIMD 操作。这里没有显式的内层循环。> 运算符会并行地应用于每个通道。
比较只是一个例子。这也可以是加法、乘法、最小值、最大值,或向量类型支持的任何其他操作。关键在于代码仍然保持相同的形态。
步骤 4:归约向量结果
现在我们有了一个布尔向量,但原始循环需要知道第一个小于等于 0xF 的值的位置。
首先,我们来处理所有值都大于 0xF 的常见情况:
if (@reduce(.And, greater_than_threshold)) continue;@reduce(.And, ...) 使用 and 将所有布尔值组合起来,并返回单个布尔值。如果每个通道都是 true,我们就 continue 并处理下一个向量。在我们的例子中,通道 3 是 false,所以 @reduce 返回 false,我们会继续往下执行,以找出究竟是哪个通道失败了。
如果有任何通道为 false,那么我们需要精确找到是哪个通道失败了:
const mask: std.meta.Int(.unsigned, lanes) = @bitCast(greater_than_threshold);
end += @ctz(~mask);
break;@bitCast 将布尔向量转换为每个通道对应一位的整数。1 位表示该值大于 0xF,0 则表示不是。我们将掩码取反,使失败的比较变为 1,然后 @ctz 会统计第一次失败之前的零位数量。这个计数就是第一个失败通道的索引。
我们将该索引加到 end 上并跳出循环,因为我们已经找到了控制字符。
使用步骤 3 中相同的值,我们可以看到每个通道的这种转换:
values: { 0x41, 0x42, 0x43, 0x0A, 0x44, 0x45, 0x46, 0x47 }
greater_than_threshold: { true, true, true, false, true, true, true, true }
mask: { 1, 1, 1, 0, 1, 1, 1, 1 }
~mask: { 0, 0, 0, 1, 0, 0, 0, 0 }@ctz(~mask) 会在第一个 1 之前数出三个零位,因此它返回 3。将 3 加到 end 上,会使其指向通道 3,其中包含 0x0A,即第一个控制字符。
这就是步骤 4:将向量结果归约为原始算法所需的任何形式。这也是在不同算法之间变化最大的步骤。求和可能会将向量累加器归约为单个数字。转换可能会将整个向量存储到输出缓冲区。我们的扫描则将向量转换为位掩码,以便找到某一个特定的通道。
步骤 5:用标量尾部收尾
在向量循环之后,我们运行与开头完全相同的标量循环:
while (end < cps.len and cps[end] > 0xF) end += 1;如果输入长度不是向量宽度的整数倍,这个循环就会处理剩余的值。例如,八通道的向量循环会给这个循环留下零到七个值。这被称为标量尾部。
这个循环还能处理 simd.lanes(u32) 返回 null 的 CPU。在这种情况下,我们会跳过所有 SIMD 代码,由标量循环处理整个输入。原始实现既是回退方案,也是尾部处理。
这就是步骤 5。它就是普通的循环。
回顾:常见形态
让我们将整个实现映射回五个步骤:
@splat(0xF)将比较值广播到每个通道。while循环一次加载lanes个值。values > threshold并行比较每个通道。@reduce、@bitCast和@ctz找到第一个失败的比较。- 原始的标量循环处理剩余部分和不支持的 CPU。
步骤 4 中的细节起初需要一些时间来理解,但整体形态是直观的。而且步骤 1、2、3 和 5 在完全不同的算法中往往看起来几乎一模一样。
每当你看到 for (byte in bytes),你就会映射到这种形态。
为什么编译器不能做这件事?
有时它确实可以!编译器可以自动向量化简单循环,尤其是没有复杂控制流的规则算术循环。在手动编写 SIMD 之前,你应该始终先以优化模式编译标量版本,看看编译器生成了什么。
但编译器在能够自动向量化的范围上受到严重限制,总体上表现也很差。自动向量化已经是编译器研究领域数十年的活跃课题,而最新研究仍然从生产级编译器经常错失向量化机会这一观察出发。这个问题我预计短期内不会消失。
更重要的是,当这个循环重要到让我在意 5 倍加速时,我希望向量化是显式且可预测的。我不希望一次不相关的代码修改或编译器更新就悄悄地把它变回标量循环。
人人都应该了解 SIMD
每一位开发者都应该能够识别这种机会,最重要的是,应该不要惧怕 SIMD。如果你看到一个热点循环在扫描、比较、计数或转换大量连续数据,你应该能够想象一次处理一个向量宽度块的方式。
本文表明,这些常见情况遵循着非常规律的模式,你很快就会习惯。有了良好的语言支持,你不需要了解任何汇编或特定于 CPU 的细节就能轻松获得提升。
每个人都应该掌握足以做到这一点的 SIMD 知识。5
脚注
像 simdutf 和 simdjson 这样非常出色的项目使用极其复杂的 SIMD 技巧来实现其目标。但这并不是我所认为的“日常 SIMD”。 ↩
C0 控制字符的范围超出了
0xF。这是 Ghostty 在这条特定代码路径中使用的截断值;ESC 及其他控制序列的处理在别处进行。 ↩通用向量消除的是特定于 CPU 的语法,而非特定于 CPU 的代码生成。Zig 仍会将这些操作降低为目标平台所启用指令集的指令。当无法选择支持的向量宽度时,Ghostty 会回退到标量代码。 ↩
比较本身是一次向量操作。加载向量、归约结果以及定位失败通道则需要额外的指令。重要的是我们正在一次进行多次比较。 ↩
本文基于我写的一条 Lobsters 评论。 ↩
随机一篇博客