Skip to main content

cubecl_common/quant/
scheme.rs

1use alloc::vec;
2use alloc::vec::Vec;
3use core::{default::Default, ops::Deref};
4use serde::{Deserialize, Serialize};
5
6/// Describes a quantization scheme/configuration.
7#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
8pub struct QuantScheme {
9    /// The logical data type of quantized input values (e.g., `QInt8`).
10    ///
11    /// This defines how values are interpreted during computation, independent of how they're stored.
12    pub value: QuantValue,
13    /// Precision used for quantization parameters (e.g., scale and biases).
14    ///
15    /// This is the only param a one-level scheme has. [`QuantLevel::BlockTensor`] adds a second one
16    /// for its per-tensor scale, so a consumer that reads this field alone will miss that factor.
17    pub param: QuantParam,
18    /// Data type used for storing quantized values.
19    pub store: QuantStore,
20    /// Granularity level of quantization (e.g., per-tensor).
21    pub level: QuantLevel,
22    /// Quantization mode (e.g., symmetric).
23    pub mode: QuantMode,
24}
25
26impl Default for QuantScheme {
27    fn default() -> Self {
28        Self {
29            value: QuantValue::Q8F,
30            param: QuantParam::F32,
31            store: QuantStore::PackedU32(0),
32            level: QuantLevel::Tensor,
33            mode: QuantMode::Symmetric,
34        }
35    }
36}
37
38impl QuantScheme {
39    /// Set the quantization level.
40    pub fn with_level(mut self, level: QuantLevel) -> Self {
41        self.level = level;
42        self
43    }
44
45    /// Set the quantization mode.
46    pub fn with_mode(mut self, mode: QuantMode) -> Self {
47        self.mode = mode;
48        self
49    }
50
51    /// Set the data type used for quantized values.
52    pub fn with_value(mut self, value: QuantValue) -> Self {
53        self.value = value;
54        self
55    }
56
57    /// Set the data type used to store quantized values.
58    pub fn with_store(mut self, store: QuantStore) -> Self {
59        self.store = store;
60        self
61    }
62
63    /// Set the precision used for quantization parameters
64    pub fn with_param(mut self, param: QuantParam) -> Self {
65        self.param = param;
66        self
67    }
68
69    /// Returns the size of the quantization storage type in bits.
70    pub fn size_bits_stored(&self) -> usize {
71        self.store.size_bits(&self.value)
72    }
73
74    /// Returns the size of the quantization storage type in bits.
75    pub fn size_bits_value(&self) -> usize {
76        self.value.size_bits()
77    }
78
79    /// Returns the number of quantized values stored in a single element.
80    pub fn num_quants(&self) -> usize {
81        self.size_bits_stored() / self.value.size_bits()
82    }
83
84    /// Returns the native packing factor for the values. When native packing > 1, the packed
85    /// representation stores `num_quants` elements grouped into packs of `native_packing` size.
86    pub fn native_packing(&self) -> usize {
87        self.value.native_packing()
88    }
89
90    /// Returns the packing dim for the store.
91    pub fn packing_dim(&self) -> Option<usize> {
92        self.store.packing_dim()
93    }
94
95    /// Swaps the packing dim if it's either of `dim0` or `dim1`.
96    /// Executes the corresponding update to `shape.swap(dim0, dim1)`.
97    pub fn swap_packing_dim(&mut self, dim0: usize, dim1: usize) {
98        if let QuantStore::PackedU32(packed_dim) | QuantStore::PackedNative(packed_dim) =
99            &mut self.store
100        {
101            if *packed_dim == dim0 {
102                *packed_dim = dim1;
103            } else if *packed_dim == dim1 {
104                *packed_dim = dim0;
105            }
106        }
107    }
108}
109
110/// Level or granularity of quantization.
111///
112/// Append new variants, never insert. Some transports serialize this with a format that encodes
113/// variants by position rather than by name, so inserting one silently reinterprets streams and
114/// stored schemes written by an older build.
115#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
116pub enum QuantLevel {
117    /// Quantize the whole tensor using a single tensor.
118    Tensor,
119    /// Quantize a tensor using multiple blocks.
120    Block(BlockSize),
121    /// Quantize a tensor using multiple blocks whose scales are themselves normalized by a single
122    /// per-tensor scale.
123    ///
124    /// See [`QuantLevel::block_tensor`] for what that buys and what it does not.
125    BlockTensor {
126        /// Size of each block. The block scales use [`QuantScheme::param`].
127        block: BlockSize,
128        /// Precision of the per-tensor scale. Only a param with more range than the block scales
129        /// use is meaningful.
130        global: QuantParam,
131    },
132}
133
134impl QuantLevel {
135    /// Converting constructor for [`QuantLevel::Block`]
136    pub fn block(values: impl AsRef<[u8]>) -> Self {
137        QuantLevel::Block(BlockSize::new(values))
138    }
139
140    /// Converting constructor for [`QuantLevel::BlockTensor`].
141    ///
142    /// The per-tensor scale absorbs the tensor's dynamic range, which is what lets the block scales
143    /// live in a narrow type. Without it a block scale has to cover that range on its own, and a
144    /// type like [`QuantParam::UE4M3`] underflows to zero for small values.
145    ///
146    /// What the block param covers is then the spread between blocks, which is still bounded. A
147    /// block whose scale falls further below the largest one than the block param can express is
148    /// stored at that param's smallest value, far too coarse for it, and every value in the block
149    /// quantizes to zero. [`QuantParam::UE4M3`] spans about 2^18 this way, from its smallest
150    /// subnormal to 448, so a tensor holding a genuine outlier can lose its ordinary values.
151    ///
152    /// The kernels in `cubecl-std` do not implement this level yet and reject it at launch rather
153    /// than reconstruct values short by the per-tensor factor.
154    pub fn block_tensor(values: impl AsRef<[u8]>, global: QuantParam) -> Self {
155        QuantLevel::BlockTensor {
156            block: BlockSize::new(values),
157            global,
158        }
159    }
160
161    /// The block size, for the levels that quantize in blocks.
162    pub fn block_size(&self) -> Option<BlockSize> {
163        match self {
164            QuantLevel::Tensor => None,
165            QuantLevel::Block(block) | QuantLevel::BlockTensor { block, .. } => Some(*block),
166        }
167    }
168
169    /// The precision of the per-tensor scale, for the levels that have one.
170    pub fn global_param(&self) -> Option<QuantParam> {
171        match self {
172            QuantLevel::Tensor | QuantLevel::Block(_) => None,
173            QuantLevel::BlockTensor { global, .. } => Some(*global),
174        }
175    }
176}
177
178impl QuantParam {
179    /// The largest finite value representable by the parameter type.
180    ///
181    /// A two-level scheme picks its per-tensor scale so that the largest block scale lands here,
182    /// which is what keeps the block scales inside the range their type can express. That recipe
183    /// only holds for a block param narrower than the scale it divides: dividing by
184    /// [`QuantParam::F32`]'s or [`QuantParam::UE8M0`]'s maximum drives the per-tensor scale
185    /// subnormal and the renormalized block scales to infinity. A two-level scheme has nothing to
186    /// gain from those params anyway, since their block scales already reach the full range.
187    pub fn max_representable(&self) -> f32 {
188        match self {
189            QuantParam::F32 => f32::MAX,
190            QuantParam::F16 => half::f16::MAX.to_f32(),
191            QuantParam::BF16 => half::bf16::MAX.to_f32(),
192            // Spelled out because `ue8m0` and `e4m3` sit behind the `fp8` feature and this
193            // function is not gated. The tests check both against those types when it is on.
194            QuantParam::UE8M0 => f32::from_bits(0x7F00_0000), // 2^127
195            QuantParam::UE4M3 => 448.0,
196        }
197    }
198
199    /// The smallest value representable by the parameter type that is not below `scale`.
200    ///
201    /// Storing a quantization scale wants this rather than the nearest value. Rounding down puts
202    /// the scale below what calibration asked for, so every value at the block maximum clips to
203    /// the quantization range; rounding up costs one step of coarseness instead. Backends have to
204    /// agree on this, or a tensor quantized on one reconstructs differently on another.
205    ///
206    /// This is not a cast. Conversion to these types rounds to nearest, which is what a cast
207    /// should do; this is the storage policy for a scale specifically.
208    ///
209    /// `scale` must not be negative. Symmetric quantization only produces non-negative scales,
210    /// and the stepping below walks away from zero for a negative input.
211    ///
212    /// [`QuantParam::UE8M0`] answers [`None`]. Its minimum is 2^-127, subnormal in f32, where the
213    /// grid below no longer holds.
214    pub fn round_up(&self, scale: f32) -> Option<f32> {
215        match self {
216            QuantParam::F32 => {
217                return Some(scale);
218            }
219            QuantParam::UE8M0 => {
220                return None;
221            }
222            _ => {}
223        }
224        if scale.is_nan() {
225            return Some(scale);
226        }
227        debug_assert!(scale >= 0.0, "a quantization scale is never negative");
228
229        // Nothing representable sits above the maximum, and converting past it yields an infinity
230        // for the params that have one, which would make every reconstructed value NaN.
231        let max = self.max_representable();
232        if scale >= max {
233            return Some(max);
234        }
235
236        let grid = self.f32_grid();
237
238        if let Some(subnormals) = grid.subnormals
239            && scale < subnormals.min_normal
240        {
241            // Below the minimum normal the spacing stops halving, so the answer is a count of steps.
242            // Qualified call: the inherent `f32::ceil` lives in std, and this crate builds no_std.
243            return Some(num_traits::Float::ceil(scale / subnormals.spacing) * subnormals.spacing);
244        }
245
246        Some(f32::from_bits(
247            (scale.to_bits() + grid.round_up_bias()) & grid.truncate_mask(),
248        ))
249    }
250
251    /// The param's grid, expressed on the f32 bit pattern. See [`F32Grid`].
252    ///
253    /// bf16 reports no subnormal range because it does not need the separate treatment: its pattern
254    /// is f32's top half all the way down, so the bit step stays right where the others stop. Its
255    /// own subnormals start at 2^-133, which is subnormal in f32 too and flushed to zero by most
256    /// backends.
257    ///
258    /// # Panics
259    ///
260    /// For [`QuantParam::F32`], which is the grid itself, and [`QuantParam::UE8M0`], which is not
261    /// yet supported.
262    pub fn f32_grid(&self) -> F32Grid {
263        /// One f32 ulp per param ulp: the mantissa bits f32 carries and the param does not.
264        const fn bit_step(mantissa_digits: u32) -> u32 {
265            1 << (f32::MANTISSA_DIGITS - mantissa_digits)
266        }
267
268        match self {
269            QuantParam::F16 => F32Grid {
270                bit_step: bit_step(half::f16::MANTISSA_DIGITS),
271                subnormals: Some(SubnormalRange {
272                    min_normal: half::f16::MIN_POSITIVE.to_f32(),
273                    spacing: half::f16::MIN_POSITIVE_SUBNORMAL.to_f32(),
274                }),
275            },
276            QuantParam::BF16 => F32Grid {
277                bit_step: bit_step(half::bf16::MANTISSA_DIGITS),
278                subnormals: None,
279            },
280            // Spelled out rather than read off `e4m3`, which sits behind the `fp8` feature while
281            // this is not gated. The tests check them against that type when it is on.
282            QuantParam::UE4M3 => F32Grid {
283                bit_step: bit_step(4),
284                subnormals: Some(SubnormalRange {
285                    min_normal: 0.015625, // 2^-6
286                    spacing: 0.001953125, // 2^-9
287                }),
288            },
289            QuantParam::F32 => {
290                unimplemented!("F32 is the grid, it has no narrower one to round onto")
291            }
292            QuantParam::UE8M0 => unimplemented!("UE8M0 scales are not yet supported"),
293        }
294    }
295}
296
297/// A narrower float format's grid, laid over the f32 bit pattern.
298///
299/// f32 carries every param this exists for exactly, so the grid can be walked there rather than
300/// through the storage type. A value representable in the param leaves the low f32 mantissa bits
301/// zero, so one param ulp is an increment at that position and the carry into the exponent falls
302/// out on its own. Working in f32 also keeps the grid available to backends with no narrow integer,
303/// and to builds without the `fp8` feature.
304#[derive(Clone, Copy, Debug, PartialEq)]
305pub struct F32Grid {
306    /// One step up in the normal range, as an increment on the f32 bit pattern.
307    pub bit_step: u32,
308    /// The param's subnormals, for the formats whose subnormals land in f32's normal range.
309    pub subnormals: Option<SubnormalRange>,
310}
311
312/// Where a format's subnormals begin and how far apart they are, in f32.
313#[derive(Clone, Copy, Debug, PartialEq)]
314pub struct SubnormalRange {
315    /// The smallest normal value, below which the spacing stops halving.
316    pub min_normal: f32,
317    /// The constant distance between neighbouring subnormals.
318    pub spacing: f32,
319}
320
321impl F32Grid {
322    /// Clears the mantissa bits the param does not carry, truncating a bit pattern onto the grid.
323    pub fn truncate_mask(&self) -> u32 {
324        !(self.bit_step - 1)
325    }
326
327    /// Added to a bit pattern before [`truncate_mask`](Self::truncate_mask) to turn that truncation
328    /// into a round up. The carry it can produce is only safe below the param's maximum, which is
329    /// why callers saturate there first.
330    pub fn round_up_bias(&self) -> u32 {
331        self.bit_step - 1
332    }
333}
334
335/// Data type used to represent quantized values.
336#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
337pub enum QuantValue {
338    /// 8-bit quantization with full range.
339    Q8F,
340    /// 8-bit floating point, e5m2 format.
341    E5M2,
342    /// 8-bit floating point, e4m3 format.
343    E4M3,
344    /// 4-bit quantization with full range.
345    Q4F,
346    /// 4-bit floating point, e2m1 format.
347    E2M1,
348    /// 2-bit quantization with full range.
349    Q2F,
350    /// 8-bit quantization with symmetric range.
351    Q8S,
352    /// 4-bit quantization with symmetric range.
353    Q4S,
354    /// 2-bit quantization with symmetric range.
355    Q2S,
356}
357
358impl QuantValue {
359    /// Returns the size of the quantization input type in bits.
360    pub fn size_bits(&self) -> usize {
361        match self {
362            QuantValue::Q8F | QuantValue::Q8S | QuantValue::E4M3 | QuantValue::E5M2 => 8,
363            QuantValue::Q4F | QuantValue::Q4S | QuantValue::E2M1 => 4,
364            QuantValue::Q2F | QuantValue::Q2S => 2,
365        }
366    }
367
368    /// Packing factor for the native representation used for intermediate values. If > 1, values
369    /// should always be processed in `native_packing` sized chunks.
370    pub fn native_packing(&self) -> usize {
371        match self {
372            QuantValue::E2M1 => 2,
373            _ => 1,
374        }
375    }
376
377    /// The possible range of values allowed by the quant value.
378    pub fn range(&self) -> (f32, f32) {
379        match self {
380            QuantValue::Q8F => (i8::MIN as f32, i8::MAX as f32),
381            QuantValue::Q4F => (-8.0, 7.0),
382            QuantValue::Q2F => (-2.0, 1.0),
383            QuantValue::Q8S => (-i8::MAX as f32, i8::MAX as f32),
384            QuantValue::Q4S => (-7.0, 7.0),
385            QuantValue::Q2S => (-1.0, 1.0),
386            QuantValue::E4M3 => (-448.0, 448.0),
387            QuantValue::E5M2 => (-57344.0, 57344.0),
388            QuantValue::E2M1 => (-6.0, 6.0), // Hardcoded because of no-std
389        }
390    }
391
392    /// If the range of values is symmetric around zero.
393    pub fn is_symmetric(&self) -> bool {
394        match self {
395            Self::Q8F | Self::Q4F | Self::Q2F | Self::E4M3 | Self::E5M2 | Self::E2M1 => false,
396            Self::Q8S | Self::Q4S | Self::Q2S => true,
397        }
398    }
399}
400
401impl QuantStore {
402    /// Returns the size of the quantization input type in bits.
403    pub fn size_bits(&self, value: &QuantValue) -> usize {
404        match self {
405            QuantStore::Native => value.size_bits(),
406            QuantStore::PackedNative(_) => value.size_bits() * value.native_packing(),
407            QuantStore::PackedU32(_) => 32,
408        }
409    }
410
411    fn packing_dim(&self) -> Option<usize> {
412        match self {
413            QuantStore::Native => None,
414            QuantStore::PackedNative(packing_dim) | QuantStore::PackedU32(packing_dim) => {
415                Some(*packing_dim)
416            }
417        }
418    }
419}
420
421/// Data type used to stored quantized values.
422#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
423pub enum QuantStore {
424    /// Native quantization doesn't require packing and unpacking.
425    Native,
426    /// Store packed quantized values in a natively supported packing format (i.e. e2m1x2).
427    /// Argument is the dimension the tensor is packed on, starting from the innermost dimension.
428    PackedNative(usize),
429    /// Store packed quantized values in a 4-byte unsigned integer.
430    /// Argument is the dimension the tensor is packed on, starting from the innermost dimension.
431    PackedU32(usize),
432    // /// Store packed quantized values in a 8-bit unsigned integer.
433    // U8,
434}
435
436/// Strategy used to quantize values.
437#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
438pub enum QuantMode {
439    /// Symmetric or scale quantization.
440    Symmetric,
441}
442
443/// Quantization floating-point precision.
444///
445/// This is used to represent the floating-point precision of quantization parameters like the scale(s)
446/// or the accumulation precision used during operations like matrix multiplication.
447#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
448pub enum QuantParam {
449    /// Full precision.
450    F32,
451    /// Half precision.
452    F16,
453    /// bfloat16 precision.
454    BF16,
455    /// unsigned floating point, e8m0 format.
456    UE8M0,
457    /// unsigned floating point, e4m3 format.
458    UE4M3,
459}
460
461const MAX_DIMS: usize = 5;
462
463/// Copyable block size, specialized version of `SmallVec`.
464#[derive(Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
465pub struct BlockSize {
466    storage: [u8; MAX_DIMS],
467    len: u8,
468}
469
470impl core::fmt::Debug for BlockSize {
471    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
472        write!(f, "BlockSize({:?})", self.as_slice())
473    }
474}
475
476impl BlockSize {
477    /// Max number of dimensions for block size
478    pub const MAX_DIMS: usize = MAX_DIMS;
479
480    /// Create a new blocksize from a set of values. The number of values must be `<= MAX_DIMS`.
481    pub fn new(values: impl AsRef<[u8]>) -> Self {
482        let values = values.as_ref();
483        debug_assert!(
484            values.len() <= MAX_DIMS,
485            "Tried creating a block size larger than the cap"
486        );
487        let len = values.len().min(MAX_DIMS);
488        let mut storage = [1; MAX_DIMS];
489        storage[..len].copy_from_slice(&values[..len]);
490        Self {
491            storage,
492            len: len as u8,
493        }
494    }
495
496    /// Create a new blocksize from a set of values. The number of values must be `<= MAX_DIMS`.
497    /// Trims any leading zeros.
498    pub fn new_trim(values: impl AsRef<[u8]>) -> Self {
499        let values = values.as_ref();
500        let first_value = values.iter().position(|s| *s != 1).unwrap_or(0);
501        Self::new(&values[first_value..])
502    }
503
504    /// Return a slice of only the initialized values
505    pub fn as_slice(&self) -> &[u8] {
506        &self.storage[..self.len as usize]
507    }
508
509    /// Return a vec of only the initialized values
510    pub fn to_vec(&self) -> Vec<u8> {
511        self.storage[..self.len as usize].to_vec()
512    }
513
514    /// Returns `N` dimensions, unsqueezing if necessary.
515    pub fn as_dim<const N: usize>(&self) -> [u8; N] {
516        let data_len = N.min(self.len as usize);
517        let data_start = N - data_len;
518        let mut out = [1; N];
519        out[data_start..].copy_from_slice(&self.storage[..data_len]);
520        out
521    }
522
523    /// Returns a vector of `len` dimensions, unsqueezing if necessary.
524    pub fn to_dim_vec(&self, len: usize) -> Vec<u8> {
525        let data_len = len.min(self.len as usize);
526        let data_start = len - data_len;
527        let mut out = vec![1; len];
528        out[data_start..].copy_from_slice(&self.storage[..data_len]);
529        out
530    }
531
532    /// Create an iterator over all stored dimensions
533    pub fn iter(&self) -> impl Iterator<Item = &u8> {
534        self.as_slice().iter()
535    }
536
537    /// Returns the total number of elements in each block
538    pub fn num_elements(&self) -> usize {
539        self.iter().map(|it| *it as usize).product()
540    }
541}
542
543impl Deref for BlockSize {
544    type Target = [u8];
545
546    fn deref(&self) -> &Self::Target {
547        self.as_slice()
548    }
549}
550
551impl<T: AsRef<[u8]>> From<T> for BlockSize {
552    fn from(value: T) -> Self {
553        BlockSize::new(value)
554    }
555}
556
557#[cfg(test)]
558mod tests {
559    use super::*;
560
561    #[test]
562    fn round_up_never_lands_below_the_scale() {
563        for param in [QuantParam::F16, QuantParam::BF16, QuantParam::UE4M3] {
564            for exp in -12..8 {
565                for step in 1..17 {
566                    let scale = (step as f32 / 16.0) * 2f32.powi(exp);
567                    let up = param.round_up(scale).unwrap();
568                    assert!(
569                        up >= scale,
570                        "{param:?}: {up} is below {scale}, which clips the block maximum"
571                    );
572                }
573            }
574        }
575    }
576
577    #[test]
578    fn round_up_saturates_rather_than_stepping_off_the_top() {
579        for param in [QuantParam::F16, QuantParam::BF16, QuantParam::UE4M3] {
580            let max = param.max_representable();
581            assert_eq!(param.round_up(max).unwrap(), max);
582            assert!(param.round_up(max * 2.0).unwrap().is_finite());
583        }
584    }
585
586    /// Every variant is dispatched somewhere, so none of them may panic here.
587    #[test]
588    fn round_up_answers_for_every_param() {
589        for param in [
590            QuantParam::F32,
591            QuantParam::F16,
592            QuantParam::BF16,
593            QuantParam::UE8M0,
594            QuantParam::UE4M3,
595        ] {
596            assert_eq!(
597                param.round_up(0.3).is_some(),
598                param != QuantParam::UE8M0,
599                "{param:?}"
600            );
601        }
602    }
603
604    #[test]
605    fn round_up_is_the_identity_for_f32() {
606        for scale in [1.0e-30, 0.1, 1.0, 12345.678, f32::MAX] {
607            assert_eq!(QuantParam::F32.round_up(scale).unwrap(), scale);
608        }
609    }
610
611    /// The checks that need the real storage types to compare against.
612    #[cfg(feature = "fp8")]
613    mod storage_types {
614        use super::*;
615
616        #[test]
617        fn round_up_is_the_nearest_representable_value_not_below() {
618            // Rounding up must not overshoot: stepping down from the answer has to land below.
619            for param in [QuantParam::F16, QuantParam::BF16, QuantParam::UE4M3] {
620                for exp in -8..6 {
621                    let scale = 1.7 * 2f32.powi(exp);
622                    let up = param.round_up(scale).unwrap();
623                    assert_eq!(
624                        up,
625                        param.round_up(up).unwrap(),
626                        "{param:?}: not idempotent at {scale}"
627                    );
628                    assert!(
629                        step(param, up, -1) < scale,
630                        "{param:?}: {up} overshoots {scale} by at least a step"
631                    );
632                }
633            }
634        }
635
636        /// `round_up` reads the grid instead of converting through the storage type, so a wrong
637        /// constant there is only visible against the type itself. Nothing else in this file would
638        /// catch one: a grid finer than the real thing still lands above the scale, still steps
639        /// down below it, and still looks idempotent.
640        #[test]
641        fn f32_grid_matches_the_storage_types() {
642            for param in [QuantParam::F16, QuantParam::BF16, QuantParam::UE4M3] {
643                let grid = param.f32_grid();
644
645                // bf16 deliberately reports no subnormal range, since its bit step covers them too.
646                if let Some(subnormals) = grid.subnormals {
647                    assert_eq!(
648                        subnormals.min_normal,
649                        min_normal(param),
650                        "{param:?}: minimum normal"
651                    );
652                    assert_eq!(
653                        subnormals.spacing,
654                        step(param, 0.0, 1),
655                        "{param:?}: subnormal spacing"
656                    );
657                }
658
659                // Walk the whole normal range: one step on the f32 pattern has to be one step in
660                // the type, at every exponent.
661                let mut value = min_normal(param);
662                let max = param.max_representable();
663                while value < max {
664                    let stepped = f32::from_bits(value.to_bits() + grid.bit_step);
665                    assert_eq!(
666                        stepped,
667                        step(param, value, 1),
668                        "{param:?}: step above {value}"
669                    );
670                    value = stepped;
671                }
672                assert_eq!(
673                    value, max,
674                    "{param:?}: the grid has to land exactly on the maximum"
675                );
676            }
677        }
678
679        #[test]
680        fn max_representable_matches_the_e4m3_type() {
681            assert_eq!(
682                QuantParam::UE4M3.max_representable(),
683                crate::e4m3::MAX.to_f32()
684            );
685        }
686
687        /// The other limit spelled out as a literal. `ue8m0` is exponent only, so its maximum is
688        /// the power of two the hex literal encodes.
689        #[test]
690        fn max_representable_matches_the_e8m0_type() {
691            assert_eq!(
692                QuantParam::UE8M0.max_representable(),
693                crate::ue8m0::MAX as f32
694            );
695        }
696
697        /// `offset` representable steps from `value` in `param`, for positive values. Counted on
698        /// the storage type's own bit pattern, so this is an oracle independent of the grid under
699        /// test.
700        fn step(param: QuantParam, value: f32, offset: i32) -> f32 {
701            match param {
702                QuantParam::F16 => half::f16::from_bits(
703                    (half::f16::from_f32(value).to_bits() as i32 + offset) as u16,
704                )
705                .to_f32(),
706                QuantParam::BF16 => half::bf16::from_bits(
707                    (half::bf16::from_f32(value).to_bits() as i32 + offset) as u16,
708                )
709                .to_f32(),
710                QuantParam::UE4M3 => crate::e4m3::from_bits(
711                    (crate::e4m3::from_f32(value).to_bits() as i32 + offset) as u8,
712                )
713                .to_f32(),
714                QuantParam::F32 | QuantParam::UE8M0 => unreachable!(),
715            }
716        }
717
718        fn min_normal(param: QuantParam) -> f32 {
719            match param {
720                QuantParam::F16 => half::f16::MIN_POSITIVE.to_f32(),
721                QuantParam::BF16 => half::bf16::MIN_POSITIVE.to_f32(),
722                QuantParam::UE4M3 => crate::e4m3::MIN_POSITIVE.to_f32(),
723                QuantParam::F32 | QuantParam::UE8M0 => unreachable!(),
724            }
725        }
726    }
727}