goosedump 0.12.5

// SPDX-License-Identifier: LGPL-2.1-or-later
// Copyright (C) Jarkko Sakkinen 2026

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[repr(transparent)]
pub struct Fp16(pub u16);

unsafe impl bytemuck::Pod for Fp16 {}
unsafe impl bytemuck::Zeroable for Fp16 {}

include!(concat!(env!("OUT_DIR"), "/generated_lut.rs"));

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[repr(transparent)]
pub struct E8m0(pub u8);

unsafe impl bytemuck::Pod for E8m0 {}
unsafe impl bytemuck::Zeroable for E8m0 {}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[repr(transparent)]
pub struct Bf16(pub u16);

unsafe impl bytemuck::Pod for Bf16 {}
unsafe impl bytemuck::Zeroable for Bf16 {}

/// Sealed-bound marker so [`QuantFloat`] cannot be implemented outside this crate.
mod private {
    pub trait Sealed {}
    impl Sealed for super::Fp16 {}
    impl Sealed for super::Bf16 {}
    impl Sealed for super::E8m0 {}
}

/// A compact quantized float that decodes to `f32`.
///
/// Sealed: only the in-crate newtypes (`Fp16`, `Bf16`, `E8m0`) implement it, so
/// adding a future format (e.g. FP8) is a single trait impl rather than a new
/// ad-hoc decode convention threaded separately through every kernel. Kernels
/// decode block scales and weights uniformly via [`QuantFloat::decode_le`] and
/// [`QuantFloat::to_f32`], instead of repeating per-type byte-pair construction
/// (`f32::from(Fp16::from_le_bytes([b0, b1]))` / `f32::from(E8m0(b0))`) inline.
pub trait QuantFloat: Copy + bytemuck::Pod + bytemuck::Zeroable + private::Sealed {
    /// Decode this quantized value to `f32`.
    #[must_use]
    fn to_f32(self) -> f32;

    /// Decode a little-endian byte slice of this type's native width into a
    /// value. `E8m0` reads one byte; `Fp16`/`Bf16` read two.
    ///
    /// # Panics
    /// Panics if `bytes` is shorter than the type's native width.
    #[must_use]
    fn decode_le(bytes: &[u8]) -> Self;
}

impl QuantFloat for Fp16 {
    #[inline]
    fn to_f32(self) -> f32 {
        FP16_TO_F32_LUT[self.0 as usize]
    }

    #[inline]
    fn decode_le(bytes: &[u8]) -> Self {
        Self(u16::from_le_bytes([bytes[0], bytes[1]]))
    }
}

impl QuantFloat for Bf16 {
    #[inline]
    fn to_f32(self) -> f32 {
        f32::from_bits(u32::from(self.0) << 16)
    }

    #[inline]
    fn decode_le(bytes: &[u8]) -> Self {
        Self(u16::from_le_bytes([bytes[0], bytes[1]]))
    }
}

impl QuantFloat for E8m0 {
    #[inline]
    fn to_f32(self) -> f32 {
        E8M0_TO_F32_LUT[self.0 as usize]
    }

    #[inline]
    fn decode_le(bytes: &[u8]) -> Self {
        Self(bytes[0])
    }
}

impl From<E8m0> for f32 {
    #[inline]
    fn from(val: E8m0) -> Self {
        val.to_f32()
    }
}

impl From<Bf16> for f32 {
    #[inline]
    fn from(val: Bf16) -> Self {
        val.to_f32()
    }
}

impl From<Fp16> for f32 {
    #[inline]
    fn from(val: Fp16) -> Self {
        val.to_f32()
    }
}

impl From<f32> for Bf16 {
    #[inline]
    fn from(val: f32) -> Self {
        Self(u16::try_from(val.to_bits() >> 16).expect("top 16 bits fit u16"))
    }
}

