Lib.rs

---

Rust的多平台高性能计算语言扩展。

TL;DR

使用CubeCL，您可以使用Rust编写GPU程序，利用零成本抽象来开发可维护、灵活和高效的计算内核。CubeCL目前完全支持函数、泛型和结构体，对特性和类型推断提供部分支持。随着项目的不断发展，我们预计还将对Rust语言原语提供更广泛的支持，同时保持最佳性能。

示例

只需使用 cube 属性注解函数，即可指示它们应在GPU上运行。

use cubecl::prelude::*;

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

#[cube]
fn gelu_scalar<F: Float>(x: F) -> F {
    x * (F::erf(x / F::sqrt(2.0.into())) + 1.0) / 2.0
}

然后，您可以使用自动生成的 gelu_array::launch 函数启动内核。

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

    gelu_array::launch::<F32, R>(
        client.clone(),
        CubeCount::Static(1, 1, 1),
        CubeDim::new(input.len() as u32, 1, 1),
        ArrayArg::new(&client.create(f32::as_bytes(input)), input.len()),
        ArrayArg::new(&output_handle, input.len()),
    );

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

    // Should be [-0.1587,  0.0000,  0.8413,  5.0000]
    println!("Executed gelu with runtime {:?} => {output:?}", R::name());
}

fn main() {
    launch::<cubecl::cuda::CudaRuntime>(&Default::default());
    launch::<cubecl::wgpu::WgpuRuntime>(&Default::default());
}

要查看其实际运行效果，请使用以下命令运行可工作的GELU示例

cargo run --example gelu --features cuda # cuda runtime
cargo run --example gelu --features wgpu # wgpu runtime

运行时

目前支持的运行时只有CUDA和WebGPU。添加更多运行时非常容易，我们打算支持很多；欢迎贡献！

• CUDA
• WebGPU
• Metal（尽管WebGPU可以原生编译为Metal）
• ROCm
• Vulkan（尽管WebGPU可以原生编译为Vulkan）

我们还计划开发一个使用SIMD指令优化的JIT CPU运行时，利用 Cranelift。

动机

CubeCL的目标是简化编写高度优化的、可跨硬件移植的计算内核的难度。当您希望获得最佳性能的同时仍然保持多平台兼容性时，目前还没有合适的解决方案。您必须为不同的硬件编写定制的内核，通常使用不同的语言，如CUDA、Metal或ROCm。为了解决这个问题，我们创建了一个即时编译器，具有三个核心功能：**自动向量化**、**编译时**和**自调优**！

这些特性对编写高性能内核的人来说非常有用，即使在可移植性不是重点的情况下。它们提高了代码的可组合性、可重用性、可测试性和可维护性，同时保持最优性能。CubeCL还附带了一种针对吞吐量优化且具有大量缓冲区重用的内存管理策略，以避免分配。

我们的目标不仅仅是提供一种优化的计算语言；我们旨在开发一个基于Rust的高性能和科学计算生态系统。为此，我们正在开发可以集成到您自己的内核中的线性代数组件。我们目前有一个高度优化的矩阵乘法模块，在支持Tensor Cores的NVIDIA硬件上运行，在不支持的情况下会优雅地回退到其他平台的基本指令。尽管还有改进的空间，尤其是在使用较新NVIDIA GPU的定制指令方面，但我们的实现已经提供了令人印象深刻的性能。

这只是开始。我们计划包括更多工具，如卷积、随机数生成、快速傅里叶变换以及其他基本算法。我们是一个小型团队，也在构建Burn，所以请不要犹豫，贡献和移植算法；这可能会比你想象的更有帮助！

它是如何工作的

CubeCL利用Rust的proc宏系统，采用独特的两步过程

1. 解析：proc宏使用syn crate解析GPU内核代码。
2. 展开：宏不是立即生成中间表示（IR），而是生成一个新的Rust函数。

生成的函数在调用时负责创建IR，这与传统的编译器不同，传统的编译器通常在解析后立即生成IR。我们的方法实现了几个关键特性

• 编译时：由于不转换原始代码，因此可以轻松地集成编译时优化。
• 自动向量化：通过简单地向量化CubeCL函数的输入，我们可以在展开过程中确定每个中间变量的向量化因子。
• Rust集成：生成的代码仍然是有效的Rust代码，因此可以无需依赖特定运行时进行捆绑。

设计

CubeCL的设计围绕着——没错——立方体！更具体地说，它基于长方体，因为不是所有轴的大小都相同。由于所有计算API都需要映射到硬件，即可以使用3D表示访问的瓦片，因此我们的拓扑可以轻松地映射到其他API的概念。

CubeCL - 拓扑

立方体由单位组成，因此3x3x3的立方体有27个单位，可以通过它们的x、y和z轴上的位置访问。同样，超立方体由立方体组成，就像立方体由单位组成一样。超立方体中的每个立方体都可以通过其在x、y和z轴上相对于超立方体的位置访问。因此，3x3x3的超立方体将包含27个立方体。在这个例子中，工作单元的总数将是27 x 27 = 729。

特殊功能

自动向量化

高性能内核应尽可能使用SIMD指令，但这会迅速变得相当复杂！使用CubeCL，您可以在启动内核时指定每个输入变量的向量化因子。在内核代码内部，您仍然只使用一种类型，该类型是动态向量化并支持自动广播的。运行时能够编译内核，并且拥有使用最佳指令所需的所有必要信息！然而，由于算法行为可能取决于向量化因子，CubeCL允许您在需要时直接在内核中访问它，而不会造成任何性能损失，使用comptime系统！

Comptime

CubeCL不仅仅是一种新的计算语言：虽然感觉您正在编写GPU内核，但实际上，您正在编写可以完全定制的编译器插件！Comptime是在编译内核的第一次编译时修改编译器IR的一种方法。

这可以在不编写许多相同内核的不同变体的同时，实现许多优化和灵活性。

自动调整

自动调整通过在运行时运行小型的基准测试来确定最佳内核及其最佳配置，从而极大地简化了内核选择；这是一个重要的可移植性功能。此功能与comptime完美结合，以测试不同comptime值对性能的影响；有时这可能会令人惊讶！

即使基准测试可能会在第一次运行应用程序时添加一些开销，但信息会被缓存到设备上并重复使用。这对于以吞吐量为导向的程序（如深度学习模型）通常是一个不言而喻的交易。您甚至可以将自动调整缓存与您的程序一起分发，以减少您对部署目标有更多控制时的冷启动时间。

资源

目前我们没有很多资源可供学习，但您可以查看线性代数库，了解CubeCL的使用方法。如果您有任何问题或想贡献力量，请毫不犹豫地加入Discord。

免责声明 & 历史

CubeCL目前处于alpha阶段。

CubeCL虽被用于Burn，但仍有许多粗糙边缘；它尚未完善。该项目最初是Burn的WebGPU唯一后端。在我们对其优化时，我们意识到我们需要一个中间表示（IR），该表示可以优化然后编译到WGSL。有了IR，就很容易支持另一个编译目标，因此我们创建了一个CUDA运行时。然而，直接在那种IR中编写内核并不容易，所以我们创建了一个使用syn crate的Rust前端。在导航CUDA和WebGPU之间的差异的同时利用这两个平台，迫使我们想出了一些通用的概念，这些概念可以在任何地方工作。因此，CubeCL诞生了！