3 篇博文 含有标签「高性能计算」

使用 CubeCL，您可以通过 Rust 编写 GPU 程序，借助零成本抽象开发出可维护、灵活且高效的计算内核。CubeCL 目前完全支持函数、泛型和结构体，并部分支持 trait、方法和类型推导。随着项目的演进，我们预计将支持更多 Rust 语言特性，同时保持最佳性能。

示例​

只需用 cube 属性标注函数，即可指定其运行在 GPU 上：

use cubecl::prelude::*;

#[cube(launch_unchecked)]
fn gelu_array<F: Float>(input: &Array<Line<F>>, output: &mut Array<Line<F>>) {
    if ABSOLUTE_POS < input.len() {
        output[ABSOLUTE_POS] = gelu_scalar(input[ABSOLUTE_POS]);
    }
}

#[cube]
fn gelu_scalar<F: Float>(x: Line<F>) -> Line<F> {
    let sqrt2 = F::new(comptime!(2.0f32.sqrt()));
    let tmp = x / Line::new(sqrt2);

    x * (Line::erf(tmp) + 1.0) / 2.0
}

通过自动生成的 gelu_array::launch_unchecked 函数即可启动内核：

pub fn launch<R: Runtime>(device: &R::Device) {
    let client = R::client(device);
    let input = &[-1., 0., 1., 5.];
    let vectorization = 4;
    let output_handle = client.empty(input.len() * core::mem::size_of::<f32>());
    let input_handle = client.create(f32::as_bytes(input));

    unsafe {
        gelu_array::launch_unchecked::<f32, R>(
            &client,
            CubeCount::Static(1, 1, 1),
            CubeDim::new(input.len() as u32 / vectorization, 1, 1),
            ArrayArg::from_raw_parts(&input_handle, input.len(), vectorization as u8),
            ArrayArg::from_raw_parts(&output_handle, input.len(), vectorization as u8),
        )
    };

    let bytes = client.read(output_handle.binding());
    let output = f32::from_bytes(&bytes);

    println!("Executed gelu with runtime {:?} => {output:?}", R::name());
}

运行以下命令即可体验 GELU 示例：

cargo run --example gelu --features cuda # 使用 CUDA 运行时
cargo run --example gelu --features wgpu # 使用 WGPU 运行时

运行时支持​

支持以下 GPU 运行时：

• WGPU：跨平台 GPU 支持（Vulkan、Metal、DirectX、WebGPU）
• CUDA：NVIDIA GPU 支持
• ROCm/HIP：AMD GPU 支持（开发中）

未来还计划开发一个使用 SIMD 指令的优化 JIT CPU 运行时， 基于 Cranelift。

项目动机​

CubeCL 的目标是简化高性能、跨硬件可移植计算内核的开发。 目前，要在不同硬件上实现最佳性能，通常需要使用不同语言（如 CUDA、Metal 或 ROCm）编写定制内核。 这种繁琐流程激发了我们开发 CubeCL 的初衷。

CubeCL 采用了以下核心特性：

1. 自动向量化：在编译时自动使用最优 SIMD 指令。
2. Comptime：在编译 GPU 内核时动态修改 IR。
3. 自动调优：在运行时选择最佳内核配置。

这些特性不仅提升了性能，还增强了代码的可组合性、可重用性和可维护性。

我们的愿景不仅是提供优化的计算语言，还包括构建一个 Rust 高性能计算生态系统。

CubeCL 已经提供了线性代数组件，并计划支持卷积、随机数生成、快速傅里叶变换等算法。

特性概览​

自动向量化​

通过 CubeCL，可以为输入变量指定向量化因子，运行时将自动使用最佳指令集。 内核代码始终保持简单，向量化处理由 CubeCL 自动完成。

Comptime 优化​

CubeCL 允许在编译时动态修改 IR，实现指令特化、循环展开、形状特化等优化。 这样可以避免为不同硬件手写多个内核变种。

自动调优​

自动调优通过运行小型基准测试，自动选择性能最佳的内核配置，并缓存结果以优化后续运行时间。

学习资源​

目前的学习资源较少，但可以参考 线性代数库 了解 CubeCL 的实际使用。

声明与历史​

CubeCL 当前处于 alpha 阶段。

最初 CubeCL 仅作为 Burn 的 WebGPU 后端。 随着优化需求的增加，开发了中间表示（IR），并支持 CUDA 编译目标。

通过 Rust 的 proc 宏，创建了一个易用的前端，最终形成了 CubeCL。

通过 CubeCL，释放 Rust 在高性能计算中的潜力，欢迎您的参与和贡献！

链接​

• CubeCL GitHub 仓库
• CubeCL 线性代数库

背景​

在深度学习和高性能计算中，矩阵乘法（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 编程的未来！