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///
31/// Integer tensors carry their values in [`elems`](Self::elems) as
32/// [`DimExpr`]s. Floating-point *scalar* constants (needed by `Range`, whose
33/// float form the CPU kernel supports) are carried additively in
34/// [`float_elems`](Self::float_elems); for those, `elems` is empty. `Eq` is not
35/// derived because `f64` is not `Eq` — `ShapeData` is only ever used as a map
36/// *value*, so structural equality is not required.
37#[derive(Clone, Debug, PartialEq)]
38pub struct ShapeData {
39 /// The element type (`Int64`/`Int32`, or `Bool` for masks; `Float32`/
40 /// `Float64` for the float-scalar side-channel).
41 pub dtype: DataType,
42 /// Static dimensions. Empty == rank-0 scalar; `[n]` == rank-1 vector.
43 pub dims: Vec<usize>,
44 /// Row-major integer elements. Length is `1` for a scalar, `dims[0]` for a
45 /// vector. Empty when this is a floating-point scalar (see `float_elems`).
46 pub elems: Vec<DimExpr>,
47 /// Row-major floating-point elements for floating-point constants. Present
48 /// only for float/double scalars captured for `Range`; `None` otherwise.
49 pub float_elems: Option<Vec<f64>>,
50}
51
52impl ShapeData {
53 /// A rank-0 scalar holding `value`.
54 pub fn scalar(dtype: DataType, value: DimExpr) -> Self {
55 Self {
56 dtype,
57 dims: Vec::new(),
58 elems: vec![value],
59 float_elems: None,
60 }
61 }
62
63 /// A rank-1 vector holding `elems`.
64 pub fn vector(dtype: DataType, elems: Vec<DimExpr>) -> Self {
65 Self {
66 dims: vec![elems.len()],
67 dtype,
68 elems,
69 float_elems: None,
70 }
71 }
72
73 /// A rank-0 floating-point scalar holding `value`.
74 pub fn float_scalar(dtype: DataType, value: f64) -> Self {
75 Self {
76 dtype,
77 dims: Vec::new(),
78 elems: Vec::new(),
79 float_elems: Some(vec![value]),
80 }
81 }
82
83 /// This value interpreted as a floating-point scalar, if it is one.
84 pub fn as_float_scalar(&self) -> Option<f64> {
85 let floats = self.float_elems.as_ref()?;
86 (self.is_scalar() && floats.len() == 1).then(|| floats[0])
87 }
88
89 /// Whether this is a rank-0 scalar.
90 pub fn is_scalar(&self) -> bool {
91 self.dims.is_empty()
92 }
93
94 /// Whether the element count is within [`MAX_SHAPE_DATA_ELEMS`].
95 pub fn within_bounds(&self) -> bool {
96 self.elems.len() <= MAX_SHAPE_DATA_ELEMS
97 }
98
99 /// This value's contents interpreted as a shape (each element is a dim).
100 ///
101 /// Used by `Reshape`/`Expand`/`ConstantOfShape` to turn a resolved shape
102 /// vector into the output shape.
103 pub fn as_shape(&self) -> Vec<DimExpr> {
104 self.elems.clone()
105 }
106
107 /// Extract shape-data from a concrete constant tensor, if it is a small
108 /// (rank ≤ 1, within [`MAX_SHAPE_DATA_ELEMS`]) integer or boolean tensor.
109 ///
110 /// Non-integer tensors, higher-rank tensors, and over-large tensors are not
111 /// tracked (returns `None`) — only the shape-computation operands matter.
112 pub fn from_tensor(dtype: DataType, dims: &[usize], data: &[u8]) -> Option<Self> {
113 if dims.len() > 1 {
114 return None;
115 }
116 let numel: usize = if dims.is_empty() {
117 1
118 } else {
119 dims.iter().product()
120 };
121 if numel > MAX_SHAPE_DATA_ELEMS {
122 return None;
123 }
124 // Floating-point *scalars* are captured additively for `Range`; higher-
125 // rank / vector float tensors (weights) are deliberately not tracked.
126 if dtype.is_float() {
127 if !dims.is_empty() {
128 return None;
129 }
130 let value = read_float_scalar(dtype, data)?;
131 return Some(Self::float_scalar(dtype, value));
132 }
133 let ints = read_ints(dtype, numel, data)?;
134 let elems: Vec<DimExpr> = ints.into_iter().map(DimExpr::constant).collect();
135 Some(Self {
136 dtype,
137 dims: dims.to_vec(),
138 elems,
139 float_elems: None,
140 })
141 }
142}
143
144/// Read a single little-endian floating-point scalar from `data`.
145///
146/// Supports `Float32` and `Float64` (the forms the `Range` kernel accepts);
147/// returns `None` for other dtypes or a too-short buffer.
148fn read_float_scalar(dtype: DataType, data: &[u8]) -> Option<f64> {
149 match dtype {
150 DataType::Float32 => {
151 let arr: [u8; 4] = data.get(..4)?.try_into().ok()?;
152 Some(f32::from_le_bytes(arr) as f64)
153 }
154 DataType::Float64 => {
155 let arr: [u8; 8] = data.get(..8)?.try_into().ok()?;
156 Some(f64::from_le_bytes(arr))
157 }
158 _ => None,
159 }
160}
161
162/// Read `numel` little-endian integer/boolean elements from `data`.
163///
164/// Returns `None` for non-integral dtypes (floating point, string) or a byte
165/// length that does not match `numel` elements.
166fn read_ints(dtype: DataType, numel: usize, data: &[u8]) -> Option<Vec<i64>> {
167 macro_rules! read {
168 ($ty:ty, $sz:expr) => {{
169 if data.len() < numel * $sz {
170 return None;
171 }
172 data.chunks_exact($sz)
173 .take(numel)
174 .map(|c| {
175 let arr: [u8; $sz] = c.try_into().ok()?;
176 Some(<$ty>::from_le_bytes(arr) as i64)
177 })
178 .collect::<Option<Vec<i64>>>()
179 }};
180 }
181 match dtype {
182 DataType::Int64 => read!(i64, 8),
183 DataType::Uint64 => read!(u64, 8),
184 DataType::Int32 => read!(i32, 4),
185 DataType::Uint32 => read!(u32, 4),
186 DataType::Int16 => read!(i16, 2),
187 DataType::Uint16 => read!(u16, 2),
188 DataType::Int8 => read!(i8, 1),
189 DataType::Uint8 | DataType::Bool => read!(u8, 1),
190 _ => None,
191 }
192}