impl From<f32> for Fp16 {
    fn from(value: f32) -> Self {
        let bits = value.to_bits();
        let sign = u16::try_from(bits >> 16).expect("top 16 bits fit u16") & 0x8000;
        let exponent = ((bits >> 23) & 0xff).cast_signed();
        let fraction = bits & 0x007f_ffff;
        if exponent == 0xff {
            return if fraction == 0 {
                Self(sign | 0x7c00)
            } else {
                Self(
                    sign | 0x7e00
                        | (u16::try_from(fraction >> 13).expect("fraction fits u16") & 0x01ff),
                )
            };
        }
        let half_exponent = exponent - 112;
        if half_exponent >= 0x1f {
            return Self(sign | 0x7c00);
        }
        if half_exponent <= 0 {
            if half_exponent < -10 {
                return Self(sign);
            }
            let significand = fraction | 0x0080_0000;
            let shift = u32::try_from(14 - half_exponent).expect("shift is non-negative");
            let mut rounded = significand >> shift;
            let remainder = significand & ((1_u32 << shift) - 1);
            let halfway = 1_u32 << (shift - 1);
            if remainder > halfway || (remainder == halfway && rounded & 1 != 0) {
                rounded += 1;
            }
            return Self(sign | u16::try_from(rounded).expect("rounded fits u16"));
        }
        let mut rounded = fraction >> 13;
        let remainder = fraction & 0x1fff;
        if remainder > 0x1000 || (remainder == 0x1000 && rounded & 1 != 0) {
            rounded += 1;
        }
        let mut encoded_exponent = u16::try_from(half_exponent).expect("half_exponent in 1..=0x1e");
        if rounded == 0x400 {
            rounded = 0;
            encoded_exponent += 1;
            if encoded_exponent == 0x1f {
                return Self(sign | 0x7c00);
            }
        }
        Self(sign | (encoded_exponent << 10) | u16::try_from(rounded).expect("rounded fits u16"))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    #[allow(clippy::float_cmp)]
    fn test_fp16_conversion_roundtrip() {
        let values = [0.0_f32, 1.0, -1.0, 0.5, -0.5, 65504.0, -65504.0];
        for &val in &values {
            let fp16 = Fp16::from(val);
            let back = f32::from(fp16);
            assert_eq!(val, back);
        }
    }

    #[test]
    #[allow(clippy::float_cmp)]
    fn test_fp16_zero_and_special() {
        assert_eq!(f32::from(Fp16(0x0000)), 0.0);
        assert_eq!(f32::from(Fp16(0x8000)), -0.0);
        assert!(f32::from(Fp16(0x7c00)).is_infinite());
        assert!(f32::from(Fp16(0x7c00)).is_sign_positive());
        assert!(f32::from(Fp16(0xfc00)).is_infinite());
        assert!(f32::from(Fp16(0xfc00)).is_sign_negative());
    }

    #[test]
    #[allow(clippy::float_cmp)]
    fn test_e8m0_conversion() {
        assert_eq!(f32::from(E8m0(0)), f32::from_bits(0x0020_0000));
        assert_eq!(f32::from(E8m0(1)), f32::from_bits(0x0040_0000));
        assert_eq!(f32::from(E8m0(2)), f32::from_bits(0x0080_0000));
        assert_eq!(f32::from(E8m0(128)), f32::from_bits(127 << 23));
    }

    #[test]
    #[allow(clippy::float_cmp)]
    fn test_bf16_conversion() {
        let val = 1.5_f32;
        let bf16 = Bf16::from(val);
        assert_eq!(f32::from(bf16), val);
    }

    #[test]
    #[allow(clippy::float_cmp)]
    fn test_quant_float_decode_le() {
        assert_eq!(Fp16::decode_le(&[0x00, 0x3c]).to_f32(), 1.0);
        assert_eq!(Bf16::decode_le(&[0x80, 0x3f]).to_f32(), 1.0);
        assert_eq!(E8m0::decode_le(&[128]).to_f32(), f32::from_bits(127 << 23));
    }
}