rlx-ir 0.2.16

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0

//! Canonical scalar f32 ↔ low-precision codec for [`ScaledFormat`].
//!
//! One place owns the bit-exact decode/encode for every native-compute
//! low-precision element format (FP8 e4m3/e5m2 OCP + AMD FNUZ, FP6 e2m3/e3m2,
//! FP4 e2m1). Hardware tensor cores do this in silicon; this module is the
//! CPU **parity oracle** every GPU backend is checked against, plus the
//! reference quantizer that produces operands for [`crate::op::Op::ScaledMatMul`].
//!
//! Decode is exact (closed-form from the field bits). Encode is
//! nearest-representable by exhaustive search of the (≤256-entry) code space —
//! simple and obviously correct for an oracle, with round-half-to-even ties and
//! saturation of overflow/inf to the largest finite magnitude (the saturating
//! quantization mode hardware uses). NaN inputs map to `0`.

use crate::quant::ScaledFormat;

/// Decode one packed low-precision code (in the low [`ScaledFormat::bit_width`]
/// bits of `code`) to f32. Returns `f32::NAN` / `±inf` for the formats that
/// encode them.
#[inline]
pub fn decode(fmt: ScaledFormat, code: u8) -> f32 {
    // FP4 has a tiny fixed table shared with the existing NVFP4 path.
    if matches!(fmt, ScaledFormat::F4E2M1) {
        return crate::nvfp4::fp4_e2m1_to_f32(code);
    }

    let (e_bits, m_bits, bias) = fmt.fields();
    let width = e_bits + m_bits; // bits excluding sign
    let sign_bit = (code >> width) & 1;
    let exp = (u32::from(code) >> m_bits) & ((1u32 << e_bits) - 1);
    let mant = u32::from(code) & ((1u32 << m_bits) - 1);
    let sign = if sign_bit == 1 { -1.0f32 } else { 1.0f32 };
    let max_exp = (1u32 << e_bits) - 1;

    // Special encodings, per family.
    if fmt.is_fnuz() {
        // AMD FNUZ: the would-be −0 slot (sign=1, exp=0, mant=0) is the
        // sole NaN; there is no infinity and no negative zero.
        if sign_bit == 1 && exp == 0 && mant == 0 {
            return f32::NAN;
        }
    } else if fmt.has_inf() {
        // OCP E5M2 (IEEE-like): max exponent → inf (mant 0) or NaN.
        if exp == max_exp {
            return if mant == 0 {
                sign * f32::INFINITY
            } else {
                f32::NAN
            };
        }
    } else if matches!(fmt, ScaledFormat::F8E4M3) {
        // OCP E4M3 (finite): only S.1111.111 is NaN; everything else finite.
        if exp == max_exp && mant == (1u32 << m_bits) - 1 {
            return f32::NAN;
        }
    }
    // FP6 e2m3 / e3m2 are all-finite — no special codes.

    let m_div = (1u32 << m_bits) as f32;
    let val = if exp == 0 {
        // Subnormal (and zero): no implicit leading 1.
        (mant as f32 / m_div) * 2f32.powi(1 - bias)
    } else {
        (1.0 + mant as f32 / m_div) * 2f32.powi(exp as i32 - bias)
    };
    sign * val
}

/// Encode an f32 to the nearest representable code (round-half-to-even,
/// saturating overflow/inf to the max finite magnitude, NaN → `0`).
#[inline]
pub fn encode(fmt: ScaledFormat, x: f32) -> u8 {
    if x.is_nan() {
        return 0;
    }
    // ±inf saturates to the largest finite magnitude of the matching sign. We
    // snap to ±max_finite directly rather than a generic huge value: a value far
    // outside the grid is equidistant from every finite code in f64, so the
    // tie-break would wrongly pick code 0.
    if x.is_infinite() {
        return encode(fmt, x.signum() * max_finite(fmt));
    }
    let n_codes: u16 = 1 << fmt.bit_width();
    let mut best: u8 = 0;
    let mut best_err = f64::INFINITY;
    let mut best_mant_lsb: u8 = 1; // prefer even mantissa LSB on ties
    let xd = x as f64;
    for c in 0..n_codes {
        let code = c as u8;
        let v = decode(fmt, code);
        if !v.is_finite() {
            continue; // never round *to* inf/NaN; overflow saturates instead
        }
        let err = (v as f64 - xd).abs();
        let mant_lsb = code & 1;
        if err < best_err || (err == best_err && mant_lsb < best_mant_lsb) {
            best_err = err;
            best = code;
            best_mant_lsb = mant_lsb;
        }
    }
    best
}

