systile 0.10.0

//! `bf16` — the brain-floating-point format that is the native compute dtype of a TPU.
//!
//! A `bf16` is the top 16 bits of an IEEE-754 `f32`: one sign bit, eight exponent bits,
//! and seven mantissa bits. It keeps the full `f32` dynamic range while throwing away
//! mantissa precision, which is exactly the trade a systolic matrix unit wants.
//!
//! This is a from-scratch software implementation. All arithmetic is performed by
//! widening to `f32`, computing in `f32`, and narrowing back with round-to-nearest-even.
//! That mirrors how a TPU accumulates bf16 products into an f32 accumulator.

use core::cmp::Ordering;
use core::fmt;
use core::iter::Sum;
use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};

/// A brain floating point number (bfloat16).
///
/// The value is stored as the raw 16 bits that would occupy the high half of the
/// corresponding `f32`.
#[derive(Clone, Copy)]
#[repr(transparent)]
pub struct Bf16(u16);

impl Bf16 {
    /// Positive zero.
    pub const ZERO: Bf16 = Bf16(0x0000);

    /// One.
    pub const ONE: Bf16 = Bf16(0x3f80);

    /// Negative one.
    pub const NEG_ONE: Bf16 = Bf16(0xbf80);

    /// Positive infinity.
    pub const INFINITY: Bf16 = Bf16(0x7f80);

    /// Negative infinity.
    pub const NEG_INFINITY: Bf16 = Bf16(0xff80);

    /// A quiet not-a-number.
    pub const NAN: Bf16 = Bf16(0x7fc0);

    /// The largest finite `bf16`.
    pub const MAX: Bf16 = Bf16(0x7f7f);

    /// The smallest (most negative) finite `bf16`.
    pub const MIN: Bf16 = Bf16(0xff7f);

    /// The smallest positive normal `bf16`.
    pub const MIN_POSITIVE: Bf16 = Bf16(0x0080);

    /// Construct a `bf16` directly from its raw 16-bit pattern.
    #[inline]
    pub const fn from_bits(bits: u16) -> Self {
        Bf16(bits)
    }

    /// Return the raw 16-bit pattern of this value.
    #[inline]
    pub const fn to_bits(self) -> u16 {
        self.0
    }

    /// Convert an `f32` to `bf16` using round-to-nearest-even.
    ///
    /// NaNs are preserved as quiet NaNs. The rounding bias is the standard
    /// "round half to even" used by hardware bf16 truncation units.
    #[inline]
    pub fn from_f32(value: f32) -> Self {
        let bits = value.to_bits();
        if value.is_nan() {
            // Force a quiet NaN with a non-zero payload so it survives narrowing.
            return Bf16((bits >> 16) as u16 | 0x0040);
        }
        // Round to nearest even: add the rounding bias derived from the bit that
        // will be dropped plus the lsb of the kept mantissa.
        let rounding_bias = 0x7fff + ((bits >> 16) & 1);
        let rounded = bits.wrapping_add(rounding_bias);
        Bf16((rounded >> 16) as u16)
    }

    /// Convert this `bf16` back to an `f32` exactly (the low 16 bits are zero).
    #[inline]
    pub fn to_f32(self) -> f32 {
        f32::from_bits((self.0 as u32) << 16)
    }

    /// True if this value is NaN.
    #[inline]
    pub fn is_nan(self) -> bool {
        (self.0 & 0x7f80) == 0x7f80 && (self.0 & 0x007f) != 0
    }

    /// True if this value is positive or negative infinity.
    #[inline]
    pub fn is_infinite(self) -> bool {
        (self.0 & 0x7fff) == 0x7f80
    }

    /// True if this value is neither NaN nor infinite.
    #[inline]
    pub fn is_finite(self) -> bool {
        (self.0 & 0x7f80) != 0x7f80
    }

    /// True if this value is exactly zero (positive or negative).
    #[inline]
    pub fn is_zero(self) -> bool {
        (self.0 & 0x7fff) == 0
    }

    /// True if the sign bit is set. Note that `-0.0` is negative under this test.
    #[inline]
    pub fn is_sign_negative(self) -> bool {
        (self.0 & 0x8000) != 0
    }

    /// Return the absolute value.
    #[inline]
    pub fn abs(self) -> Self {
        Bf16(self.0 & 0x7fff)
    }

    /// Return a number with the magnitude of `self` and the sign of `sign`.
    #[inline]
    pub fn copysign(self, sign: Self) -> Self {
        Bf16((self.0 & 0x7fff) | (sign.0 & 0x8000))
    }

    /// Return the larger of two values, ignoring NaN where possible.
    #[inline]
    pub fn max(self, other: Self) -> Self {
        if self.is_nan() {
            return other;
        }
        if other.is_nan() {
            return self;
        }
        if self.to_f32() >= other.to_f32() {
            self
        } else {
            other
        }
    }

    /// Return the smaller of two values, ignoring NaN where possible.
    #[inline]
    pub fn min(self, other: Self) -> Self {
        if self.is_nan() {
            return other;
        }
        if other.is_nan() {
            return self;
        }
        if self.to_f32() <= other.to_f32() {
            self
        } else {
            other
        }
    }

