onnx_runtime_ir/
layout.rs

1//! Physical strided layout on tensor values (see `docs/ORT2.md` §5).
2//!
3//! Unlike upstream ONNX / `onnx-ir`, every [`crate::Value`] carries a
4//! [`TensorLayout`]. This lets optimization passes track non-contiguous
5//! (transposed / broadcast) layouts and eliminate copies at EP boundaries.
6
7use crate::dtype::DataType;
8use crate::error::IrError;
9
10/// Compute row-major (C-order) contiguous strides, in **elements**, for a shape.
11pub fn compute_contiguous_strides(shape: &[usize]) -> Vec<i64> {
12    let n = shape.len();
13    let mut strides = vec![1i64; n];
14    for i in (0..n.saturating_sub(1)).rev() {
15        strides[i] = strides[i + 1] * shape[i + 1] as i64;
16    }
17    strides
18}
19
20/// Whether `strides` describe a row-major contiguous layout for `shape`.
21pub fn is_contiguous(shape: &[usize], strides: &[i64]) -> bool {
22    strides == compute_contiguous_strides(shape).as_slice()
23}
24
25/// Compute the output shape of a numpy-style broadcast of `a` and `b`.
26pub fn broadcast_shapes(a: &[usize], b: &[usize]) -> Result<Vec<usize>, IrError> {
27    let max_ndim = a.len().max(b.len());
28    let mut result = Vec::with_capacity(max_ndim);
29    for i in 0..max_ndim {
30        let da = if i < a.len() { a[a.len() - 1 - i] } else { 1 };
31        let db = if i < b.len() { b[b.len() - 1 - i] } else { 1 };
32        if da == db || db == 1 {
33            result.push(da);
34        } else if da == 1 {
35            result.push(db);
36        } else {
37            return Err(IrError::BroadcastIncompatible {
38                a: a.to_vec(),
39                b: b.to_vec(),
40            });
41        }
42    }
43    result.reverse();
44    Ok(result)
45}
46
47/// Memory-format hint used to pick vectorized kernels.
48#[derive(Clone, Debug, PartialEq, Eq, Default)]
49pub enum MemoryFormat {
50    /// Standard row-major.
51    #[default]
52    Contiguous,
53    /// NHWC channels-last.
54    ChannelsLast,
55    /// Blocked/tiled format with the given block width (e.g. 16 for VNNI/AMX).
56    Blocked(usize),
57    /// An arbitrary strided layout that matches none of the named formats.
58    Custom,
59}
60
61/// First-class strided layout for a value.
62///
63/// `strides == None` means "contiguous row-major for the value's shape"; this
64/// is the common case and avoids materializing strides for every value.
65#[derive(Clone, Debug, PartialEq)]
66pub struct TensorLayout {
67    /// Physical strides in **elements**. `None` == contiguous row-major.
68    pub strides: Option<Vec<i64>>,
69    /// Memory-format hint.
70    pub format: MemoryFormat,
71    /// Required alignment in bytes for the backing allocation.
72    pub alignment: usize,
73}
74
75/// Default alignment (bytes) — 64 covers AVX-512 / cache-line requirements.
76pub const DEFAULT_ALIGNMENT: usize = 64;
77
78impl Default for TensorLayout {
79    fn default() -> Self {
80        Self {
81            strides: None,
82            format: MemoryFormat::Contiguous,
83            alignment: DEFAULT_ALIGNMENT,
84        }
85    }
86}
87
88impl TensorLayout {
89    /// A contiguous row-major layout (strides implied by shape).
90    pub fn contiguous() -> Self {
91        Self::default()
92    }
93
94    /// A layout with explicit strides (marked [`MemoryFormat::Custom`]).
95    pub fn strided(strides: Vec<i64>) -> Self {
96        Self {
97            strides: Some(strides),
98            format: MemoryFormat::Custom,
99            alignment: DEFAULT_ALIGNMENT,
100        }
101    }
102
103    /// Whether this layout is contiguous row-major for `shape`.
104    pub fn is_contiguous(&self, shape: &[usize]) -> bool {
105        match &self.strides {
106            None => true,
107            Some(s) => is_contiguous(shape, s),
108        }
109    }
110
111    /// The strides for `shape` under this layout, materializing the implied
112    /// contiguous strides when `strides == None`.
113    pub fn resolved_strides(&self, shape: &[usize]) -> Vec<i64> {
114        self.strides
115            .clone()
116            .unwrap_or_else(|| compute_contiguous_strides(shape))
117    }
118
119    /// Reorder axes without copying data (a lazy transpose).
120    pub fn transpose(&self, shape: &[usize], perm: &[usize]) -> Self {
121        let base = self.resolved_strides(shape);
122        let strides = perm.iter().map(|&p| base[p]).collect();
123        Self {
124            strides: Some(strides),
125            format: MemoryFormat::Custom,
126            alignment: self.alignment,
127        }
128    }
129
130    /// Total backing storage size in bytes: the largest byte offset reachable
131    /// via the strides, plus one element. Handles negative strides.
132    pub fn storage_size(&self, shape: &[usize], dtype: DataType) -> usize {
133        let elem = dtype.byte_size().max(1);
134        match &self.strides {
135            None => shape.iter().product::<usize>() * elem,
136            Some(strides) => {
137                let max_offset: i64 = shape
138                    .iter()
139                    .zip(strides.iter())
140                    .map(|(&dim, &stride)| dim.saturating_sub(1) as i64 * stride.abs())
141                    .sum();
142                (max_offset as usize + 1) * elem
143            }
144        }
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    #[test]
153    fn contiguous_strides_row_major() {
154        assert_eq!(compute_contiguous_strides(&[2, 3, 4]), vec![12, 4, 1]);
155        assert_eq!(compute_contiguous_strides(&[5]), vec![1]);
156        assert_eq!(compute_contiguous_strides(&[]), Vec::<i64>::new());
157    }
158
159    #[test]
160    fn is_contiguous_check() {
161        assert!(is_contiguous(&[2, 3], &[3, 1]));
162        assert!(!is_contiguous(&[2, 3], &[1, 2]));
163    }
164
165    #[test]
166    fn broadcast_basic() {
167        assert_eq!(broadcast_shapes(&[3, 1], &[1, 4]).unwrap(), vec![3, 4]);
168        assert_eq!(broadcast_shapes(&[5], &[3, 5]).unwrap(), vec![3, 5]);
169        assert_eq!(broadcast_shapes(&[], &[2, 2]).unwrap(), vec![2, 2]);
170    }
171
172    #[test]
173    fn broadcast_incompatible() {
174        assert!(matches!(
175            broadcast_shapes(&[3], &[4]),
176            Err(IrError::BroadcastIncompatible { .. })
177        ));
178    }
179
180    #[test]
181    fn transpose_swaps_strides() {
182        let l = TensorLayout::contiguous();
183        let t = l.transpose(&[2, 3], &[1, 0]);
184        // contiguous [2,3] -> strides [3,1]; transposed -> [1,3]
185        assert_eq!(t.strides, Some(vec![1, 3]));
186        assert!(!t.is_contiguous(&[3, 2]));
187    }
188
189    #[test]
190    fn storage_size_contiguous_and_strided() {
191        let l = TensorLayout::contiguous();
192        assert_eq!(l.storage_size(&[2, 3], DataType::Float32), 24);
193        // transposed view still covers the same 6 elements
194        let t = l.transpose(&[2, 3], &[1, 0]);
195        assert_eq!(t.storage_size(&[3, 2], DataType::Float32), 24);
196    }
197}
onnx_runtime_ir/layout.rs

onnx_runtime_ir/
layout.rs
