无畏的 SIMD 之路:七年后的回顾与展望

Raph Levien, 2025年3月29日

七年前,我写了一篇博文 Towards fearless SIMD,概述了将 Rust 打造成编写快速 SIMD 程序的引人注目的语言的愿景。 现在我们进展如何了呢?

不幸的是,现在在 Rust 中编写 SIMD 的体验仍然相当粗糙,尽管已经取得了一些进展,并且正在进行一些有希望的努力。 与之前的文章一样,本文将概述一个可能的愿景。

到目前为止,Linebender 项目还没有使用 SIMD,但这种情况正在改变。 当我们致力于 CPU/GPU 混合渲染技术时,很明显我们需要 SIMD 才能最大限度地提高 CPU 端的性能。 我们还在更快的颜色转换和加速的 2D 几何图元中看到了机会。

这篇博文也是我最近与 André Popovitch 录制的播客的补充。 该播客是对 SIMD 概念的一个很好的介绍,而这篇博文则更多地关注未来的方向。

一个简单的例子

作为一个运行示例,我们将计算一个包含 4 个值的向量的 sigmoid 函数。 标量版本如下:

fn sigmoid(x: [f32; 4]) -> [f32; 4] {
  x.map(|y| y / (1.0 + y * y).sqrt())
}

我不认为有什么理由不能自动向量化,但根据 Godbolt 来看,它是优化不良的标量代码。

安全性

在 Rust 中编写 SIMD 的最大问题之一是,所有暴露的 SIMD intrinsics 都被标记为 unsafe,即使在它们可以安全使用的情况下也是如此。 原因是 SIMD 特性的支持差异很大,并且在不支持它的 CPU 上执行 SIMD 指令是未定义的行为——芯片可能会崩溃、忽略该指令或执行一些意想不到的操作。 为了安全使用,必须有其他机制来确定 CPU 是否支持该特性。

这是手写的 intrinsic 代码中的运行示例,展示了访问 SIMD intrinsics 时需要编写 unsafe

#[cfg(target_arch = "aarch64")]
fn sigmoid_neon(x: [f32; 4]) -> [f32; 4] {
use core::arch::aarch64::*;
unsafe {
let x_simd = core::mem::transmute(x);
let x_squared = vmulq_f32(x_simd, x_simd);
let ones = vdupq_n_f32(1.0);
let sum = vaddq_f32(ones, x_squared);
let sqrt = vsqrtq_f32(sum);
let ratio = vdivq_f32(x_simd, sqrt);
    core::mem::transmute(ratio)
  }
}
#[cfg(target_arch = "x86_64")]
fn sigmoid_sse2(x: [f32; 4]) -> [f32; 4] {
use core::arch::x86_64::*;
unsafe {
let x_simd = core::mem::transmute(x);
let x_squared = _mm_mul_ps(x_simd, x_simd);
let ones = _mm_set1_ps(1.0);
let sum = _mm_add_ps(ones, x_squared);
let sqrt = _mm_sqrt_ps(sum);
let ratio = _mm_div_ps(x_simd, sqrt);
    core::mem::transmute(ratio)
  }
}

这只是一个简化的例子。 首先,SIMD 宽度固定为 4 lanes (128 bits)。 最有可能的是,在实践中,你会在一个更大的 slice 上迭代,获取等于自然 SIMD 宽度的块。

多版本 (Multiversioning)

对于 SIMD 来说,一个核心问题是多版本和运行时分发。 在某些情况下,你确切地知道 CPU 目标,例如,当你编译一个只在你机器上运行的二进制文件时(在这种情况下,target-cpu=native 是合适的)。 但是,当更广泛地分发软件时,可能存在一系列能力。 为了获得最高的性能,有必要编译代码的多个版本,并进行运行时检测以分发到硬件可以运行的最佳 SIMD 代码。 这个问题在最初的 fearless SIMD 博文中已经表达过,并且自那时以来,Rust 语言级别没有取得重大进展。

在 C++ 世界中,Highway 库为非常广泛的目标提供了出色的 SIMD 支持,并且还解决了多版本问题。 其他用途包括 JPEG-XL 图像格式的编解码器。 这样的编解码器是 SIMD 编程的理想用例,并且在浏览器中发布它们需要一个良好的多版本解决方案。 Highway 对他们处理多版本的方法有一个非常好的解释。 仔细研究它,看看他们是如何解决各种问题的,将会很有用。 我想看到的可以简洁地表达为“Rust 的 Highway”。

