onnx_runtime_shape_inference/
shape_data.rs

1//! Shape-DATA propagation: [`ShapeData`].
2//!
3//! The single most important idea carried over from the reference
4//! implementation ([`justinchuby/onnx-shape-inference`]). Standard shape
5//! inference propagates the *shape* of each tensor; it cannot resolve a
6//! `Reshape` whose target vector is *computed* at runtime from a `Shape` op.
7//! Shape-data propagation additionally tracks the known *element values* of the
8//! small integer tensors that make up those shape-computation subgraphs —
9//! `Shape → Slice → Concat → Gather → Unsqueeze → Reshape` — so the reshape
10//! target resolves symbolically without executing the graph. This is exactly
11//! what lets transformer graphs (BERT and friends) infer statically.
12//!
13//! A [`ShapeData`] models a rank-0 (scalar) or rank-1 (vector) integer tensor
14//! whose elements are [`DimExpr`]s (so they may be concrete *or* symbolic —
15//! e.g. the `batch`/`seq_len` dims read out of a `Shape` op). Only small
16//! integer tensors are tracked; everything else has no shape-data.
17//!
18//! [`justinchuby/onnx-shape-inference`]: https://github.com/justinchuby/onnx-shape-inference
19
20use onnx_runtime_ir::DataType;
21
22use crate::dim_expr::DimExpr;
23/// Upper bound on the element count a [`ShapeData`] will hold. Shape-vectors are
24/// tiny (a handful of dims); this keeps propagation bounded and prevents an
25/// accidental attempt to track a large weight tensor's contents.
26pub const MAX_SHAPE_DATA_ELEMS: usize = 1024;
27
28/// The known element values of a rank-0 or rank-1 integer tensor flowing
29/// through a shape-computation subgraph.
30#[derive(Clone, Debug, PartialEq, Eq)]
31pub struct ShapeData {
32    /// The integer element type (`Int64`/`Int32`, or `Bool` for masks).
33    pub dtype: DataType,
34    /// Static dimensions. Empty == rank-0 scalar; `[n]` == rank-1 vector.
35    pub dims: Vec<usize>,
36    /// Row-major elements. Length is `1` for a scalar, `dims[0]` for a vector.
37    pub elems: Vec<DimExpr>,
38}
39
40impl ShapeData {
41    /// A rank-0 scalar holding `value`.
42    pub fn scalar(dtype: DataType, value: DimExpr) -> Self {
43        Self {
44            dtype,
45            dims: Vec::new(),
46            elems: vec![value],
47        }
48    }
49
50    /// A rank-1 vector holding `elems`.
51    pub fn vector(dtype: DataType, elems: Vec<DimExpr>) -> Self {
52        Self {
53            dims: vec![elems.len()],
54            dtype,
55            elems,
56        }
57    }
58
59    /// Whether this is a rank-0 scalar.
60    pub fn is_scalar(&self) -> bool {
61        self.dims.is_empty()
62    }
63
64    /// Whether the element count is within [`MAX_SHAPE_DATA_ELEMS`].
65    pub fn within_bounds(&self) -> bool {
66        self.elems.len() <= MAX_SHAPE_DATA_ELEMS
67    }
68
69    /// This value's contents interpreted as a shape (each element is a dim).
70    ///
71    /// Used by `Reshape`/`Expand`/`ConstantOfShape` to turn a resolved shape
72    /// vector into the output shape.
73    pub fn as_shape(&self) -> Vec<DimExpr> {
74        self.elems.clone()
75    }
76
77    /// Extract shape-data from a concrete constant tensor, if it is a small
78    /// (rank ≤ 1, within [`MAX_SHAPE_DATA_ELEMS`]) integer or boolean tensor.
79    ///
80    /// Non-integer tensors, higher-rank tensors, and over-large tensors are not
81    /// tracked (returns `None`) — only the shape-computation operands matter.
82    pub fn from_tensor(dtype: DataType, dims: &[usize], data: &[u8]) -> Option<Self> {
83        if dims.len() > 1 {
84            return None;
85        }
86        let numel: usize = if dims.is_empty() {
87            1
88        } else {
89            dims.iter().product()
90        };
91        if numel > MAX_SHAPE_DATA_ELEMS {
92            return None;
93        }
94        let ints = read_ints(dtype, numel, data)?;
95        let elems: Vec<DimExpr> = ints.into_iter().map(DimExpr::constant).collect();
96        Some(Self {
97            dtype,
98            dims: dims.to_vec(),
99            elems,
100        })
101    }
102}
103
104/// Read `numel` little-endian integer/boolean elements from `data`.
105///
106/// Returns `None` for non-integral dtypes (floating point, string) or a byte
107/// length that does not match `numel` elements.
108fn read_ints(dtype: DataType, numel: usize, data: &[u8]) -> Option<Vec<i64>> {
109    macro_rules! read {
110        ($ty:ty, $sz:expr) => {{
111            if data.len() < numel * $sz {
112                return None;
113            }
114            data.chunks_exact($sz)
115                .take(numel)
116                .map(|c| {
117                    let arr: [u8; $sz] = c.try_into().ok()?;
118                    Some(<$ty>::from_le_bytes(arr) as i64)
119                })
120                .collect::<Option<Vec<i64>>>()
121        }};
122    }
123    match dtype {
124        DataType::Int64 => read!(i64, 8),
125        DataType::Uint64 => read!(u64, 8),
126        DataType::Int32 => read!(i32, 4),
127        DataType::Uint32 => read!(u32, 4),
128        DataType::Int16 => read!(i16, 2),
129        DataType::Uint16 => read!(u16, 2),
130        DataType::Int8 => read!(i8, 1),
131        DataType::Uint8 | DataType::Bool => read!(u8, 1),
132        _ => None,
133    }
134}
onnx_runtime_shape_inference/shape_data.rs

onnx_runtime_shape_inference/
shape_data.rs
