cubecl_ir/types/
fp8.rs

1use cubecl_common::{e4m3, e5m2};
2use enumset::EnumSetType;
3use pliron::{context::Context, r#type::TypeHandle};
4
5use crate::types::scalar::{Float8E4M3Type, Float8E5M2Type};
6
7/// Bit mask of the seven magnitude bits of an fp8 code.
8const FP8_MAGNITUDE_MASK: u32 = 0x7F;
9
10/// The two IEEE-style fp8 formats, described by the field constants their host codecs publish.
11#[derive(Debug, Hash, EnumSetType)]
12pub enum Fp8Format {
13    E4M3,
14    E5M2,
15}
16
17impl Fp8Format {
18    pub const fn exponent_bits(self) -> u32 {
19        7 - self.mantissa_bits()
20    }
21
22    pub const fn mantissa_bits(self) -> u32 {
23        match self {
24            Fp8Format::E4M3 => e4m3::MANTISSA_DIGITS - 1,
25            Fp8Format::E5M2 => e5m2::MANTISSA_DIGITS - 1,
26        }
27    }
28
29    pub const fn bias(self) -> u32 {
30        let min_exp = match self {
31            Fp8Format::E4M3 => e4m3::MIN_EXP,
32            Fp8Format::E5M2 => e5m2::MIN_EXP,
33        };
34        (2 - min_exp) as u32
35    }
36
37    pub const fn max_value(self) -> f32 {
38        match self {
39            Fp8Format::E4M3 => e4m3::MAX.to_f32(),
40            Fp8Format::E5M2 => e5m2::MAX.to_f32(),
41        }
42    }
43
44    pub const fn max_code(self) -> u32 {
45        match self {
46            Fp8Format::E4M3 => e4m3::MAX.to_bits() as u32,
47            Fp8Format::E5M2 => e5m2::MAX.to_bits() as u32,
48        }
49    }
50
51    /// Both formats have several NaN encodings and the payload is not carried across: the codecs
52    /// truncate the source mantissa, the hardware paths do as they please. This is the code for a
53    /// canonical NaN, which is what the codec returns for `f32::NAN`. Reading `NAN` instead would
54    /// give e5m2 a payload no conversion produces.
55    pub const fn nan_code(self) -> u32 {
56        let nan = match self {
57            Fp8Format::E4M3 => e4m3::from_f32(f32::NAN).to_bits(),
58            Fp8Format::E5M2 => e5m2::from_f32(f32::NAN).to_bits(),
59        };
60        nan as u32 & FP8_MAGNITUDE_MASK
61    }
62
63    pub const fn has_infinity(self) -> bool {
64        self.decode(self.max_code() + 1).is_infinite()
65    }
66
67    pub const fn min_normal(self) -> f32 {
68        match self {
69            Fp8Format::E4M3 => e4m3::MIN_POSITIVE.to_f32(),
70            Fp8Format::E5M2 => e5m2::MIN_POSITIVE.to_f32(),
71        }
72    }
73
74    pub const fn subnormal_step(self) -> f32 {
75        self.decode(1)
76    }
77
78    const fn decode(self, code: u32) -> f32 {
79        match self {
80            Fp8Format::E4M3 => e4m3::from_bits(code as u8).to_f32(),
81            Fp8Format::E5M2 => e5m2::from_bits(code as u8).to_f32(),
82        }
83    }
84
85    pub fn of_type(ctx: &Context, ty: TypeHandle) -> Option<Self> {
86        let ty = ty.deref(ctx);
87        if ty.is::<Float8E4M3Type>() {
88            Some(Fp8Format::E4M3)
89        } else if ty.is::<Float8E5M2Type>() {
90            Some(Fp8Format::E5M2)
91        } else {
92            None
93        }
94    }
95}
cubecl_ir/types/fp8.rs

cubecl_ir/types/
fp8.rs
