Rust 可移植 SIMD 在 GPU 上原生运行：warp 即向量单元

线程之下的并行性

在将 Rust 线程带到 GPU时，我们将每个 std::thread 映射到 GPU 的一个 warp。这让我们可以在 GPU 上运行许多并发线程，但并没有使用每个线程/warp 内部的并行 lane。

在 CPU 上，线程内并行的抽象是 SIMD。一条指令对打包在向量单元中的多个数据元素进行操作：标量代码将两个数相加，而 SIMD 加法则取两个向量，比如各含八个 f32 值，一次性产生八个和。这种数据并行性位于单个线程_内部_，处于操作系统调度任何内容的层级之下。

CPU thread
SIMD op
0 1 2 N ⋯
SIMD lanes
CPU thread

Rust 的可移植 SIMD

在过去，用 Rust 编写 SIMD 意味着使用 core::arch 中特定于架构的厂商内建函数，例如 x86-64 上的 _mm256_add_ps 或 Arm 上的 vaddq_f32。这些内建函数特定于单一指令集，因此，一个要在多种架构上运行的程序需要为每种架构分别实现。

Rust 的可移植 SIMD 则是在这些内建函数之上增加了一层抽象。它提供了一个通用的类型 Simd<T, N>，表示 N 个 T 类型元素的向量。程序针对 Simd 编写一次算术、比较、归约和 lane 重排操作，编译器会将其降级为目标 CPU 所拥有的任何向量指令。

在 VectorWare，我们意识到 GPU 不过是可移植 SIMD 可以面向的又一种向量硬件。 额外的好处是，可移植 SIMD 位于 core 而非 std 中，它甚至不需要我们为 GPU 带来的 std 支持。

SIMT 就是 SIMD

GPU 以 NVIDIA 所称的 SIMT（单指令多线程）模型执行。一个 warp 发出一条指令，它的 32 个 lane 各自对该指令操作自己的数据。一条指令对多个数据元素进行操作_正是_ SIMD 的含义，而 SIMT 添加的逐 lane 寻址并不会改变这一点。warp 就是一个宽的向量单元，可移植 SIMD 向量直接映射到该单元上。

CPU thread 0 1 2 N ⋯ SIMD lanes
≈ GPU warp 0 1 2 N ⋯ warp lanes

例如，一个 Simd<i16, 32> 为 warp 的 32 个 lane 各提供一个 i16 元素，将两个这样的向量相加会编译成一条 warp 指令，每个 lane 同时将自己的元素相加。

CPU
let a: Simd<i16, 32> = [1, 1, 1, ..., 1];
let b: Simd<i16, 32> = [2, 2, 2, ..., 2];
let c = a + b;

compiles to
vpaddw %zmm2, %zmm1, %zmm0
a0+b0 lane 0
a1+b1 lane 1
a2+b2 lane 2
...
a31+b31 lane 31

println!("{c:?}");

GPU
let a: Simd<i16, 32> = [1, 1, 1, ..., 1];
let b: Simd<i16, 32> = [2, 2, 2, ..., 2];
let c = a + b;

compiles to
add.s16 %rs3, %rs1, %rs2;
a0+b0 lane 0
a1+b1 lane 1
a2+b2 lane 2
...
a31+b31 lane 31

println!("{c:?}");

这种新的映射补全了我们早期工作中的并行性层级。在 CPU 上，线程包含 SIMD lane；在 GPU 上，我们的 std::thread 就是一个 warp，其硬件 lane 扮演同样的角色。在两种情况下，core::simd 都驱动着这些 lane。

CPU ... thread 0 [0 1 2 N ⋯] thread 1 [0 1 2 N ⋯] thread N [0 1 2 N ⋯] SIMD lanes
≈ GPU ... warp 0 [0 1 2 N ⋯] warp 1 [0 1 2 N ⋯] warp N [0 1 2 N ⋯] warp lanes

世界首创：GPU 上的 core::simd

与我们之前的文章一样，这在视觉上很难展示，因为代码就是普通的 Rust。相同的 core::simd 类型在笔记本电脑上降级为 x86-64 SIMD，在 GPU 上则降级为 warp 操作，源代码无需任何更改。

这里我们定义了一个小的可移植 SIMD 例程并从 main 中调用它。它用到了该模型的核心特性：逐元素算术、产生 lane 掩码的比较、由该掩码驱动的 select，以及跨 lane 的水平归约。

