nxpu_ir/
types.rs

1//! Type system for the NxPU IR.
2
3use crate::arena::Handle;
4
5/// Width of a scalar type in bytes.
6pub type Bytes = u8;
7
8/// The kind of a scalar type.
9#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
10pub enum ScalarKind {
11    /// Boolean.
12    Bool,
13    /// Signed integer.
14    Sint,
15    /// Unsigned integer.
16    Uint,
17    /// Floating point.
18    Float,
19    /// Brain floating point.
20    BFloat,
21}
22
23/// A scalar type: kind + byte width.
24#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
25pub struct Scalar {
26    /// The kind of scalar (bool, int, float, etc.).
27    pub kind: ScalarKind,
28    /// Width of the scalar in bytes.
29    pub width: Bytes,
30}
31
32impl Scalar {
33    pub const BOOL: Self = Self {
34        kind: ScalarKind::Bool,
35        width: 1,
36    };
37    pub const I32: Self = Self {
38        kind: ScalarKind::Sint,
39        width: 4,
40    };
41    pub const U32: Self = Self {
42        kind: ScalarKind::Uint,
43        width: 4,
44    };
45    pub const F16: Self = Self {
46        kind: ScalarKind::Float,
47        width: 2,
48    };
49    pub const F32: Self = Self {
50        kind: ScalarKind::Float,
51        width: 4,
52    };
53    pub const I8: Self = Self {
54        kind: ScalarKind::Sint,
55        width: 1,
56    };
57    pub const U8: Self = Self {
58        kind: ScalarKind::Uint,
59        width: 1,
60    };
61    pub const BF16: Self = Self {
62        kind: ScalarKind::BFloat,
63        width: 2,
64    };
65}
66
67/// Number of components in a vector.
68#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
69pub enum VectorSize {
70    /// 2 components.
71    Bi = 2,
72    /// 3 components.
73    Tri = 3,
74    /// 4 components.
75    Quad = 4,
76}
77
78/// Size of an array.
79#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
80pub enum ArraySize {
81    /// Fixed-size array.
82    Constant(u32),
83    /// Runtime-sized array.
84    Dynamic,
85}
86
87/// A single tensor dimension: either a fixed size, a named symbolic parameter,
88/// or fully dynamic (unknown at compile time).
89#[derive(Clone, Debug, Hash, Eq, PartialEq)]
90pub enum Dimension {
91    /// Fixed (static) size known at compile time.
92    Fixed(u32),
93    /// Dynamic (unknown) dimension with an optional name (e.g. "batch").
94    Dynamic(Option<String>),
95    /// Symbolic dimension with a named constraint (e.g. "batch", "seq_len").
96    ///
97    /// Unlike `Dynamic`, symbolic dimensions carry a semantic constraint:
98    /// all dimensions sharing the same symbolic name must have the same
99    /// runtime value. This enables shape validation across tensor operands.
100    Symbolic(String),
101}
102
103impl Dimension {
104    /// Returns `true` if this dimension is statically known.
105    pub fn is_fixed(&self) -> bool {
106        matches!(self, Self::Fixed(_))
107    }
108
109    /// Returns `true` if this dimension is dynamic (either unnamed or named).
110    pub fn is_dynamic(&self) -> bool {
111        matches!(self, Self::Dynamic(_))
112    }
113
114    /// Returns `true` if this dimension is symbolic (named constraint).
115    pub fn is_symbolic(&self) -> bool {
116        matches!(self, Self::Symbolic(_))
117    }
118
119    /// Returns the fixed size, if known.
120    pub fn fixed_size(&self) -> Option<u32> {
121        match self {
122            Self::Fixed(n) => Some(*n),
123            Self::Dynamic(_) | Self::Symbolic(_) => None,
124        }
125    }
126
127    /// Returns the symbolic name, if any.
128    pub fn symbolic_name(&self) -> Option<&str> {
129        match self {
130            Self::Symbolic(name) => Some(name),
131            Self::Dynamic(Some(name)) => Some(name),
132            _ => None,
133        }
134    }
135}
136
137/// A multi-dimensional tensor shape supporting mixed static/dynamic dimensions.
138///
139/// For example, `[batch, 224, 224, 3]` where `batch` is dynamic and spatial
140/// dimensions are fixed.
141#[derive(Clone, Debug, Hash, Eq, PartialEq)]
142pub struct TensorShape {
143    /// The dimensions of the tensor.
144    pub dims: Vec<Dimension>,
145}
146
147impl TensorShape {
148    /// Create a shape where all dimensions are fixed.
149    pub fn fixed(sizes: &[u32]) -> Self {
150        Self {
151            dims: sizes.iter().map(|&s| Dimension::Fixed(s)).collect(),
152        }
153    }
154
155    /// Create a shape where all dimensions are dynamic (unnamed).
156    pub fn all_dynamic(rank: usize) -> Self {
157        Self {
158            dims: (0..rank).map(|_| Dimension::Dynamic(None)).collect(),
159        }
160    }
161
162    /// Returns the number of dimensions (rank).
163    pub fn rank(&self) -> usize {
164        self.dims.len()
165    }
166
167    /// Returns `true` if all dimensions are statically known.
168    pub fn is_fully_static(&self) -> bool {
169        self.dims.iter().all(|d| d.is_fixed())
170    }
171
172    /// Returns `true` if all dimensions are dynamic (unnamed).
173    pub fn is_fully_dynamic(&self) -> bool {
174        self.dims.iter().all(|d| d.is_dynamic())
175    }
176
177    /// Returns `true` if any dimension is not statically known
178    /// (either dynamic or symbolic).
179    pub fn has_dynamic_dims(&self) -> bool {
180        self.dims.iter().any(|d| !d.is_fixed())
181    }
182
183    /// Returns `true` if the shape contains a mix of static and dynamic/symbolic dims.
184    pub fn is_mixed(&self) -> bool {
185        !self.is_fully_static() && !self.is_fully_dynamic()
186    }
187}
188
189/// Memory layout for tensor data.
190///
191/// Different NPU hardware expects tensors in specific memory formats.
192/// This annotation allows the compiler to insert layout conversions
193/// only when necessary.
194#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
195pub enum MemoryLayout {
196    /// Row-major (C-style) — last dimension varies fastest.
197    RowMajor,
198    /// Column-major (Fortran-style) — first dimension varies fastest.
199    ColMajor,
200    /// Channels-last: (N, H, W, C) — used by TFLite, Arm Ethos.
201    Nhwc,
202    /// Channels-first: (N, C, H, W) — used by ONNX, Intel NPU.
203    Nchw,
204}
205
206impl MemoryLayout {
207    /// Returns a human-readable name for the layout.
208    pub fn name(self) -> &'static str {
209        match self {
210            Self::RowMajor => "row_major",
211            Self::ColMajor => "col_major",
212            Self::Nhwc => "nhwc",
213            Self::Nchw => "nchw",
214        }
215    }
216}
217
218/// A member of a struct type.
219#[derive(Clone, Debug, Hash, Eq, PartialEq)]
220pub struct StructMember {
221    /// Optional member name.
222    pub name: Option<String>,
223    /// The type of this member.
224    pub ty: Handle<Type>,
225    /// Byte offset within the struct.
226    pub offset: u32,
227}
228
229/// A named type.
230#[derive(Clone, Debug, Hash, Eq, PartialEq)]
231pub struct Type {
232    /// Optional human-readable name.
233    pub name: Option<String>,
234    /// The concrete type shape.
235    pub inner: TypeInner,
236}
237
238/// The concrete shape of a type.
239#[derive(Clone, Debug, Hash, Eq, PartialEq)]
240pub enum TypeInner {
241    /// A single scalar value.
242    Scalar(Scalar),
243    /// A vector of scalars.
244    Vector { size: VectorSize, scalar: Scalar },
245    /// A matrix of column vectors.
246    Matrix {
247        columns: VectorSize,
248        rows: VectorSize,
249        scalar: Scalar,
250    },
251    /// An atomic scalar.
252    Atomic(Scalar),
253    /// A pointer to a value in a given address space.
254    Pointer {
255        base: Handle<Type>,
256        space: crate::AddressSpace,
257    },
258    /// A fixed-size or runtime-sized array.
259    Array {
260        base: Handle<Type>,
261        size: ArraySize,
262        stride: u32,
263    },
264    /// A composite struct type.
265    Struct {
266        members: Vec<StructMember>,
267        span: u32,
268    },
269    /// A multi-dimensional tensor with element type and shape.
270    ///
271    /// Supports mixed static/dynamic dimensions for production ML models
272    /// (e.g. dynamic batch with fixed spatial dimensions).
273    Tensor {
274        /// Element scalar type (e.g. F32, F16, I8).
275        scalar: Scalar,
276        /// Shape with mixed static/dynamic dimensions.
277        shape: TensorShape,
278    },
279}
280
281#[cfg(test)]
282mod tests {
283    use super::*;
284    use crate::arena::UniqueArena;
285
286    #[test]
287    fn scalar_constants() {
288        assert_eq!(Scalar::F32.kind, ScalarKind::Float);
289        assert_eq!(Scalar::F32.width, 4);
290        assert_eq!(Scalar::U32.kind, ScalarKind::Uint);
291        assert_eq!(Scalar::U32.width, 4);
292        assert_eq!(Scalar::BOOL.width, 1);
293        assert_eq!(Scalar::F16.width, 2);
294    }
295
296    #[test]
297    fn type_dedup() {
298        let mut types = UniqueArena::new();
299        let t0 = types.insert(Type {
300            name: None,
301            inner: TypeInner::Scalar(Scalar::F32),
302        });
303        let t1 = types.insert(Type {
304            name: None,
305            inner: TypeInner::Scalar(Scalar::F32),
306        });
307        assert_eq!(t0, t1);
308        assert_eq!(types.len(), 1);
309    }
310
311    #[test]
312    fn different_types_not_deduped() {
313        let mut types = UniqueArena::new();
314        let t0 = types.insert(Type {
315            name: None,
316            inner: TypeInner::Scalar(Scalar::F32),
317        });
318        let t1 = types.insert(Type {
319            name: None,
320            inner: TypeInner::Scalar(Scalar::I32),
321        });
322        assert_ne!(t0, t1);
323        assert_eq!(types.len(), 2);
324    }
325
326    #[test]
327    fn vector_type() {
328        let ty = TypeInner::Vector {
329            size: VectorSize::Tri,
330            scalar: Scalar::F32,
331        };
332        if let TypeInner::Vector { size, scalar } = ty {
333            assert_eq!(size, VectorSize::Tri);
334            assert_eq!(scalar, Scalar::F32);
335        } else {
336            panic!("expected Vector");
337        }
338    }
339
340    #[test]
341    fn vector_size_values() {
342        assert_eq!(VectorSize::Bi as u32, 2);
343        assert_eq!(VectorSize::Tri as u32, 3);
344        assert_eq!(VectorSize::Quad as u32, 4);
345    }
346
347    #[test]
348    fn dimension_fixed() {
349        let d = Dimension::Fixed(224);
350        assert!(d.is_fixed());
351        assert!(!d.is_dynamic());
352        assert_eq!(d.fixed_size(), Some(224));
353    }
354
355    #[test]
356    fn dimension_dynamic() {
357        let d = Dimension::Dynamic(Some("batch".into()));
358        assert!(!d.is_fixed());
359        assert!(d.is_dynamic());
360        assert!(!d.is_symbolic());
361        assert_eq!(d.fixed_size(), None);
362        assert_eq!(d.symbolic_name(), Some("batch"));
363    }
364
365    #[test]
366    fn dimension_symbolic() {
367        let d = Dimension::Symbolic("batch".into());
368        assert!(!d.is_fixed());
369        assert!(!d.is_dynamic());
370        assert!(d.is_symbolic());
371        assert_eq!(d.fixed_size(), None);
372        assert_eq!(d.symbolic_name(), Some("batch"));
373    }
374
375    #[test]
376    fn tensor_shape_has_dynamic_dims() {
377        let static_shape = TensorShape::fixed(&[1, 224, 224, 3]);
378        assert!(!static_shape.has_dynamic_dims());
379
380        let mixed_shape = TensorShape {
381            dims: vec![Dimension::Symbolic("batch".into()), Dimension::Fixed(224)],
382        };
383        assert!(mixed_shape.has_dynamic_dims());
384
385        let dynamic_shape = TensorShape::all_dynamic(3);
386        assert!(dynamic_shape.has_dynamic_dims());
387    }
388
389    #[test]
390    fn tensor_shape_fixed() {
391        let shape = TensorShape::fixed(&[1, 224, 224, 3]);
392        assert_eq!(shape.rank(), 4);
393        assert!(shape.is_fully_static());
394        assert!(!shape.is_fully_dynamic());
395        assert!(!shape.is_mixed());
396    }
397
398    #[test]
399    fn tensor_shape_all_dynamic() {
400        let shape = TensorShape::all_dynamic(3);
401        assert_eq!(shape.rank(), 3);
402        assert!(shape.is_fully_dynamic());
403        assert!(!shape.is_fully_static());
404        assert!(!shape.is_mixed());
405    }
406
407    #[test]
408    fn tensor_shape_mixed() {
409        let shape = TensorShape {
410            dims: vec![
411                Dimension::Dynamic(Some("batch".into())),
412                Dimension::Fixed(224),
413                Dimension::Fixed(224),
414                Dimension::Fixed(3),
415            ],
416        };
417        assert_eq!(shape.rank(), 4);
418        assert!(shape.is_mixed());
419        assert!(!shape.is_fully_static());
420        assert!(!shape.is_fully_dynamic());
421    }
422
423    #[test]
424    fn tensor_type_inner() {
425        let mut types = UniqueArena::new();
426        let t = types.insert(Type {
427            name: Some("image".into()),
428            inner: TypeInner::Tensor {
429                scalar: Scalar::F32,
430                shape: TensorShape {
431                    dims: vec![
432                        Dimension::Dynamic(Some("batch".into())),
433                        Dimension::Fixed(224),
434                        Dimension::Fixed(224),
435                        Dimension::Fixed(3),
436                    ],
437                },
438            },
439        });
440        let ty = &types[t];
441        match &ty.inner {
442            TypeInner::Tensor { scalar, shape } => {
443                assert_eq!(*scalar, Scalar::F32);
444                assert_eq!(shape.rank(), 4);
445                assert!(shape.is_mixed());
446            }
447            _ => panic!("expected Tensor"),
448        }
449    }
450}
nxpu_ir/types.rs

nxpu_ir/
types.rs
