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//! This library provides the core types that define how Burn tensor data is represented, stored, and interpreted.
6
7#[macro_use]
8extern crate derive_new;
9
10extern crate alloc;
11
12/// [`Backend`] trait and required types.
13pub mod backend;
14pub use backend::*;
15
16// Re-exported types
17pub use burn_std::reader::*; // Useful so that backends don't have to add `burn_std` as a dependency.
18pub use burn_std::{
19    AllocationProperty, BoolDType, BoolStore, Bytes, DType, DataError, DeviceHandle, Distribution,
20    DistributionSampler, DistributionSamplerKind, Element, ElementAdd, ElementConversion,
21    ElementEq, ElementOrdered, ElementRandom, FloatDType, IntDType, Scalar, SplitPolicy,
22    TensorData, Tolerance, bf16, distribution, element, f16, stream::StreamId,
23};
24
25/// Shape definition.
26pub mod shape {
27    pub use burn_std::shape::*;
28}
29pub use shape::*;
30
31/// Slice utilities.
32pub mod slice {
33    pub use burn_std::{s, slice::*};
34}
35pub use slice::*;
36
37/// Indexing utilities.
38pub mod indexing {
39    pub use burn_std::indexing::*;
40}
41pub use indexing::*;
42
43mod alias;
44pub use alias::*;
45
46/// Quantization data representation.
47pub mod quantization;
48
49/// CubeCL inter-operation helpers (gated by the `cubecl` feature).
50///
51/// Provides plain conversion functions between burn's [`DType`] and cubecl's
52/// `ElemType` / `StorageType`. They are intentionally exposed as named
53/// functions rather than `From`/`Into` impls so the cubecl type tree does not
54/// leak into `burn-std`'s public surface.
55#[cfg(feature = "cubecl")]
56pub mod cubecl;
57
58#[cfg(any(
59    feature = "cubecl-wgpu",
60    feature = "cubecl-metal",
61    feature = "cubecl-vulkan",
62    feature = "cubecl-webgpu"
63))]
64mod cube_wgpu {
65    use crate::backend::DeviceOps;
66    use burn_std::{BoolStore, DType, DeviceSettings};
67    use cubecl::wgpu::WgpuDevice;
68
69    impl DeviceOps for WgpuDevice {
70        #[cfg(not(any(feature = "cubecl-metal", feature = "cubecl-vulkan")))]
71        fn defaults(&self) -> DeviceSettings {
72            DeviceSettings::new(
73                DType::F32,
74                DType::I32,
75                DType::Bool(BoolStore::U32),
76                Default::default(),
77            )
78        }
79
80        #[cfg(any(feature = "cubecl-metal", feature = "cubecl-vulkan"))]
81        fn defaults(&self) -> DeviceSettings {
82            DeviceSettings::new(
83                DType::F32,
84                DType::I32,
85                DType::Bool(BoolStore::U8),
86                Default::default(),
87            )
88        }
89    }
90}
91
92#[cfg(feature = "cubecl-cuda")]
93mod cube_cuda {
94    use crate::backend::DeviceOps;
95    use burn_std::{BoolStore, DType, DeviceSettings};
96    use cubecl::cuda::CudaDevice;
97
98    impl DeviceOps for CudaDevice {
99        fn defaults(&self) -> DeviceSettings {
100            DeviceSettings::new(
101                DType::F32,
102                DType::I32,
103                DType::Bool(BoolStore::U8),
104                Default::default(),
105            )
106        }
107    }
108}
109
110#[cfg(feature = "cubecl-cpu")]
111mod cube_cpu {
112    use crate::backend::DeviceOps;
113    use burn_std::{BoolStore, DType, DeviceSettings};
114    use cubecl::cpu::CpuDevice;
115
116    impl DeviceOps for CpuDevice {
117        fn defaults(&self) -> DeviceSettings {
118            DeviceSettings::new(
119                DType::F32,
120                DType::I32,
121                DType::Bool(BoolStore::U8),
122                Default::default(),
123            )
124        }
125    }
126}
127
128#[cfg(feature = "cubecl-hip")]
129mod cube_hip {
130    use crate::backend::DeviceOps;
131    use burn_std::{BoolStore, DType, DeviceSettings};
132    use cubecl::hip::AmdDevice;
133
134    impl DeviceOps for AmdDevice {
135        fn defaults(&self) -> DeviceSettings {
136            DeviceSettings::new(
137                DType::F32,
138                DType::I32,
139                DType::Bool(BoolStore::U8),
140                Default::default(),
141            )
142        }
143    }
144}
145
146/// Convenience macro to link to the `burn-tensor` docs for this crate version.
147///
148/// Usage:
149/// ```rust,ignore
150/// # use burn_backend::doc_tensor;
151/// doc_tensor!();        // Links to `Tensor` struct
152/// doc_tensor!("zeros"); // Links to `Tensor::zeros` method
153/// ```
154#[macro_export]
155macro_rules! doc_tensor {
156    () => {
157        concat!(
158            "[`Tensor`](https://docs.rs/burn-tensor/",
159            env!("CARGO_PKG_VERSION"),
160            "/burn_tensor/struct.Tensor.html)"
161        )
162    };
163
164    ($method:literal) => {
165        concat!(
166            "[`Tensor::",
167            $method,
168            "`](",
169            "https://docs.rs/burn-tensor/",
170            env!("CARGO_PKG_VERSION"),
171            "/burn_tensor/struct.Tensor.html#method.",
172            $method,
173            ")"
174        )
175    };
176}
burn_backend/lib.rs

burn_backend/
lib.rs