#![feature(portable_simd)]

use core::simd::cmp::SimdPartialOrd;
use core::simd::num::SimdFloat;
use core::simd::{Select, Simd};

// Portable SIMD. This exact function also compiles and runs on the CPU,
// where it lowers to x86-64, Arm, or scalar code depending on the target.
// 可移植 SIMD。这个完全相同的函数也能在 CPU 上编译和运行，
// 根据目标平台降级为 x86-64、Arm 或标量代码。
fn relu_dot(a: Simd<f32, 32>, b: Simd<f32, 32>) -> f32 {
    // Elementwise multiply: 32 products computed at once.
    // 逐元素乘法：一次计算出 32 个乘积。
    let products = a * b;

    // Per-lane comparison produces a mask, one boolean per lane.
    // 逐 lane 比较产生一个掩码，每个 lane 一个布尔值。
    let positive = products.simd_gt(Simd::splat(0.0));

    // Keep the positive products, replace the rest with zero.
    // 保留正的乘积，其余替换为零。
    let clamped = positive.select(products, Simd::splat(0.0));

    // Horizontal add across all lanes down to a single scalar.
    // 跨所有 lane 的水平加法，归约为单个标量。
    clamped.reduce_sum()
}

fn main() {
    // Two 32-wide vectors, built with ordinary Rust.
    // 两个宽度为 32 的向量，用普通 Rust 构建。
    let a = Simd::<f32, 32>::splat(2.0);
    let b = Simd::<f32, 32>::from_array(std::array::from_fn(|i| i as f32 - 16.0));

    // Elementwise ops, a comparison mask, a select, and a reduction:
    // all ordinary portable SIMD, all running on the GPU.
    // 逐元素运算、比较掩码、select 和归约：
    // 全部是普通的可移植 SIMD，全部运行在 GPU 上。
    let result = relu_dot(a, b);

    // Printed from the GPU using our std support.
    // 使用我们的 std 支持从 GPU 打印输出。
    println!("relu_dot = {result}");
}

入口点是一个普通的 fn main，没有任何 GPU 专用注解。我们的工具链将其编译为 GPU kernel，结果通过我们的 std 支持从设备端打印出来。

下面是该程序在 GPU 上运行的录制视频，产生的结果与在 CPU 上运行完全相同。

实现

如前所述，这种映射建立在一个观察之上：warp 就是一个向量单元，其 lane 是各自独立寻址的。一旦 Simd<T, N> 按 lane 分布好，每一类操作都有直接的 warp 级对应物。

SIMD 逐元素操作是简单的情况。加法、乘法、比较以及其他逐 lane 运算符来自 Simd 上诸如 Add 的普通 Rust trait 实现。GPU 可以原生执行它们。

SIMD 归约，例如 reduce_sum 和 reduce_max，将每个 lane 合并为一个标量。这些操作使用 GPU 的 warp shuffle 指令在 lane 之间交换和合并数值，在每个 lane 中产生相同的标量结果。

SIMD 跨 lane 重排，例如 simd_swizzle! 和循环移位，在 lane 之间移动元素。由于 SIMD lane 就是 GPU warp lane，这些操作映射到相同的 warp shuffle 原语上——正是这些原语让 GPU lane 在数据交换方面如此出色。

SIMD 掩码的映射同样简洁。一个 Mask<T, N> 为每个 SIMD lane 提供一个谓词。Mask::select 在每个 warp lane 中执行一次选择。水平掩码查询，如 any 和 all，使用 GPU 的 vote 和 ballot 指令。

周围代码中的标量值，如循环计数器或常量，由每个 lane 以相同方式计算，因此就像在普通 CUDA 中一样在 warp 内简单复制。这与 ISPC 等数据并行语言所明确区分的 uniform 与 varying 是同一个区别，只不过在这里它是从 Rust 自身的类型中自然得出的：一个普通的 f32 是 uniform，而 Simd<f32, 32> 是 varying。

处理 lane

抽象与硬件不一致的地方只有一处：lane 数量。在 CPU 上，Simd<T, N> 允许 N 取 1 到 64 之间的任何值，但 GPU 硬件具有固定的宽度：NVIDIA 为 32 个 lane，AMD 为 32 或 64 个。只有当 N 与该宽度匹配时，映射才是一对一的。较小的 N 会使一些 lane 闲置，而较大的 N 则会让部分或全部 lane 处理多个元素。