一种可能的方法是一个名为 multiversion 的 crate,它使用宏来复制多个版本的代码。 最近基于宏的方法是 rust-target-feature-dispatch。 它通常与 multiversion 类似,并且具体差异在该 crate 的 README 中列出。

正如我在 2018 年的博文中首次提倡的那样,另一种方法是编写在表示 SIMD 能力的零大小类型上多态的函数,然后依靠单态化 (monomorphization) 来创建各种版本。 这种方法的一个动机是在 Rust 的类型系统中编码安全性。 拥有零大小的 token 是底层 CPU 具有一定 SIMD 能力的证明,因此调用这些 intrinsics 是安全的。 使用这种方法的主要库是 pulp,它也为 faer 线性代数库提供支持。

我开始把运行的例子放在一起,但是遇到了一个直接的问题,它缺乏一个 sqrt intrinsic(然而,这很容易添加)。 它的工作方式也有点不同,因为它只支持自然宽度的向量,而不支持固定宽度的向量。 对于一般的线性代数来说,这很好,但对于其他一些应用程序来说,它会增加摩擦,例如带有 alpha 的颜色自然是 4 个标量的块。 要查看 pulp 代码的示例以及一些讨论,请参阅此 Zulip 线程

fearless_simd#2 中,我提出了一个相当符合人体工程学的 SIMD 多版本的原型。 像最初的 fearless_simd 原型一样,向量数据类型在 SIMD 级别上是多态的。 新原型在几个重要方面超越了这一点。 首先,std::ops 中的算术 traits 是为向量类型实现的,因此可以将两个向量加在一起,将向量乘以标量等等。

这是该原型中的运行示例:

#[inline(always)]
fn sigmoid_impl<S: Simd>(simd: S, x: [f32; 4]) -> [f32; 4] {
let x_simd: f32x4<S> = x.simd_into(simd);
  (x_simd / (1.0 + x_simd * x_simd).sqrt()).into()
}
simd_dispatch!(sigmoid(level, rgba: [f32; 4]) -> [f32; 4] = sigmoid_impl);

fearless_simd#2 原型优于 pulp 的一个优点是基于 SIMD 级别进行向下转换 (downcasting) 的功能,因此可以编写针对不同芯片优化的不同代码。 有关更多详细信息,请参阅该 pull request 中的 srgb example。 尽管有明显的优势,但在这一点上,我不确定这是否是前进的方向。 构建所有需要的类型和操作将需要大量工作,并且库中可能会出现大量重复的样板代码,这反过来可能会导致编译时问题。 另一个可能的方向是一个更智能的、类似编译器的 proc macro,它根据源程序中的类型和操作按需合成 SIMD intrinsics。

Rust 的另一个考虑因素是运行时特性检测的实现比它应该的慢。 因此,特性检测和分发不应在每次函数调用时完成。 一个好的工作解决方案是在程序启动时进行一次特性检测,然后将该 token 通过函数调用传递下去。 这是可行的,但绝对是一种人体工程学上的不足 (ergonomic paper cut)。

FP16 和 AVX-512

并行计算的一个普遍趋势(实际上是由 AI 工作负载推动的)是具有更高吞吐量的小标量。 虽然在 x86_64 上还不常见,但 FP16 扩展在所有 Apple Silicon 桌面 CPU 和大多数最新的高端 ARM 手机上都受支持。 由于 Neon 只有 128 bits 宽,因此拥有 8 lanes 是受欢迎的。 我发现 f16 格式对于像素值特别有用,因为它可以编码颜色值,其精度足以避免视觉伪影(8 bits 不太够,但对于某些应用程序来说已经足够了,只要你不尝试使用 HDR)。

