onnx_runtime_ir/
tensor.rs

1//! Constant tensor storage, weight references, and ONNX type descriptors.
2
3use std::path::PathBuf;
4
5use crate::dtype::DataType;
6use crate::shape::Shape;
7
8mod sealed {
9    pub trait Sealed {}
10}
11
12/// A primitive numeric type that can be decoded from little-endian bytes.
13pub trait FromLeBytes: sealed::Sealed + Sized {
14    const BYTE_SIZE: usize;
15
16    fn from_le_bytes(bytes: &[u8]) -> Self;
17}
18
19macro_rules! impl_from_le_bytes {
20    ($($type:ty),+ $(,)?) => {
21        $(
22            impl sealed::Sealed for $type {}
23
24            impl FromLeBytes for $type {
25                const BYTE_SIZE: usize = size_of::<Self>();
26
27                fn from_le_bytes(bytes: &[u8]) -> Self {
28                    let mut array = [0_u8; size_of::<Self>()];
29                    array.copy_from_slice(bytes);
30                    Self::from_le_bytes(array)
31                }
32            }
33        )+
34    };
35}
36
37impl_from_le_bytes!(i32, i64, f32, f64);
38
39/// An invalid byte length for little-endian numeric decoding.
40#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
41pub enum RawBytesError {
42    #[error(
43        "invalid byte length for little-endian {type_name} scalar: expected {expected}, got {actual}"
44    )]
45    ScalarLength {
46        type_name: &'static str,
47        expected: usize,
48        actual: usize,
49    },
50    #[error(
51        "invalid byte length for little-endian {type_name} vector: {actual} is not a multiple of {element_size}"
52    )]
53    VectorLength {
54        type_name: &'static str,
55        element_size: usize,
56        actual: usize,
57    },
58}
59
60/// Decode one primitive numeric value from an exact-length little-endian byte slice.
61pub fn read_scalar_le<T: FromLeBytes>(bytes: &[u8]) -> Result<T, RawBytesError> {
62    if bytes.len() != T::BYTE_SIZE {
63        return Err(RawBytesError::ScalarLength {
64            type_name: std::any::type_name::<T>(),
65            expected: T::BYTE_SIZE,
66            actual: bytes.len(),
67        });
68    }
69    Ok(T::from_le_bytes(bytes))
70}
71
72/// Decode primitive numeric values from a little-endian byte slice.
73// `as_chunks::<N>()` needs a const generic argument, and `T::BYTE_SIZE` is an
74// associated const of a generic parameter, which is not permitted in that
75// position on stable. The chunk size here is not statically known per
76// monomorphisation site the way the lint assumes, so `chunks_exact` stays.
77#[allow(
78    clippy::chunks_exact_to_as_chunks,
79    reason = "chunk size is an associated const of a generic parameter"
80)]
81pub fn read_vec_le<T: FromLeBytes>(bytes: &[u8]) -> Result<Vec<T>, RawBytesError> {
82    if !bytes.len().is_multiple_of(T::BYTE_SIZE) {
83        return Err(RawBytesError::VectorLength {
84            type_name: std::any::type_name::<T>(),
85            element_size: T::BYTE_SIZE,
86            actual: bytes.len(),
87        });
88    }
89    Ok(bytes
90        .chunks_exact(T::BYTE_SIZE)
91        .map(T::from_le_bytes)
92        .collect())
93}
94
95/// A concrete constant tensor held inline (e.g. an attribute value or a small
96/// initializer). Element bytes are stored little-endian and densely packed.
97///
98/// Large model weights are referenced lazily via [`WeightRef`] instead.
99#[derive(Clone, Debug, PartialEq)]
100pub struct TensorData {
101    pub name: Option<String>,
102    pub dtype: DataType,
103    /// Static dimensions (constants always have a fully known shape).
104    pub dims: Vec<usize>,
105    /// Raw little-endian element bytes. Sub-byte values are densely packed
106    /// (two 4-bit or four 2-bit elements per byte); for [`DataType::String`]
107    /// this is empty and `strings` is used instead.
108    pub data: Vec<u8>,
109    /// String payloads for [`DataType::String`] tensors.
110    pub strings: Vec<String>,
111}
112
113impl TensorData {
114    /// A numeric tensor from raw little-endian bytes.
115    pub fn from_raw(dtype: DataType, dims: Vec<usize>, data: Vec<u8>) -> Self {
116        Self {
117            name: None,
118            dtype,
119            dims,
120            data,
121            strings: Vec::new(),
122        }
123    }
124
125    /// Number of elements (product of dims; `1` for a scalar).
126    pub fn numel(&self) -> usize {
127        self.checked_numel().expect("tensor element count overflow")
128    }
129
130    /// Number of elements, or `None` when the dimensions overflow `usize`.
131    pub fn checked_numel(&self) -> Option<usize> {
132        checked_numel(&self.dims)
133    }
134
135    /// Expected byte length for `numel` elements of `dtype`, accounting for
136    /// sub-byte packing.
137    pub fn expected_bytes(&self) -> usize {
138        self.checked_expected_bytes()
139            .expect("tensor byte count overflow")
140    }
141
142    /// Expected byte length, or `None` when the element or byte count overflows.
143    pub fn checked_expected_bytes(&self) -> Option<usize> {
144        checked_expected_bytes(self.dtype, &self.dims)
145    }
146}
147
148/// Number of elements in `dims`, or `None` when their product overflows.
149pub fn checked_numel(dims: &[usize]) -> Option<usize> {
150    dims.iter()
151        .try_fold(1usize, |product, &dimension| product.checked_mul(dimension))
152}
153
154/// Dense storage size for `dtype` and `dims`, or `None` on geometry overflow.
155pub fn checked_expected_bytes(dtype: DataType, dims: &[usize]) -> Option<usize> {
156    let element_count = checked_numel(dims)?;
157    if dtype == DataType::Undefined {
158        return None;
159    }
160    if dtype.is_sub_byte() {
161        let elements_per_byte = 8 / dtype.bit_size();
162        return (element_count / elements_per_byte).checked_add(usize::from(
163            !element_count.is_multiple_of(elements_per_byte),
164        ));
165    }
166    element_count.checked_mul(dtype.byte_size())
167}
168
169/// A sparse constant tensor in COO form.
170#[derive(Clone, Debug, PartialEq)]
171pub struct SparseTensorData {
172    /// Non-zero values.
173    pub values: TensorData,
174    /// Indices of the non-zero values (int64), shape `[nnz, rank]` or `[nnz]`.
175    pub indices: TensorData,
176    /// Dense shape.
177    pub dims: Vec<usize>,
178}
179
180/// An ONNX `TypeProto`: the type of a value, which may be a tensor or a
181/// container of tensors (see `docs/architecture/ORT2.md` §3.2).
182#[derive(Clone, Debug, PartialEq)]
183pub enum TypeProto {
184    Tensor {
185        dtype: DataType,
186        shape: Shape,
187    },
188    Sequence(Box<TypeProto>),
189    Optional(Box<TypeProto>),
190    Map {
191        key: DataType,
192        value: Box<TypeProto>,
193    },
194    SparseTensor {
195        dtype: DataType,
196        shape: Shape,
197    },
198}
199
200/// A reference to initializer (weight) data.
201///
202/// Small weights may be inlined; large weights are memory-mapped from an
203/// external file at load time (see `docs/architecture/ORT2.md` §12). The IR only stores the
204/// *reference*; the loader/`onnx-runtime-memory` crate performs the mmap.
205#[derive(Clone, Debug, PartialEq)]
206pub enum WeightRef {
207    /// Weight bytes stored inline in the model.
208    Inline(TensorData),
209    /// Weight bytes located in an external file at `[offset, offset+length)`.
210    External {
211        path: PathBuf,
212        offset: usize,
213        length: usize,
214        dtype: DataType,
215        dims: Vec<usize>,
216    },
217}
218
219impl WeightRef {
220    /// The element type of the referenced weight.
221    pub fn dtype(&self) -> DataType {
222        match self {
223            WeightRef::Inline(t) => t.dtype,
224            WeightRef::External { dtype, .. } => *dtype,
225        }
226    }
227
228    /// The static dimensions of the referenced weight.
229    pub fn dims(&self) -> &[usize] {
230        match self {
231            WeightRef::Inline(t) => &t.dims,
232            WeightRef::External { dims, .. } => dims,
233        }
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240
241    #[test]
242    fn tensor_numel_and_bytes() {
243        let t = TensorData::from_raw(DataType::Float32, vec![2, 3], vec![0u8; 24]);
244        assert_eq!(t.numel(), 6);
245        assert_eq!(t.expected_bytes(), 24);
246    }
247
248    #[test]
249    fn sub_byte_expected_bytes() {
250        let t = TensorData::from_raw(DataType::Int4, vec![3], vec![0u8; 2]);
251        assert_eq!(t.numel(), 3);
252        assert_eq!(t.expected_bytes(), 2); // 3 packed nibbles -> 2 bytes
253    }
254
255    #[test]
256    fn checked_geometry_rejects_overflow() {
257        let tensor = TensorData::from_raw(DataType::Float32, vec![usize::MAX, 2], Vec::new());
258        assert_eq!(tensor.checked_numel(), None);
259        assert_eq!(tensor.checked_expected_bytes(), None);
260
261        let byte_overflow =
262            TensorData::from_raw(DataType::Float64, vec![usize::MAX / 4], Vec::new());
263        assert_eq!(byte_overflow.checked_numel(), Some(usize::MAX / 4));
264        assert_eq!(byte_overflow.checked_expected_bytes(), None);
265    }
266
267    #[test]
268    fn weight_ref_accessors() {
269        let w = WeightRef::External {
270            path: PathBuf::from("weights.bin"),
271            offset: 128,
272            length: 4096,
273            dtype: DataType::Float16,
274            dims: vec![64, 32],
275        };
276        assert_eq!(w.dtype(), DataType::Float16);
277        assert_eq!(w.dims(), &[64, 32]);
278    }
279
280    #[test]
281    fn read_little_endian_values() {
282        assert_eq!(read_scalar_le::<i32>(&42_i32.to_le_bytes()), Ok(42));
283
284        let bytes = [1.5_f32.to_le_bytes(), (-2.0_f32).to_le_bytes()].concat();
285        assert_eq!(read_vec_le::<f32>(&bytes), Ok(vec![1.5, -2.0]));
286        assert_eq!(read_vec_le::<i64>(&[]), Ok(Vec::new()));
287    }
288
289    #[test]
290    fn read_little_endian_values_rejects_wrong_lengths() {
291        assert!(matches!(
292            read_scalar_le::<i64>(&[0; 7]),
293            Err(RawBytesError::ScalarLength {
294                expected: 8,
295                actual: 7,
296                ..
297            })
298        ));
299        assert!(matches!(
300            read_vec_le::<i32>(&[0; 5]),
301            Err(RawBytesError::VectorLength {
302                element_size: 4,
303                actual: 5,
304                ..
305            })
306        ));
307    }
308}
onnx_runtime_ir/tensor.rs

onnx_runtime_ir/
tensor.rs
