burn_backend/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2#![warn(missing_docs)]
3#![cfg_attr(docsrs, feature(doc_cfg))]
4
5//! The contract between Burn's tensor API and the backends that execute it.
6//!
7//! Application code does not use this crate directly: it works with `burn::tensor::Tensor`
8//! and `Device`, and dispatch picks a backend at runtime. This crate is for code below that
9//! boundary, such as backend implementations, backend decorators and backend extensions.
10//!
11//! - [`BackendTypes`] names a backend's tensor primitives and device type.
12//! - [`Backend`] and the operation traits in [`ops`] define every tensor operation a backend
13//!   must implement. [`AutodiffBackend`] adds gradient support on top.
14//! - [`DeviceOps`] describes a backend device and its default dtypes.
15//! - [`TensorData`], [`DType`], [`Shape`] and the element traits describe tensor data
16//!   independently of any backend.
17//!
18//! Backend operations report device failures as [`ExecutionError`] instead of panicking.
19//!
20//! # Feature flags
21//!
22//! - `std` (default): standard library support. Without it the crate is `no_std` with `alloc`.
23//! - `cubecl`: conversions between Burn and CubeCL types.
24//! - `cubecl-device` and the `cubecl-<runtime>` features: implement [`DeviceOps`] for
25//!   `cubecl::Device`.
26//! - `tracing`: instrument operations with the `tracing` crate.
27
28#[macro_use]
29extern crate derive_new;
30
31extern crate alloc;
32
33/// [`Backend`] trait and required types.
34pub mod backend;
35pub use backend::*;
36
37// Re-exported types
38pub use burn_std::reader::*; // Useful so that backends don't have to add `burn_std` as a dependency.
39pub use burn_std::{
40    AllocationProperty, BoolDType, BoolStore, Bytes, DType, DataError, DeviceHandle, Distribution,
41    DistributionSampler, DistributionSamplerKind, Element, ElementAdd, ElementConversion,
42    ElementEq, ElementOrdered, ElementRandom, FloatDType, IntDType, Scalar, SplitPolicy,
43    TensorData, Tolerance, bf16, distribution, element, f16, stream::StreamId,
44};
45
46/// Shape definition.
47pub mod shape {
48    pub use burn_std::shape::*;
49}
50pub use shape::*;
51
52/// Slice utilities.
53pub mod slice {
54    pub use burn_std::{s, slice::*};
55}
56pub use slice::*;
57
58/// Indexing utilities.
59pub mod indexing {
60    pub use burn_std::indexing::*;
61}
62pub use indexing::*;
63
64mod alias;
65pub use alias::*;
66
67/// Quantization data representation.
68pub mod quantization;
69
70/// CubeCL inter-operation helpers (gated by the `cubecl` feature).
71///
72/// Provides plain conversion functions between burn's [`DType`] and cubecl's
73/// `ElemType` / `StorageType`. They are intentionally exposed as named
74/// functions rather than `From`/`Into` impls so the cubecl type tree does not
75/// leak into `burn-std`'s public surface.
76#[cfg(feature = "cubecl")]
77pub mod cubecl;
78
79// Not gated on the `cubecl-*` runtime features: a build gets its cubecl runtime from whichever
80// crate asked for one, which need not be this one — `burn-cubecl` compiles with no runtime feature
81// of its own and still needs this impl. So the impl is its own feature, and a crate that needs
82// `cubecl::Device` to be a burn device says so.
83#[cfg(feature = "cubecl-device")]
84mod cube_device {
85    use crate::backend::DeviceOps;
86    use burn_std::{BoolStore, DType, DeviceSettings};
87    use cubecl::{Device, RuntimeId};
88    use cubecl::{
89        features::TypeUsage,
90        ir::{ElemType, UIntKind},
91    };
92
93    impl DeviceOps for Device {
94        fn defaults(&self) -> DeviceSettings {
95            // Cargo features make native compilers available, but automatic devices can still
96            // fall back to WGSL. Only use byte-sized bools when this device can store and convert
97            // them; WGSL needs a word. Other runtimes continue to use a byte.
98            let bool_store = match self.runtime() {
99                RuntimeId::Wgpu => {
100                    let usage = self
101                        .client()
102                        .properties()
103                        .type_usage(ElemType::UInt(UIntKind::U8));
104                    if usage.is_superset(TypeUsage::Buffer | TypeUsage::Conversion) {
105                        BoolStore::U8
106                    } else {
107                        BoolStore::U32
108                    }
109                }
110                _ => BoolStore::U8,
111            };
112
113            DeviceSettings::new(
114                DType::F32,
115                DType::I32,
116                DType::Bool(bool_store),
117                Default::default(),
118            )
119        }
120    }
121}
122
123/// Convenience macro to link to the `burn-tensor` docs for this crate version.
124///
125/// Usage:
126/// ```rust,ignore
127/// # use burn_backend::doc_tensor;
128/// doc_tensor!();        // Links to `Tensor` struct
129/// doc_tensor!("zeros"); // Links to `Tensor::zeros` method
130/// ```
131#[macro_export]
132macro_rules! doc_tensor {
133    () => {
134        concat!(
135            "[`Tensor`](https://docs.rs/burn-tensor/",
136            env!("CARGO_PKG_VERSION"),
137            "/burn_tensor/struct.Tensor.html)"
138        )
139    };
140
141    ($method:literal) => {
142        concat!(
143            "[`Tensor::",
144            $method,
145            "`](",
146            "https://docs.rs/burn-tensor/",
147            env!("CARGO_PKG_VERSION"),
148            "/burn_tensor/struct.Tensor.html#method.",
149            $method,
150            ")"
151        )
152    };
153}
burn_backend/lib.rs

burn_backend/
lib.rs
