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 的高性能和科学计算生态系统。为了实现这一目标，我们正在开发线性代数组件，您可以将它们集成到自己的内核中。我们目前有一个高度优化的矩阵乘法模块，当可用时利用 NVIDIA 硬件上的 Tensor 核心，同时在其他平台上优雅地回退到基本指令。虽然还有改进的空间，尤其是在使用更新型 NVIDIA GPU 上的自定义指令方面，但我们的实现已经展示了令人印象深刻的表现。

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

它是如何工作的

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

1. 解析：proc 宏使用 syn 包解析 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包创建了一个Rust前端。在利用两个平台的同时，处理CUDA和WebGPU之间的差异，迫使我们想出适用于所有场合的通用概念。因此，CubeCL应运而生！