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#[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#[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 type FloatElem = f32;
103
104 type IntTensorPrimitive = HostTensor;
105 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 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 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 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#[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 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}