tracel-ai/cubecl

What it solves

CubeCL provides a way to write high-performance compute kernels once in Rust and run them across multiple hardware platforms (NVIDIA, AMD, Apple, Vulkan, WebGPU, and CPU) without having to write separate code in different shader languages like CUDA or WGSL.

How it works

It acts as a Rust language extension and Just-in-Time (JIT) compiler. Using a #[cube] attribute, it transforms Rust functions into an intermediate representation (IR) that is then compiled on-demand to the target platform's native language (e.g., CUDA C++, HIP, Metal, or SPIR-V).

To maintain peak performance across diverse hardware, it uses a four-axis parallelism model (Vector, Plane, CubeDim, and CubeCount) that allows kernels to be adaptive rather than hardcoded to specific hardware constants like warp sizes.

Who it’s for

It is designed for developers building scientific computing libraries, deep learning frameworks (such as Burn), and high-performance compute kernels who want the safety and composability of Rust without sacrificing the raw performance of low-level GPU/CPU programming.

Highlights

• Multi-platform Support: Compiles to CUDA, HIP, Metal, SPIR-V, WGSL, and CPU SIMD.
• Comptime: Allows modifying the compiler IR at the time of first compilation for instruction specialization and loop unrolling.
• Autotuning: Automatically benchmarks different kernel configurations at runtime to find the most efficient version for the current hardware.
• Automatic Vectorization: Lowers high-level vector types to the platform's native SIMD instructions with automatic broadcasting.
• JIT Compilation: Only compiles the specific kernel variants actually launched, reducing binary size compared to ahead-of-time compilation.