How tracel-ai/cubecl works
1. This is a Rust workspace for a JIT-compiled compute language, not a single application. The top-level cubecl crate is a thin public entrypoint that re-exports feature-gated pieces of the stack, while most real architecture lives in the sibling crates.
2. The main user-facing flow starts with the #[cube] proc macro in cubecl-macros. It parses Rust kernel code, rewrites it into ordinary Rust that builds IR at call time, and supports related DSL features such as comptime, autotune annotations, kernel launches, and trait/type generation.
3. cubecl-core is the frontend and kernel model. It defines the high-level language constructs for elements, tensors, containers, control flow, runtime options, validation, and post-processing. This crate is where kernels are represented before backend lowering.
4. cubecl-ir is the backend-neutral intermediate representation. It models operations such as arithmetic, atomics, branching, barriers, memory, tensor indexing, synchronization, metadata, and runtime properties. Everything downstream consumes this IR, so it is the key architecture boundary between frontend semantics and target codegen.
5. cubecl-opt is an IR optimization layer. It runs analysis-driven passes such as liveness, dominance, uniformity, range analysis, and GVN, plus transforms like dead-code elimination, expression merging, disaggregation, and strength reduction. In the architecture diagram, this sits between IR construction and backend code emission.
6. cubecl-runtime is the generic execution/runtime layer. It owns compilation caching, autotuning, logging/profiling, stream scheduling, memory management, validation, client/server coordination, and storage abstractions. This crate is the bridge between compiled kernels and actual device execution, independent of any specific backend.
7. Backend-specific runtime crates provide the device-facing implementations:
- cubecl-cuda for NVIDIA/CUDA
- cubecl-hip for AMD/HIP
- cubecl-metal for Apple/Metal
- cubecl-wgpu for WebGPU/Vulkan/Metal-backed execution
- cubecl-cpu for CPU execution and SIMD
These crates each contain their own device, compute, stream, storage, and server layers, reflecting that scheduling, synchronization, and memory handling differ materially by platform.
8. Code generation is split by target family. cubecl-cpp contains shared C++-style codegen helpers for CUDA/HIP/Metal paths, while cubecl-wgpu has a WGSL compiler and backend adapters, and cubecl-spirv provides a SPIR-V compiler/IR lowering path. The architecture is therefore “one frontend, multiple code generators.”
9. cubecl-common is the shared foundation crate for low-level utilities used across the workspace: byte buffers, caching, device handles, profiling, readers, hash/format helpers, quantization, float formats, and assorted infrastructure. It is cross-cutting support rather than a public compute layer.
10. cubecl-std is the higher-level standard library on top of the core/runtime stack. It provides reusable tensor/layout/view abstractions, fast math, quantization helpers, swizzle/reinterpret utilities, events, and test coverage for these higher-level operations. This is the library boundary most users would build on after the kernel language itself.
11. cubecl-zspace appears to provide shared shape/stride/indexing utilities for tensor-like layouts. It is a supporting domain library for coordinate transforms and layout validation, which feeds into the tensor and view machinery in cubecl-std and cubecl-core.
12. The execution path is roughly: user Rust kernel -> proc macro expansion -> IR construction in cubecl-core -> optimization in cubecl-opt / post-processing -> backend-specific lowering -> runtime compilation/cache/autotune -> device dispatch through the selected backend crate. That is the central data flow to diagram.
13. The repository also includes examples/ and cubecl-book/ as usage and documentation surfaces, plus xtask/ for workspace maintenance tasks like build, test, docs, validation, profiling, and book generation. These are important for developer workflow but are not part of the runtime execution architecture.
Main entry points: src/lib.rs Runtime · src/lib.rs Core frontend · src/lib.rs IR · src/lib.rs Optimizer
Components 21
Workspace
Top-level entrypoints, docs, examples, and maintenance tooling.
- cubecl · Thin workspace entrypoint that re-exports feature-gated pieces.
crates/cubecl/src/lib.rs - Examples · Workspace examples demonstrating kernels, sharing, fusion, normalization, and tracing.
examples - Book · User and developer documentation for features, usage, and advanced topics.
cubecl-book/src/overview.md - xtask · Maintenance commands for build, test, docs, validation, profiling, and book generation.
xtask/src/main.rs
Frontend
User-facing kernel language, macro expansion, and frontend IR construction.
- Macros · Parses and rewrites #[cube] kernels, comptime, autotune, launch, and type/trait generation.
crates/cubecl-macros/src/lib.rs - Macro internals · Shared parsing and code-generation helpers for macro expansion.
crates/cubecl-macros-internal/src/lib.rs - Core frontend · Kernel language model, element/container abstractions, validation, options, and post-processing.
crates/cubecl-core/src/lib.rs - Codegen hooks · Frontend-side compilation metadata and integration points for lowering kernels.
crates/cubecl-core/src/codegen/mod.rs
IR
Backend-neutral intermediate representation and optimization.
Runtime
Generic execution services shared across backends.
- Runtime · Compilation cache, autotuning, logging, streams, memory management, validation, server/client coordination, and storage.
crates/cubecl-runtime/src/lib.rs
Backends
Backend-specific lowering and device runtime implementations.
- C++ codegen · Shared C++-style lowering helpers for CUDA, HIP, and Metal paths.
crates/cubecl-cpp/src/lib.rs - SPIR-V codegen · SPIR-V lowering and instruction translation path.
crates/cubecl-spirv/src/lib.rs - WGPU backend · WGSL compiler plus backend adapters for WebGPU/Vulkan/Metal-backed execution.
crates/cubecl-wgpu/src/lib.rs - CPU backend · CPU execution path with compiler, scheduling, streams, and memory handling.
crates/cubecl-cpu/src/lib.rs - CUDA backend · NVIDIA device runtime with compute, storage, streams, sync, and server layers.
crates/cubecl-cuda/src/lib.rs - HIP backend · AMD device runtime with compute, storage, streams, fencing, and server layers.
crates/cubecl-hip/src/lib.rs - Metal backend · Apple device runtime with compute, memory, and stream layers.
crates/cubecl-metal/src/lib.rs
Libraries
Shared utility and higher-level tensor support crates.
- Common · Cross-cutting utilities for buffers, cache, device handles, profiling, hashing, formatting, quantization, and float formats.
crates/cubecl-common/src/lib.rs - Std · Reusable tensor, layout, view, quantization, fast-math, swizzle, and event abstractions on top of core/runtime.
crates/cubecl-std/src/lib.rs - Z-space · Shared shape, stride, indexing, and layout validation utilities for tensor-like coordinates.
crates/cubecl-zspace/src/lib.rs
Connections 35
- cubecl → Macros: re-exports · Public entrypoint surfaces the proc-macro-driven kernel API.
- cubecl → Core frontend: re-exports · Facade exposes the core frontend and language model.
- cubecl → Runtime: re-exports · Facade exposes runtime services needed by consumers.
- Macros → Macro internals: uses · Macro crate delegates parsing and generation helpers.
- Macros → Core frontend: expands into · Kernel syntax is rewritten into frontend/kernel-building Rust.
- Macros → Runtime: launches via · Expanded kernels call into runtime launch and tuning paths.
- Core frontend → IR: lowers to · Frontend kernel model is represented as backend-neutral IR.
- Core frontend → Codegen hooks: integrates · Frontend carries codegen metadata and compilation integration.
- Core frontend → Z-space: depends on · Tensor/layout semantics rely on shared shape and indexing utilities.
- Codegen hooks → IR: emits · Codegen layer produces or prepares IR-facing compilation artifacts.
- IR → Optimizer: optimized by · IR flows through analysis and transformation passes before code emission.
- Optimizer → SPIR-V codegen: feeds · Optimized IR is lowered toward SPIR-V code generation.
- Optimizer → C++ codegen: feeds · Optimized IR is lowered through shared C++-style backends.
- Optimizer → WGPU backend: feeds · Optimized IR is compiled into WGSL/backend adapters.
- Optimizer → CPU backend: feeds · Optimized IR is compiled into CPU execution code.
- Runtime → Common: depends on · Runtime builds on shared cache, device, profiling, and storage primitives.
- Runtime → CPU backend: dispatches to · Runtime coordinates execution through the selected backend runtime.
- Runtime → CUDA backend: dispatches to · Runtime coordinates execution through the selected backend runtime.
- Runtime → HIP backend: dispatches to · Runtime coordinates execution through the selected backend runtime.
- Runtime → Metal backend: dispatches to · Runtime coordinates execution through the selected backend runtime.
- Runtime → WGPU backend: dispatches to · Runtime coordinates execution through the selected backend runtime.
- Std → Core frontend: builds on · Std library layers tensor/view abstractions over the core language model.
- Std → Runtime: builds on · Std library uses runtime services for execution and memory handling.
- Std → Z-space: uses · Layout/view operations use shared coordinate and stride helpers.
- C++ codegen → IR: consumes · Target-family codegen lowers IR into C++-style dialects.
- WGPU backend → IR: consumes · WGSL compiler and adapters translate IR into web GPU execution artifacts.
- SPIR-V codegen → IR: consumes · SPIR-V compiler lowers backend-neutral IR into SPIR-V instructions.
- CPU backend → Common: uses · CPU runtime shares common buffer, profiling, and device utilities.
- CUDA backend → Common: uses · CUDA runtime shares common buffer, profiling, and device utilities.
- HIP backend → Common: uses · HIP runtime shares common buffer, profiling, and device utilities.
- Metal backend → Common: uses · Metal runtime shares common buffer, profiling, and device utilities.
- Examples → cubecl: demonstrates · Examples show end-user usage of the public facade and kernels.
- Book → Std: documents · Documentation explains higher-level tensor and language features.
- xtask → Examples: supports · Workspace tooling builds and validates examples.
- xtask → Book: supports · Workspace tooling generates and checks the documentation book.