CubeCL：用于 CUDA、ROCm 和 WGPU 的 Rust GPU Kernel

tracel-ai / cubecl Public

用于 Rust 的多平台高性能计算语言扩展。 burn.dev

License

Apache-2.0, MIT licenses found

Licenses found

Apache-2.0 LICENSE-APACHE MIT LICENSE-MIT 988 stars 54 forks Branches Tags Activity

tracel-ai/cubecl

main BranchesTags

Go to file Code

Folders and files

Name| Name| Last commit message| Last commit date ---|---|---|---

Latest commit

History

780 Commits .cargo [.cargo] .github [.github] .vscode [.vscode] assets [assets] crates [crates] cubecl-book [cubecl-book] examples [examples] xtask [xtask] .gitignore [.gitignore] CONTRIBUTING.md [CONTRIBUTING.md] Cargo.toml [Cargo.toml] LICENSE-APACHE [LICENSE-APACHE] LICENSE-MIT [LICENSE-MIT] NOTICE.md [NOTICE.md] README.md [README.md] View all files

Repository files navigation

Discord Current Crates.io Version Minimum Supported Rust Version Test Status license NVIDIA AMD WGPU

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

简介

使用 CubeCL，你可以使用 Rust 编程你的 GPU，利用零成本抽象来开发可维护、灵活和高效的计算 Kernel。CubeCL 目前完全支持函数、泛型和结构体，部分支持 traits、方法和类型推断。随着项目的演进，我们预计将更广泛地支持 Rust 语言原语，同时保持最佳性能。

示例

只需使用 cube 属性注解函数，以表明它们应该在 GPU 上运行。

use cubecl::prelude::*;
#[cube(launch_unchecked)]
/// A [Line] represents a contiguous series of elements where SIMD operations may be available.
/// The runtime will automatically use SIMD instructions when possible for improved performance.
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> {
  // Execute the sqrt function at comptime.
  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 函数启动 Kernel。

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::<f32>(&input_handle, input.len(), vectorization as u8),
      ArrayArg::from_raw_parts::<f32>(&output_handle, input.len(), vectorization as u8),
    )
  };
  let bytes = client.read_one(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());
}

要查看它的实际效果，请使用以下命令运行 GELU 示例：

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

运行时

我们支持以下 GPU 运行时：

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

动机

CubeCL 的目标是减轻编写可在硬件之间移植的高度优化计算 Kernel 的痛苦。当你想要最佳性能同时又是多平台时，目前没有足够的解决方案。你必须为不同的硬件编写自定义 Kernel，通常使用不同的语言，例如 CUDA、Metal 或 ROCm。为了解决这个问题，我们创建了一个具有三个核心功能的 Just-in-Time 编译器：自动向量化、comptime 和 autotune！

这些功能对于任何编写高性能 Kernel 的人来说都非常有用，即使可移植性不是问题。它们提高了代码的可组合性、可重用性、可测试性和可维护性，同时保持最佳状态。CubeCL 还提供了一种内存管理策略，该策略针对具有大量缓冲区重用的吞吐量进行了优化，以避免分配。

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

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

工作原理

CubeCL 在 Rust 的 proc macro 系统中利用了独特的两步过程：

  1. 解析：proc macro 使用 syn crate 解析 GPU Kernel 代码。
  2. 扩展：宏不是立即生成中间表示 (IR)，而是生成一个新的 Rust 函数。

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

设计

CubeCL 的设计围绕 - 你猜对了 - Cubes！更具体地说，它基于长方体，因为并非所有轴的大小都相同。由于所有计算 API 都需要映射到硬件，硬件是可以使用 3D 表示访问的 tiles，因此我们的拓扑可以轻松映射到其他 API 的概念。

CubeCL - 拓扑

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

拓扑等价 👇

由于所有拓扑变量在 Kernel 入口点内都是常量，因此我们选择使用带有大写字母的 Rust 常量语法。通常在创建 Kernel 时，我们并不总是关心单元在立方体内沿每个轴的相对位置，但我们通常只关心它的一般位置。因此，每种变量也有其自身的轴独立变量，这在其他语言中通常不存在。

