cubecl_ir/types/
fp8.rs

1use crate::EnumSetType;
2use cubecl_common::{e4m3, e5m2, ue8m0};
3use pliron::{context::Context, r#type::TypeHandle};
4
5use crate::types::scalar::{Float8E4M3Type, Float8E5M2Type, Float8E8M0Type};
6
7/// Bit mask of the seven magnitude bits of an fp8 code.
8const FP8_MAGNITUDE_MASK: u32 = 0x7F;
9
10/// The 8-bit float formats, described by the field constants their host codecs publish.
11///
12/// [`Fp8Format::UE8M0`] is the odd one: it carries no sign and no mantissa, so it is a bare
13/// exponent rather than an IEEE-style minifloat. It belongs here anyway because the question this
14/// enum answers — which 8-bit formats does the backend convert natively, and which need the
15/// software path — is the same question for it, and it is the scale type MXFP4 is defined on.
16/// The generic field arithmetic below does not serve it; the codec branches on it instead.
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, EnumSetType)]
18pub enum Fp8Format {
19    E4M3,
20    E5M2,
21    UE8M0,
22}
23
24impl Fp8Format {
25    pub const fn exponent_bits(self) -> u32 {
26        match self {
27            // Eight bits and no sign, so the usual `7 - mantissa` does not describe it.
28            Fp8Format::UE8M0 => 8,
29            _ => 7 - self.mantissa_bits(),
30        }
31    }
32
33    pub const fn mantissa_bits(self) -> u32 {
34        match self {
35            Fp8Format::E4M3 => e4m3::MANTISSA_DIGITS - 1,
36            Fp8Format::E5M2 => e5m2::MANTISSA_DIGITS - 1,
37            Fp8Format::UE8M0 => 0,
38        }
39    }
40
41    pub const fn bias(self) -> u32 {
42        let min_exp = match self {
43            Fp8Format::E4M3 => e4m3::MIN_EXP,
44            Fp8Format::E5M2 => e5m2::MIN_EXP,
45            // Code 127 is 2^0, so the bias is 127. Spelled out because `ue8m0` publishes no
46            // `MIN_EXP`: it has no subnormals for one to describe.
47            Fp8Format::UE8M0 => return 127,
48        };
49        (2 - min_exp) as u32
50    }
51
52    pub const fn max_value(self) -> f32 {
53        match self {
54            Fp8Format::E4M3 => e4m3::MAX.to_f32(),
55            Fp8Format::E5M2 => e5m2::MAX.to_f32(),
56            // 2^127. Written as a bit pattern because `to_f32` on `ue8m0` is not `const`.
57            Fp8Format::UE8M0 => f32::from_bits(0x7F00_0000),
58        }
59    }
60
61    pub const fn max_code(self) -> u32 {
62        match self {
63            Fp8Format::E4M3 => e4m3::MAX.to_bits() as u32,
64            Fp8Format::E5M2 => e5m2::MAX.to_bits() as u32,
65            Fp8Format::UE8M0 => ue8m0::MAX.to_bits() as u32,
66        }
67    }
68
69    /// Both formats have several NaN encodings and the payload is not carried across: the codecs
70    /// truncate the source mantissa, the hardware paths do as they please. This is the code for a
71    /// canonical NaN, which is what the codec returns for `f32::NAN`. Reading `NAN` instead would
72    /// give e5m2 a payload no conversion produces.
73    pub const fn nan_code(self) -> u32 {
74        let nan = match self {
75            Fp8Format::E4M3 => e4m3::from_f32(f32::NAN).to_bits(),
76            Fp8Format::E5M2 => e5m2::from_f32(f32::NAN).to_bits(),
77            // The all-ones code, one past the maximum. Unsigned, so no magnitude mask applies.
78            Fp8Format::UE8M0 => return 0xFF,
79        };
80        nan as u32 & FP8_MAGNITUDE_MASK
81    }
82
83    pub const fn has_infinity(self) -> bool {
84        match self {
85            // Its one code past the maximum is the NaN, so there is no infinity to reach.
86            Fp8Format::UE8M0 => false,
87            _ => self.decode(self.max_code() + 1).is_infinite(),
88        }
89    }
90
91    pub const fn min_normal(self) -> f32 {
92        match self {
93            Fp8Format::E4M3 => e4m3::MIN_POSITIVE.to_f32(),
94            Fp8Format::E5M2 => e5m2::MIN_POSITIVE.to_f32(),
95            // Every `ue8m0` code is normal in its own terms, so its smallest normal is the bottom
96            // of its range: 2^-127, written as a bit pattern because it is subnormal in f32.
97            Fp8Format::UE8M0 => f32::from_bits(0x0040_0000),
98        }
99    }
100
101    pub const fn subnormal_step(self) -> f32 {
102        match self {
103            // No subnormal ladder to step along — the codec never takes that arm for `ue8m0`.
104            Fp8Format::UE8M0 => 0.0,
105            _ => self.decode(1),
106        }
107    }
108
109    const fn decode(self, code: u32) -> f32 {
110        match self {
111            Fp8Format::E4M3 => e4m3::from_bits(code as u8).to_f32(),
112            Fp8Format::E5M2 => e5m2::from_bits(code as u8).to_f32(),
113            // `ue8m0::to_f32` is not `const`; the code is the f32 exponent field outright.
114            Fp8Format::UE8M0 => f32::from_bits(code << 23),
115        }
116    }
117
118    pub fn of_type(ctx: &Context, ty: TypeHandle) -> Option<Self> {
119        let ty = ty.deref(ctx);
120        if ty.is::<Float8E4M3Type>() {
121            Some(Fp8Format::E4M3)
122        } else if ty.is::<Float8E5M2Type>() {
123            Some(Fp8Format::E5M2)
124        } else if ty.is::<Float8E8M0Type>() {
125            Some(Fp8Format::UE8M0)
126        } else {
127            None
128        }
129    }
130}
cubecl_ir/types/fp8.rs

cubecl_ir/types/
fp8.rs
