2 posts tagged with "CubeCL"

2024年标志着Burn架构的重大演变。 传统的深度学习框架常常要求开发者在性能、可移植性和灵活性之间做出妥协；而我们的目标是超越这些权衡。 展望2025年，我们致力于将这一理念应用于整个计算栈，从嵌入式设备到数据中心，涵盖所有领域。

2024年回顾：突破硬件限制​

重新定义内核开发​

今年之初，我们面临一个限制：我们的WGPU后端依赖于基础的WGSL模板，限制了我们的适应能力。这个挑战促使我们创建了CubeCL [1]，这是我们统一内核开发的解决方案。这项任务非常复杂——设计一个抽象层，适用于各种不同的硬件，同时保持顶级性能。我们的结果证明了这一策略的有效性，在大多数基准测试中，性能现在已匹配甚至超过LibTorch。

多后端架构​

后端生态系统现已包括CUDA [2]、HIP/ROCm [3]以及支持WebGPU和Vulkan的先进WGPU实现 [4]。迄今为止最显著的成就是在相同硬件上实现不同后端的性能平衡。例如，无论是在CUDA还是Vulkan上执行矩阵乘法操作，性能几乎相同，这直接反映了我们平台无关优化的策略。

我们还引入了新的Router和HTTP后端：Router后端支持多后端的动态混合，而HTTP后端则支持跨多台机器的分布式处理。为了解决内存管理挑战，我们实施了池化和检查点机制，即使在反向传播期间也能实现操作融合。

硬件无关加速​

我们的硬件加速策略标志着一个重要的技术里程碑。我们并不依赖于特定平台的库，如cuBLAS [5]或rocBLAS [6]，而是开发了一套编译器栈，利用每个平台的最佳特性，同时确保跨平台的兼容性。这涉及克服代码生成和优化中的复杂挑战，尤其是对于矩阵乘法等操作，必须高效利用各种硬件架构的张量核心。

2025年路线图：拥抱极端​

在2025年，我们将解决深度学习部署中的两个基本挑战。

小规模：量化​

量化对于资源有限的计算至关重要。 我们的方法使用复杂操作的融合，通过“读取时融合”功能，实现如归约等任务在计算管道中的无缝集成。 这种融合策略自动处理操作的打包和解包，确保量化操作高效运行，无需手动调整。 结果是什么？高性能的量化操作在保持精度的同时，降低了资源需求。

大规模：可扩展的分布式计算​

在另一端，是分布式计算。 通过利用我们的Router和HTTP后端构建强大的分布式训练基础设施，我们旨在创建一个流畅的分布式计算体验， 使工作负载能够在不同硬件和后端配置之间轻松流动，同时优化异构计算环境中的资源利用。

为了支持这种普遍兼容性的愿景，我们正在扩展我们的后端生态系统，包括：

• 开发Metal后端，充分利用Apple Silicon的能力，超越当前WGPU的功能；
• 在Rust中实现一个即时向量化的CPU后端，以增强CPU性能；
• 开启新的后端可能性，如FPGA支持，确保Burn能够适应任何计算环境。

我们还将大量投资于开发者体验，提供全面的CubeCL文档，并推动Burn API的稳定化。这些改进将使开发者更容易利用我们跨平台能力的全部潜力。

在2024年，我们证明了跨平台性能不需要妥协。 展望2025年，我们将这一原则扩展到整个计算领域——从微控制器到服务器农场。 通过解决两个极端的技术挑战，我们致力于使深度学习在任何规模或硬件限制下都更加高效和易用。

参考文献​

• Going Big and Small for 2025
• CubeCL: Multi-platform high-performance compute language extension for Rust
• CUDA Toolkit
• AMD ROCm Software
• Rust implementation of SPIR-V
• cuBLAS: CUDA Basic Linear Algebra Subroutine Library
• rocBLAS: ROCm Basic Linear Algebra Subprograms Library
• Integer Quantization for Deep Learning Inference: Principles and Empirical Evaluation
• Rust Bindings for Metal

使用 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 线性代数库