Skip to main content

ruda_tensor_host/
backend.rs

1use alloc::string::String;
2use core::marker::PhantomData;
3
4use ruda_tensor::{Backend, BackendTypes, DType, DTypeUsage, DTypeUsageSet, DeviceId, DeviceOps};
5use ruda_tensor::graph::{BackendIr, HandleKind, TensorHandle};
6use ruda_core::device::Device;
7
8use crate::qtensor::HostQTensor;
9use crate::tensor::HostTensor;
10
11pub use rurand_host::HostRng as HostRng;
12
13/// CPU device for the Host backend.
14///
15/// Unit struct since there's only one CPU device.
16#[derive(Clone, Copy, Default, PartialEq, Eq, Hash)]
17pub struct HostDevice;
18
19impl Device for HostDevice {
20    fn to_id(&self) -> DeviceId {
21        DeviceId::new(0, 0)
22    }
23
24    fn from_id(_id: DeviceId) -> Self {
25        Self
26    }
27}
28
29impl DeviceOps for HostDevice {}
30
31impl core::fmt::Display for HostDevice {
32    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
33        write!(f, "Cpu")
34    }
35}
36
37impl core::fmt::Debug for HostDevice {
38    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
39        core::fmt::Display::fmt(self, f)
40    }
41}
42
43/// The Host backend, a fast, portable CPU backend for Ruda.
44///
45/// The `E` and `I` type parameters exist purely to match the shape of other Ruda
46/// backends (e.g. `NdArray<E, I, Q>`) so `Host` slots into `ruda-dispatch`'s
47/// generic dispatch macros. The body of `Host` uses runtime `DType` dispatch, so
48/// both parameters are phantom and unused at runtime.
49///
50/// # Limitations of the phantom generics
51///
52/// The `Backend` impl is provided only for the default instantiation
53/// `Host<f32, i32>`. Writing `Host` (with no arguments) resolves to the default
54/// and works exactly as before. Writing `Host<f64, i64>` or any other non-default
55/// combination is a valid Rust type but will not satisfy trait bounds requiring
56/// `Backend`, producing errors like:
57///
58/// ```text
59/// the trait bound `Host<f64, i64>: Backend` is not satisfied
60/// ```
61///
62/// This is a deliberate compromise for the initial migration: making `Host`
63/// generic over element types at the trait-impl level is a follow-up that would
64/// require rewriting all `impl FooOps<Host> for Host` blocks plus internal
65/// `Host::method()` calls (tracked in
66/// [#4762](https://github.com/shuqi2077/RUDA/blob/main/THIRD_PARTY_NOTICES.md)). Until then, treat
67/// the generic parameters as opaque shape placeholders; real element-type
68/// selection happens at runtime via `DType`.
69///
70/// The bound is locked in by a compile-fail doctest so that if someone later
71/// makes the `Backend` impl generic over `E`/`I`, this documentation gets
72/// flagged as out of date:
73///
74/// ```compile_fail
75/// use ruda_tensor::Backend;
76/// use ruda_tensor_host::Host;
77/// fn requires_backend<B: Backend>() {}
78/// requires_backend::<Host<f64, i64>>();
79/// ```
80#[derive(Clone, Copy, Default)]
81pub struct Host<E = f32, I = i32> {
82    _e: PhantomData<E>,
83    _i: PhantomData<I>,
84}
85
86impl<E: core::fmt::Debug, I: core::fmt::Debug> core::fmt::Debug for Host<E, I> {
87    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
88        f.debug_struct("Flex")
89            .field("_e", &self._e)
90            .field("_i", &self._i)
91            .finish()
92    }
93}
94
95impl BackendTypes for Host {
96    type Device = HostDevice;
97
98    type FloatTensorPrimitive = HostTensor;
99    /// Default float element type. Determines the dtype for `.float()` conversions and
100    /// `Tensor::from_data` when no explicit dtype is provided.
101    /// Prefer explicit dtypes via `(&device, DType::F32)`.
102    type FloatElem = f32;
103
104    type IntTensorPrimitive = HostTensor;
105    /// Default int element type. Determines the dtype for `.int()` conversions and
106    /// `Tensor::from_data` when no explicit dtype is provided.
107    /// Set to i32 to match ruda's ecosystem default (test suite, record settings, ruda-tensor-remote).
108    /// Prefer explicit dtypes via `(&device, DType::I32)`.
109    type IntElem = i32;
110
111    type BoolTensorPrimitive = HostTensor;
112    type BoolElem = bool;
113
114    type QuantizedTensorPrimitive = HostQTensor;
115}
116
117impl Backend for Host {
118    fn flush(_device: &Self::Device) -> Result<(), ruda_tensor::ExecutionError> {
119        Ok(())
120    }
121
122    fn name(_device: &Self::Device) -> String {
123        "flex".into()
124    }
125
126    fn seed(_device: &Self::Device, seed: u64) {
127        rurand_host::seed(seed)
128    }
129
130    fn rng_state(_device: &Self::Device) -> Result<alloc::vec::Vec<u8>, ruda_tensor::ExecutionError> {
131        Ok(rurand_host::rng_state())
132    }
133
134    fn restore_rng_state(_device: &Self::Device, state: &[u8]) -> Result<(), ruda_tensor::ExecutionError> {
135        rurand_host::restore_rng_state(state)
136            .map_err(|reason| ruda_tensor::ExecutionError::WithContext { reason: reason.into() })
137    }
138
139    fn device_count(_type_id: u16) -> usize {
140        1
141    }
142
143    fn dtype_usage(_device: &Self::Device, dtype: DType) -> DTypeUsageSet {
144        match dtype {
145            // Full support for standard types
146            DType::F64 | DType::F32 | DType::F16 | DType::BF16 => {
147                DTypeUsage::Storage | DTypeUsage::Arithmetic
148            }
149            DType::I64 | DType::I32 | DType::I16 | DType::I8 => {
150                DTypeUsage::Storage | DTypeUsage::Arithmetic
151            }
152            DType::U64 | DType::U32 | DType::U16 | DType::U8 => {
153                DTypeUsage::Storage | DTypeUsage::Arithmetic
154            }
155            // Bool storage: flex stores bools as 1 byte per element, so Native and
156            // U8 are both supported (they share the same layout, only the tag
157            // differs). Bool(U32) would require 4-byte-per-element storage
158            // throughout the backend and is not yet implemented.
159            DType::Bool(ruda_tensor::BoolStore::Native | ruda_tensor::BoolStore::U8) => {
160                DTypeUsage::Storage | DTypeUsage::Arithmetic
161            }
162            DType::Bool(ruda_tensor::BoolStore::U32) => DTypeUsageSet::empty(),
163            // Quantized types: storage only for now
164            DType::QFloat(_) => DTypeUsage::Storage.into(),
165            _ => DTypeUsageSet::empty(),
166        }
167    }
168}
169
170impl BackendIr for Host {
171    type Handle = HandleKind<Self>;
172
173    fn float_tensor(handle: TensorHandle<Self::Handle>) -> HostTensor {
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>) -> HostTensor {
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>) -> HostTensor {
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>) -> HostQTensor {
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: HostTensor) -> Self::Handle {
202        HandleKind::Float(tensor)
203    }
204
205    fn int_tensor_handle(tensor: HostTensor) -> Self::Handle {
206        HandleKind::Int(tensor)
207    }
208
209    fn bool_tensor_handle(tensor: HostTensor) -> Self::Handle {
210        HandleKind::Bool(tensor)
211    }
212
213    fn quantized_tensor_handle(tensor: HostQTensor) -> 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 ruda_tensor::{Backend, DType};
223    use ruda_tensor::BoolStore;
224
225    use super::*;
226
227    #[test]
228    fn supports_bool_native() {
229        let device = HostDevice;
230        assert!(Host::supports_dtype(
231            &device,
232            DType::Bool(BoolStore::Native)
233        ));
234    }
235
236    #[test]
237    fn supports_bool_u8() {
238        let device = HostDevice;
239        assert!(Host::supports_dtype(&device, DType::Bool(BoolStore::U8)));
240    }
241
242    #[test]
243    fn does_not_support_bool_u32() {
244        let device = HostDevice;
245        assert!(
246            !Host::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 ruda_tensor::ops::BoolTensorOps;
254        let shape = ruda_tensor::Shape::from(alloc::vec![3]);
255        let t = Host::bool_empty(shape, &HostDevice, ruda_tensor::BoolDType::Native);
256        assert_eq!(t.dtype(), DType::Bool(BoolStore::Native));
257    }
258
259    #[test]
260    fn bool_empty_preserves_u8_dtype() {
261        use ruda_tensor::ops::BoolTensorOps;
262        let shape = ruda_tensor::Shape::from(alloc::vec![3]);
263        let t = Host::bool_empty(shape, &HostDevice, ruda_tensor::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!("{:?}", HostDevice), "Cpu");
271        assert_eq!(format!("{}", HostDevice), "Cpu");
272    }
273
274    #[test]
275    fn comparison_preserves_out_dtype_native() {
276        let lhs = HostTensor::from_data(ruda_tensor::TensorData::from([1.0f32, 2.0, 3.0]));
277        let rhs = HostTensor::from_data(ruda_tensor::TensorData::from([2.0f32, 2.0, 1.0]));
278        let result = crate::ops::comparison::greater(lhs, rhs, ruda_tensor::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 = HostTensor::from_data(ruda_tensor::TensorData::from([1.0f32, 2.0, 3.0]));
285        let rhs = HostTensor::from_data(ruda_tensor::TensorData::from([2.0f32, 2.0, 1.0]));
286        let result = crate::ops::comparison::greater(lhs, rhs, ruda_tensor::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 = HostTensor::from_data(ruda_tensor::TensorData::from([1.0f32, 2.0]));
294        let rhs = HostTensor::from_data(ruda_tensor::TensorData::from([2.0f32, 1.0]));
295        let _ = crate::ops::comparison::greater(lhs, rhs, ruda_tensor::BoolDType::U32);
296    }
297
298    #[test]
299    fn bool_not_preserves_u8_dtype() {
300        use ruda_tensor::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            ruda_tensor::Shape::from(alloc::vec![3]),
307            ruda_tensor::BoolDType::U8,
308        );
309        let result = Host::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}