burn_tensor/tensor/backend/base.rs
1use alloc::string::String;
2
3use crate::TensorMetadata;
4use crate::tensor::Element;
5use crate::{ops::*, quantization::QTensorPrimitive};
6
7use super::DeviceOps;
8
9/// This trait defines all types and functions needed for a backend to be used with burn.
10///
11/// ## Design
12///
13/// This trait aims to be as unopinionated as possible and allows implementations to define
14/// their own types and patterns. Therefore, there are few pre-defined abstractions baked
15/// into this trait.
16///
17/// Backends must define their own tensor types for each data type: `float`, `int`, and `bool`.
18/// Since we minimize assumptions, we chose to separate these types, as they are used in
19/// different contexts. However, some backends may have a generic tensor type that is used
20/// for all data types.
21///
22/// ### Eager Mode
23///
24/// Because burn supports dynamic graphs, the backend trait is designed around kernel
25/// implementations that can be called without any mutable context or graph. This may not be
26/// ideal for backends that want to configure their computational graphs and execute them
27/// multiple times.
28///
29/// To implement this kind of backend, channels could be used to communicate with a backend
30/// server thread to build the computation graphs and re-execute the ones that are repeated,
31/// with some form of cache. Once that pattern has matured, a graph mode backend trait could
32/// be extracted from it, allowing other backends of the same kind to be quickly integrated
33/// with burn. This pattern could also be used to create an operation fusion trait, which
34/// allows backends to define what kind of graph structures can be fused into one operation.
35///
36/// ### Multi-Threaded
37///
38/// Backend tensor types are all `Clone` + `Send`, which allows them to be safely
39/// sent between threads. It is recommended to wrap tensors with [Arc](alloc::sync::Arc),
40/// which avoids copying the tensor's buffer. Note that it is still possible to mutate and
41/// reuse tensors' buffer without locking; see the next section on the Mutable API.
42///
43/// ### Mutable API
44///
45/// There is no mutable or inplace operation API to implement, but that does not mean that
46/// backends cannot support them. Using [try_unwrap](alloc::sync::Arc::try_unwrap) and
47/// [get_mut](alloc::sync::Arc::get_mut) allows backends to have access to an owned or mutable
48/// reference to their tensor buffer data structure if the tensor is not shared. In that case,
49/// backends can dispatch to their owned inplace operations for better performance.
50///
51/// ## Documentation
52///
53/// Most of the documentation for each function can be found on the user API [tensor struct](crate::Tensor).
54/// For modules, public functions are often created, which can be used by `burn-core` modules.
55pub trait Backend:
56 FloatTensorOps<Self>
57 + BoolTensorOps<Self>
58 + IntTensorOps<Self>
59 + ModuleOps<Self>
60 + ActivationOps<Self>
61 + QTensorOps<Self>
62 + TransactionOps<Self>
63 + Clone
64 + Default
65 + Sized
66 + Send
67 + Sync
68 + core::fmt::Debug
69 + 'static
70{
71 /// Device type.
72 type Device: DeviceOps;
73
74 /// Tensor primitive to be used for all float operations.
75 type FloatTensorPrimitive: TensorMetadata + 'static;
76 /// Default float element type.
77 type FloatElem: Element;
78
79 /// Tensor primitive to be used for all int operations.
80 type IntTensorPrimitive: TensorMetadata + 'static;
81 /// Int element type.
82 type IntElem: Element;
83
84 /// Tensor primitive to be used for all bool operations.
85 type BoolTensorPrimitive: TensorMetadata + 'static;
86 /// Tensor primitive to be used for all bool operations.
87 type BoolElem: Element;
88
89 /// Tensor primitive to be used for all quantized operations.
90 type QuantizedTensorPrimitive: TensorMetadata + QTensorPrimitive + 'static;
91 /// Quantized tensor encoding type.
92 type QuantizedEncoding: Element;
93
94 /// If autodiff is enabled.
95 fn ad_enabled() -> bool {
96 false
97 }
98
99 /// Name of the backend.
100 fn name(device: &Self::Device) -> String;
101
102 /// Seed the backend.
103 fn seed(seed: u64);
104
105 /// Sync the backend, ensure that all computation are finished.
106 fn sync(_device: &Self::Device) {}
107}
108
109/// Trait that allows a backend to support autodiff.
110pub trait AutodiffBackend: Backend {
111 /// The inner backend type.
112 type InnerBackend: Backend<Device = Self::Device, FloatElem = Self::FloatElem, IntElem = Self::IntElem>;
113
114 /// Gradients type.
115 type Gradients: Send;
116
117 /// Backward pass.
118 ///
119 /// # Arguments
120 ///
121 /// * `tensor` - The tensor is the last node of computational graph where the gradients are computed.
122 ///
123 /// # Returns
124 ///
125 /// The gradients.
126 fn backward(tensor: FloatTensor<Self>) -> Self::Gradients;
127
128 /// Returns the gradients of a tensor.
129 ///
130 /// # Arguments
131 ///
132 /// * `tensor` - The tensor to extract the gradients from.
133 ///
134 /// # Returns
135 ///
136 /// An optional tensor containing the gradient.
137 fn grad(
138 tensor: &FloatTensor<Self>,
139 grads: &Self::Gradients,
140 ) -> Option<FloatTensor<Self::InnerBackend>>;
141
142 /// Pops the gradients of a tensor and returns them.
143 ///
144 /// # Arguments
145 ///
146 /// * `tensor` - The tensor to pop the gradients from.
147 /// * `grads` - The gradients.
148 ///
149 /// # Returns
150 ///
151 /// An optional tensor containing the given gradients.
152 fn grad_remove(
153 tensor: &FloatTensor<Self>,
154 grads: &mut Self::Gradients,
155 ) -> Option<FloatTensor<Self::InnerBackend>>;
156
157 /// Replace the gradients of a tensor with the one provided.
158 ///
159 /// If no gradient existed for the provided tensor, register it.
160 ///
161 /// # Arguments
162 ///
163 /// * `tensor` - The tensor to pop the gradients from.
164 /// * `grads` - The gradients.
165 /// * `grad` - The updated grad tensor.
166 fn grad_replace(
167 tensor: &FloatTensor<Self>,
168 grads: &mut Self::Gradients,
169 grad: FloatTensor<Self::InnerBackend>,
170 );
171
172 /// Returns the tensor with inner backend type.
173 ///
174 /// # Arguments
175 ///
176 /// * `tensor` - The tensor to get the inner backend tensor for.
177 ///
178 /// # Returns
179 ///
180 /// The inner backend tensor.
181 fn inner(tensor: FloatTensor<Self>) -> FloatTensor<Self::InnerBackend>;
182
183 /// Returns the tensor with inner backend type.
184 ///
185 /// # Arguments
186 ///
187 /// * `tensor` - The tensor to get the inner backend tensor for.
188 ///
189 /// # Returns
190 ///
191 /// The inner backend tensor.
192 fn int_inner(tensor: IntTensor<Self>) -> IntTensor<Self::InnerBackend>;
193
194 /// Returns the tensor with inner backend type.
195 ///
196 /// # Arguments
197 ///
198 /// * `tensor` - The tensor to get the inner backend tensor for.
199 ///
200 /// # Returns
201 ///
202 /// The inner backend tensor.
203 fn bool_inner(tensor: BoolTensor<Self>) -> BoolTensor<Self::InnerBackend>;
204
205 /// Returns the tensor with inner backend type.
206 ///
207 /// # Arguments
208 ///
209 /// * `tensor` - The tensor to get the inner backend tensor for.
210 ///
211 /// # Returns
212 ///
213 /// The inner backend tensor.
214 fn q_inner(tensor: QuantizedTensor<Self>) -> QuantizedTensor<Self::InnerBackend>;
215
216 /// Converts the inner backend tensor to the autodiff backend tensor.
217 ///
218 /// # Arguments
219 ///
220 /// * `tensor` - The inner backend tensor to convert.
221 ///
222 ///
223 /// # Returns
224 ///
225 /// The autodiff backend tensor.
226 fn from_inner(tensor: FloatTensor<Self::InnerBackend>) -> FloatTensor<Self>;
227
228 /// Converts the inner backend tensor to the autodiff backend tensor.
229 ///
230 /// # Arguments
231 ///
232 /// * `tensor` - The inner backend tensor to convert.
233 ///
234 ///
235 /// # Returns
236 ///
237 /// The autodiff backend tensor.
238 fn int_from_inner(tensor: IntTensor<Self::InnerBackend>) -> IntTensor<Self>;
239
240 /// Converts the inner backend tensor to the autodiff backend tensor.
241 ///
242 /// # Arguments
243 ///
244 /// * `tensor` - The inner backend tensor to convert.
245 ///
246 ///
247 /// # Returns
248 ///
249 /// The autodiff backend tensor.
250 fn bool_from_inner(tensor: BoolTensor<Self::InnerBackend>) -> BoolTensor<Self>;
251
252 /// Converts the inner backend tensor to the autodiff backend tensor.
253 ///
254 /// # Arguments
255 ///
256 /// * `tensor` - The inner backend tensor to convert.
257 ///
258 ///
259 /// # Returns
260 ///
261 /// The autodiff backend tensor.
262 fn q_from_inner(tensor: QuantizedTensor<Self::InnerBackend>) -> QuantizedTensor<Self>;
263}