native_onnx/
tensor.rs

1use crate::Result;
2use anyhow::{bail, ensure};
3pub use half::{bf16, f16};
4
5/// Dense, contiguous, row-major tensor metadata.
6#[derive(Clone, Debug, PartialEq, Eq)]
7pub struct TensorSpec {
8    /// Input or output name from the ONNX graph.
9    pub name: String,
10    /// Required element type.
11    pub dtype: DType,
12    /// `None` means unknown rank; `Some(vec![])` is a scalar.
13    /// Within a known rank, `None` is a dynamic dimension and zero is an empty dimension.
14    pub shape: Option<Vec<Option<usize>>>,
15}
16
17impl TensorSpec {
18    /// Concrete dimensions, or `None` if any dimension is dynamic.
19    pub fn fixed_shape(&self) -> Option<Vec<usize>> {
20        self.shape.as_ref()?.iter().copied().collect()
21    }
22    /// Check a tensor's name, element type, dimensions, and storage length.
23    pub fn validate(&self, tensor: &TensorView<'_>) -> Result<()> {
24        ensure!(
25            tensor.name == self.name,
26            "Tensor name mismatch: {}",
27            tensor.name
28        );
29        ensure!(
30            tensor.data.dtype() == self.dtype,
31            "Tensor {} expects {:?}, got {:?}",
32            tensor.name,
33            self.dtype,
34            tensor.data.dtype()
35        );
36        ensure!(
37            self.shape.as_ref().is_none_or(|shape| {
38                tensor.shape.len() == shape.len()
39                    && shape
40                        .iter()
41                        .zip(tensor.shape)
42                        .all(|(expected, actual)| expected.is_none_or(|v| v == *actual))
43            }),
44            "Tensor {} expects shape {:?}, got {:?}",
45            tensor.name,
46            self.shape,
47            tensor.shape
48        );
49        ensure!(
50            element_count(tensor.shape)? == tensor.data.len(),
51            "Tensor length mismatch: {}",
52            tensor.name
53        );
54        Ok(())
55    }
56}
57
58/// Checked product of a tensor's dimensions.
59/// Scalars (`[]`) have one element; any zero dimension makes the tensor empty.
60/// Returns an error when a nonempty tensor's element count overflows `usize`.
61pub fn element_count(shape: &[usize]) -> Result<usize> {
62    if shape.contains(&0) {
63        return Ok(0);
64    }
65    shape.iter().try_fold(1usize, |n, &d| {
66        n.checked_mul(d)
67            .ok_or_else(|| anyhow::anyhow!("Tensor element count overflow"))
68    })
69}
70
71macro_rules! tensor_types {
72    ($($variant:ident: $ty:ty),+ $(,)?) => {
73        /// Supported dense tensor element types.
74        #[derive(Clone, Copy, Debug, PartialEq, Eq)]
75        pub enum DType {
76            $(#[doc = concat!("Elements of type `", stringify!($ty), "`.")]
77            $variant),+
78        }
79
80        /// Borrowed, contiguous tensor elements.
81        #[derive(Clone, Copy, Debug)]
82        pub enum TensorData<'a> {
83            $(#[doc = concat!("Borrowed `", stringify!($ty), "` elements.")]
84            $variant(&'a [$ty])),+
85        }
86
87        /// Exclusively borrowed, contiguous tensor elements.
88        #[derive(Debug)]
89        pub enum TensorDataMut<'a> {
90            $(#[doc = concat!("Mutable `", stringify!($ty), "` elements.")]
91            $variant(&'a mut [$ty])),+
92        }
93
94        /// Owned, contiguous tensor elements stored on the CPU.
95        #[derive(Clone, Debug, PartialEq)]
96        pub enum TensorBuffer {
97            $(#[doc = concat!("Owned `", stringify!($ty), "` elements.")]
98            $variant(Vec<$ty>)),+
99        }
100
101        impl TensorData<'_> {
102            /// Element type of this slice.
103            pub fn dtype(&self) -> DType {
104                match self { $(Self::$variant(_) => DType::$variant),+ }
105            }
106            /// Number of elements, not bytes.
107            pub fn len(&self) -> usize {
108                match self { $(Self::$variant(x) => x.len()),+ }
109            }
110            /// Whether the slice contains no elements.
111            pub fn is_empty(&self) -> bool {
112                self.len() == 0
113            }
114            /// Copy the elements into an owned CPU buffer.
115            pub fn to_owned(self) -> TensorBuffer {
116                match self { $(Self::$variant(x) => TensorBuffer::$variant(x.to_vec())),+ }
117            }
118        }
119
120        impl TensorDataMut<'_> {
121            /// Element type of this slice.
122            pub fn dtype(&self) -> DType {
123                match self { $(Self::$variant(_) => DType::$variant),+ }
124            }
125            /// Number of elements, not bytes.
126            pub fn len(&self) -> usize {
127                match self { $(Self::$variant(x) => x.len()),+ }
128            }
129            /// Whether the slice contains no elements.
130            pub fn is_empty(&self) -> bool {
131                self.len() == 0
132            }
133            /// Copy elements after checking that types and lengths match.
134            pub fn copy_from(&mut self, source: TensorData<'_>) -> Result<()> {
135                ensure!(self.len() == source.len(), "Output buffer length mismatch");
136                match (self, source) {
137                    $((Self::$variant(dst), TensorData::$variant(src)) => dst.copy_from_slice(src)),+,
138                    _ => bail!("Output buffer dtype mismatch"),
139                }
140                Ok(())
141            }
142        }
143
144        impl TensorBuffer {
145            /// Element type of this buffer.
146            pub fn dtype(&self) -> DType {
147                match self { $(Self::$variant(_) => DType::$variant),+ }
148            }
149            /// Number of elements, not bytes.
150            pub fn len(&self) -> usize {
151                match self { $(Self::$variant(x) => x.len()),+ }
152            }
153            /// Whether the buffer contains no elements.
154            pub fn is_empty(&self) -> bool {
155                self.len() == 0
156            }
157            /// Borrow the elements without copying.
158            pub fn view(&self) -> TensorData<'_> {
159                match self { $(Self::$variant(x) => TensorData::$variant(x)),+ }
160            }
161            /// Borrow the elements for mutation without copying.
162            pub fn view_mut(&mut self) -> TensorDataMut<'_> {
163                match self { $(Self::$variant(x) => TensorDataMut::$variant(x)),+ }
164            }
165        }
166    };
167}
168tensor_types!(F32: f32, F64: f64, F16: f16, BF16: bf16, I64: i64, I32: i32, I16: i16, I8: i8, U64: u64, U32: u32, U16: u16, U8: u8, Bool: bool);
169
170/// Named tensor borrowing contiguous, row-major CPU storage.
171/// Creating a view does not validate it; inference validates it against the model.
172#[derive(Clone, Copy, Debug)]
173pub struct TensorView<'a> {
174    /// Input or output name from the ONNX graph.
175    pub name: &'a str,
176    /// Concrete dimensions; use `[]` for a scalar.
177    pub shape: &'a [usize],
178    /// Borrowed elements in row-major order.
179    pub data: TensorData<'a>,
180}
181impl<'a> TensorView<'a> {
182    /// Construct a view over FP32 elements without copying or validating them.
183    pub fn f32(name: &'a str, shape: &'a [usize], data: &'a [f32]) -> Self {
184        Self {
185            name,
186            shape,
187            data: TensorData::F32(data),
188        }
189    }
190    /// Copy the name, shape, and elements into an owned tensor.
191    pub fn to_owned(self) -> Tensor {
192        Tensor {
193            name: self.name.into(),
194            shape: self.shape.to_vec(),
195            data: self.data.to_owned(),
196        }
197    }
198    /// Borrow the FP32 elements, or return an error for another element type.
199    pub fn as_f32(&self) -> Result<&'a [f32]> {
200        match self.data {
201            TensorData::F32(x) => Ok(x),
202            _ => bail!("Not an FP32 tensor"),
203        }
204    }
205}
206
207/// Named tensor exclusively borrowing contiguous, row-major CPU storage.
208#[derive(Debug)]
209pub struct TensorViewMut<'a> {
210    /// Input or output name from the ONNX graph.
211    pub name: &'a str,
212    /// Concrete dimensions; use `[]` for a scalar.
213    pub shape: &'a [usize],
214    /// Exclusively borrowed elements in row-major order.
215    pub data: TensorDataMut<'a>,
216}
217impl<'a> TensorViewMut<'a> {
218    /// Construct a mutable FP32 view without copying or validating the elements.
219    pub fn f32(name: &'a str, shape: &'a [usize], data: &'a mut [f32]) -> Self {
220        Self {
221            name,
222            shape,
223            data: TensorDataMut::F32(data),
224        }
225    }
226    /// Mutably borrow FP32 elements, or return an error for another element type.
227    pub fn as_f32_mut(&mut self) -> Result<&mut [f32]> {
228        match &mut self.data {
229            TensorDataMut::F32(x) => Ok(x),
230            _ => bail!("Not an FP32 tensor"),
231        }
232    }
233}
234
235/// Owned, named tensor with contiguous, row-major CPU storage.
236#[derive(Clone, Debug, PartialEq)]
237pub struct Tensor {
238    /// Input or output name from the ONNX graph.
239    pub name: String,
240    /// Concrete dimensions; an empty shape denotes a scalar.
241    pub shape: Vec<usize>,
242    /// Owned elements in row-major order.
243    pub data: TensorBuffer,
244}
245impl Tensor {
246    /// Borrow the tensor without copying its metadata or elements.
247    pub fn view(&self) -> TensorView<'_> {
248        TensorView {
249            name: &self.name,
250            shape: &self.shape,
251            data: self.data.view(),
252        }
253    }
254    /// Borrow the tensor for element mutation without copying.
255    pub fn view_mut(&mut self) -> TensorViewMut<'_> {
256        TensorViewMut {
257            name: &self.name,
258            shape: &self.shape,
259            data: self.data.view_mut(),
260        }
261    }
262}
263
264pub(crate) fn copy_outputs(source: &[Tensor], targets: &mut [TensorViewMut<'_>]) -> Result<()> {
265    anyhow::ensure!(source.len() == targets.len(), "Output count mismatch");
266    for (i, target) in targets.iter().enumerate() {
267        anyhow::ensure!(
268            !targets[..i].iter().any(|x| x.name == target.name),
269            "Duplicate output {}",
270            target.name
271        );
272        let value = source
273            .iter()
274            .find(|x| x.name == target.name)
275            .ok_or_else(|| anyhow::anyhow!("Unknown output {}", target.name))?;
276        anyhow::ensure!(
277            target.shape == value.shape
278                && target.data.dtype() == value.data.dtype()
279                && target.data.len() == value.data.len(),
280            "Output shape/type/length mismatch: {}",
281            target.name
282        );
283    }
284    for target in targets {
285        let source = source
286            .iter()
287            .find(|x| x.name == target.name)
288            .expect("validated output");
289        target.data.copy_from(source.data.view())?;
290    }
291    Ok(())
292}
native_onnx/tensor.rs

native_onnx/
tensor.rs
