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。为了解决这个问题，我们创建了一个即时编译器，具有三个核心特性：自动向量化、comptime和autotune！

这些功能对于编写高性能内核的人来说非常有用，即使不考虑可移植性也是如此。它们提高了代码的组成性、可重用性、可测试性和可维护性，同时保持最优。CubeCL 还附带了一种针对吞吐量优化的内存管理策略，通过大量重复使用缓冲区来避免分配。

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

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

工作原理

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

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

生成的函数，与原始函数在语义上相似，负责在调用时创建 IR。这种方法与传统编译器不同，传统编译器通常在解析后直接生成 IR。我们的方法实现了几个关键特性

• 编译时：由于不转换原始代码，因此可以非常容易地集成编译时优化。
• 自动向量化：通过简单地向量化 CubeCL 函数的输入，我们可以在扩展过程中确定每个中间变量的向量化因子。
• Rust 集成：生成的代码仍然是有效的 Rust 代码，允许它捆绑而没有对特定运行时的任何依赖。

设计

CubeCL 是围绕——没错——立方体设计的！更具体地说，它是基于长方体，因为所有轴的尺寸并不相同。由于所有计算 API 都需要映射到硬件，这些硬件是可以通过三维表示访问的瓦片，因此我们的拓扑可以轻松地映射到其他 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诞生了！