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}