onnx_runtime_ir/
shape.rs

1//! Shapes with static and symbolic (dynamic) dimensions (see `docs/ORT2.md`
2//! §3.2 and §11).
3
4/// Unique identifier for a symbolic dimension.
5///
6/// Two dimensions that share the same protobuf dim-param name (e.g.
7/// `"batch_size"`) are interned to the same `SymbolId` by the loader so that
8/// shape reasoning can equate them.
9#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
10pub struct SymbolId(pub u32);
11
12/// A single dimension: either a concrete size or a symbol.
13#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
14pub enum Dim {
15    /// A statically known extent.
16    Static(usize),
17    /// A dynamic extent identified by a [`SymbolId`].
18    Symbolic(SymbolId),
19}
20
21impl Dim {
22    /// The concrete size, if statically known.
23    pub fn as_static(self) -> Option<usize> {
24        match self {
25            Dim::Static(n) => Some(n),
26            Dim::Symbolic(_) => None,
27        }
28    }
29
30    /// Whether this dimension is statically known.
31    pub fn is_static(self) -> bool {
32        matches!(self, Dim::Static(_))
33    }
34}
35
36impl From<usize> for Dim {
37    fn from(n: usize) -> Self {
38        Dim::Static(n)
39    }
40}
41
42impl From<SymbolId> for Dim {
43    fn from(s: SymbolId) -> Self {
44        Dim::Symbolic(s)
45    }
46}
47
48/// A tensor shape: an ordered list of [`Dim`]s.
49///
50/// A rank-0 (scalar) shape is the empty vector. The IR `Shape` therefore always
51/// denotes a *known* rank; it does not carry a distinct "unknown rank"
52/// inhabitant (ONNX `tensor_type.shape` entirely absent). The loader maps an
53/// absent `TensorShapeProto` to an empty `Shape`, so unknown-rank values are
54/// currently indistinguishable from rank-0 scalars — an accepted limitation for
55/// the dense fp / quantized models the runtime targets, where value types carry
56/// shapes. Distinguishing the two would require making `Value::shape` optional
57/// across the frozen IR contract and is deliberately out of scope here. Known,
58/// per-axis unknown *dimensions* (of a known rank) are fully modeled via
59/// [`Dim::Symbolic`].
60pub type Shape = Vec<Dim>;
61
62/// Convenience constructor for a fully-static shape.
63pub fn static_shape(dims: impl IntoIterator<Item = usize>) -> Shape {
64    dims.into_iter().map(Dim::Static).collect()
65}
66
67/// Returns `Some(dims)` if every dimension of `shape` is static.
68pub fn as_static_shape(shape: &[Dim]) -> Option<Vec<usize>> {
69    shape.iter().map(|d| d.as_static()).collect()
70}
71
72/// Whether every dimension of `shape` is static.
73pub fn is_fully_static(shape: &[Dim]) -> bool {
74    shape.iter().all(|d| d.is_static())
75}
76
77/// Constraints on a symbolic dimension, used by shape bucketing, kernel
78/// specialization, and tiling (see `docs/ORT2.md` §11).
79#[derive(Clone, Debug, PartialEq, Eq, Default)]
80pub struct SymbolConstraints {
81    pub id: Option<SymbolId>,
82    /// Human-readable name, e.g. `"batch_size"`, `"seq_len"`.
83    pub name: Option<String>,
84    /// Inclusive minimum value.
85    pub min: Option<usize>,
86    /// Inclusive maximum value.
87    pub max: Option<usize>,
88    /// The value must be a multiple of this (for tiling / vectorization).
89    pub divisible_by: Option<usize>,
90}
91
92impl SymbolConstraints {
93    /// A fresh constraint set for `id` with an optional `name`.
94    pub fn new(id: SymbolId, name: Option<String>) -> Self {
95        Self {
96            id: Some(id),
97            name,
98            ..Default::default()
99        }
100    }
101
102    /// Whether `value` satisfies all present constraints.
103    pub fn accepts(&self, value: usize) -> bool {
104        self.min.is_none_or(|lo| value >= lo)
105            && self.max.is_none_or(|hi| value <= hi)
106            && self
107                .divisible_by
108                .is_none_or(|m| m != 0 && value.is_multiple_of(m))
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115
116    #[test]
117    fn static_shape_helpers() {
118        let s = static_shape([2, 3, 4]);
119        assert_eq!(s.len(), 3);
120        assert!(is_fully_static(&s));
121        assert_eq!(as_static_shape(&s), Some(vec![2, 3, 4]));
122    }
123
124    #[test]
125    fn symbolic_shape_is_not_static() {
126        let s = vec![Dim::Symbolic(SymbolId(0)), Dim::Static(768)];
127        assert!(!is_fully_static(&s));
128        assert_eq!(as_static_shape(&s), None);
129        assert_eq!(s[1].as_static(), Some(768));
130    }
131
132    #[test]
133    fn dim_conversions() {
134        assert_eq!(Dim::from(5usize), Dim::Static(5));
135        assert_eq!(Dim::from(SymbolId(7)), Dim::Symbolic(SymbolId(7)));
136    }
137
138    #[test]
139    fn constraints_accept() {
140        let c = SymbolConstraints {
141            id: Some(SymbolId(0)),
142            name: Some("seq".into()),
143            min: Some(1),
144            max: Some(2048),
145            divisible_by: Some(8),
146        };
147        assert!(c.accepts(8));
148        assert!(c.accepts(2048));
149        assert!(!c.accepts(0)); // below min
150        assert!(!c.accepts(4096)); // above max
151        assert!(!c.accepts(12)); // not divisible by 8
152    }
153}
onnx_runtime_ir/shape.rs

onnx_runtime_ir/
shape.rs