/// Largest finite magnitude representable in `fmt`.
pub fn max_finite(fmt: ScaledFormat) -> f32 {
    let n_codes: u16 = 1 << fmt.bit_width();
    (0..n_codes)
        .map(|c| decode(fmt, c as u8).abs())
        .filter(|v| v.is_finite())
        .fold(0.0f32, f32::max)
}

/// Decode a whole buffer of codes (1 code per byte) into f32.
pub fn decode_slice(fmt: ScaledFormat, codes: &[u8], out: &mut [f32]) {
    debug_assert_eq!(codes.len(), out.len());
    for (o, &c) in out.iter_mut().zip(codes) {
        *o = decode(fmt, c);
    }
}

/// Round a positive f32 scale **up** to the nearest power of two and return its
/// OCP **E8M0** byte (biased exponent, bias 127; `0xFF` is NaN, `0x00` is
/// 2^−127). Used for `ScaleLayout::BlockMxE8M0`.
#[inline]
pub fn f32_to_e8m0(scale: f32) -> u8 {
    if !scale.is_finite() || scale <= 0.0 {
        return 0; // degenerate block → smallest scale
    }
    // ceil(log2(scale)) as a biased E8M0 exponent.
    let mut exp = scale.log2().ceil() as i32 + 127;
    exp = exp.clamp(0, 254);
    exp as u8
}

