2 posts tagged with "SPIR-V"

概述​

Burn 0.15.0 带来了显著的性能改进，特别是在矩阵乘法和卷积操作方面。

此外，此版本还引入了以下重要更新：

• 实验性支持：新增 ROCm/HIP 和 SPIR-V 支持，通过 CubeCL 运行时实现。
• 多后端兼容性：奠定多后端支持的基础。
• 新特性：增加了量化操作支持。
• ONNX 支持扩展：包括更多的算子支持和错误修复，以提升覆盖率。

除此之外，Burn 0.15.0 还包含多项错误修复、性能优化、新的张量操作，以及改进的文档支持。

模块与张量相关更新​

• 移除：对常量泛型模块的拷贝限制。
• 新增：deform_conv2d（实现于 torchvision）、Softmin、round、floor、ceil 等浮点操作。
• 增强：为张量同步增加支持，添加 tensor.one_hot 整数操作。
• 更改：LR 调度器调整为首次调用 .step() 时返回初始学习率。

ONNX 支持扩展​

• 支持多维索引的 gather 操作。
• 增强张量形状跟踪能力。
• 新增 ConvTranspose1d 和 trilu 操作支持。
• 修复 where 操作在标量输入下的行为。

后端改进​

• 支持 CudaDevice 和 MetalDevice，避免重复创建设备。
• 新增 SPIR-V 编译器支持 (burn-wgpu) 和 HIP 支持 (burn-hip)。
• 引入 BackendRouter，为分布式后端处理铺路。
• 修复自动微分相关的内存泄漏和 NaN 问题。

文档与示例​

• 新增自定义 cubecl 内核的文档。
• 改进了回归任务的示例和 burn-tch 文档。
• 修复了多个 Burn Book 的链接及 Raspberry Pi 示例的编译问题。

性能与优化​

• 性能提升：增强了切片内核的性能，改进了 conv2d 和 conv_transpose2d 的自动调优。
• 数据局部性优化：为隐式 GEMM 提供更好的性能支持，并新增边界检查以支持任意输入形状。

Miscellaneous 更新​

• 工具链：更新了 CI 工作流及工具，修复编译器设置的多处问题。
• 兼容性：确保最小支持 Rust 版本为 1.81。

参考​

• Burn 官方 GitHub Release 页面
• Burn 文档
• 原文链接
• CubeCL 文档

通过 Burn 0.15.0，深度学习开发者可以更高效地利用 GPU 加速和量化技术，同时享受多后端支持带来的灵活性。欢迎尝试新版本并加入我们的社区，共同推动 Rust 生态的技术进步！

引言​

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 编程的未来！