burn_cubecl/
lib.rs

1#![warn(missing_docs)]
2#![cfg_attr(docsrs, feature(doc_cfg))]
3// TODO: remove when fixed in cubecl
4#![allow(semicolon_in_expressions_from_non_local_macros)]
5
6//! The Burn backend for every [CubeCL](https://github.com/tracel-ai/cubecl) runtime.
7//!
8//! [`CubeBackend`] implements Burn's tensor operations as CubeCL kernels, compiled just in time
9//! for the device they run on. CUDA, ROCm, Metal, Vulkan, WebGPU, wgpu and the CubeCL CPU
10//! runtime all share this one backend type: a tensor's [`CubeDevice`] says which runtime it
11//! uses. [`Cube`] is the type dispatch uses, wrapped in `burn_fusion::Fusion` when the `fusion`
12//! feature is on.
13//!
14//! Applications reach this backend through a `burn` feature such as `cuda`, `wgpu` or `cpu`
15//! and a `Device` constructor; the runtime crates (`burn-cuda`, `burn-wgpu`, `burn-rocm`,
16//! `burn-cpu`) are thin wrappers that select a runtime. Use this crate directly to write
17//! custom kernels: [`kernel`] and [`ops`] hold the building blocks, and [`cubecl`] is
18//! re-exported so kernels use the same CubeCL version.
19//!
20//! # Feature flags
21//!
22//! - `cuda`, `hip`, `wgpu`, `metal`, `vulkan`, `webgpu`, `cpu`: compile in a CubeCL runtime.
23//! - `fusion`: kernel fusion through `burn-fusion`.
24//! - `autotune`: benchmark kernel variants at runtime and keep the fastest.
25//! - `fft`: FFT kernels.
26//! - `template`: launch hand-written, non-JIT kernels.
27//! - `tracing`: instrument operations with the `tracing` crate.
28
29#[macro_use]
30extern crate derive_new;
31extern crate alloc;
32
33/// Utilities for implementing JIT kernels
34pub mod ops;
35
36/// Kernel module
37pub mod kernel;
38/// Tensor module.
39pub mod tensor;
40
41/// Elements for JIT backend
42pub mod element;
43
44pub use element::{BoolElement, CubeElement, FloatElement, IntElement};
45
46mod backend;
47
48pub use backend::*;
49
50// Re-export cubecl.
51pub use cubecl;
52
53mod tune_key;
54pub use tune_key::CubeAutotuneKey;
55
56#[cfg(any(feature = "fusion", test))]
57/// Module for interacting with fusion
58pub mod fusion;
59
60#[cfg(feature = "template")]
61/// Module for compiling custom non-jit kernels
62pub mod template;
63
64/// The device a cube tensor lives on.
65///
66/// One type across every runtime: which runtime a tensor runs on is what its
67/// device *says*, not what its type is.
68pub use cubecl::Device as CubeDevice;
69
70pub use cubecl::CubeTuneId;
71
72/// The tensor backend for every cubecl runtime.
73///
74/// CUDA, ROCm, Metal, Vulkan, WebGPU, wgpu and the CPU runtime are all this one
75/// type; which of them a tensor runs on is what its [`CubeDevice`] says. Fusion
76/// wraps it when the `fusion` feature is on.
77#[cfg(not(feature = "fusion"))]
78pub type Cube = CubeBackend;
79
80/// The tensor backend for every cubecl runtime, fusing operations across
81/// streams. See [`CubeBackend`] for the unfused type.
82#[cfg(feature = "fusion")]
83pub type Cube = burn_fusion::Fusion<CubeBackend>;
burn_cubecl/lib.rs

burn_cubecl/
lib.rs