/// Decode an OCP E8M0 scale byte to f32 (`2^(byte-127)`).
#[inline]
pub fn e8m0_to_f32(byte: u8) -> f32 {
    if byte == 0xFF {
        return f32::NAN;
    }
    2f32.powi(byte as i32 - 127)
}

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

    const ALL: [ScaledFormat; 7] = [
        F8E4M3, F8E5M2, F8E4M3Fnuz, F8E5M2Fnuz, F6E2M3, F6E3M2, F4E2M1,
    ];

    #[test]
    fn round_trip_every_representable_code() {
        // decode → encode → decode must be a fixed point for every finite code.
        for fmt in ALL {
            let n = 1u16 << fmt.bit_width();
            for c in 0..n {
                let v = decode(fmt, c as u8);
                if !v.is_finite() {
                    continue;
                }
                let re = decode(fmt, encode(fmt, v));
                assert_eq!(
                    re, v,
                    "{fmt}: code {c:#04x} decoded {v}, re-encoded to {re}"
                );
            }
        }
    }

    #[test]
    fn known_max_finite() {
        assert_eq!(max_finite(F8E4M3), 448.0);
        assert_eq!(max_finite(F8E5M2), 57344.0);
        assert_eq!(max_finite(F8E4M3Fnuz), 240.0);
        assert_eq!(max_finite(F8E5M2Fnuz), 57344.0);
        assert_eq!(max_finite(F6E2M3), 7.5);
        assert_eq!(max_finite(F6E3M2), 28.0);
        assert_eq!(max_finite(F4E2M1), 6.0);
    }

    #[test]
    fn unity_decodes_to_one() {
        // 1.0 = (1 + 0) * 2^0 → exp == bias, mant == 0.
        for fmt in ALL {
            let (e, m, bias) = fmt.fields();
            let code = ((bias as u32) << m) as u8;
            let _ = e;
            assert_eq!(decode(fmt, code), 1.0, "{fmt}: 1.0 code {code:#04x}");
            // and encode(1.0) round-trips
            assert_eq!(decode(fmt, encode(fmt, 1.0)), 1.0, "{fmt}");
        }
    }

    #[test]
    fn e4m3_has_one_nan_no_inf() {
        assert!(decode(F8E4M3, 0x7F).is_nan());
        assert!(decode(F8E4M3, 0xFF).is_nan());
        // every other code is finite
        let finite = (0..256u16)
            .filter(|&c| decode(F8E4M3, c as u8).is_finite())
            .count();
        assert_eq!(finite, 254);
    }

    #[test]
    fn e5m2_has_inf_and_nan() {
        assert_eq!(decode(F8E5M2, 0x7C), f32::INFINITY); // 0.11111.00
        assert_eq!(decode(F8E5M2, 0xFC), f32::NEG_INFINITY);
        assert!(decode(F8E5M2, 0x7D).is_nan());
    }

    #[test]
    fn fnuz_single_nan_no_neg_zero() {
        assert!(decode(F8E4M3Fnuz, 0x80).is_nan());
        assert!(decode(F8E5M2Fnuz, 0x80).is_nan());
        // +0 only
        assert_eq!(decode(F8E4M3Fnuz, 0x00), 0.0);
        // no infinities anywhere
        assert!(!(0..256u16).any(|c| decode(F8E4M3Fnuz, c as u8).is_infinite()));
    }

    #[test]
    fn saturates_on_overflow() {
        // Values past the max finite saturate (never become inf for finite fmts).
        for fmt in ALL {
            let big = 1.0e9;
            let v = decode(fmt, encode(fmt, big));
            assert!(v.is_finite(), "{fmt} overflow not finite");
            assert_eq!(v, max_finite(fmt), "{fmt} did not saturate to max");
            // ±inf saturates too (never encodes to an inf/NaN slot, even for E5M2
            // which has a representable inf code).
            assert_eq!(
                decode(fmt, encode(fmt, f32::INFINITY)),
                max_finite(fmt),
                "{fmt} +inf did not saturate"
            );
            assert_eq!(
                decode(fmt, encode(fmt, f32::NEG_INFINITY)),
                -max_finite(fmt),
                "{fmt} -inf did not saturate"
            );
        }
    }

    #[test]
    fn e8m0_round_trips_powers_of_two() {
        for p in -10..=10i32 {
            let s = 2f32.powi(p);
            assert_eq!(e8m0_to_f32(f32_to_e8m0(s)), s, "2^{p}");
        }
    }

    // --- Parameterized `Custom { exp_bits, mant_bits, bias }` minifloats ---

    #[test]
    fn custom_f4e3m0_value_table() {
        // f4e3m0 = 1 sign + 3 exp + 0 mant, IEEE bias 3, all-finite. The 8
        // non-negative codes are {±0, 0.25, 0.5, 1, 2, 4, 8, 16}; setting the
        // sign bit (bit 3) negates.
        let f = ScaledFormat::custom(3, 0);
        assert_eq!(f.bit_width(), 4);
        assert_eq!(f.fields(), (3, 0, 3));
        let pos = [0.0f32, 0.25, 0.5, 1.0, 2.0, 4.0, 8.0, 16.0];
        for (code, &want) in pos.iter().enumerate() {
            assert_eq!(decode(f, code as u8), want, "code {code}");
            assert_eq!(decode(f, code as u8 | 0x8), -want, "code {}", code + 8);
        }
        assert_eq!(max_finite(f), 16.0);
        // No NaN/inf slots — every code is finite.
        assert!((0..16u8).all(|c| decode(f, c).is_finite()));
    }

    #[test]
    fn custom_e2m1_matches_named_f4e2m1() {
        // A Custom built with E2M1's fields must decode bit-identically to the
        // named F4E2M1 (LUT fast-path) for all 16 codes — proving the generic
        // formula reproduces the hand-written table.
        let custom = ScaledFormat::custom(2, 1); // bias 1
        assert_eq!(custom.fields(), ScaledFormat::F4E2M1.fields());
        for c in 0..16u8 {
            let named = decode(ScaledFormat::F4E2M1, c);
            let generic = decode(custom, c);
            assert!(
                named.to_bits() == generic.to_bits(),
                "e2m1 code {c:#04x}: named {named} vs custom {generic}"
            );
        }
        assert_eq!(max_finite(custom), max_finite(ScaledFormat::F4E2M1));
    }

    #[test]
    fn custom_encode_is_nearest_and_saturates() {
        let f = ScaledFormat::custom(3, 0); // f4e3m0, max finite 16
        for v in [0.25f32, 0.5, 1.0, 2.0, 8.0, 16.0, -4.0, -0.5] {
            assert_eq!(decode(f, encode(f, v)), v, "exact grid point {v}");
        }
        assert_eq!(
            decode(f, encode(f, 100.0)),
            16.0,
            "finite overflow saturates"
        );
        assert_eq!(
            decode(f, encode(f, f32::INFINITY)),
            16.0,
            "+inf saturates to max"
        );
        assert_eq!(
            decode(f, encode(f, f32::NEG_INFINITY)),
            -16.0,
            "-inf saturates"
        );
        assert_eq!(encode(f, f32::NAN), 0, "NaN encodes to 0");
    }

    #[test]
    fn custom_round_trip_every_representable_code() {
        // decode → encode → decode is a fixed point for a spread of exp/mant
        // splits, including a 6-bit and an 8-bit custom.
        for f in [
            ScaledFormat::custom(3, 0), // f4e3m0
            ScaledFormat::custom(2, 2), // f5e2m2
            ScaledFormat::custom(4, 3), // f8e4m3-shaped (finite, no special codes)
            ScaledFormat::custom(3, 4), // f8e3m4
        ] {
            let n = 1u16 << f.bit_width();
            for c in 0..n {
                let v = decode(f, c as u8);
                assert!(v.is_finite(), "{f}: code {c:#04x} unexpectedly non-finite");
                assert_eq!(decode(f, encode(f, v)), v, "{f}: code {c:#04x} ({v})");
            }
        }
    }
}