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

burn_flex/
backend.rs
