software

使用 Rust 和 wgpu 的高性能 GPGPU

发布: (2025年12月14日 GMT+8 22:46)
6 min read
原文: Dev.to

Source: Dev.to

计算应用的架构

GPGPU 应用与传统渲染循环有显著差异。在图形上下文中，管线复杂，涉及顶点着色器、片段着色器、光栅化和深度缓冲。相比之下，计算管线则清爽得多，主要由数据缓冲区和计算着色器组成。工作流程包括初始化 GPU 设备、加载着色器代码、创建 GPU 可访问的内存缓冲区，以及分发 工作组 来执行逻辑。

wgpu 的核心抽象包括 Instance、Adapter、Device 和 Queue。

  • Instance – API 的入口点。
  • Adapter – 表示物理硬件。
  • Device – 逻辑连接，允许你创建资源。
  • Queue – 用于提交命令缓冲区以供执行的地方。

与需要窗口表面的图形渲染不同，计算上下文可以完全 无头（headless）运行，非常适合后台处理工具或服务器端应用。

用 WGSL 编写内核

在 GPU 上执行的逻辑使用 WebGPU 着色语言（WGSL）编写。这种语言感觉像是 Rust 与 GLSL 的混合体。对于计算着色器，我们定义一个带有 @compute 属性的入口点，并指定工作组大小。GPU 会在三维网格上并行执行此函数。

// shader.wgsl
@group(0) @binding(0)
var data: array;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) global_id: vec3) {
    let index = global_id.x;
    // Guard against out‑of‑bounds access if the array size
    // isn't a perfect multiple of the workgroup size
    if (index < arrayLength(&data)) {
        data[index] = data[index] * data[index];
    }
}

工作组大小设为 64。当我们从 Rust 端分发工作时，需要计算多少个 64 的组才能覆盖整个数据数组。函数内部的逻辑很简单，但硬件会同时执行成千上万的实例。

缓冲区管理与绑定组

内存管理是 GPGPU 编程中最关键的环节。CPU 与 GPU 往往拥有不同的内存空间。为了解决这个问题，wgpu 使用 缓冲区。对于计算操作，我们通常需要一个 存储缓冲区（Storage Buffer），它允许着色器读写任意数据。然而，CPU 直接读取 GPU 内存要么很慢，要么根本不可能，所以常用 暂存缓冲区（Staging Buffer） 策略：

  1. 创建一个驻留在 GPU 上的缓冲区用于处理。
  2. 创建一个可以映射供 CPU 读取的独立缓冲区。

缓冲区创建完毕后，需要告诉着色器它们的位置。这通过 绑定组（Bind Groups） 完成。绑定组布局（Bind Group Layout） 描述接口——例如，绑定槽 0 是一个存储缓冲区。绑定组（Bind Group） 本身则将实际的 wgpu::Buffer 对象连接到该槽位。这种分离使 wgpu 能在 GPU 看到任何命令之前验证资源使用，避免了许多低层图形 API 常见的崩溃。

分发工作

管线创建并且数据已上传后，我们对命令进行编码：

let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
    label: None,
});

{
    let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
        label: None,
        timestamp_writes: None,
    });
    cpass.set_pipeline(&compute_pipeline);
    cpass.set_bind_group(0, &bind_group, &[]);
    // Example: 1024 elements, workgroup size 64 → 16 workgroups on X axis
    cpass.dispatch_workgroups(data_size / 64, 1, 1);
}

分发完毕后，如果我们想把结果读回 CPU，需要发出一次拷贝命令，将 GPU 常驻的存储缓冲区的数据转移到可映射读取的暂存缓冲区。最后，完成编码器并将命令缓冲区提交到队列。

异步读取

wgpu 是异步的。向队列提交工作会立即返回，但 GPU 会在稍后处理这些命令。要读取数据，我们必须映射暂存缓冲区，这会返回一个 Future。应用需要轮询设备，例如：

device.poll(wgpu::Maintain::Wait);

这会阻塞主线程，直到 GPU 操作完成并且映射回调触发。缓冲区映射后，我们可以把原始字节转换回 Rust 切片，复制到本地向量，并解除映射，从而创建一个同步点，保证在 CPU 访问结果之前 GPU 已经完成工作。

结论

wgpu 生态为 GPGPU 编程提供了一个稳健的基础，兼顾安全性和可移植性，同时不牺牲硬件的原始并行算力。通过标准化 WGSL 与 WebGPU 资源模型，开发者可以编写在桌面、移动端和网页上无缝运行的计算内核。虽然设置管线和管理内存缓冲区的样板代码比高级 CPU 线程更冗长，但其回报是能够并行处理海量数据，释放出 CPU 单独难以实现的性能潜能。

Back to Blog

相关文章

阅读更多 »