什么Rust已经进军GPU优化了吗？C++游戏统治地位要崩塌吗？

傻梦兽
2,862 阅读13分钟

什么Rust已经进军GPU优化了吗？C++游戏统治地位要崩塌吗？

大家好，我是梦兽。一个 WEB 全栈开发和 Rust 爱好者。如果你对 Rust 非常感兴趣，可以关注梦兽编程公众号获取群，进入和梦兽一起交流。

学习输出知识不容易，能否给我一个赞。

我最近拜读了 Zach Nussbaum 的精彩文章《Optimizing a WebGPU Matmul Kernel for 1TFLOP+ Performance》，文章内容深入探讨了如何优化一个 WebGPU 矩阵乘法内核，从而实现出色的 1TFLOP+ 性能。这让我萌生了一个想法——为什么不尝试用 Rust GPU 来重新实现这个过程呢？不仅会很好玩，也能带来一些有趣的技术探索。

在接下来的过程中，我会紧密参考 Zach 的原作，从实现到优化全程保持一致。同时，我还计划对比使用 Rust 与原文所采用的 WGSL 和 TypeScript 时，在哪些关键点上存在差异，从实际开发效率到性能表现，一一展开讨论。

更值得期待的是，最后我会分享在 GPU 开发中选择 Rust 所带来的一些独特优势。这不仅是一次实现上的实验，也是一次技术上的深度剖析。

在这里，我要特别感谢 Zach 大方地允许我基于他的博客内容开展这个尝试！

什么是 Rust GPU?

Rust GPU 是一个允许使用 Rust 编程语言为 GPU 编写代码的项目。传统上GPU 通常使用诸如 WGSL、GLSL、MSL 或 HLSL 等专门的语言进行编程。Rust GPU 的出现改变了这一现状，让开发者可以使用 Rust 来编写 GPU 程序（通常称为“着色器”或“内核”）。

这些 Rust GPU 程序会被编译为 SPIR-V，这是大多数 GPU 能够理解的一种低级格式。由于 SPIR-V 是 Vulkan 使用的格式，Rust GPU 因而能够将基于 Rust 的 GPU 程序集成到任何兼容 Vulkan 的工作流程中。

想了解更多信息，可以访问 Rust GPU 的官网或者 GitHub 仓库。

Rust GPU 是如何工作的？

简单来说，Rust GPU 的核心目标是将你的 Rust 代码编译成 SPIR-V 格式，这正是 GPU 执行的代码形式。这里需要注意的是，Rust GPU 并不会干涉你如何处理 CPU 和 GPU 之间的通信或数据传输——在这方面，你有完全可控得选择任何适合你项目需求的主机 CPU 库，不管它是用什么语言编写的。

在 Rust 生态中，有几款热门的库可以帮助你完成这类操作：

  • ash：Rust 的低级 Vulkan 绑定库，它能为开发者提供最高的操作灵活性和控制力。
  • vulkano：这是一个更高级的 Vulkan 库，可以帮助开发者简化一些常见任务。
  • wgpu：一个跨平台的 GPU 操作抽象库，它支持 Vulkan、DirectX、Metal 和 WebGPU 等多种后端。

不过需要强调的是，就算你在 GPU 端使用的是 Rust 代码，CPU 端的语言不一定也非得是 Rust——任何语言都可以，只要符合你的需求就行。

我们将使用什么？

在 Zach 的文章中，他使用 WGSL 来编写 GPU 程序。这些程序及其数据通过 TypeScript 与浏览器内置的 WebGPU CPU 代码进行交互，在 CPU 和 GPU 之间进行传递。

我们将采取一种不同的方法：通过 Rust GPU 用 Rust 编写 GPU 程序，同时管理包括 CPU 端代码在内的所有内容。无论是 GPU 程序还是控制它们的代码，都用同一种语言编写。如果你熟悉 Web 编程，可以把我们的方式类比为 JavaScript 同时运行在服务器和客户端上的概念。

在 CPU 和 GPU 层面都使用 Rust 有很多优势，例如统一的开发工具和可共享的代码库。但这也要求我们明确区分代码在哪一端运行。我尽量确保这一点易于理解和跟随。

