#![cfg(feature = "inference")]
use g_math::fixed_point::OverflowDetected;
use g_math::tq19::bits::{decompose, to_q64_raw, to_raw, to_tq19_raw, HalfKind, WeightBits};
use g_math::tq19::quantize::quantize_tq19;
#[allow(dead_code)]
mod refs {
use g_math::tq19::bits::HalfKind;
include!("data/weight_bits_refs.rs");
}
const FNV_OFFSET: u64 = 0xcbf29ce484222325;
const FNV_PRIME: u64 = 0x100000001b3;
fn fnv(mut h: u64, data: &[u8]) -> u64 {
for &b in data {
h = (h ^ b as u64).wrapping_mul(FNV_PRIME);
}
h
}
fn fold(h: u64, r: Result<i128, OverflowDetected>) -> u64 {
match r {
Ok(v) => fnv(fnv(h, &[0]), &v.to_le_bytes()),
Err(OverflowDetected::InvalidInput) => fnv(h, &[0xE1]),
Err(OverflowDetected::TierOverflow) => fnv(h, &[0xE2]),
Err(e) => panic!("unexpected error {e:?}"),
}
}
#[test]
fn to_raw_matches_the_exact_model_on_every_pattern() {
let expected: [(HalfKind, [u64; 8]); 2] = [
(HalfKind::Binary16, [
refs::BINARY16_RAW_F2, refs::BINARY16_RAW_F10, refs::BINARY16_RAW_F16, refs::BINARY16_RAW_F24,
refs::BINARY16_RAW_F30, refs::BINARY16_RAW_F32, refs::BINARY16_RAW_F64, refs::BINARY16_RAW_F100,
]),
(HalfKind::BFloat16, [
refs::BFLOAT16_RAW_F2, refs::BFLOAT16_RAW_F10, refs::BFLOAT16_RAW_F16, refs::BFLOAT16_RAW_F24,
refs::BFLOAT16_RAW_F30, refs::BFLOAT16_RAW_F32, refs::BFLOAT16_RAW_F64, refs::BFLOAT16_RAW_F100,
]),
];
for (kind, hashes) in expected {
for (i, &f) in refs::RAW_SPLITS.iter().enumerate() {
let mut h = FNV_OFFSET;
for bits in 0u16..=0xFFFF {
h = fold(h, to_raw(bits, kind, f));
}
assert_eq!(h, hashes[i], "{kind:?} at {f} fractional bits");
}
}
}
#[test]
fn to_tq19_raw_matches_the_exact_model_on_every_pattern() {
for (kind, want) in [(HalfKind::Binary16, refs::BINARY16_TQ19), (HalfKind::BFloat16, refs::BFLOAT16_TQ19)] {
let mut h = FNV_OFFSET;
for bits in 0u16..=0xFFFF {
h = fold(h, to_tq19_raw(bits, kind).map(|v| v as i128));
}
assert_eq!(h, want, "{kind:?}");
}
}
#[test]
fn spot_values_and_structure() {
for &(kind, bits, f, want) in refs::RAW_SPOTS {
assert_eq!(to_raw(bits, kind, f), Ok(want), "{kind:?} {bits:#06x} F={f}");
}
assert_eq!(decompose(0x3F80, HalfKind::BFloat16), Ok((false, 0x80, -7)));
assert_eq!(decompose(0x3C00, HalfKind::Binary16), Ok((false, 0x400, -10)));
assert_eq!(to_tq19_raw(0x3F80, HalfKind::BFloat16), Ok(19_683));
assert_eq!(to_tq19_raw(0xBC00, HalfKind::Binary16), Ok(-19_683));
for kind in [HalfKind::Binary16, HalfKind::BFloat16] {
assert_eq!(to_raw(0x0000, kind, 64), Ok(0));
assert_eq!(to_raw(0x8000, kind, 64), Ok(0));
assert_eq!(to_q64_raw(0x8000, kind), Ok(0));
}
assert_eq!(to_raw(0x7C00, HalfKind::Binary16, 10), Err(OverflowDetected::InvalidInput));
assert_eq!(to_raw(0x7E00, HalfKind::Binary16, 10), Err(OverflowDetected::InvalidInput));
assert_eq!(to_raw(0x7F80, HalfKind::BFloat16, 10), Err(OverflowDetected::InvalidInput));
assert_eq!(to_tq19_raw(0xFFC0, HalfKind::BFloat16), Err(OverflowDetected::InvalidInput));
assert_eq!(to_raw(0x3DCD, HalfKind::BFloat16, 10), Ok(102));
assert_eq!(to_raw(0xBDCD, HalfKind::BFloat16, 10), Ok(-102));
}
#[cfg(any(table_format = "q16_16", table_format = "q32_32"))]
#[test]
fn storage_raw_is_to_raw_at_the_build_split() {
use g_math::fixed_point::imperative::BinaryStorage;
use g_math::tq19::bits::{to_fixed, to_storage_raw};
#[cfg(table_format = "q16_16")]
let f = g_math::fixed_point::frac_config::FRAC_BITS;
#[cfg(table_format = "q32_32")]
let f = 32;
let (mut ok, mut refused) = (0u32, 0u32);
for kind in [HalfKind::Binary16, HalfKind::BFloat16] {
for bits in 0u16..=0xFFFF {
let want = match to_raw(bits, kind, f) {
Ok(v) if v.abs() <= BinaryStorage::MAX as i128 => Ok(v as BinaryStorage),
Ok(_) => Err(OverflowDetected::TierOverflow),
Err(e) => Err(e),
};
assert_eq!(to_storage_raw(bits, kind), want, "{kind:?} {bits:#06x}");
assert_eq!(to_fixed(bits, kind).map(|v| v.raw()), want);
if want.is_ok() { ok += 1 } else { refused += 1 }
}
}
assert!(ok > 60_000 && refused > 250, "ok {ok} refused {refused}");
}
fn matrix(kind: HalfKind, bits: &[u16]) -> WeightBits {
let bytes: Vec<u8> = bits.iter().flat_map(|b| b.to_le_bytes()).collect();
WeightBits::from_le_bytes(refs::Q_ROWS, refs::Q_COLS, kind, &bytes).unwrap()
}
#[test]
fn weight_bits_from_le_bytes() {
let w = matrix(HalfKind::BFloat16, refs::BFLOAT16_BITS);
assert_eq!(w.bits, refs::BFLOAT16_BITS);
assert_eq!((w.rows, w.cols), (refs::Q_ROWS, refs::Q_COLS));
assert!(WeightBits::from_le_bytes(2, 2, HalfKind::Binary16, &[0u8; 7]).is_err());
assert_eq!(WeightBits::from_le_bytes(1, 1, HalfKind::Binary16, &[0x00, 0x3C, 0xFF]).unwrap().bits, vec![0x3C00]);
}
#[test]
fn quantize_tq19_matches_the_exact_model() {
for (kind, bits, big, want) in [
(HalfKind::BFloat16, refs::BFLOAT16_BITS, refs::BFLOAT16_BIG_BITS, refs::BFLOAT16_TQ19_DATA),
(HalfKind::Binary16, refs::BINARY16_BITS, refs::BINARY16_BIG_BITS, refs::BINARY16_TQ19_DATA),
] {
let q = quantize_tq19(&matrix(kind, bits)).expect("in range");
assert_eq!(q.data(), want, "{kind:?}");
assert!(quantize_tq19(&matrix(kind, big)).is_none(), "{kind:?} big");
}
let mut bad = refs::BFLOAT16_BITS.to_vec();
bad[7] = 0x7F80;
assert!(quantize_tq19(&matrix(HalfKind::BFloat16, &bad)).is_none());
}
#[cfg(any(table_format = "q16_16", table_format = "q32_32"))]
#[test]
fn quantize_tq19_rowscaled_matches_the_exact_model() {
use g_math::tq19::quantize::quantize_tq19_rowscaled;
for (kind, bits, data, scales) in [
(HalfKind::BFloat16, refs::BFLOAT16_BITS, refs::BFLOAT16_RS19_DATA, refs::BFLOAT16_RS19_SCALES),
(HalfKind::BFloat16, refs::BFLOAT16_BIG_BITS, refs::BFLOAT16_BIG_RS19_DATA, refs::BFLOAT16_BIG_RS19_SCALES),
(HalfKind::Binary16, refs::BINARY16_BITS, refs::BINARY16_RS19_DATA, refs::BINARY16_RS19_SCALES),
(HalfKind::Binary16, refs::BINARY16_BIG_BITS, refs::BINARY16_BIG_RS19_DATA, refs::BINARY16_BIG_RS19_SCALES),
] {
let q = quantize_tq19_rowscaled(&matrix(kind, bits)).unwrap();
assert_eq!(q.data(), data, "{kind:?}");
assert_eq!(q.scales_q32(), scales, "{kind:?}");
assert_eq!(q.scales_q32()[2], 0);
}
}
#[cfg(table_format = "q16_16")]
#[test]
fn quantize_tq5_rowscaled_matches_the_exact_model() {
use g_math::tq19::quantize::quantize_tq5_rowscaled;
for (kind, bits, data, scales) in [
(HalfKind::BFloat16, refs::BFLOAT16_BITS, refs::BFLOAT16_TQ5_DATA, refs::BFLOAT16_TQ5_SCALES),
(HalfKind::BFloat16, refs::BFLOAT16_BIG_BITS, refs::BFLOAT16_BIG_TQ5_DATA, refs::BFLOAT16_BIG_TQ5_SCALES),
(HalfKind::Binary16, refs::BINARY16_BITS, refs::BINARY16_TQ5_DATA, refs::BINARY16_TQ5_SCALES),
(HalfKind::Binary16, refs::BINARY16_BIG_BITS, refs::BINARY16_BIG_TQ5_DATA, refs::BINARY16_BIG_TQ5_SCALES),
] {
let q = quantize_tq5_rowscaled(&matrix(kind, bits)).unwrap();
assert_eq!(q.data(), data, "{kind:?}");
assert_eq!(q.scales_q32(), scales, "{kind:?}");
for r in [0usize, 1, 3, 4] {
assert!(q.data()[r * refs::Q_COLS..(r + 1) * refs::Q_COLS].iter().any(|c| c.abs() == 121), "row {r}");
}
}
}