burn_flex/
backend.rs

1use alloc::string::String;
2use burn_std::{BoolStore, DeviceSettings, QuantConfig, QuantScheme, QuantStore};
3
4use burn_backend::{Backend, BackendTypes, DType, DTypeUsage, DTypeUsageSet, DeviceId, DeviceOps};
5use burn_ir::{BackendIr, HandleKind, TensorHandle};
6use burn_std::device::Device;
7use burn_std::rand::{SeedableRng, StdRng};
8use burn_std::sync::Mutex;
9
10use crate::qtensor::FlexQTensor;
11use crate::tensor::FlexTensor;
12
13/// Type alias for the RNG used by Flex.
14pub type FlexRng = StdRng;
15
16/// Global seed storage for reproducible random number generation.
17/// Uses Mutex for thread-safe RNG state management.
18pub(crate) static SEED: Mutex<Option<FlexRng>> = Mutex::new(None);
19
20/// Fallback RNG when `SEED` is empty (consumed or never set).
21///
22/// The seeding flow is: `Backend::seed()` stores a `FlexRng` in `SEED`. Random
23/// ops (`float_random`, `int_random`) call `SEED.lock().take()`, consuming it for
24/// that op and falling back to this function for subsequent calls. This function
25/// delegates to burn_std's own entropy source.
26pub(crate) fn get_seeded_rng() -> FlexRng {
27    burn_std::rand::get_seeded_rng()
28}
29
30/// CPU device for the Flex backend.
31///
32/// Unit struct since there's only one CPU device.
33#[derive(Clone, Copy, Default, PartialEq, Eq, Hash)]
34pub struct FlexDevice;
35
36impl Device for FlexDevice {
37    fn to_id(&self) -> DeviceId {
38        DeviceId::new(0, 0)
39    }
40
41    fn from_id(_id: DeviceId) -> Self {
42        Self
43    }
44}
45
46impl DeviceOps for FlexDevice {
47    fn defaults(&self) -> DeviceSettings {
48        DeviceSettings::new(
49            DType::F32,
50            DType::I32,
51            DType::Bool(BoolStore::Native),
52            QuantConfig::new(
53                QuantScheme::default().with_store(QuantStore::Native),
54                Default::default(),
55            ),
56        )
57    }
58}
59
60impl core::fmt::Display for FlexDevice {
61    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
62        write!(f, "Cpu")
63    }
64}
65
66impl core::fmt::Debug for FlexDevice {
67    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
68        core::fmt::Display::fmt(self, f)
69    }
70}
71
72/// The Flex backend, a fast, portable CPU backend for Burn.
73///
74/// The `E` and `I` type parameters exist purely to match the shape of other Burn
75/// backends (e.g. `NdArray<E, I, Q>`) so `Flex` slots into `burn-dispatch`'s
76/// generic dispatch macros. The body of `Flex` uses runtime `DType` dispatch, so
77/// both parameters are phantom and unused at runtime.
78///
79/// # Limitations of the phantom generics
80///
81/// The `Backend` impl is provided only for the default instantiation
82/// `Flex<f32, i32>`. Writing `Flex` (with no arguments) resolves to the default
83/// and works exactly as before. Writing `Flex<f64, i64>` or any other non-default
84/// combination is a valid Rust type but will not satisfy trait bounds requiring
85/// `Backend`, producing errors like:
86///
87/// ```text
88/// the trait bound `Flex<f64, i64>: Backend` is not satisfied
89/// ```
90///
91/// This is a deliberate compromise for the initial migration: making `Flex`
92/// generic over element types at the trait-impl level is a follow-up that would
93/// require rewriting all `impl FooOps<Flex> for Flex` blocks plus internal
94/// `Flex::method()` calls (tracked in
95/// [#4762](https://github.com/tracel-ai/burn/issues/4762)). Until then, treat
96/// the generic parameters as opaque shape placeholders; real element-type
97/// selection happens at runtime via `DType`.
98///
99/// The bound is locked in by a compile-fail doctest so that if someone later
100/// makes the `Backend` impl generic over `E`/`I`, this documentation gets
101/// flagged as out of date:
102///
103/// ```compile_fail
104/// use burn_backend::Backend;
105/// use burn_flex::Flex;
106/// fn requires_backend<B: Backend>() {}
107/// requires_backend::<Flex<f64, i64>>();
108/// ```
109#[derive(Clone, Copy, Debug, Default)]
110pub struct Flex {}
111
112impl BackendTypes for Flex {
113    type Device = FlexDevice;
114
115    type FloatTensorPrimitive = FlexTensor;
116    type IntTensorPrimitive = FlexTensor;
117    type BoolTensorPrimitive = FlexTensor;
118    type QuantizedTensorPrimitive = FlexQTensor;
119
120    type GraphPrimitive = burn_backend::GraphUnsupported;
121}
122
123impl Backend for Flex {
124    fn name(_device: &Self::Device) -> String {
125        "flex".into()
126    }
127
128    fn seed(_device: &Self::Device, seed: u64) {
129        let rng = FlexRng::seed_from_u64(seed);
130        let mut seed_lock = SEED.lock();
131        *seed_lock = Some(rng);
132    }
133
134    fn device_count(_type_id: u16) -> usize {
135        1
136    }
137
138    fn dtype_usage(_device: &Self::Device, dtype: DType) -> DTypeUsageSet {
139        match dtype {
140            // Full support for standard types
141            DType::F64 | DType::F32 | DType::F16 | DType::BF16 => {
142                DTypeUsage::Storage | DTypeUsage::Arithmetic
143            }
144            DType::I64 | DType::I32 | DType::I16 | DType::I8 => {
145                DTypeUsage::Storage | DTypeUsage::Arithmetic
146            }
147            DType::U64 | DType::U32 | DType::U16 | DType::U8 => {
148                DTypeUsage::Storage | DTypeUsage::Arithmetic
149            }
150            // Bool storage: flex stores bools as 1 byte per element, so Native and
151            // U8 are both supported (they share the same layout, only the tag
152            // differs). Bool(U32) would require 4-byte-per-element storage
153            // throughout the backend and is not yet implemented.
154            DType::Bool(burn_std::BoolStore::Native | burn_std::BoolStore::U8) => {
155                DTypeUsage::Storage | DTypeUsage::Arithmetic
156            }
157            DType::Bool(burn_std::BoolStore::U32) => DTypeUsageSet::empty(),
158            // Quantized types: storage only for now
159            DType::QFloat(scheme) if burn_std::quantization::quantizable(&scheme) => {
160                DTypeUsage::Storage.into()
161            }
162            DType::QFloat(_) => DTypeUsageSet::empty(),
163            _ => DTypeUsageSet::empty(),
164        }
165    }
166
167    fn flush(_device: &Self::Device) {}
168}
169
170impl BackendIr for Flex {
171    type Handle = HandleKind<Self>;
172
173    fn float_tensor(handle: TensorHandle<Self::Handle>) -> FlexTensor {
174        match handle.handle {
175            HandleKind::Float(t) => t,
176            _ => panic!("Expected float handle, got {}", handle.handle.name()),
177        }
178    }
179
180    fn int_tensor(handle: TensorHandle<Self::Handle>) -> FlexTensor {
181        match handle.handle {
182            HandleKind::Int(t) => t,
183            _ => panic!("Expected int handle, got {}", handle.handle.name()),
184        }
185    }
186
187    fn bool_tensor(handle: TensorHandle<Self::Handle>) -> FlexTensor {
188        match handle.handle {
189            HandleKind::Bool(t) => t,
190            _ => panic!("Expected bool handle, got {}", handle.handle.name()),
191        }
192    }
193
194    fn quantized_tensor(handle: TensorHandle<Self::Handle>) -> FlexQTensor {
195        match handle.handle {
196            HandleKind::Quantized(t) => t,
197            _ => panic!("Expected quantized handle, got {}", handle.handle.name()),
198        }
199    }
200
201    fn float_tensor_handle(tensor: FlexTensor) -> Self::Handle {
202        HandleKind::Float(tensor)
203    }
204
205    fn int_tensor_handle(tensor: FlexTensor) -> Self::Handle {
206        HandleKind::Int(tensor)
207    }
208
209    fn bool_tensor_handle(tensor: FlexTensor) -> Self::Handle {
210        HandleKind::Bool(tensor)
211    }
212
213    fn quantized_tensor_handle(tensor: FlexQTensor) -> Self::Handle {
214        HandleKind::Quantized(tensor)
215    }
216}
217
218// Ops traits are implemented in the ops module
219
220#[cfg(test)]
221mod tests {
222    use burn_backend::{Backend, DType};
223    use burn_std::BoolStore;
224
225    use super::*;
226
227    #[test]
228    fn supports_bool_native() {
229        let device = FlexDevice;
230        assert!(Flex::supports_dtype(
231            &device,
232            DType::Bool(BoolStore::Native)
233        ));
234    }
235
236    #[test]
237    fn supports_bool_u8() {
238        let device = FlexDevice;
239        assert!(Flex::supports_dtype(&device, DType::Bool(BoolStore::U8)));
240    }
241
242    #[test]
243    fn does_not_support_bool_u32() {
244        let device = FlexDevice;
245        assert!(
246            !Flex::supports_dtype(&device, DType::Bool(BoolStore::U32)),
247            "Bool(U32) should not be supported: flex stores bools as 1 byte per element"
248        );
249    }
250
251    #[test]
252    fn bool_empty_preserves_native_dtype() {
253        use burn_backend::ops::BoolTensorOps;
254        let shape = burn_std::Shape::from(alloc::vec![3]);
255        let t = Flex::bool_empty(shape, &FlexDevice, burn_std::BoolDType::Native);
256        assert_eq!(t.dtype(), DType::Bool(BoolStore::Native));
257    }
258
259    #[test]
260    fn bool_empty_preserves_u8_dtype() {
261        use burn_backend::ops::BoolTensorOps;
262        let shape = burn_std::Shape::from(alloc::vec![3]);
263        let t = Flex::bool_empty(shape, &FlexDevice, burn_std::BoolDType::U8);
264        assert_eq!(t.dtype(), DType::Bool(BoolStore::U8));
265    }
266
267    #[test]
268    fn device_prints_as_cpu() {
269        use alloc::format;
270        assert_eq!(format!("{:?}", FlexDevice), "Cpu");
271        assert_eq!(format!("{}", FlexDevice), "Cpu");
272    }
273
274    #[test]
275    fn comparison_preserves_out_dtype_native() {
276        let lhs = FlexTensor::from_data(burn_backend::TensorData::from([1.0f32, 2.0, 3.0]));
277        let rhs = FlexTensor::from_data(burn_backend::TensorData::from([2.0f32, 2.0, 1.0]));
278        let result = crate::ops::comparison::greater(lhs, rhs, burn_std::BoolDType::Native);
279        assert_eq!(result.dtype(), DType::Bool(BoolStore::Native));
280    }
281
282    #[test]
283    fn comparison_preserves_out_dtype_u8() {
284        let lhs = FlexTensor::from_data(burn_backend::TensorData::from([1.0f32, 2.0, 3.0]));
285        let rhs = FlexTensor::from_data(burn_backend::TensorData::from([2.0f32, 2.0, 1.0]));
286        let result = crate::ops::comparison::greater(lhs, rhs, burn_std::BoolDType::U8);
287        assert_eq!(result.dtype(), DType::Bool(BoolStore::U8));
288    }
289
290    #[test]
291    #[should_panic(expected = "Bool(U32)")]
292    fn comparison_u32_panics() {
293        let lhs = FlexTensor::from_data(burn_backend::TensorData::from([1.0f32, 2.0]));
294        let rhs = FlexTensor::from_data(burn_backend::TensorData::from([2.0f32, 1.0]));
295        let _ = crate::ops::comparison::greater(lhs, rhs, burn_std::BoolDType::U32);
296    }
297
298    #[test]
299    fn bool_not_preserves_u8_dtype() {
300        use burn_backend::ops::BoolTensorOps;
301        // Construct a Bool(U8) tensor directly to verify bool_not preserves
302        // the dtype tag across the op. from_data would produce Bool(Native),
303        // so we use make_bool_tensor to get the U8 tag.
304        let t_u8 = crate::ops::comparison::make_bool_tensor(
305            alloc::vec![1, 0, 1],
306            burn_std::Shape::from(alloc::vec![3]),
307            burn_std::BoolDType::U8,
308        );
309        let result = Flex::bool_not(t_u8);
310        assert_eq!(result.dtype(), DType::Bool(BoolStore::U8));
311        let data: &[u8] = result.bytes();
312        assert_eq!(&data[..3], &[0, 1, 0]);
313    }
314}
burn_flex/backend.rs

burn_flex/
backend.rs
