CubeCL
lastmod 2025-04-28created 2025-04-281 edits
Multi-platform high-performance compute language extension for Rust.
Rust-based GPU kernels for CUDA, ROCm, and WGPU.
About
- CubeCL is a high-performance multi-platform language extension that lets you write GPU kernels in Rust
- Fully supports functions, generics, and structs, with partial support for traits, methods, and type inference
- Supports WGPU-, CUDA-, and ROCm-based runtimes, and an optimized JIT CPU runtime using SIMD instructions is also in development
- Provides automatic vectorization, compile-time optimization (Comptime), and autotune features to maximize code performance
- Uses a unique approach that leverages Rust's procedural macro system to parse and expand GPU kernel code
Introduction to CubeCL
- CubeCL is a high-performance computing language extension that extends the Rust language to support writing GPU kernels directly
- Leverages zero-cost abstractions to enable efficient, maintainable kernel development
- Currently fully supports functions, generics, and structs, with partial support for traits, methods, and type inference
Simple usage example
- The #[cube] attribute marks functions to run on the GPU
- The gelu_array and gelu_scalar functions automatically use SIMD instructions to improve performance
- Kernels can be launched with the launch_unchecked function
Supported runtimes
- WGPU: cross-platform GPU support via Vulkan, Metal, DirectX, and WebGPU
- CUDA: optimized support for NVIDIA GPUs
- ROCm/HIP: AMD GPU support (in development)
- JIT CPU runtime: an optimized Cranelift-based JIT CPU runtime using SIMD instructions is in development
Goals of CubeCL
- Supports writing high-performance computing kernels independent of hardware
- Improves code composability, reusability, testability, and maintainability through automatic vectorization, comptime, and autotuning
- Provides optimized buffer reuse for high throughput through memory management strategies
- Aims to build a high-performance and scientific computing ecosystem by developing linear algebra components
How CubeCL works
- Uses a unique approach that leverages Rust's procedural macro system to parse and expand GPU kernel code
- Comptime: by not transforming the original code, compile-time optimizations can be integrated easily
- Automatic vectorization: vectorizes the inputs of a CubeCL function to determine the vectorization factor of each intermediate variable
- Rust integration: the generated code remains valid Rust code and can be bundled without depending on a specific runtime
Design and special features
- CubeCL is designed around the cube, and all compute APIs map to the hardware and can use a 3D representation
- Automatic vectorization: can improve performance using SIMD instructions, and the vectorization factor can be accessed directly within the kernel
- Comptime: modifies the compiler IR at runtime to provide optimization and flexibility
- Autotune: runs small benchmarks at runtime to select the best kernel and configuration
Caveats and history
- CubeCL is currently in alpha, with some rough edges
- Originally started as the WebGPU backend of the Burn project, but
- expanded into a general-purpose IR and Rust frontend with the addition of the CUDA runtime
- Developed general-purpose concepts to overcome the differences between CUDA and WebGPU