rustnn/
runtime_checks.rs

1//! Runtime shape checks for tensors bound to a graph with static or dynamic dimensions.
2//!
3//! [`RuntimeShapeState`] validates that bound tensors match the graph's descriptors: same
4//! rank, static dimensions equal, dynamic dimensions within `max_size`, and dimensions that
5//! share a name agree across all inputs and outputs. Used by `MLContext::dispatch` and the
6//! `*_checked` executor functions.
7
8use std::collections::HashMap;
9
10use crate::error::GraphError;
11use crate::graph::{Dimension, OperandDescriptor};
12
13/// Whether a binding is a graph input or output; selects the error wording.
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum TensorKind {
16    /// Graph input.
17    Input,
18    /// Graph output.
19    Output,
20}
21
22impl TensorKind {
23    fn as_str(self) -> &'static str {
24        match self {
25            TensorKind::Input => "input",
26            TensorKind::Output => "output",
27        }
28    }
29}
30
31#[derive(Debug, Clone)]
32struct BoundDynamicDim {
33    value: usize,
34}
35
36/// Tracks the values bound to named dynamic dimensions during one validation pass.
37#[derive(Debug, Default, Clone)]
38pub struct RuntimeShapeState {
39    bound_dims: HashMap<String, BoundDynamicDim>,
40}
41
42impl RuntimeShapeState {
43    /// Empty state; one per dispatch.
44    pub fn new() -> Self {
45        Self::default()
46    }
47
48    /// Check that `actual_shapes` covers exactly the named `descriptors` and that every shape
49    /// passes [`Self::validate_shape`].
50    #[tracing::instrument(
51        name = "validate_shapes",
52        skip_all,
53        err(level = "debug"),
54        level = "debug",
55        fields(kind = ?kind, tensors = actual_shapes.len())
56    )]
57    pub fn validate_named_shapes(
58        &mut self,
59        actual_shapes: &HashMap<String, Vec<usize>>,
60        descriptors: &HashMap<String, OperandDescriptor>,
61        kind: TensorKind,
62    ) -> Result<(), GraphError> {
63        for name in descriptors.keys() {
64            if !actual_shapes.contains_key(name) {
65                return Err(GraphError::RuntimeTensorMissing {
66                    kind: kind.as_str().to_string(),
67                    name: name.clone(),
68                });
69            }
70        }
71
72        for name in actual_shapes.keys() {
73            if !descriptors.contains_key(name) {
74                return Err(GraphError::RuntimeTensorUnexpected {
75                    kind: kind.as_str().to_string(),
76                    name: name.clone(),
77                });
78            }
79        }
80
81        for (name, actual_shape) in actual_shapes {
82            let descriptor = descriptors.get(name).expect("checked above");
83            self.validate_shape(name, actual_shape, descriptor, kind)?;
84        }
85
86        Ok(())
87    }
88
89    /// Check rank, static dimensions, dynamic bounds and named-dimension consistency of one
90    /// tensor.
91    pub fn validate_shape(
92        &mut self,
93        name: &str,
94        actual_shape: &[usize],
95        descriptor: &OperandDescriptor,
96        kind: TensorKind,
97    ) -> Result<(), GraphError> {
98        if actual_shape.len() != descriptor.shape.len() {
99            return Err(GraphError::RuntimeTensorRankMismatch {
100                kind: kind.as_str().to_string(),
101                name: name.to_string(),
102                expected_rank: descriptor.shape.len(),
103                actual_rank: actual_shape.len(),
104            });
105        }
106
107        for (axis, (actual, expected_dim)) in actual_shape
108            .iter()
109            .copied()
110            .zip(&descriptor.shape)
111            .enumerate()
112        {
113            match expected_dim {
114                Dimension::Static(expected) => {
115                    if actual != *expected as usize {
116                        return Err(GraphError::RuntimeStaticDimensionMismatch {
117                            kind: kind.as_str().to_string(),
118                            name: name.to_string(),
119                            axis,
120                            expected: *expected,
121                            actual,
122                        });
123                    }
124                }
125                Dimension::Dynamic(dynamic) => {
126                    if actual > dynamic.max_size as usize {
127                        return Err(GraphError::RuntimeDynamicDimensionExceeded {
128                            kind: kind.as_str().to_string(),
129                            name: name.to_string(),
130                            axis,
131                            dim_name: dynamic.name.clone(),
132                            max_size: dynamic.max_size,
133                            actual,
134                        });
135                    }
136
137                    if let Some(bound) = self.bound_dims.get(&dynamic.name) {
138                        if bound.value != actual {
139                            return Err(GraphError::RuntimeDynamicDimensionNameMismatch {
140                                dim_name: dynamic.name.clone(),
141                                expected: bound.value,
142                                actual,
143                            });
144                        }
145                    } else {
146                        self.bound_dims
147                            .insert(dynamic.name.clone(), BoundDynamicDim { value: actual });
148                    }
149                }
150            }
151        }
152
153        Ok(())
154    }
155}
156
157/// Check that `data_len` elements match the element count of `shape`.
158#[tracing::instrument(
159    name = "validate_data_length",
160    skip_all,
161    err,
162    level = "debug",
163    fields(name = %name, shape = ?shape, elements = data_len)
164)]
165pub fn validate_shape_data_length(
166    name: &str,
167    shape: &[usize],
168    data_len: usize,
169) -> Result<(), GraphError> {
170    let expected = shape
171        .iter()
172        .try_fold(1usize, |acc, &dim| acc.checked_mul(dim))
173        .ok_or_else(|| GraphError::RuntimeTensorShapeOverflow {
174            name: name.to_string(),
175            shape: shape.to_vec(),
176        })?;
177
178    if data_len != expected {
179        return Err(GraphError::RuntimeTensorDataLengthMismatch {
180            name: name.to_string(),
181            expected,
182            actual: data_len,
183        });
184    }
185
186    Ok(())
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192    use crate::graph::{DataType, Dimension, DynamicDimension, OperandDescriptor};
193
194    fn input_desc(shape: Vec<Dimension>) -> OperandDescriptor {
195        OperandDescriptor {
196            data_type: DataType::Float32,
197            shape,
198            pending_permutation: vec![],
199        }
200    }
201
202    #[test]
203    fn validates_static_shape() {
204        let mut state = RuntimeShapeState::new();
205        let desc = input_desc(vec![Dimension::Static(2), Dimension::Static(3)]);
206        state
207            .validate_shape("x", &[2, 3], &desc, TensorKind::Input)
208            .unwrap();
209    }
210
211    #[test]
212    fn rejects_static_mismatch() {
213        let mut state = RuntimeShapeState::new();
214        let desc = input_desc(vec![Dimension::Static(2), Dimension::Static(3)]);
215        let err = state
216            .validate_shape("x", &[2, 4], &desc, TensorKind::Input)
217            .unwrap_err();
218        assert!(matches!(
219            err,
220            GraphError::RuntimeStaticDimensionMismatch { axis: 1, .. }
221        ));
222    }
223
224    #[test]
225    fn rejects_dynamic_greater_than_max_size() {
226        let mut state = RuntimeShapeState::new();
227        let desc = input_desc(vec![Dimension::Dynamic(DynamicDimension {
228            name: "batch".to_string(),
229            max_size: 8,
230        })]);
231        let err = state
232            .validate_shape("x", &[9], &desc, TensorKind::Input)
233            .unwrap_err();
234        assert!(matches!(
235            err,
236            GraphError::RuntimeDynamicDimensionExceeded { .. }
237        ));
238    }
239
240    #[test]
241    fn enforces_same_named_dynamic_dimensions_across_inputs_and_outputs() {
242        let mut state = RuntimeShapeState::new();
243        let input = input_desc(vec![
244            Dimension::Dynamic(DynamicDimension {
245                name: "batch".to_string(),
246                max_size: 16,
247            }),
248            Dimension::Static(64),
249        ]);
250        let output = input_desc(vec![
251            Dimension::Dynamic(DynamicDimension {
252                name: "batch".to_string(),
253                max_size: 16,
254            }),
255            Dimension::Static(10),
256        ]);
257
258        state
259            .validate_shape("x", &[4, 64], &input, TensorKind::Input)
260            .unwrap();
261
262        let err = state
263            .validate_shape("y", &[5, 10], &output, TensorKind::Output)
264            .unwrap_err();
265        assert!(matches!(
266            err,
267            GraphError::RuntimeDynamicDimensionNameMismatch { .. }
268        ));
269    }
270
271    #[test]
272    fn validates_named_shape_sets() {
273        let mut state = RuntimeShapeState::new();
274        let mut descs = HashMap::new();
275        descs.insert(
276            "x".to_string(),
277            input_desc(vec![Dimension::Static(2), Dimension::Static(3)]),
278        );
279        let mut actual = HashMap::new();
280        actual.insert("x".to_string(), vec![2, 3]);
281
282        state
283            .validate_named_shapes(&actual, &descs, TensorKind::Input)
284            .unwrap();
285    }
286
287    #[test]
288    fn rejects_missing_named_tensor() {
289        let mut state = RuntimeShapeState::new();
290        let mut descs = HashMap::new();
291        descs.insert("x".to_string(), input_desc(vec![Dimension::Static(1)]));
292        let actual = HashMap::new();
293
294        let err = state
295            .validate_named_shapes(&actual, &descs, TensorKind::Input)
296            .unwrap_err();
297        assert!(matches!(err, GraphError::RuntimeTensorMissing { .. }));
298    }
299
300    #[test]
301    fn rejects_unexpected_named_tensor() {
302        let mut state = RuntimeShapeState::new();
303        let descs = HashMap::new();
304        let mut actual = HashMap::new();
305        actual.insert("x".to_string(), vec![1]);
306
307        let err = state
308            .validate_named_shapes(&actual, &descs, TensorKind::Input)
309            .unwrap_err();
310        assert!(matches!(err, GraphError::RuntimeTensorUnexpected { .. }));
311    }
312
313    #[test]
314    fn validates_shape_data_length() {
315        validate_shape_data_length("x", &[2, 3], 6).unwrap();
316        validate_shape_data_length("x", &[], 1).unwrap();
317    }
318
319    #[test]
320    fn rejects_shape_data_length_mismatch() {
321        let err = validate_shape_data_length("x", &[2, 3], 7).unwrap_err();
322        assert!(matches!(
323            err,
324            GraphError::RuntimeTensorDataLengthMismatch { .. }
325        ));
326    }
327}
rustnn/runtime_checks.rs

rustnn/
runtime_checks.rs
