onnx_runtime_ir/
dtype.rs

1//! Element data types, matching `onnx.TensorProto.DataType`.
2
3/// Supported tensor element types.
4///
5/// The discriminant values match ONNX `TensorProto.DataType` so that the
6/// loader can cast the protobuf integer directly (see [`DataType::from_onnx`]).
7#[repr(u8)]
8#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
9pub enum DataType {
10    Undefined = 0,
11    Float32 = 1,
12    Uint8 = 2,
13    Int8 = 3,
14    Uint16 = 4,
15    Int16 = 5,
16    Int32 = 6,
17    Int64 = 7,
18    String = 8,
19    Bool = 9,
20    Float16 = 10,
21    Float64 = 11,
22    Uint32 = 12,
23    Uint64 = 13,
24    Complex64 = 14,
25    Complex128 = 15,
26    BFloat16 = 16,
27    Float8E4M3FN = 17,
28    Float8E4M3FNUZ = 18,
29    Float8E5M2 = 19,
30    Float8E5M2FNUZ = 20,
31    Uint4 = 21,
32    Int4 = 22,
33    Float4E2M1 = 23,
34    Float8E8M0 = 24,
35    Uint2 = 25,
36    Int2 = 26,
37}
38
39impl DataType {
40    /// Byte size per element. Sub-byte (packed) and variable-width types
41    /// return `0`; use [`DataType::bit_size`] for those.
42    pub fn byte_size(self) -> usize {
43        match self {
44            Self::Float32 | Self::Int32 | Self::Uint32 => 4,
45            Self::Float64 | Self::Int64 | Self::Uint64 | Self::Complex64 => 8,
46            Self::Complex128 => 16,
47            Self::Float16 | Self::BFloat16 | Self::Int16 | Self::Uint16 => 2,
48            Self::Int8
49            | Self::Uint8
50            | Self::Bool
51            | Self::Float8E4M3FN
52            | Self::Float8E4M3FNUZ
53            | Self::Float8E5M2
54            | Self::Float8E5M2FNUZ
55            | Self::Float8E8M0 => 1,
56            Self::Int4 | Self::Uint4 | Self::Float4E2M1 | Self::Int2 | Self::Uint2 => 0,
57            Self::String | Self::Undefined => 0, // variable-width / no concrete storage
58        }
59    }
60
61    /// Bit size per element. Sub-byte types return their true width (e.g. 4).
62    pub fn bit_size(self) -> usize {
63        match self {
64            Self::Int4 | Self::Uint4 | Self::Float4E2M1 => 4,
65            Self::Int2 | Self::Uint2 => 2,
66            other => other.byte_size() * 8,
67        }
68    }
69
70    /// Whether this is a floating-point type (any precision).
71    pub fn is_float(self) -> bool {
72        matches!(
73            self,
74            Self::Float32
75                | Self::Float64
76                | Self::Float16
77                | Self::BFloat16
78                | Self::Float8E4M3FN
79                | Self::Float8E4M3FNUZ
80                | Self::Float8E5M2
81                | Self::Float8E5M2FNUZ
82                | Self::Float8E8M0
83                | Self::Float4E2M1
84        )
85    }
86
87    /// Whether this is a signed or unsigned integer type (excludes `Bool`).
88    pub fn is_int(self) -> bool {
89        matches!(
90            self,
91            Self::Uint8
92                | Self::Int8
93                | Self::Uint16
94                | Self::Int16
95                | Self::Int32
96                | Self::Int64
97                | Self::Uint32
98                | Self::Uint64
99                | Self::Int4
100                | Self::Uint4
101                | Self::Int2
102                | Self::Uint2
103        )
104    }
105
106    /// Whether elements contain real and imaginary floating-point components.
107    pub fn is_complex(self) -> bool {
108        matches!(self, Self::Complex64 | Self::Complex128)
109    }
110
111    /// Whether elements are packed multiple-per-byte.
112    pub fn is_sub_byte(self) -> bool {
113        matches!(
114            self,
115            Self::Int4 | Self::Uint4 | Self::Float4E2M1 | Self::Int2 | Self::Uint2
116        )
117    }
118
119    /// Number of bytes needed to store `count` elements, accounting for
120    /// sub-byte packing.
121    ///
122    /// Panics on `usize` overflow of the `count * byte_size` product. Callers
123    /// that size heap allocations from an untrusted shape MUST use
124    /// [`DataType::checked_storage_bytes`] instead so a crafted static shape
125    /// cannot wrap the multiply and under-allocate.
126    pub fn storage_bytes(self, count: usize) -> usize {
127        self.checked_storage_bytes(count)
128            .expect("storage_bytes overflow")
129    }
130
131    /// Number of bytes needed to store `count` elements, returning `None` when
132    /// the `count * byte_size` product overflows `usize`.
133    ///
134    /// This is the overflow-safe counterpart of [`DataType::storage_bytes`]:
135    /// even though an element *count* may fit in `usize`, the element-count →
136    /// bytes multiply can still wrap for a fixed-width dtype (e.g. `2^61`
137    /// elements of an 8-byte type). Allocation sites route through this so a
138    /// wrapped product becomes a clean error instead of a 1-byte allocation
139    /// followed by an out-of-bounds access.
140    pub fn checked_storage_bytes(self, count: usize) -> Option<usize> {
141        if self == Self::Undefined {
142            None
143        } else if self.is_sub_byte() {
144            let elements_per_byte = 8 / self.bit_size();
145            Some(count / elements_per_byte + usize::from(!count.is_multiple_of(elements_per_byte)))
146        } else {
147            count.checked_mul(self.byte_size())
148        }
149    }
150
151    /// Convert from the raw ONNX `TensorProto.DataType` integer.
152    ///
153    /// Returns `None` for `UNDEFINED` (0) and any out-of-range or future value
154    /// the runtime does not model. The discriminants below mirror the vendored
155    /// `onnx.proto3` `TensorProto.DataType` enum verbatim.
156    pub fn from_onnx(raw: i32) -> Option<Self> {
157        Some(match raw {
158            1 => Self::Float32,
159            2 => Self::Uint8,
160            3 => Self::Int8,
161            4 => Self::Uint16,
162            5 => Self::Int16,
163            6 => Self::Int32,
164            7 => Self::Int64,
165            8 => Self::String,
166            9 => Self::Bool,
167            10 => Self::Float16,
168            11 => Self::Float64,
169            12 => Self::Uint32,
170            13 => Self::Uint64,
171            14 => Self::Complex64,
172            15 => Self::Complex128,
173            16 => Self::BFloat16,
174            17 => Self::Float8E4M3FN,
175            18 => Self::Float8E4M3FNUZ,
176            19 => Self::Float8E5M2,
177            20 => Self::Float8E5M2FNUZ,
178            21 => Self::Uint4,
179            22 => Self::Int4,
180            23 => Self::Float4E2M1,
181            24 => Self::Float8E8M0,
182            25 => Self::Uint2,
183            26 => Self::Int2,
184            _ => return None,
185        })
186    }
187
188    /// The raw ONNX `TensorProto.DataType` integer for this type.
189    pub fn to_onnx(self) -> i32 {
190        self as i32
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197
198    #[test]
199    fn byte_and_bit_sizes() {
200        assert_eq!(DataType::Undefined.checked_storage_bytes(1), None);
201        assert_eq!(DataType::Float32.byte_size(), 4);
202        assert_eq!(DataType::Float32.bit_size(), 32);
203        assert_eq!(DataType::Int64.byte_size(), 8);
204        assert_eq!(DataType::Float16.byte_size(), 2);
205        assert_eq!(DataType::Bool.byte_size(), 1);
206        assert_eq!(DataType::Int4.byte_size(), 0);
207        assert_eq!(DataType::Int4.bit_size(), 4);
208        assert_eq!(DataType::Int2.bit_size(), 2);
209    }
210
211    #[test]
212    fn sub_byte_storage_rounds_up() {
213        assert_eq!(DataType::Int4.storage_bytes(4), 2);
214        assert_eq!(DataType::Int4.storage_bytes(5), 3);
215        assert_eq!(DataType::Int2.storage_bytes(5), 2);
216        assert_eq!(DataType::Float32.storage_bytes(4), 16);
217    }
218
219    #[test]
220    fn checked_storage_bytes_normal_counts() {
221        assert_eq!(DataType::Float32.checked_storage_bytes(4), Some(16));
222        assert_eq!(DataType::Float64.checked_storage_bytes(3), Some(24));
223        assert_eq!(DataType::Int4.checked_storage_bytes(5), Some(3));
224        assert_eq!(DataType::Uint2.checked_storage_bytes(5), Some(2));
225        assert_eq!(DataType::Bool.checked_storage_bytes(0), Some(0));
226    }
227
228    #[test]
229    fn checked_storage_bytes_detects_byte_overflow() {
230        // The element *count* fits in usize but count * byte_size wraps: this
231        // is the exploited path (a `[2^61]`-of-8-byte shape passes the numel
232        // check yet `2^61 * 8` wraps to 0 → 1-byte allocation).
233        assert_eq!(
234            DataType::Float64.checked_storage_bytes(usize::MAX / 4),
235            None
236        );
237        assert_eq!(DataType::Int64.checked_storage_bytes(usize::MAX), None);
238        // Sub-byte packing can never overflow (div_ceil shrinks the count).
239        assert_eq!(
240            DataType::Int4.checked_storage_bytes(usize::MAX),
241            Some(usize::MAX / 2 + 1)
242        );
243        assert_eq!(
244            DataType::Int2.checked_storage_bytes(usize::MAX),
245            Some(usize::MAX / 4 + 1)
246        );
247    }
248
249    #[test]
250    fn classification() {
251        assert!(DataType::Float16.is_float());
252        assert!(!DataType::Int32.is_float());
253        assert!(DataType::Int32.is_int());
254        assert!(!DataType::Bool.is_int());
255        assert!(DataType::Complex64.is_complex());
256        assert!(DataType::Complex128.is_complex());
257        assert!(!DataType::Float32.is_complex());
258        assert!(DataType::Uint4.is_sub_byte());
259        // float8 / float4 variants are floats but not ints.
260        for dt in [
261            DataType::Float8E4M3FN,
262            DataType::Float8E4M3FNUZ,
263            DataType::Float8E5M2,
264            DataType::Float8E5M2FNUZ,
265            DataType::Float8E8M0,
266            DataType::Float4E2M1,
267        ] {
268            assert!(dt.is_float(), "{dt:?} should be float");
269            assert!(!dt.is_int(), "{dt:?} should not be int");
270        }
271        // 4-bit types (incl. Float4E2M1) are sub-byte; float8s are full bytes.
272        assert!(DataType::Int4.is_sub_byte());
273        assert!(DataType::Float4E2M1.is_sub_byte());
274        assert!(DataType::Int2.is_sub_byte());
275        assert!(DataType::Uint2.is_sub_byte());
276        assert!(!DataType::Float8E5M2.is_sub_byte());
277        assert!(!DataType::Float8E4M3FNUZ.is_sub_byte());
278    }
279
280    #[test]
281    fn float4_sub_byte_storage_rounds_up() {
282        assert_eq!(DataType::Float4E2M1.byte_size(), 0);
283        assert_eq!(DataType::Float4E2M1.bit_size(), 4);
284        assert_eq!(DataType::Float4E2M1.storage_bytes(4), 2);
285        assert_eq!(DataType::Float4E2M1.storage_bytes(5), 3);
286        assert_eq!(
287            DataType::Float4E2M1.checked_storage_bytes(usize::MAX),
288            Some(usize::MAX / 2 + 1)
289        );
290    }
291
292    #[test]
293    fn float8_variants_are_one_byte() {
294        for dt in [
295            DataType::Float8E4M3FN,
296            DataType::Float8E4M3FNUZ,
297            DataType::Float8E5M2,
298            DataType::Float8E5M2FNUZ,
299            DataType::Float8E8M0,
300        ] {
301            assert_eq!(dt.byte_size(), 1, "{dt:?} should be 1 byte");
302            assert_eq!(dt.bit_size(), 8, "{dt:?} should be 8 bits");
303            assert!(!dt.is_sub_byte(), "{dt:?} is not packed");
304        }
305    }
306
307    /// Pins every `DataType` variant to its authoritative ONNX
308    /// `TensorProto.DataType` integer (vendored `onnx.proto3`). A regression
309    /// here silently corrupts weights, so the table is exhaustive.
310    #[test]
311    fn onnx_discriminants_match_spec() {
312        assert_eq!(DataType::Undefined.to_onnx(), 0);
313        let table: &[(DataType, i32)] = &[
314            (DataType::Float32, 1),
315            (DataType::Uint8, 2),
316            (DataType::Int8, 3),
317            (DataType::Uint16, 4),
318            (DataType::Int16, 5),
319            (DataType::Int32, 6),
320            (DataType::Int64, 7),
321            (DataType::String, 8),
322            (DataType::Bool, 9),
323            (DataType::Float16, 10),
324            (DataType::Float64, 11),
325            (DataType::Uint32, 12),
326            (DataType::Uint64, 13),
327            (DataType::Complex64, 14),
328            (DataType::Complex128, 15),
329            (DataType::BFloat16, 16),
330            (DataType::Float8E4M3FN, 17),
331            (DataType::Float8E4M3FNUZ, 18),
332            (DataType::Float8E5M2, 19),
333            (DataType::Float8E5M2FNUZ, 20),
334            (DataType::Uint4, 21),
335            (DataType::Int4, 22),
336            (DataType::Float4E2M1, 23),
337            (DataType::Float8E8M0, 24),
338            (DataType::Uint2, 25),
339            (DataType::Int2, 26),
340        ];
341        for &(dt, raw) in table {
342            assert_eq!(dt.to_onnx(), raw, "{dt:?} to_onnx mismatch");
343            assert_eq!(
344                DataType::from_onnx(raw),
345                Some(dt),
346                "from_onnx({raw}) mismatch"
347            );
348        }
349    }
350
351    #[test]
352    fn onnx_roundtrip() {
353        for dt in [
354            DataType::Float32,
355            DataType::Int64,
356            DataType::BFloat16,
357            DataType::Uint4,
358            DataType::Float4E2M1,
359            DataType::Float8E5M2FNUZ,
360            DataType::Float8E8M0,
361            DataType::Uint2,
362            DataType::Int2,
363        ] {
364            assert_eq!(DataType::from_onnx(dt.to_onnx()), Some(dt));
365        }
366    }
367
368    /// Unsupported / unknown raw values must return `None` (clean error path),
369    /// never a wrong variant or a panic.
370    #[test]
371    fn unknown_raw_values_return_none() {
372        // UNDEFINED and out-of-range/future values.
373        for raw in [0, 27, 100, 9999, -1, i32::MAX, i32::MIN] {
374            assert_eq!(DataType::from_onnx(raw), None, "raw {raw} should be None");
375        }
376    }
377}
onnx_runtime_ir/dtype.rs

onnx_runtime_ir/
dtype.rs