为了实现 CPU 和 GPU 之间的通信，我们将使用 wgpu。wgpu 是一个高级的 Rust 库，它实现了 WebGPU API。在 Web 环境下，它直接与浏览器的 WebGPU 实现交互；而在原生平台上，它会将 API 调用转换为对应平台的 GPU API（如 Vulkan、DirectX 或 Metal）。通过这种方式，我们可以在包括 Windows、Linux、macOS2、iOS3、Android 和 Web4 在内的多个平台上运行同一套代码。

选择 Rust GPU 和 wgpu，让我们拥有了一个干净、可移植的开发环境，并用Rust 完整地编写所有内容。这种方式巧妙地结合了高性能和跨平台支持，为我们的项目奠定了坚实的技术基础。

GPU程序基础

最小的执行单位是线程，它负责执行GPU程序。

线程的集合被称为“工作组”，它们被组织在一起并以并行的方式运行（在CUDA中被称为线程块）。工作组内的线程能够访问相同的共享内存。

可以同时调度多个这样的工作组。在CUDA中，这种多工作组的结构被称为网格（由线程块组成）。

工作组和工作组的调度方式是以三维形式定义的。工作组的大小通过compute(threads((x, y, z)))来表示，其中工作组内线程的总数等于x * y * z。

编写内核

内核 1：基础内核
计算矩阵 A 和 B 的点积并将结果写入矩阵 C 的最简单方法是，对于矩阵 A（形状为 M）的每一行，遍历 A（形状为 K）的列，并与矩阵 B 中对应的值相乘。

在这里，我们与 Zach 的文章有了第一个不同之处。在 WGSL 中，您必须在顶层作用域中定义输入：

struct Dimensions {
  M: u32,
  K: u32,
  N: u32,
}

@group(0) @binding(0) var<uniform> dimensions: Dimensions;
@group(0) @binding(1) var<storage, read> a: array<f32>;
@group(0) @binding(2) var<storage, read> b: array<f32>;
@group(0) @binding(3) var<storage, read_write> result: array<f32>;

然后编写你的内核：

 @compute @workgroup_size(1)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
  let index = global_id.x;
  let row = index / dimensions.N;
  let col = index % dimensions.N;

  if (index < dimensions.M * dimensions.N) {
    var sum = 0.0;
    for (var i: u32 = 0u; i < dimensions.K; i = i + 1u) {
      sum = sum + a[row * dimensions.K + i] * b[i * dimensions.N + col];
    }
    result[row * dimensions.N + col] = sum;
  }
}

在使用 Rust GPU 时，我们通过将输入作为内核的参数来进行指定，并借助过程宏对其进行配置。这种方式不仅简化了开发流程，还能够提升代码的可读性和维护性，为开发者提供了更高效的 GPU 编程体验。

#![no_std]

use settings::Dimensions;
use spirv_std::glam::UVec3;
use spirv_std::spirv;

#[spirv(compute(threads(1)))]
pub fn matmul(
    #[spirv(global_invocation_id)] global_id: UVec3,
    #[spirv(uniform, descriptor_set = 0, binding = 0)] dimensions: &Dimensions,
    #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] a: &[f32],
    #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] b: &[f32],
    #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] result: &mut [f32],
) {
    let index = global_id.x;
    let row = index / dimensions.n;
    let col = index % dimensions.n;

    if index < dimensions.m * dimensions.n {
        let mut sum = 0.0;

        for i in 0..dimensions.k {
            let a_val = a[(row * dimensions.k + i) as usize];
            let b_val = b[(i * dimensions.n + col) as usize];
            sum += a_val * b_val;
        }

        result[(row * dimensions.n + col) as usize] = sum;
    }
}

这段代码看起来与普通的 Rust 代码无异，但运行时却完全依赖于 GPU。

在讨论这种 Rust 实现时，有几个关键点需要注意：