当工作量超出 warp 的宽度时，我们需要一种方式来表达哪些 lane 做什么。将 warp 视作一个独立的"小机器"会有所帮助：一组固定的原语用于跨 lane 移动和合并数据，再加上关于哪些 lane 处于活动状态以及每个 lane 持有多少数据的约束。"对它进行编程"意味着在这些规则内将工作放置到 lane 上。

在 VectorWare，我们为那台机器提供了一个 IR。 它不是独立的数据结构，而是用 Rust 的类型系统，通过类型、泛型、const 泛型和 trait 约束来编码。程序由类型化的操作组成：ballot、shuffle、归约、扫描、gather、scatter、原子操作，以及针对比 warp 更宽的向量的分段循环。操作数、执行形状和容量也都是类型化的。由于操作在类型中携带其形状，许多无效程序根本无法构造出来。

该 IR 在 GPU 上不需要解释器。每个操作直接降级为对应的指令，与手写 PTX 相比零开销。同样的类型也让我们能在 CPU 上运行它。我们构建了一个引用解释器，可以确定性地执行该 IR，它就像是为 warp lane 编程准备的 Miri。我们在 GPU 上模拟代码时使用它，也用它进行差分测试。

我们的工作目前面向 NVIDIA，但这里没有任何东西是 CUDA 特有的。AMD 的 wavefront 和 Vulkan 的 subgroup 暴露了类似的原语和语义。该 IR 本身是与架构无关的 Rust。

优势

相同的源代码可以在 CPU 和 GPU 上运行。已经使用可移植 SIMD 的代码和库无需重写即可成为 GPU 执行的候选。

未经修改的 CPU 代码可以使用 GPU 的 lane 级并行性。GPU 感知的代码还可以更进一步，使用直接映射到 PTX 的 core::arch 内建函数。

Simd<T, N> 是一个普通的拥有所有权的值。借用检查器、生命周期和类型检查对它的作用方式与在 CPU 上完全相同。我们不是在添加一个 GPU 专用的向量类型或一套新的注解。我们是将 Rust 现有的可移植 SIMD 映射到 GPU 的原生执行模型上。在 VectorWare，我们正在让 GPU 表现得像一个普通的 Rust 平台。

不足

可移植 SIMD 在 Rust 中仍然不稳定。它需要 nightly 的 #![feature(portable_simd)]，且在其稳定之前接口可能会发生变化。

比 warp 窄的向量会让 lane 闲置，比 warp 宽的向量则会将每个操作变成更多条指令。只有当向量宽度与 warp lane 数量匹配时，该抽象才是零成本的。

并非每个跨 lane 操作都能映射为一条高效的 warp 指令。与硬件支持的 pattern 匹配的 shuffle 代价很低，但任意的排列可能需要多条指令或经由共享内存。归约和 all/any 等水平操作还会在 warp 内充当同步点，这限制了调度器在重叠工作时的自由度。

我们必须修改编译器，才能使该抽象在与 Rust 其他特性交互时保持健全。由于这是未经探索的领域，我们还不能确信已经覆盖了所有情况。

未来工作

随着 SIMD、线程和 async 都映射到了 GPU 上，自然的下一步是将它们组合起来：线程将工作分散到各个 warp，core::simd 将数据分散到每个 warp 内的 lane 上，async 则组织它们之间的并发。

我们还对将矩阵形状的 SIMD 降级到 GPU 的 tensor core 感兴趣，也希望能将普通的标量 Rust 循环自动向量化为 Simd 操作，这样代码无需针对 core::simd 编写就能获得 warp 级并行性。作为 Rust 编译器团队成员，我们热衷于探索其中有多少可以在编译器自身中完成。

CPU 和 GPU 共享的向量表示是有价值的，尽管目前尚不清楚今天的可移植 SIMD 类型是否是正确的基石。另一方面，它们在 core 和 std API 中很大程度上自成一体。还需要做更多的探索。

VectorWare 只关注 Rust 吗？

我们在 GPU 上取得进展的速度证明了 Rust 抽象和生态系统的强大。

作为一家公司，我们理解并非每个人都使用 Rust。我们未来的产品将支持多种编程语言和运行时。然而，我们相信 Rust 特别适合构建高性能、可靠的 GPU 原生应用，这正是我们最兴奋的地方。

• 原文链接： vectorware.com/blog/simd...
• 鸿途知科网 AI 助手，为大家转译优秀英文文章，如有翻译不通的地方，还请包涵～