2 篇博文 含有标签「矩阵乘法优化」

背景​

在深度学习和高性能计算中，矩阵乘法（Matmul）是核心操作之一， 也是现代 AI 模型如 GPT 和 Transformer 的基础计算单元。

随着 WebGPU 的发展，我们可以在浏览器中高效运行 GPU 计算，为前端机器学习应用带来了更多可能。

本文将通过五个阶段，从基础内核出发，逐步优化 WebGPU 矩阵乘法内核， 最终达到 超过 1TFLOPS 的性能，并探讨 WebGPU 与 CUDA 的区别及应用场景。

什么是 WebGPU？​

WebGPU 是为浏览器设计的下一代 GPU 编程接口，原生支持计算着色器（Compute Shader）， 通过使用 WGSL（WebGPU Shading Language） 编写 GPU 代码， 并支持多种硬件平台（如 Vulkan 和 Metal）。

优势​

1. 跨平台：兼容 Vulkan、Metal 和 DirectX。
2. 高性能：原生支持并行计算，如矩阵乘法和深度学习。
3. 便捷性：无需传统 WebGL 的复杂 hack，可直接进行机器学习计算。

WebGPU 与 CUDA 的区别​

WebGPU 计算着色器基础​

1. 线程（Thread）：最小的并行执行单元。
2. 工作组（Workgroup）：线程的集合，支持组内存共享。
3. 网格（Grid）：多个工作组组成的并行执行结构。

示例：@workgroup_size(x, y, z) 定义每个工作组的线程数量为 (x \times y \times z)。

矩阵乘法优化的五个阶段​

阶段 1：基础实现​

Python 示例：​

def matmul(a, b, c):
    m, k, n = len(a), len(a[0]), len(b[0])
    for i in range(m):
        for j in range(n):
            c[i][j] = sum(a[i][l] * b[l][j] for l in range(k))

WGSL 实现：​

@compute @workgroup_size(1)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
  let row = global_id.x / dimensions.N;
  let col = global_id.x % dimensions.N;
  if (row < dimensions.M && col < dimensions.N) {
    var sum = 0.0;
    for (var i: u32 = 0; i < dimensions.K; i++) {
      sum += a[row * dimensions.K + i] * b[i * dimensions.N + col];
    }
    result[row * dimensions.N + col] = sum;
  }
}

存在问题：​

• 每个线程仅计算一个结果，导致大量工作组启动开销高。
• 每个工作组重复加载数据，没有利用缓存。

阶段 2：增加线程数量​

通过提高每个工作组的线程数（如 @workgroup_size(256)），显著减少工作组的数量，从而降低启动开销。

阶段 3：二维工作组优化​

通过将工作组从一维扩展到二维（如 (16 \times 16)），使每个工作组能够并行计算更多结果。

@compute @workgroup_size(16, 16)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
  let row = global_id.y;
  let col = global_id.x;
  ...
}

阶段 4：内核平铺（Tiling）​

采用平铺策略，每个线程一次计算多个结果（如 (1 \times 4)），进一步提升性能。

阶段 5：循环展开（Unrolling）​

通过手动展开循环，减少 GPU 在运行时的循环控制开销，并利用指令级并行，性能大幅提升。

优化成果​

• 性能提升 超过 1000 倍，达到 1TFLOPS 的运算强度。
• 有效利用 WebGPU 的多线程并行与缓存机制。
• 实现了更高效的矩阵乘法内核，适用于前端高性能计算场景。

参考资料​

• WebGPU 官方文档
• WGSL 语言规范
• CUDA 编程指南
• Optimizing a WebGPU Matmul Kernel for 1TFLOP+ Performance

引言​

GPU 编程通常依赖于如 WGSL、GLSL 或 HLSL 等语言。然而，Rust GPU 项目开辟了新的可能，允许开发者直接使用 Rust 编程语言 编写 GPU 内核代码，结合强大的类型安全性和性能优化能力。

本文基于 Zach Nussbaum 的文章《Optimizing a WebGPU Matmul Kernel for 1TFLOP+ Performance》，详细探讨如何在 Rust GPU 中实现矩阵乘法（matmul）内核优化，逐步探索 Rust 在 GPU 编程中的独特优势。

什么是 Rust GPU？​

