cubecl_common/float/fp4/
fp4_e2m1.rs

1use core::{
2    fmt::{Debug, Display},
3    ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign},
4};
5
6use bytemuck::{Pod, Zeroable};
7use num_traits::{NumCast, ToPrimitive};
8
9use super::cvt::{f64_to_fp4, fp4_to_f64};
10
11/// A 4-bit floating point type with 2 exponent bits and 1 mantissa bit.
12///
13/// [`Minifloat`]: https://en.wikipedia.org/wiki/Minifloat
14#[allow(non_camel_case_types)]
15#[repr(transparent)]
16#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
17#[derive(Clone, Copy, Default, Zeroable, PartialEq, PartialOrd)]
18pub struct e2m1(u8);
19
20/// A 4-bit floating point type with 2 exponent bits and 1 mantissa bit. Packed with two elements
21/// per value, to allow for conversion to/from bytes. Care must be taken to ensure the shape is
22/// adjusted appropriately.
23///
24/// [`Minifloat`]: https://en.wikipedia.org/wiki/Minifloat
25#[allow(non_camel_case_types)]
26#[repr(transparent)]
27#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
28#[derive(Clone, Copy, Default, Zeroable, Pod, PartialEq, PartialOrd)]
29pub struct e2m1x2(u8);
30
31impl e2m1 {
32    /// Maximum representable value
33    pub const MAX: e2m1 = e2m1(0x7);
34    /// Minimum representable value
35    pub const MIN: e2m1 = e2m1(0xf);
36
37    /// Constructs a [`e2m1`] value from the raw bits.
38    #[inline]
39    #[must_use]
40    pub const fn from_bits(bits: u8) -> e2m1 {
41        e2m1(bits)
42    }
43
44    /// Constructs a [`e2m1`] value from a 32-bit floating point value.
45    ///
46    /// This operation is lossy: values past ±6 saturate to ±6, NaN becomes +6, and every other
47    /// value rounds to the nearest representable value, ties to even.
48    #[inline]
49    #[must_use]
50    pub const fn from_f32(value: f32) -> e2m1 {
51        Self::from_f64(value as f64)
52    }
53
54    /// Constructs a [`e2m1`] value from a 64-bit floating point value.
55    ///
56    /// This operation is lossy: values past ±6 saturate to ±6, NaN becomes +6, and every other
57    /// value rounds to the nearest representable value, ties to even.
58    #[inline]
59    #[must_use]
60    pub const fn from_f64(value: f64) -> e2m1 {
61        e2m1(f64_to_fp4(value))
62    }
63
64    /// Converts a [`e2m1`] into the underlying bit representation.
65    #[inline]
66    #[must_use]
67    pub const fn to_bits(self) -> u8 {
68        self.0
69    }
70
71    /// Converts a [`e2m1`] value into an [`f32`] value.
72    ///
73    /// This conversion is lossless as all values can be represented exactly in [`f32`].
74    #[inline]
75    #[must_use]
76    pub fn to_f32(self) -> f32 {
77        self.to_f64() as f32
78    }
79
80    /// Converts a [`e2m1`] value into an [`f64`] value.
81    ///
82    /// This conversion is lossless as all values can be represented exactly in [`f64`].
83    #[inline]
84    #[must_use]
85    pub fn to_f64(self) -> f64 {
86        fp4_to_f64(self.0)
87    }
88}
89
90impl Neg for e2m1 {
91    type Output = Self;
92
93    fn neg(self) -> Self::Output {
94        Self::from_f32(self.to_f32().neg())
95    }
96}
97
98impl Mul for e2m1 {
99    type Output = Self;
100
101    fn mul(self, rhs: Self) -> Self::Output {
102        Self::from_f32(self.to_f32() * rhs.to_f32())
103    }
104}
105
106impl MulAssign for e2m1 {
107    fn mul_assign(&mut self, rhs: Self) {
108        *self = *self * rhs;
109    }
110}
111
112impl Div for e2m1 {
113    type Output = Self;
114
115    fn div(self, rhs: Self) -> Self::Output {
116        Self::from_f32(self.to_f32() / rhs.to_f32())
117    }
118}
119
120impl DivAssign for e2m1 {
121    fn div_assign(&mut self, rhs: Self) {
122        *self = *self / rhs;
123    }
124}
125
126impl Add for e2m1 {
127    type Output = Self;
128
129    fn add(self, rhs: Self) -> Self::Output {
130        Self::from_f32(self.to_f32() + rhs.to_f32())
131    }
132}
133
134impl AddAssign for e2m1 {
135    fn add_assign(&mut self, rhs: Self) {
136        *self = *self + rhs;
137    }
138}
139
140impl Sub for e2m1 {
141    type Output = Self;
142
143    fn sub(self, rhs: Self) -> Self::Output {
144        Self::from_f32(self.to_f32() - rhs.to_f32())
145    }
146}
147
148impl SubAssign for e2m1 {
149    fn sub_assign(&mut self, rhs: Self) {
150        *self = *self - rhs;
151    }
152}
153
154impl ToPrimitive for e2m1 {
155    fn to_i64(&self) -> Option<i64> {
156        Some(e2m1::to_f32(*self) as i64)
157    }
158
159    fn to_u64(&self) -> Option<u64> {
160        Some(e2m1::to_f64(*self) as u64)
161    }
162
163    fn to_f32(&self) -> Option<f32> {
164        Some(e2m1::to_f32(*self))
165    }
166
167    fn to_f64(&self) -> Option<f64> {
168        Some(e2m1::to_f64(*self))
169    }
170}
171
172impl NumCast for e2m1 {
173    fn from<T: num_traits::ToPrimitive>(n: T) -> Option<Self> {
174        Some(Self::from_f32(n.to_f32()?))
175    }
176}
177
178impl Display for e2m1 {
179    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
180        write!(f, "{}", e2m1::to_f32(*self))
181    }
182}
183
184impl Debug for e2m1 {
185    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
186        write!(f, "{self}")
187    }
188}
189
190impl e2m1x2 {
191    /// Create a new e2m1x2 from bits
192    pub fn from_bits(bits: u8) -> Self {
193        Self(bits)
194    }
195
196    /// Retrieve the stored bits from an e2m1x2
197    pub fn to_bits(self) -> u8 {
198        self.0
199    }
200
201    /// Create a slice of packed fp4 values from a slice of f32s
202    pub fn from_f32_slice(f32s: &[f32]) -> alloc::vec::Vec<e2m1x2> {
203        let mut out = alloc::vec![e2m1x2(0); f32s.len().div_ceil(2)];
204        for (i, chunk) in f32s.chunks(2).enumerate() {
205            let mut chunk = chunk.iter().copied();
206            let a = chunk.next().unwrap_or_default();
207            let b = chunk.next().unwrap_or_default();
208
209            let a = e2m1::from_f32(a).0 & 0x0F;
210            let b = (e2m1::from_f32(b).0 << 4) & 0xF0;
211            out[i] = e2m1x2(a | b);
212        }
213        out
214    }
215}
216
217impl Debug for e2m1x2 {
218    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
219        let a = e2m1::from_bits(self.0 & 0xF).to_f32();
220        let b = e2m1::from_bits((self.0 >> 4) & 0xF).to_f32();
221        f.debug_tuple("e2m1x2").field(&a).field(&b).finish()
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228
229    #[test]
230    fn every_code_converts_to_its_value_and_back() {
231        for code in 0..16u8 {
232            let value = e2m1::from_bits(code).to_f32();
233            assert_eq!(
234                e2m1::from_f32(value).to_bits(),
235                code,
236                "{code:#x} is {value}"
237            );
238        }
239    }
240
241    #[test]
242    fn values_round_to_the_nearest_code_ties_to_even_and_saturate() {
243        for (value, expected) in [
244            (0.25, 0.0),
245            (0.75, 1.0),
246            (1.25, 1.0),
247            (1.75, 2.0),
248            (2.5, 2.0),
249            (3.5, 4.0),
250            (5.0, 4.0),
251            (0.3, 0.5),
252            (2.4, 2.0),
253            (-0.75, -1.0),
254            (-2.5, -2.0),
255            (100.0, 6.0),
256            (-100.0, -6.0),
257            (f32::NAN, 6.0),
258        ] {
259            assert_eq!(e2m1::from_f32(value).to_f32(), expected, "{value}");
260        }
261    }
262
263    #[test]
264    fn a_negative_too_small_for_any_code_keeps_its_sign() {
265        assert_eq!(e2m1::from_f32(-0.0).to_bits(), 0x8);
266        assert_eq!(e2m1::from_f32(-0.1).to_bits(), 0x8);
267    }
268}
cubecl_common/float/fp4/fp4_e2m1.rs

cubecl_common/float/fp4/
fp4_e2m1.rs
