CubeCL Documentation

CubeCL is a multi-platform high-performance compute language extension for Rust. It acts as a Just-in-Time (JIT) compiler and provides a set of runtimes for writing high-performance compute kernels.

By using the #[cube] macro, a single Rust function can be compiled on demand to multiple backends including:

• CUDA (NVIDIA)
• HIP (AMD)
• Metal (Apple)
• SPIR-V (Vulkan)
• WGSL (WebGPU)
• CPU SIMD

This allows developers to write kernels once in regular Rust (benefiting from type-checking, borrow-checking, and composability) while achieving peak performance on each specific hardware target.

The cubecl-runtime crate provides the infrastructure for creating high-performance asynchronous runtimes. It is designed to manage the lifecycle of kernel executions and memory in a way that optimizes performance for compute-intensive tasks.

Key capabilities include:

• Asynchronous kernel executions: Managing the scheduling and execution of kernels across different backends.
• Memory allocation management: Handling memory lifecycle and allocation strategies.
• Autotuning: Providing mechanisms to optimize kernel performance through automated tuning.

The cubecl-common package is a shared utility crate designed to host code that must be accessible across multiple CubeCL packages. It is architected to support both std and no_std environments, ensuring compatibility for various runtime requirements.

Note for developers: To maintain compatibility with restricted environments, this package must be able to build using cargo build --no-default-features.

CubeCL simplifies high-performance kernel development by allowing you to specify a vectorization factor for input variables during kernel launch. Instead of manually writing SIMD code, you use a single type within the kernel, which the runtime then dynamically vectorizes.

Key features include:

• Automatic Broadcasting: The runtime handles broadcasting for vectorized types.
• Hardware Optimization: Runtimes use the provided vectorization factor to compile kernels using the best available hardware instructions.
• Algorithmic Control: If your algorithm's logic depends on the vectorization factor, you can access it directly inside the kernel using CubeCL's comptime system without incurring performance penalties.

You can mark struct fields with the #[cube(comptime)] attribute. Values in these fields are known at kernel compilation time, allowing you to use the comptime! macro for specialization within your kernel logic. This is useful for conditional branching that doesn't incur runtime overhead on the GPU.

#[derive(CubeType, CubeLaunch)]
pub struct TaggedSlice {
    pub array: Box<[f32]>,
    #[cube(comptime)]
    pub tag: String,
}

#[cube(launch_unchecked)]
pub fn kernel_with_tag(output: &mut TaggedSlice) {
    if UNIT_POS == 0 {
        if comptime! {&output.tag == "zero"} {
            output.array[0] = 0.0;
        } else {
            output.array[0] = 1.0;
        }
    }
}

CubeCL is part of a modular ecosystem designed with separated concerns:

• cubecl: The core repository containing the language, JIT compiler, IR, and per-platform runtimes.
• cubek: A library of high-level kernels (matrix multiplication, convolutions, attention, etc.) built on top of CubeCL.
• Burn: A deep learning framework that uses CubeCL as its production proving ground.

In CubeCL, GPUs are organized into groups of threads that operate in lockstep called planes. These are equivalent to CUDA warps, Vulkan subgroups, or Metal SIMD groups.

CubeCL provides the following built-in constants to manage work within a plane:

• PLANE_DIM: Returns the actual dimension (size) of the current plane as determined by the driver at runtime. This is useful because plane sizes can vary based on hardware or factors like register pressure.
• PLANE_POS: Returns the position of the current plane within the larger cube. This is used to facilitate dividing work across multiple planes.
• UNIT_POS_PLANE: Returns the position of the current unit (thread) within its plane, ranging from 0 to PLANE_DIM. This is primarily used for cooperative plane operations like plane_broadcast and plane_shuffle to calculate relative positions.

CubeCL provides several high-level abstractions to manage low-level GPU/CPU complexity:

• Automatic Vectorization: You specify a vectorization factor during kernel launch. Inside the kernel, you use a single type that is dynamically vectorized and supports automatic broadcasting, which the runtime lowers to native SIMD instructions.
• Comptime: Allows modifying the compiler IR at the time of the first compilation. This enables instruction specialization, adaptive parallelism (reading PLANE_DIM, CUBE_DIM, and vector width), loop unrolling, and shape specialization without runtime cost.
• Autotuning: Automatically selects the best kernel and configuration for the current hardware by running small benchmarks at runtime. Results are cached on the device to avoid repeated overhead.