f16 类型的原生 Rust 支持尚未落地(在 rust#125440 中跟踪),这使得这种标量大小的使用更加困难。 但是,在 half 库中以及在 fearless_simd#2 原型中也提供了一些支持,该原型通过内联汇编导出了许多 FP16 Neon 指令。 当真正的 f16 支持落地时,就可以切换到 intrinsics,这将具有更好的优化和人体工程学(例如,相同的方法将 splat 在编译时转换为 f16 的常量和在运行时转换为 f32 的变量)。

AVX-512 是一种有些争议的 SIMD 功能。 它首先出现在命运不济的 Larrabee 项目中,该项目从 2010 年开始以有限的数量作为 Xeon Phi 出售,并且此后出现在零星的 Intel CPU 中,但做出了妥协。 特别是,即使在程序中少量使用 AVX-512 代码也可能导致降频,从而降低所有工作负载的性能(有关更多详细信息,请参阅 Stack Overflow thread on throttling)。 如今,获得具有 AVX-512 的 CPU 的最可能方式是 AMD Zen 4 或 Zen 5; 正是由于它们的实力,AVX-512 在 Steam 硬件调查中约占计算机的 16%。

增加的宽度不是对 AVX-512 感到兴奋的主要原因。 事实上,在 Zen 4 和大多数 Zen 5 芯片上,数据通道是 256 bits,因此完整的 512 bit 指令是“double pumped”。 最令人兴奋的方面是基于掩码的predication,这是 GPU 上的常见实现技术。 特别是,当掩码位为零时,内存加载和存储操作是安全的,这对于在字符串上有效地使用 SIMD 特别有帮助。 如果没有 predication,一种常见的技术是编写两个循环,第一个循环仅处理 SIMD 宽度的偶数倍,第二个循环通常编写为标量,用于处理奇数大小的“尾部”。 这有很多问题——代码膨胀、更差的分支预测、无法利用 SIMD 处理略小于自然 SIMD 宽度的块(随着 SIMD 变得更宽,这种情况会变得更糟),以及两个循环的行为不完全相同的风险。

展望未来,Intel 已经提出了 AVX10,并且有望在未来几年内发布 AVX 10.2 芯片。 此扩展具有 AVX-512 的几乎所有功能,并进行了一些清理和新功能(直到最近,AVX10 被定义为具有 256 bit 的基本宽度,可选 512,但 512 现在是基线)。 此外,AVX10.2 将包括 16-bit floats(目前仅在 Sapphire Rapids 高端服务器和工作站芯片中可用)。

关于 std::simd

“portable SIMD” 工作已经进行了多年,目前位于 nightly std::simd。 虽然我认为它在许多应用程序中非常有用,但我个人对我的应用程序不是很兴奋。 首先,因为它强调可移植性,所以它鼓励采用“最低公分母”方法,而我相信,对于某些用例,调整算法以最好地利用不同 SIMD 实现的特定怪癖非常重要。 其次,std::simd 本身并不能解决多版本问题。 从我的角度来看,最好将其视为 autovectorization 的增强版本。

语言演进

Rust 对 SIMD 的开箱即用支持仍然相当粗糙,尤其是需要广泛使用 unsafe。 虽然一些差距可以用库来填补,但可以说语言本身的目标应该是支持安全的 SIMD 代码。 在这方面已经取得了一些进展。

首先,原始版本的 target_feature 需要 unsafe 才能调用 任何#[target_feature] 注释的函数。 一项放宽这一限制的提案,以便已经处于 target_feature gate 下的函数可以安全地调用另一个具有相同 gate 的函数,称为 “target_feature 1.1" 并且计划在 1.86 中发布。 密切相关的是,一旦进入合适的 target_feature gate,大多数 SIMD intrinsics(广义上讲,那些不通过指针进行内存访问的 intrinsics)应被编译器视为安全的,并且该特性(core::arch 中的安全 intrinsics)也在进行中。

可以做更多的事情来帮助 Rust 编译器识别 SIMD 的使用何时是安全的,特别是当作为函数参数传入 SIMD 级别的具体 witness 时允许 target_features。 “struct target_features” 提案 (RFC 3525) 启用了在这种情况下启用 target_feature,并且是在拟议的 Rust 项目目标 Nightly support for ergonomic SIMD multiversioning 中考虑的提案之一。

总的来说,改进 Rust SIMD 支持将需要库和 Rust 语言中的支持。 库级别的不同方法可能表明不同的语言特性能够最好地支持它们。

展望未来

我提出这些原型以及撰写这些博客文章的主要目标是引发关于如何在 Rust 中最好地支持 SIMD 编程的讨论。 如果做得好,这对该语言来说是一个绝佳的机会,并且符合其对性能和可移植性的关注。

当我们构建 Vello hybrid CPU/GPU renderer 时,CPU 组件的性能将很大程度上依赖于 SIMD,因此我们需要投入大量精力来编写 SIMD 代码。 最保守的方法是为所有目标手动编写基于不安全 intrinsics 的代码,但这需要大量工作,并且使用不安全的代码并不吸引人。 我希望 Rust 生态系统能够团结起来并构建良好的基础设施,与 Highway 竞争。 现在,我认为是时候仔细考虑设计空间并尝试就它应该是什么样子达成共识了。