proxima_primitives/
float_layout.rs

1//! A float layout is a tuple of exponent width, mantissa width and special-value policy; the
2//! scalar codecs `proxima-gguf/src/quant/bf8_e5m2.rs` and `bf4_e2m1.rs` will instantiate it
3//! instead of each hand-coding the same arithmetic.
4
5/// How a layout treats the all-ones exponent and finite overflow.
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
7pub enum Specials {
8    /// No inf or nan encodings; the all-ones exponent is a normal value and overflow saturates
9    /// to the signed max normal (E3M2, E2M1).
10    None,
11    /// IEEE inf and nan encodings are reserved; finite overflow and inf input clamp to the max
12    /// normal, nan still encodes as nan (E5M2_SATURATING, the GGUF bf8 storage policy).
13    SaturateOnOverflow,
14    /// IEEE inf and nan encodings are reserved; overflow rounds to inf (binary32, BF16).
15    InfinityOnOverflow,
16}
17
18/// How the sign of a stored number is encoded.
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
20pub enum SignEncoding {
21    /// No sign bit: the layout is unsigned and `width()` has no sign bit (`uint`).
22    Absent,
23    /// A sign bit above the other fields selects the sign of the magnitude, so zero has two
24    /// encodings (every IEEE-family float; sign-magnitude integers).
25    Magnitude,
26    /// A sign bit of weight `-2^(mantissa_bits)` over the mantissa field: the range is
27    /// asymmetric and zero has one encoding (`int8` and the other native integers).
28    TwosComplement,
29}
30
31/// A sign encoding, `exponent_bits`, `mantissa_bits`; at most 32 bits wide. An integer is the
32/// layout with `exponent_bits` 0 and a `bias_override` of `1 - mantissa_bits`, which makes its
33/// scale exactly 1.
34#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
35pub struct FloatLayout {
36    /// Width of the exponent field; 3 for E3M2, 8 for binary32.
37    pub exponent_bits: u8,
38    /// Width of the stored mantissa field, excluding the implicit leading one; 2 for E3M2, 23
39    /// for binary32.
40    pub mantissa_bits: u8,
41    /// Policy for the all-ones exponent and for overflow; `Specials::None` for E3M2.
42    pub specials: Specials,
43    /// Replaces the IEEE bias `2^(exponent_bits - 1) - 1` when set; `Some(4)` on E3M2 shifts
44    /// its max normal from 28.0 to 14.0.
45    pub bias_override: Option<i16>,
46    /// How the sign is encoded; `Magnitude` for IEEE floats, `TwosComplement` for `int8`.
47    pub sign: SignEncoding,
48}
49
50const fn pow2(exponent: i32) -> f32 {
51    match exponent {
52        -126..=127 => f32::from_bits(((exponent + 127) as u32) << 23),
53        -149..=-127 => f32::from_bits(1 << (exponent + 149)),
54        ..=-150 => 0.0,
55        _ => f32::INFINITY,
56    }
57}
58
59impl FloatLayout {
60    pub const E5M2: Self = Self::ieee_like(5, 2, Specials::InfinityOnOverflow);
61    /// E5M2 widths with saturation: the GGUF bf8 storage policy, which keeps infinities out of stored weights.
62    pub const E5M2_SATURATING: Self = Self::ieee_like(5, 2, Specials::SaturateOnOverflow);
63    /// IEEE 754 binary16: bias 15, inf and nan reserved, max normal 65504.0.
64    pub const F16: Self = Self::ieee_like(5, 10, Specials::InfinityOnOverflow);
65    pub const E2M1: Self = Self::ieee_like(2, 1, Specials::None);
66    /// Proxima cold-storage E2M2: bias 1, no inf or nan, max normal 7.0.
67    pub const E2M2: Self = Self::ieee_like(2, 2, Specials::None);
68    /// OCP MX v1.0 FP6 E3M2: bias 3, no inf or nan, max normal 28.0 (page 12, section 5.3.2, Table 4).
69    pub const E3M2: Self = Self::ieee_like(3, 2, Specials::None);
70    pub const E2M3: Self = Self::ieee_like(2, 3, Specials::None);
71    pub const BF32: Self = Self::ieee_like(8, 23, Specials::InfinityOnOverflow);
72    pub const BF16: Self = Self::ieee_like(8, 7, Specials::InfinityOnOverflow);
73    pub const BF8: Self = Self::E5M2;
74    pub const BF4: Self = Self::E2M1;
75    pub const INT8: Self = Self::integer(8, SignEncoding::TwosComplement);
76    pub const UINT8: Self = Self::integer(8, SignEncoding::Absent);
77    pub const INT16: Self = Self::integer(16, SignEncoding::TwosComplement);
78    pub const UINT16: Self = Self::integer(16, SignEncoding::Absent);
79    pub const INT32: Self = Self::integer(32, SignEncoding::TwosComplement);
80    pub const UINT32: Self = Self::integer(32, SignEncoding::Absent);
81    pub const UINT1: Self = Self::integer(1, SignEncoding::Absent);
82    pub const UINT2: Self = Self::integer(2, SignEncoding::Absent);
83    pub const UINT3: Self = Self::integer(3, SignEncoding::Absent);
84    pub const UINT4: Self = Self::integer(4, SignEncoding::Absent);
85    pub const UINT5: Self = Self::integer(5, SignEncoding::Absent);
86    pub const UINT6: Self = Self::integer(6, SignEncoding::Absent);
87    /// Sign-magnitude eight bits: range -127..=127 with a negative zero.
88    pub const SIGN_MAGNITUDE_INT8: Self = Self::integer(8, SignEncoding::Magnitude);
89
90    const fn ieee_like(exponent_bits: u8, mantissa_bits: u8, specials: Specials) -> Self {
91        Self {
92            exponent_bits,
93            mantissa_bits,
94            specials,
95            bias_override: None,
96            sign: SignEncoding::Magnitude,
97        }
98    }
99
100    const fn integer(width: u8, sign: SignEncoding) -> Self {
101        let mantissa_bits = match sign {
102            SignEncoding::Absent => width,
103            SignEncoding::Magnitude | SignEncoding::TwosComplement => width - 1,
104        };
105        Self {
106            exponent_bits: 0,
107            mantissa_bits,
108            specials: Specials::None,
109            bias_override: Some(1 - mantissa_bits as i16),
110            sign,
111        }
112    }
113
114    #[must_use]
115    pub const fn bias(&self) -> i16 {
116        match self.bias_override {
117            Some(bias) => bias,
118            None => (1i16 << self.exponent_bits.saturating_sub(1)) - 1,
119        }
120    }
121
122    #[must_use]
123    pub const fn width(&self) -> u8 {
124        let sign_bits = match self.sign {
125            SignEncoding::Absent => 0,
126            SignEncoding::Magnitude | SignEncoding::TwosComplement => 1,
127        };
128        sign_bits + self.exponent_bits + self.mantissa_bits
129    }
130
131    /// Whether two layouts are the same format: equal exponent and mantissa widths, equal
132    /// specials, and equal effective bias. `==` also compares the raw `bias_override`, so a
133    /// layout spelled with `Some(15)` and one with `None` differ under `==` and agree here.
134    #[must_use]
135    pub fn same_format(&self, other: &Self) -> bool {
136        self.exponent_bits == other.exponent_bits
137            && self.mantissa_bits == other.mantissa_bits
138            && self.specials == other.specials
139            && self.bias() == other.bias()
140            && self.sign == other.sign
141    }
142
143    /// Bytes `elements` values occupy packed densely at `width()` bits each, plus one `scale` code
144    /// per `block` elements when `scale_block` is `Some((scale, block))`, rounded up to whole
145    /// bytes: 32 elements of E3M2 under an f16 scale are 32 * 6 + 16 = 208 bits, 26 bytes.
146    #[must_use]
147    pub const fn packed_bytes(
148        &self,
149        elements: u64,
150        scale_block: Option<(FloatLayout, u64)>,
151    ) -> u64 {
152        let element_bits = elements * self.width() as u64;
153        let scale_bits = match scale_block {
154            Some((scale, block)) if block > 0 => elements.div_ceil(block) * scale.width() as u64,
155            _ => 0,
156        };
157        (element_bits + scale_bits).div_ceil(8)
158    }
159
160    #[must_use]
161    pub const fn max_normal(&self) -> f32 {
162        let exponent_codes = (1i32 << self.exponent_bits) - 1;
163        let top_code = match self.specials {
164            Specials::None => exponent_codes,
165            Specials::SaturateOnOverflow | Specials::InfinityOnOverflow => exponent_codes - 1,
166        };
167        let significand = 2.0 - pow2(-(self.mantissa_bits as i32));
168        significand * pow2(top_code - self.bias() as i32)
169    }
170
171    #[must_use]
172    pub const fn min_normal(&self) -> f32 {
173        pow2(1 - self.bias() as i32)
174    }
175
176    #[must_use]
177    pub const fn min_subnormal(&self) -> f32 {
178        pow2(1 - self.bias() as i32 - self.mantissa_bits as i32)
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185
186    const CASES: [(FloatLayout, i16, u8, f32, f32, f32); 6] = [
187        (
188            FloatLayout::E5M2,
189            15,
190            8,
191            57344.0,
192            6.103_515_6e-5,
193            1.525_878_9e-5,
194        ),
195        (FloatLayout::E2M1, 1, 4, 6.0, 1.0, 0.5),
196        (FloatLayout::E3M2, 3, 6, 28.0, 0.25, 0.0625),
197        (FloatLayout::E2M3, 1, 6, 7.5, 1.0, 0.125),
198        (
199            FloatLayout::BF16,
200            127,
201            16,
202            3.389_531_4e38,
203            f32::MIN_POSITIVE,
204            f32::from_bits(1 << 16),
205        ),
206        (
207            FloatLayout::BF32,
208            127,
209            32,
210            f32::MAX,
211            f32::MIN_POSITIVE,
212            1.0e-45,
213        ),
214    ];
215
216    #[test]
217    fn float_layout_derived_bias_per_layout() {
218        let biases = CASES.map(|case| case.0.bias());
219        assert_eq!(biases, CASES.map(|case| case.1));
220    }
221
222    #[test]
223    fn float_layout_derived_width_per_layout() {
224        let widths = CASES.map(|case| case.0.width());
225        assert_eq!(widths, CASES.map(|case| case.2));
226    }
227
228    #[test]
229    fn float_layout_derived_max_normal_per_layout() {
230        let maxima = CASES.map(|case| case.0.max_normal());
231        assert_eq!(maxima, CASES.map(|case| case.3));
232    }
233
234    #[test]
235    fn float_layout_derived_min_normal_per_layout() {
236        let minima = CASES.map(|case| case.0.min_normal());
237        assert_eq!(minima, CASES.map(|case| case.4));
238    }
239
240    #[test]
241    fn float_layout_derived_min_subnormal_per_layout() {
242        let minima = CASES.map(|case| case.0.min_subnormal());
243        assert_eq!(minima, CASES.map(|case| case.5));
244    }
245
246    #[test]
247    fn float_layout_derived_bias_override_shifts_range() {
248        let shifted = FloatLayout {
249            bias_override: Some(4),
250            ..FloatLayout::E3M2
251        };
252
253        assert_eq!(shifted.bias(), 4);
254        assert_eq!(shifted.max_normal(), 14.0);
255    }
256
257    #[test]
258    fn float_layout_constants_bf8_equals_e5m2() {
259        assert_eq!(FloatLayout::BF8, FloatLayout::E5M2);
260    }
261
262    #[test]
263    fn float_layout_constants_bf4_equals_e2m1() {
264        assert_eq!(FloatLayout::BF4, FloatLayout::E2M1);
265    }
266
267    #[test]
268    fn float_layout_constants_bf32_is_e8m23() {
269        let layout = FloatLayout::BF32;
270
271        assert_eq!((layout.exponent_bits, layout.mantissa_bits), (8, 23));
272        assert_eq!(layout.max_normal(), f32::MAX);
273        assert_eq!(layout.min_subnormal().to_bits(), 1);
274    }
275
276    #[test]
277    fn float_layout_constants_f16_is_e5m10() {
278        let layout = FloatLayout::F16;
279
280        assert_eq!((layout.exponent_bits, layout.mantissa_bits), (5, 10));
281        assert_eq!(layout.max_normal(), 65504.0);
282        assert_eq!(layout.min_subnormal(), 2.0f32.powi(-24));
283    }
284
285    #[test]
286    fn float_layout_constants_bf16_is_e8m7() {
287        let layout = FloatLayout::BF16;
288
289        assert_eq!((layout.exponent_bits, layout.mantissa_bits), (8, 7));
290        assert_eq!(layout.max_normal().to_bits(), 0x7f7f_0000);
291    }
292
293    #[test]
294    fn float_layout_constants_specials_and_widths_per_layout() {
295        let observed = [
296            FloatLayout::E5M2,
297            FloatLayout::BF8,
298            FloatLayout::E5M2_SATURATING,
299            FloatLayout::BF16,
300            FloatLayout::BF32,
301            FloatLayout::E3M2,
302            FloatLayout::E2M3,
303            FloatLayout::E2M1,
304            FloatLayout::BF4,
305        ]
306        .map(|layout| (layout.specials, layout.width()));
307
308        let saturate = Specials::SaturateOnOverflow;
309        let infinity = Specials::InfinityOnOverflow;
310        let expected = [
311            (infinity, 8),
312            (infinity, 8),
313            (saturate, 8),
314            (infinity, 16),
315            (infinity, 32),
316            (Specials::None, 6),
317            (Specials::None, 6),
318            (Specials::None, 4),
319            (Specials::None, 4),
320        ];
321        assert_eq!(observed, expected);
322    }
323}
proxima_primitives/float_layout.rs

proxima_primitives/
float_layout.rs