首先，内核使用了 Rust 的 #![no_std] 属性。这是运行于 GPU 上的必要条件，因为 GPU 上无法访问 Rust 的标准库 (std)。取而代之的是使用 core 和 spirv_std 提供类似 std 的功能，从而满足代码的基本需求。

其次，库的引入方式与普通 Rust 一样，通过 use 进行导入。模块系统的工作方式没有任何变化，非常直观和简洁。

此外，我们还引入了一个 vendored 副本的 glam 库，它就是来自 crates.io 上的 glam crate，没有任何魔改，让开发者感觉熟悉、友好。

关于循环部分，在内层循环 for i in 0..dimensions.k 中，直接使用了 Rust 中的 for 语法和范围表达式。这种抽象层次比 WGSL、GLSL 或 HLSL 等着色器语言中需要手动操作索引的实现要高级得多，既简洁，也更符合惯用模式。

最后，值得一提的是，代码中对于只读输入数据，可以通过不可变引用 (&Dimensions / &[f32]) 表示；而对于可写输出数据，则使用可变引用 (&mut [f32])。这与普通 Rust 代码的写法非常相似，对于熟悉 Rust 的开发者来说，上手没有任何难度。

这种模式不仅逻辑清晰，还保留了 Rust 一贯的表达能力和安全性，为开发 GPU 应用创造了便利。

为什么总是出现那么多usize？

在Rust中，usize被定义为与运行代码的硬件的本地指针宽度一致的数据类型。这一点非常重要，因为Rust使用usize来索引切片，确保访问是适当的指针对齐状态。

在大多数GPU硬件上，usize实际上等同于u32。但Rust编译器不会直接假定这一点，因为这样可能导致问题，例如如果代码运行在usize为u64的硬件上。Rust不会让你隐式地将u32视为usize使用，你必须显式地进行类型转换，从而明确告诉编译器“我知道在目标硬件上这样做是安全的”。

虽然这种明确性看起来可能有些繁琐，但它正是Rust防止隐蔽性错误的一种方式。它迫使你考虑自己对硬件对齐和指针大小的假设是否正确，从而使代码更加具有可移植性和可靠性。

矩阵乘法是一个具有大量索引操作以及行列计算的特殊情况。大多数Rust GPU代码并不会像这些示例一样频繁地进行usize类型转换。据我了解目前大多数的WebGPU是32位的

调度工作组