Rust GPU 是一个专为 GPU 编程设计的项目，通过将 Rust 代码编译为 GPU 可识别的 SPIR-V 格式，使其能够无缝集成到 Vulkan 等兼容的 GPU 编程生态中。

核心特点​

• Rust 编程支持：无需依赖 WGSL 等传统 GPU 专用语言。
• 生态兼容性：与 Vulkan、DirectX 和 Metal 集成。
• 安全与高效：Rust 的类型系统和零开销抽象为 GPU 开发提供更高的稳定性。

Rust GPU 的工作原理​

Rust GPU 专注于将 Rust 代码编译为 SPIR-V，而 CPU 与 GPU 的通信通常通过其他库（如 wgpu、vulkano 或 ash）实现。

在本文中，我们使用 wgpu 库来管理 CPU 和 GPU 的交互，确保通信的高效性和跨平台支持。

核心概念：线程与工作组​

GPU 的并行计算由以下核心概念构成：

1. 线程（Thread）：最小执行单元，运行 GPU 内核代码。
2. 工作组（Workgroup）：线程的集合，能够共享组内存并协作计算。
3. 网格（Grid）：由多个工作组组成，适合大规模任务的并行执行。

工作组维度可通过 (x, y, z) 三维定义，如下所示：

#[spirv(compute(threads(x, y, z)))]
pub fn kernel(...) { ... }

Rust GPU 的实现：从简单到优化​

以下是矩阵乘法内核优化的四个阶段。

阶段 1：基础矩阵乘法内核​

我们从最基础的矩阵乘法实现开始，为矩阵 (A) 和 (B) 计算结果矩阵 (C)。以下是 Rust GPU 的实现代码：

#![no_std]

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 {
            sum += a[(row * dimensions.k + i) as usize] * b[(i * dimensions.n + col) as usize];
        }
        result[(row * dimensions.n + col) as usize] = sum;
    }
}

问题：​

• 每个线程仅计算一个结果，导致启动大量工作组，增加开销。
• 矩阵数据重复加载，未充分利用缓存。

阶段 2：增加线程数量​

通过提高工作组线程数（如 compute(threads(256))），可以显著减少工作组的数量，降低启动开销。

阶段 3：二维工作组​

为支持更大的矩阵，将工作组扩展为二维（如 (16 \times 16)），使每个工作组可以处理更多矩阵元素。

#[spirv(compute(threads(16, 16)))]
pub fn matmul(...) { ... }

阶段 4：内核平铺（Tiling）​

通过平铺策略，每个线程一次计算多个矩阵元素，进一步减少启动开销。

#[spirv(compute(threads(16, 16)))]
pub fn matmul(...) {
    let row = global_id.y * TILE_M;
    let col = global_id.x * TILE_N;

    let mut sums = [[0.0; TILE_N as usize]; TILE_M as usize];

    for k in 0..dimensions.k as usize {
        for i in 0..TILE_M as usize {
            let a_elem = a.get(row + i).unwrap_or(&0.0);
            for j in 0..TILE_N as usize {
                let b_elem = b.get(col + j).unwrap_or(&0.0);
                sums[i][j] += a_elem * b_elem;
            }
        }
    }

    for i in 0..TILE_M as usize {
        for j in 0..TILE_N as usize {
            let output_row = row + i;
            let output_col = col + j;
            if output_row < dimensions.m as usize && output_col < dimensions.n as usize {
                result[output_row * dimensions.n as usize + output_col] = sums[i][j];
            }
        }
    }
}

Rust GPU 的独特优势​

1. 共享代码：Rust 模块化设计可让 CPU 和 GPU 使用相同数据结构，避免重复定义。
2. 条件编译与 CPU 调试：支持在 CPU 上运行 GPU 内核，方便调试和验证。
3. 生态系统支持：Rust 的 no_std 和现有库（如 spirv_std）提供了丰富的功能复用能力。
4. 泛型与零开销抽象：通过特性（Traits）和泛型优化代码的扩展性与可维护性。

总结​

Rust GPU 结合 Rust 的安全性与性能优势，为 GPU 编程提供了强大支持。通过本文的四阶段优化，从基础实现到高级平铺技术，展示了如何有效提升矩阵乘法内核性能。

Rust GPU 不仅提升了 GPU 编程的开发体验，更为跨平台高性能计算带来了新的可能性。欢迎开发者加入 Rust GPU 项目，探索 GPU 编程的未来！