clawhdf5/crates/clawhdf5-format/src/float16.rs

//! IEEE-754 half precision (binary16) conversions.
//!
//! Pure integer bit manipulation, so it works under `no_std` and needs no
//! `libm`. The writer ([`crate::type_builders::DatasetBuilder::with_f16_data`]),
//! the reader and `clawhdf5-agent`'s half-precision embedding store all use
//! these two functions, so a value rounded in memory is bit-for-bit the value
//! that reads back from the file.

/// Largest finite half-precision value. Anything larger in magnitude rounds
/// to infinity.
pub const F16_MAX: f32 = 65504.0;

/// Convert an `f32` to the bit pattern of the nearest half-precision value,
/// rounding ties to even (the IEEE default, and what numpy and the `half`
/// crate do).
///
/// Values beyond ±[`F16_MAX`] become ±infinity, values too small for a
/// subnormal become signed zero, and NaN stays NaN (quiet, payload
/// truncated).
pub fn f32_to_f16_bits(value: f32) -> u16 {
    let x = value.to_bits();
    let sign = (x >> 16) & 0x8000;
    let exp = x & 0x7F80_0000;
    let man = x & 0x007F_FFFF;

    // Infinity and NaN.
    if exp == 0x7F80_0000 {
        let quiet_nan = if man == 0 { 0 } else { 0x0200 };
        return (sign | 0x7C00 | quiet_nan | (man >> 13)) as u16;
    }

    let half_exp = ((exp >> 23) as i32) - 127 + 15;

    // Too large: infinity.
    if half_exp >= 0x1F {
        return (sign | 0x7C00) as u16;
    }

    // Subnormal half, or zero.
    if half_exp <= 0 {
        if 14 - half_exp > 24 {
            return sign as u16;
        }
        let man = man | 0x0080_0000; // implicit leading bit
        let shift = (14 - half_exp) as u32;
        let mut half_man = man >> shift;
        let round_bit = 1u32 << (shift - 1);
        // Round half to even: up if above half, or exactly half and odd.
        if (man & round_bit) != 0 && (man & (3 * round_bit - 1)) != 0 {
            half_man += 1;
        }
        return (sign | half_man) as u16;
    }

    // Normal half. A mantissa carry correctly rolls into the exponent (and
    // from the largest finite value into infinity).
    let half = sign | ((half_exp as u32) << 10) | (man >> 13);
    let round_bit = 0x0000_1000;
    if (man & round_bit) != 0 && (man & (3 * round_bit - 1)) != 0 {
        (half + 1) as u16
    } else {
        half as u16
    }
}

/// Convert the bit pattern of a half-precision value to `f32` (exact: every
/// half value is representable as an `f32`).
pub fn f16_bits_to_f32(h: u16) -> f32 {
    let h = h as u32;
    let sign = (h & 0x8000) << 16;
    let exp = (h >> 10) & 0x1f;
    let mant = h & 0x3ff;
    let bits = if exp == 0 {
        if mant == 0 {
            sign // signed zero
        } else {
            // Subnormal: normalize into an f32 normal.
            let mut e: i32 = -1;
            let mut m = mant;
            loop {
                e += 1;
                m <<= 1;
                if m & 0x400 != 0 {
                    break;
                }
            }
            let m = m & 0x3ff;
            sign | (((127 - 15 - e) as u32) << 23) | (m << 13)
        }
    } else if exp == 0x1f {
        sign | 0x7f80_0000 | (mant << 13) // inf / NaN
    } else {
        sign | ((exp + 127 - 15) << 23) | (mant << 13)
    };
    f32::from_bits(bits)
}

/// Round an `f32` to the nearest half-precision value, returned as `f32`.
pub fn round_to_f16(value: f32) -> f32 {
    f16_bits_to_f32(f32_to_f16_bits(value))
}

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

    #[test]
    fn every_half_value_round_trips() {
        for bits in 0..=u16::MAX {
            let v = f16_bits_to_f32(bits);
            if v.is_nan() {
                assert!(f16_bits_to_f32(f32_to_f16_bits(v)).is_nan(), "{bits:#06x}");
            } else {
                assert_eq!(f32_to_f16_bits(v), bits, "{bits:#06x} -> {v}");
            }
        }
    }

    #[test]
    fn matches_the_half_crate() {
        // Every 257th f32 bit pattern (~16.7M values) covers every exponent,
        // the subnormal range, both signs, ties and the overflow boundary.
        let mut bits: u32 = 0;
        loop {
            let v = f32::from_bits(bits);
            let ours = f32_to_f16_bits(v);
            let theirs = half::f16::from_f32(v);
            if v.is_nan() {
                assert!(theirs.is_nan() && f16_bits_to_f32(ours).is_nan());
            } else {
                assert_eq!(ours, theirs.to_bits(), "{bits:#010x} ({v:e})");
                assert_eq!(f16_bits_to_f32(ours).to_bits(), theirs.to_f32().to_bits());
            }
            match bits.checked_add(257) {
                Some(b) => bits = b,
                None => break,
            }
        }
    }

    #[test]
    fn rounds_ties_to_even_and_saturates_to_infinity() {
        // 1 + 2^-11 is exactly halfway between 1.0 and the next half (1 + 2^-10).
        assert_eq!(round_to_f16(1.0 + 2f32.powi(-11)), 1.0);
        assert_eq!(
            round_to_f16(1.0 + 3.0 * 2f32.powi(-11)),
            1.0 + 2.0 * 2f32.powi(-10)
        );
        assert_eq!(round_to_f16(F16_MAX), F16_MAX);
        assert_eq!(round_to_f16(65520.0), f32::INFINITY); // halfway to 2^16 rounds up
        assert_eq!(round_to_f16(-1e9), f32::NEG_INFINITY);
        assert_eq!(round_to_f16(1e-9).to_bits(), 0);
        assert_eq!(round_to_f16(-1e-9).to_bits(), (-0.0f32).to_bits());
    }
}