burn_flex/
backend.rs

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

burn_flex/
backend.rs
