ruccl/tensor_device/
element.rs

1use crate::rank::ElementType;
2use ruda_tensor::{Element, bf16, f16};
3
4mod sealed {
5    pub trait Wire: Sized {
6        fn append_le(self, bytes: &mut Vec<u8>);
7        fn read_le(bytes: &[u8]) -> Self;
8    }
9}
10
11/// Storage types supported by the shared collective device adapter.
12///
13/// The sealed implementations preserve F32/F16/BF16/I32 storage bits on the
14/// little-endian rank transport; they do not convert half values through F32.
15pub trait TensorElement: Element + sealed::Wire {
16    const ELEMENT_TYPE: ElementType;
17}
18
19macro_rules! native_element {
20    ($ty:ty, $kind:ident) => {
21        impl TensorElement for $ty {
22            const ELEMENT_TYPE: ElementType = ElementType::$kind;
23        }
24        impl sealed::Wire for $ty {
25            fn append_le(self, bytes: &mut Vec<u8>) {
26                bytes.extend_from_slice(&self.to_le_bytes());
27            }
28            fn read_le(bytes: &[u8]) -> Self {
29                Self::from_le_bytes(bytes.try_into().expect("complete wire element"))
30            }
31        }
32    };
33}
34
35macro_rules! half_element {
36    ($ty:ty, $kind:ident) => {
37        impl TensorElement for $ty {
38            const ELEMENT_TYPE: ElementType = ElementType::$kind;
39        }
40        impl sealed::Wire for $ty {
41            fn append_le(self, bytes: &mut Vec<u8>) {
42                bytes.extend_from_slice(&self.to_bits().to_le_bytes());
43            }
44            fn read_le(bytes: &[u8]) -> Self {
45                Self::from_bits(u16::from_le_bytes(bytes.try_into().expect("complete wire element")))
46            }
47        }
48    };
49}
50
51native_element!(f32, F32);
52native_element!(i32, I32);
53half_element!(f16, F16);
54half_element!(bf16, BF16);
55
56pub(super) fn encode<T: TensorElement>(values: &[T]) -> Vec<u8> {
57    let mut bytes = Vec::with_capacity(std::mem::size_of_val(values));
58    for value in values {
59        value.append_le(&mut bytes);
60    }
61    bytes
62}
63
64pub(super) fn decode<T: TensorElement>(bytes: &[u8]) -> Result<Vec<T>, super::TensorDeviceError> {
65    let width = std::mem::size_of::<T>();
66    if !bytes.len().is_multiple_of(width) {
67        return Err(super::TensorDeviceError::InvalidBuffer("incomplete wire element"));
68    }
69    Ok(bytes.chunks_exact(width).map(T::read_le).collect())
70}
ruccl/tensor_device/element.rs

ruccl/tensor_device/
element.rs