    /// Clamp to the inclusive range `[lo, hi]`.
    #[inline]
    pub fn clamp(self, lo: Self, hi: Self) -> Self {
        self.max(lo).min(hi)
    }

    /// The reciprocal `1.0 / self`.
    #[inline]
    pub fn recip(self) -> Self {
        Bf16::from_f32(self.to_f32().recip())
    }

    /// `self * a + b` computed in f32 and narrowed once, like a fused multiply-add.
    #[inline]
    pub fn mul_add(self, a: Self, b: Self) -> Self {
        Bf16::from_f32(self.to_f32().mul_add(a.to_f32(), b.to_f32()))
    }

    /// `1.0`, `0.0`, or `-1.0` according to the sign; NaN maps to NaN.
    #[inline]
    pub fn signum(self) -> Self {
        if self.is_nan() {
            Bf16::NAN
        } else if self.is_zero() {
            self
        } else if self.is_sign_negative() {
            Bf16::NEG_ONE
        } else {
            Bf16::ONE
        }
    }

    /// True if this value is a normal (not zero, subnormal, infinite, or NaN) number.
    #[inline]
    pub fn is_normal(self) -> bool {
        let exp = self.0 & 0x7f80;
        exp != 0 && exp != 0x7f80
    }

    /// The little-endian byte representation of the raw bits.
    #[inline]
    pub fn to_le_bytes(self) -> [u8; 2] {
        self.0.to_le_bytes()
    }

    /// Reconstruct a value from its little-endian byte representation.
    #[inline]
    pub fn from_le_bytes(bytes: [u8; 2]) -> Self {
        Bf16(u16::from_le_bytes(bytes))
    }
}

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

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

impl From<f64> for Bf16 {
    #[inline]
    fn from(value: f64) -> Self {
        Bf16::from_f32(value as f32)
    }
}

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

impl From<i8> for Bf16 {
    #[inline]
    fn from(value: i8) -> Self {
        Bf16::from_f32(value as f32)
    }
}

impl From<u8> for Bf16 {
    #[inline]
    fn from(value: u8) -> Self {
        Bf16::from_f32(value as f32)
    }
}

impl From<i16> for Bf16 {
    #[inline]
    fn from(value: i16) -> Self {
        Bf16::from_f32(value as f32)
    }
}

impl Default for Bf16 {
    #[inline]
    fn default() -> Self {
        Bf16::ZERO
    }
}

impl PartialEq for Bf16 {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        if self.is_nan() || other.is_nan() {
            return false;
        }
        // +0.0 and -0.0 must compare equal.
        if self.is_zero() && other.is_zero() {
            return true;
        }
        self.0 == other.0
    }
}

impl PartialOrd for Bf16 {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.to_f32().partial_cmp(&other.to_f32())
    }
}

impl Neg for Bf16 {
    type Output = Bf16;
    #[inline]
    fn neg(self) -> Self {
        Bf16(self.0 ^ 0x8000)
    }
}

impl Add for Bf16 {
    type Output = Bf16;
    #[inline]
    fn add(self, rhs: Self) -> Self {
        Bf16::from_f32(self.to_f32() + rhs.to_f32())
    }
}

impl Sub for Bf16 {
    type Output = Bf16;
    #[inline]
    fn sub(self, rhs: Self) -> Self {
        Bf16::from_f32(self.to_f32() - rhs.to_f32())
    }
}

impl Mul for Bf16 {
    type Output = Bf16;
    #[inline]
    fn mul(self, rhs: Self) -> Self {
        Bf16::from_f32(self.to_f32() * rhs.to_f32())
    }
}

impl Div for Bf16 {
    type Output = Bf16;
    #[inline]
    fn div(self, rhs: Self) -> Self {
        Bf16::from_f32(self.to_f32() / rhs.to_f32())
    }
}

impl AddAssign for Bf16 {
    #[inline]
    fn add_assign(&mut self, rhs: Self) {
        *self = *self + rhs;
    }
}

impl SubAssign for Bf16 {
    #[inline]
    fn sub_assign(&mut self, rhs: Self) {
        *self = *self - rhs;
    }
}

impl MulAssign for Bf16 {
    #[inline]
    fn mul_assign(&mut self, rhs: Self) {
        *self = *self * rhs;
    }
}

impl DivAssign for Bf16 {
    #[inline]
    fn div_assign(&mut self, rhs: Self) {
        *self = *self / rhs;
    }
}

impl Sum for Bf16 {
    #[inline]
    fn sum<I: Iterator<Item = Bf16>>(iter: I) -> Self {
        // Accumulate in f32 to mirror a TPU's f32 accumulator, then narrow once.
        Bf16::from_f32(iter.map(|x| x.to_f32()).sum())
    }
}

impl fmt::Debug for Bf16 {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Bf16({} /* 0x{:04x} */)", self.to_f32(), self.0)
    }
}

impl fmt::Display for Bf16 {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&self.to_f32(), f)
    }
}