由于每个工作组仅对应一个线程 (#[spirv(compute(threads(1)))]), 因此每个工作组仅处理一个结果元素 result[i, j]。

为了计算完整的矩阵，我们需要启动与矩阵 m * n 条目数量相同的工作组。这里我们具体指定为 (Uvec3::new(m * n, 1, 1) 在 CPU 上：

impl GridComputation for Naive {
    fn workgroup(&self) -> UVec3 {
        UVec3::new(1, 1, 1)
    }

    fn dispatch_count(&self, m: u32, n: u32) -> UVec3 {
        UVec3::new(m * n, 1, 1)
    }
}

dispatch_count() 函数在 CPU 上运行，用于通过 CPU 到 GPU 的接口（在我们的案例中是 wgpu）来配置并分配工作到 GPU 上。

let dispatch_count = <T as GridComputation>::dispatch_count(&self.variant, m, n);
...
compute_pass.dispatch_workgroups(dispatch_count.x, dispatch_count.y, dispatch_count.z);

这段代码看上去比实际需要的更复杂。我在写博客文章时，为了能够轻松替换不同的内核及其设置，使用了泛型和特性对与GPU交互的CPU端代码进行了抽象。

其实，你也可以直接硬编码值来简化代码。

内核2：更多线程！

通过增加线程数量来提升性能。以下是具体的改进点：

  1. 线程分布：每个线程负责计算结果矩阵 C 的一个元素。
  2. 工作组划分：将线程分组为工作组（workgroups），每个工作组中的线程可以共享内存。
  3. 并行化计算：通过 Rust GPU 的 spirv 宏定义线程数，并利用 GPU 的并行能力处理更多数据。 在 Rust 中，内核代码如下：
#[spirv(compute(threads(256)))] // 256 个线程
pub fn matmul(...) {
    // 每个线程计算一个矩阵元素
    let index = global_id.x;
    let row = index / dimensions.n;
    let col = index % dimensions.n;
    ...
}

相比第一版内核（Kernel 1），这种方法显著提升了性能，因为更多的线程可以同时工作。

内核 3：使用二维工作组进行计算

在 GPU 优化中，使用 2D 工作组（2D Workgroups） 是一种常见的策略，用于优化计算任务，尤其是矩阵乘法（matmul）这样具有二维数据结构的任务。以下是 “Calculating with 2D workgroups” 的优化原理和方式：

  1. 什么是 2D 工作组？ GPU 的线程组织可以是 1D、2D 或 3D 的形式。对于矩阵乘法来说，矩阵本身是二维的，因此使用 2D 工作组 更符合任务的结构特点。

  2. 1D 工作组：线程只按一维排列，适合线性任务。

  3. 2D 工作组：线程按二维网格排列，每个线程负责矩阵中的一个元素或一个小块。

  4. 3D 工作组：线程按三维网格排列，适合三维数据（如体积渲染）。 在 2D 工作组中，线程的二维索引（如 (x, y)）直接映射到矩阵的行和列，从而简化了线程任务的分配。

这是数学模型问题，具体内容需要读者自己补课。这里不过多解释

#[spirv(compute(threads(16, 16)))] // 定义工作组为 16x16
pub fn matmul_2d_workgroup(
    #[spirv(global_invocation_id)] global_id: UVec3, // 全局线程 ID
    #[spirv(workgroup_id)] workgroup_id: UVec3,     // 工作组 ID
    #[spirv(local_invocation_id)] local_id: UVec3,  // 工作组内线程 ID
    #[spirv(uniform, descriptor_set = 0, binding = 0)] dimensions: &Dimensions, // 矩阵维度
    #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] a: &[f32], // 矩阵 A
    #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] b: &[f32], // 矩阵 B
    #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] result: &mut [f32], // 矩阵 C
    #[spirv(workgroup)] shared_memory: &mut [f32], // 共享内存
) {
    let row = global_id.y; // 当前线程对应结果矩阵的行
    let col = global_id.x; // 当前线程对应结果矩阵的列

    let local_row = local_id.y; // 工作组内线程的行索引
    let local_col = local_id.x; // 工作组内线程的列索引

    let mut sum = 0.0;

    // 分块计算
    for tile in 0..(dimensions.k / 16) {
        // 加载 A 和 B 的子块到共享内存
        let a_index = row * dimensions.k + (tile * 16 + local_col);
        let b_index = (tile * 16 + local_row) * dimensions.n + col;

        shared_memory[local_row * 16 + local_col] = a[a_index];
        shared_memory[256 + local_row * 16 + local_col] = b[b_index];

        // 等待所有线程加载完共享内存
        unsafe { spirv_std::arch::workgroup_barrier() };

        // 计算子块的结果
        for k in 0..16 {
            let a_val = shared_memory[local_row * 16 + k];
            let b_val = shared_memory[256 + k * 16 + local_col];
            sum += a_val * b_val;
        }

        // 等待所有线程完成计算
        unsafe { spirv_std::arch::workgroup_barrier() };
    }

    // 写入结果矩阵
    if row < dimensions.m && col < dimensions.n {
        result[row * dimensions.n + col] = sum;
    }
}

总结

这篇文章展示了 Rust 在 GPU 编程中的强大潜力，以及如何通过 Rust GPU 和 wgpu 优化矩阵乘法内核。相比传统的 GPU 编程语言，Rust 提供了更安全、更高效的开发体验，同时实现了跨平台支持。

注意

感谢您阅读至此。如果您觉得这篇文章有用，请点赞并分享。也许其他人也会觉得它有用。💖

创建和维护这个博客以及相关的库带来了十分庞大的工作量，即便我十分热爱它们，仍然需要你们的支持。或者转发文章。通过赞助我，可以让我有能投入更多时间与精力在创造新内容，开发新功能上。赞助我最好的办法是微信公众号看看广告。

avatar
文章
阅读
粉丝