CubeCL | CUDA | WebGPU | Metal ---|---|---|--- CUBE_COUNT | N/A | N/A | N/A CUBE_COUNT_X | gridDim.x | num_workgroups.x | threadgroups_per_grid.x CUBE_COUNT_Y | gridDim.y | num_workgroups.y | threadgroups_per_grid.y CUBE_COUNT_Z | gridDim.z | num_workgroups.z | threadgroups_per_grid.z CUBE_POS | N/A | N/A | N/A CUBE_POS_X | blockIdx.x | workgroup_id.x | threadgroup_position_in_grid.x CUBE_POS_Y | blockIdx.y | workgroup_id.y | threadgroup_position_in_grid.y CUBE_POS_Z | blockIdx.z | workgroup_id.z | threadgroup_position_in_grid.z CUBE_DIM | N/A | N/A | N/A CUBE_DIM_X | blockDim.x | workgroup_size.x | threads_per_threadgroup.x CUBE_DIM_Y | blockDim.y | workgroup_size.y | threads_per_threadgroup.y CUBE_DIM_Z | blockDim.z | workgroup_size.z | threads_per_threadgroup.z UNIT_POS | N/A | local_invocation_index | thread_index_in_threadgroup UNIT_POS_X | threadIdx.x | local_invocation_id.x | thread_position_in_threadgroup.x UNIT_POS_Y | threadIdx.y | local_invocation_id.y | thread_position_in_threadgroup.y UNIT_POS_Z | threadIdx.z | local_invocation_id.z | thread_position_in_threadgroup.z PLANE_POS | N/A | subgroup_id | simdgroup_index_in_threadgroup PLANE_DIM | warpSize | subgroup_size | threads_per_simdgroup UNIT_POS_PLANE | N/A | subgroup_invocation_id | thread_index_in_simdgroup ABSOLUTE_POS | N/A | N/A | N/A ABSOLUTE_POS_X | N/A | global_id.x | thread_position_in_grid.x ABSOLUTE_POS_Y | N/A | global_id.y | thread_position_in_grid.y ABSOLUTE_POS_Z | N/A | global_id.z | thread_position_in_grid.z

特殊功能

自动向量化

高性能 Kernel 应尽可能依赖 SIMD 指令，但这样做很快就会变得非常复杂！使用 CubeCL，你可以在启动 Kernel 时指定每个输入变量的向量化因子。在 Kernel 代码中，你仍然只使用一种类型，该类型是动态向量化的，并支持自动广播。运行时能够编译 Kernel，并具有使用最佳指令的所有必要信息！但是，由于算法行为可能取决于向量化因子，因此 CubeCL 允许你在需要时直接在 Kernel 中访问它，而不会有任何性能损失，使用 comptime 系统！

Comptime

CubeCL 不仅仅是一种新的计算语言：虽然感觉你正在编写 GPU Kernel，但实际上你正在编写可以完全自定义的编译器插件！Comptime 是一种在首次编译 Kernel 时在运行时修改编译器 IR 的方法。

这可以实现大量的优化和灵活性，而无需编写许多相同的 Kernel 的单独变体，以确保最佳性能。

功能 | 描述 ---|--- 指令专业化 | 并非所有指令都可在所有硬件上使用，但当存在专用指令时，应使用简单的 if 语句启用它。 自动向量化 | 当你可以使用 SIMD 指令时，你应该这样做！但是由于并非所有硬件都支持相同的向量化因子，因此可以在运行时注入它！ 循环展开 | 你可能想要相同 Kernel 的多个版本，其中循环展开仅适用于特定范围的值。可以使用 Comptime 轻松配置它。 形状专业化 | 对于深度学习 Kernel，通常必须依赖于不同输入大小的不同 Kernel；你可以通过将形状信息作为 Comptime 值传递来实现。 编译时计算 | 通常，你可以使用 Rust 运行时属性计算一个常量，并将其注入到 Kernel 中，以避免在每次执行期间重新计算它。

Autotuning

Autotuning 通过在运行时运行小型基准测试来简化 Kernel 选择，以找出在当前硬件上运行的最佳 Kernel 以及最佳配置；这是可移植性的基本功能。此功能与 comptime 完美结合，可以测试不同 comptime 值对性能的影响；有时这可能会令人惊讶！

即使基准测试在首次运行应用程序时可能会增加一些开销，但信息会被缓存在设备上并将被重用。对于面向吞吐量的程序（例如深度学习模型），这通常是一个明智的权衡。你甚至可以将 autotune 缓存与你的程序一起发布，从而缩短在你可以更好地控制部署目标时的冷启动时间。

资源

目前我们没有很多学习资源，但你可以查看 线性代数库 了解如何使用 CubeCL。如果你有任何问题或想做出贡献，请随时加入 Discord。

免责声明 & 历史

CubeCL 目前处于 alpha 阶段。

虽然 CubeCL 用于 Burn，但仍然有很多粗糙的地方；它尚未完善。该项目最初是 Burn 的一个仅限 WebGPU 的后端。当我们对其进行优化时，我们意识到我们需要一个可以优化然后编译为 WGSL 的中间表示 (IR)。拥有 IR 可以轻松支持另一个编译目标，因此我们制作了一个 CUDA 运行时。但是，直接在该 IR 中编写 Kernel 并不容易，因此我们使用 syn crate 创建了一个 Rust 前端。在利用这两个平台的同时，了解 CUDA 和 WebGPU 之间的差异迫使我们提出适用于所有地方的通用概念。因此，CubeCL 诞生了！

关于

用于 Rust 的多平台高性能计算语言扩展。 burn.dev

Topics

rust gpu cuda jit gpgpu webgpu linalg

Resources

Readme

License

Apache-2.0, MIT licenses found

Licenses found

Apache-2.0 LICENSE-APACHE MIT LICENSE-MIT Activity Custom properties

Stars

988 stars

Watchers

14 watching

Forks

54 forks

Releases 3

v0.5.0 Latest Apr 23, 2025 + 2 releases

Packages 0

No packages published

Used by 295

Contributors 25

+ 11 contributors

Languages