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

onnx_runtime_ir/
dtype.rs
