GPU 上的 Rust SIMD

在 VectorWare，我们正在打造第一家 GPU 原生软件公司。今天，我们很高兴地宣布我们可以在 GPU 上成功使用 Rust 的便携式 SIMD ( core::simd )。这一里程碑标志着我们朝着让开发人员能够使用熟悉的 Rust 抽象来编写复杂的高性能应用程序，充分利用 GPU 硬件的全部功能的愿景迈出了重要一步。线程下方的并行性当我们将 Rust 线​​程引入 GPU 时，我们将每个 std::thread 映射到 GPU warp 。这让我们可以在 GPU 上运行许多并发线程，但不使用每个线程/扭曲中的并行通道。在 CPU 上，线程内并行性的抽象是 SIMD 。一条指令对打包到向量单元中的多个数据元素进行操作：其中标量代码将两个数字相加，SIMD 加法采用两个向量（例如，八个 f32 值）并一次生成八个和。这种数据并行性位于单个线程内，低于操作系统调度任何内容的级别。 CPU 线程 SIMD op 0 1 2 N ⋯ SIMD 通道 CPU 线程 Rust 的便携式 SIMD 从历史上看，在 Rust 中编写 SIMD 意味着在 core::arch 中获取特定于架构的供应商内在函数，例如 x86-64 上的 _mm256_add_ps 或 Arm 上的 vaddq_f32 。这些内在函数特定于单个指令集，因此在多个体系结构上运行的程序需要为每个体系结构单独实现。相反，Rust 的可移植 SIMD 在这些内在函数之上添加了一个抽象层。它提供了一个通用类型 Simd ，表示由 N 个 T 类型元素组成的向量。程序针对 Simd 编写一次算术、比较、归约和通道洗牌，然后编译器将它们降低为目标 CPU 具有的任何向量指令。在 VectorWare，我们意识到 GPU 只是便携式 SIMD 目标的又一个矢量硬件。额外的好处是，便携式 SIMD 存在于核心而不是 std 中，它甚至不需要我们为 GPU 带来的 std 支持。 SIMT 是 SIMD GPU 在 NVIDIA 称为 SIMT（即单指令多线程）的模型中执行的。 Warp 发出一条指令，其 32 个通道中的每一个都根据自己的数据运行该指令。对许多数据元素进行操作的一条指令正是 SIMD 的含义，SIMT 添加的每通道寻址不会改变这一点。扭曲是一个宽向量单元，便携式 SIMD 向量直接映射到该单元。 CPU 线程 0 1 2 N ⋯ SIMD 通道 ≈ GPU 扭曲 0 1 2 N ⋯ 扭曲通道 例如，Simd 为扭曲的 32 个通道中的每一个提供一个 i16 元素，并且添加两个这样的向量编译为单个扭曲指令，其中每个通道一次添加其元素。 CPU 让 a: Simd = [ 1 , 1 , 1 , ... , 1 ];让 b: Simd = [ 2 , 2 , 2 , ... , 2 ];令 c = a + b；编译为 vpaddw %zmm2, %zmm1, %zmm0 a0+b0 Lane 0 a1+b1 Lane 1 a2+b2 Lane 2 a31+b31 Lane 31 ⋯ println!("{c:?}"); GPU 让 a: Simd = [ 1 , 1 , 1 , ... , 1 ];让 b: Simd = [ 2 , 2 , 2 , ... , 2 ];令 c = a + b；编译为 add.s16 %rs3, %rs1, %rs2; a0+b0 车道 0 a1+b1 车道 1 a2+b2 车道 2 a31+b31 车道 31 ⋯ println!("{c:?}");这个新的映射完善了我们早期工作中的并行层次结构。在 CPU 上，线程包含 SIMD 通道，而在 GPU 上，我们的 std::thread 是一个扭曲，其硬件通道起着相同的作用。在这两种情况下，core::simd 驱动这些通道。 CPU ⋯ 线程 0 0 1 2 N ⋯ 线程 1 0 1 2 N ⋯ 线程 N 0 1 2 N ⋯ SIMD 通道 ≈ GPU ⋯ 扭曲 0 0 1 2 N ⋯ 扭曲 1 0 1 2 N ⋯ 扭曲 N 0 1 2 N ⋯ 扭曲通道 世界第一：GPU 上的 core::simd 与我们之前的文章一样，这很难直观地展示，因为代码是普通的铁锈。相同的 core::simd 类型在笔记本电脑上降低到 x86-64 SIMD，降低到 GPU 上的扭曲操作，而无需更改源。这里我们定义了一个小型的可移植 SIMD 例程并从 main 中调用它。它运用了模型的核心功能：元素运算、生成通道掩码的比较、由该掩码驱动的选择以及跨通道的水平缩减。 #![feature(portable_simd)] 使用 core :: simd :: cmp :: SimdPartialOrd ；使用 core::simd::num:SimdFloat ；使用 core :: simd :: { Select , Simd }; // 便携式SIMD。这个确切的函数也在 CPU 上编译和运行， // 根据目标，它会降低到 x86-64、Arm 或标量代码。 fn relu_dot (a : Simd < f32 , 32>, b : Simd < f32 , 32>) -> f32 { // 元素相乘：一次计算 32 个乘积。令产品 = a * b； // 每通道比较产生一个掩码，每通道一个布尔值。设正数=产品。 simd_gt ( Simd::splat( 0.0 )); // 保留正积，其余的用零代替。让夹紧=正。选择（产品，Simd :: splat（0.0））； // 将所有通道水平相加至单个标量。夹住。 reduce_sum () } fn main () { // 两个 32 宽的向量，用普通 Rust 构建。让 a = Simd :: < f32 , 32 > :: splat ( 2.0 );让 b = Simd :: < f32 , 32 > :: from_array ( std :: array :: from_fn ( | i | i as f32 - 16.0 )); // 艾尔