uor-matmul-codec 0.1.0

//! The E8 lattice codebook (S4).
//!
//! 256 codewords of eight `i8` each: 2048 bytes of read-only data, which on an
//! embedded target lives in flash and is reached by a pointer.
//!
//! # What this table is, and is not
//!
//! It is the E8 lattice: the `D8` shell of vectors with even coordinate sum,
//! together with the half-integer coset, scaled so the codewords land in `i8`,
//! sorted by norm, and truncated to the 256 shortest. Code `0` is the origin,
//! so a zeroed code stream decodes to the alphabet's zero --- which is what
//! makes zero padding of a coded operand exact (S8).
//!
//! It is **not** a claim about reconstruction quality. Whether these 256 points
//! are a good codebook for any particular weight distribution is measured per
//! (model, codebook) and reported `open`, never asserted (N3). What is claimed
//! here is only that the table is the lattice it says it is, which
//! `e8_is_the_lattice_it_claims` checks by re-deriving the defining property
//! rather than by comparing against a copy.

use uor_matmul_core::{Alphabet, Bound, IntegerElement};

use crate::codecs::Book;

/// Codewords in the table.
pub const ENTRIES: usize = 256;
/// Alphabet elements per codeword.
pub const BLOCK: usize = 8;

/// The raw codewords.
pub static E8_I8: [[i8; BLOCK]; ENTRIES] = [
    [0, 0, 0, 0, 0, 0, 0, 0],
    [-32, -32, 0, 0, 0, 0, 0, 0],
    [-32, 0, -32, 0, 0, 0, 0, 0],
    [-32, 0, 0, -32, 0, 0, 0, 0],
    [-32, 0, 0, 0, -32, 0, 0, 0],
    [-32, 0, 0, 0, 0, -32, 0, 0],
    [-32, 0, 0, 0, 0, 0, -32, 0],
    [-32, 0, 0, 0, 0, 0, 0, -32],
    [-32, 0, 0, 0, 0, 0, 0, 32],
    [-32, 0, 0, 0, 0, 0, 32, 0],
    [-32, 0, 0, 0, 0, 32, 0, 0],
    [-32, 0, 0, 0, 32, 0, 0, 0],
    [-32, 0, 0, 32, 0, 0, 0, 0],
    [-32, 0, 32, 0, 0, 0, 0, 0],
    [-32, 32, 0, 0, 0, 0, 0, 0],
    [-16, -16, -16, -16, -16, -16, -16, -16],
    [-16, -16, -16, -16, -16, -16, 16, 16],
    [-16, -16, -16, -16, -16, 16, -16, 16],
    [-16, -16, -16, -16, -16, 16, 16, -16],
    [-16, -16, -16, -16, 16, -16, -16, 16],
    [-16, -16, -16, -16, 16, -16, 16, -16],
    [-16, -16, -16, -16, 16, 16, -16, -16],
    [-16, -16, -16, -16, 16, 16, 16, 16],
    [-16, -16, -16, 16, -16, -16, -16, 16],
    [-16, -16, -16, 16, -16, -16, 16, -16],
    [-16, -16, -16, 16, -16, 16, -16, -16],
    [-16, -16, -16, 16, -16, 16, 16, 16],
    [-16, -16, -16, 16, 16, -16, -16, -16],
    [-16, -16, -16, 16, 16, -16, 16, 16],
    [-16, -16, -16, 16, 16, 16, -16, 16],
    [-16, -16, -16, 16, 16, 16, 16, -16],
    [-16, -16, 16, -16, -16, -16, -16, 16],
    [-16, -16, 16, -16, -16, -16, 16, -16],
    [-16, -16, 16, -16, -16, 16, -16, -16],
    [-16, -16, 16, -16, -16, 16, 16, 16],
    [-16, -16, 16, -16, 16, -16, -16, -16],
    [-16, -16, 16, -16, 16, -16, 16, 16],
    [-16, -16, 16, -16, 16, 16, -16, 16],
    [-16, -16, 16, -16, 16, 16, 16, -16],
    [-16, -16, 16, 16, -16, -16, -16, -16],
    [-16, -16, 16, 16, -16, -16, 16, 16],
    [-16, -16, 16, 16, -16, 16, -16, 16],
    [-16, -16, 16, 16, -16, 16, 16, -16],
    [-16, -16, 16, 16, 16, -16, -16, 16],
    [-16, -16, 16, 16, 16, -16, 16, -16],
    [-16, -16, 16, 16, 16, 16, -16, -16],
    [-16, -16, 16, 16, 16, 16, 16, 16],
    [-16, 16, -16, -16, -16, -16, -16, 16],
    [-16, 16, -16, -16, -16, -16, 16, -16],
    [-16, 16, -16, -16, -16, 16, -16, -16],
    [-16, 16, -16, -16, -16, 16, 16, 16],
    [-16, 16, -16, -16, 16, -16, -16, -16],
    [-16, 16, -16, -16, 16, -16, 16, 16],
    [-16, 16, -16, -16, 16, 16, -16, 16],
    [-16, 16, -16, -16, 16, 16, 16, -16],
    [-16, 16, -16, 16, -16, -16, -16, -16],
    [-16, 16, -16, 16, -16, -16, 16, 16],
    [-16, 16, -16, 16, -16, 16, -16, 16],
    [-16, 16, -16, 16, -16, 16, 16, -16],
    [-16, 16, -16, 16, 16, -16, -16, 16],
    [-16, 16, -16, 16, 16, -16, 16, -16],
    [-16, 16, -16, 16, 16, 16, -16, -16],
    [-16, 16, -16, 16, 16, 16, 16, 16],
    [-16, 16, 16, -16, -16, -16, -16, -16],
    [-16, 16, 16, -16, -16, -16, 16, 16],
    [-16, 16, 16, -16, -16, 16, -16, 16],
    [-16, 16, 16, -16, -16, 16, 16, -16],
    [-16, 16, 16, -16, 16, -16, -16, 16],
    [-16, 16, 16, -16, 16, -16, 16, -16],
    [-16, 16, 16, -16, 16, 16, -16, -16],
    [-16, 16, 16, -16, 16, 16, 16, 16],
    [-16, 16, 16, 16, -16, -16, -16, 16],
    [-16, 16, 16, 16, -16, -16, 16, -16],
    [-16, 16, 16, 16, -16, 16, -16, -16],
    [-16, 16, 16, 16, -16, 16, 16, 16],
    [-16, 16, 16, 16, 16, -16, -16, -16],
    [-16, 16, 16, 16, 16, -16, 16, 16],
    [-16, 16, 16, 16, 16, 16, -16, 16],
    [-16, 16, 16, 16, 16, 16, 16, -16],
    [0, -32, -32, 0, 0, 0, 0, 0],
    [0, -32, 0, -32, 0, 0, 0, 0],
    [0, -32, 0, 0, -32, 0, 0, 0],
    [0, -32, 0, 0, 0, -32, 0, 0],
    [0, -32, 0, 0, 0, 0, -32, 0],
    [0, -32, 0, 0, 0, 0, 0, -32],
    [0, -32, 0, 0, 0, 0, 0, 32],
    [0, -32, 0, 0, 0, 0, 32, 0],
    [0, -32, 0, 0, 0, 32, 0, 0],
    [0, -32, 0, 0, 32, 0, 0, 0],
    [0, -32, 0, 32, 0, 0, 0, 0],
    [0, -32, 32, 0, 0, 0, 0, 0],
    [0, 0, -32, -32, 0, 0, 0, 0],
    [0, 0, -32, 0, -32, 0, 0, 0],
    [0, 0, -32, 0, 0, -32, 0, 0],
    [0, 0, -32, 0, 0, 0, -32, 0],
    [0, 0, -32, 0, 0, 0, 0, -32],
    [0, 0, -32, 0, 0, 0, 0, 32],
    [0, 0, -32, 0, 0, 0, 32, 0],
    [0, 0, -32, 0, 0, 32, 0, 0],
    [0, 0, -32, 0, 32, 0, 0, 0],
    [0, 0, -32, 32, 0, 0, 0, 0],
    [0, 0, 0, -32, -32, 0, 0, 0],
    [0, 0, 0, -32, 0, -32, 0, 0],
    [0, 0, 0, -32, 0, 0, -32, 0],
    [0, 0, 0, -32, 0, 0, 0, -32],
    [0, 0, 0, -32, 0, 0, 0, 32],
    [0, 0, 0, -32, 0, 0, 32, 0],
    [0, 0, 0, -32, 0, 32, 0, 0],
    [0, 0, 0, -32, 32, 0, 0, 0],
    [0, 0, 0, 0, -32, -32, 0, 0],
    [0, 0, 0, 0, -32, 0, -32, 0],
    [0, 0, 0, 0, -32, 0, 0, -32],
    [0, 0, 0, 0, -32, 0, 0, 32],
    [0, 0, 0, 0, -32, 0, 32, 0],
    [0, 0, 0, 0, -32, 32, 0, 0],
    [0, 0, 0, 0, 0, -32, -32, 0],
    [0, 0, 0, 0, 0, -32, 0, -32],
    [0, 0, 0, 0, 0, -32, 0, 32],
    [0, 0, 0, 0, 0, -32, 32, 0],
    [0, 0, 0, 0, 0, 0, -32, -32],
    [0, 0, 0, 0, 0, 0, -32, 32],
    [0, 0, 0, 0, 0, 0, 32, -32],
    [0, 0, 0, 0, 0, 0, 32, 32],
    [0, 0, 0, 0, 0, 32, -32, 0],
    [0, 0, 0, 0, 0, 32, 0, -32],
    [0, 0, 0, 0, 0, 32, 0, 32],
    [0, 0, 0, 0, 0, 32, 32, 0],
    [0, 0, 0, 0, 32, -32, 0, 0],
    [0, 0, 0, 0, 32, 0, -32, 0],
    [0, 0, 0, 0, 32, 0, 0, -32],
    [0, 0, 0, 0, 32, 0, 0, 32],
    [0, 0, 0, 0, 32, 0, 32, 0],
    [0, 0, 0, 0, 32, 32, 0, 0],
    [0, 0, 0, 32, -32, 0, 0, 0],
    [0, 0, 0, 32, 0, -32, 0, 0],
    [0, 0, 0, 32, 0, 0, -32, 0],
    [0, 0, 0, 32, 0, 0, 0, -32],
    [0, 0, 0, 32, 0, 0, 0, 32],
    [0, 0, 0, 32, 0, 0, 32, 0],
    [0, 0, 0, 32, 0, 32, 0, 0],
    [0, 0, 0, 32, 32, 0, 0, 0],
    [0, 0, 32, -32, 0, 0, 0, 0],
    [0, 0, 32, 0, -32, 0, 0, 0],
    [0, 0, 32, 0, 0, -32, 0, 0],
    [0, 0, 32, 0, 0, 0, -32, 0],
    [0, 0, 32, 0, 0, 0, 0, -32],
    [0, 0, 32, 0, 0, 0, 0, 32],
    [0, 0, 32, 0, 0, 0, 32, 0],
    [0, 0, 32, 0, 0, 32, 0, 0],
    [0, 0, 32, 0, 32, 0, 0, 0],
    [0, 0, 32, 32, 0, 0, 0, 0],
    [0, 32, -32, 0, 0, 0, 0, 0],
    [0, 32, 0, -32, 0, 0, 0, 0],
    [0, 32, 0, 0, -32, 0, 0, 0],
    [0, 32, 0, 0, 0, -32, 0, 0],
    [0, 32, 0, 0, 0, 0, -32, 0],
    [0, 32, 0, 0, 0, 0, 0, -32],
    [0, 32, 0, 0, 0, 0, 0, 32],
    [0, 32, 0, 0, 0, 0, 32, 0],
    [0, 32, 0, 0, 0, 32, 0, 0],
    [0, 32, 0, 0, 32, 0, 0, 0],
    [0, 32, 0, 32, 0, 0, 0, 0],
    [0, 32, 32, 0, 0, 0, 0, 0],
    [16, -16, -16, -16, -16, -16, -16, 16],
    [16, -16, -16, -16, -16, -16, 16, -16],
    [16, -16, -16, -16, -16, 16, -16, -16],
    [16, -16, -16, -16, -16, 16, 16, 16],
    [16, -16, -16, -16, 16, -16, -16, -16],
    [16, -16, -16, -16, 16, -16, 16, 16],
    [16, -16, -16, -16, 16, 16, -16, 16],
    [16, -16, -16, -16, 16, 16, 16, -16],
    [16, -16, -16, 16, -16, -16, -16, -16],
    [16, -16, -16, 16, -16, -16, 16, 16],
    [16, -16, -16, 16, -16, 16, -16, 16],
    [16, -16, -16, 16, -16, 16, 16, -16],
    [16, -16, -16, 16, 16, -16, -16, 16],
    [16, -16, -16, 16, 16, -16, 16, -16],
    [16, -16, -16, 16, 16, 16, -16, -16],
    [16, -16, -16, 16, 16, 16, 16, 16],
    [16, -16, 16, -16, -16, -16, -16, -16],
    [16, -16, 16, -16, -16, -16, 16, 16],
    [16, -16, 16, -16, -16, 16, -16, 16],
    [16, -16, 16, -16, -16, 16, 16, -16],
    [16, -16, 16, -16, 16, -16, -16, 16],
    [16, -16, 16, -16, 16, -16, 16, -16],
    [16, -16, 16, -16, 16, 16, -16, -16],
    [16, -16, 16, -16, 16, 16, 16, 16],
    [16, -16, 16, 16, -16, -16, -16, 16],
    [16, -16, 16, 16, -16, -16, 16, -16],
    [16, -16, 16, 16, -16, 16, -16, -16],
    [16, -16, 16, 16, -16, 16, 16, 16],
    [16, -16, 16, 16, 16, -16, -16, -16],
    [16, -16, 16, 16, 16, -16, 16, 16],
    [16, -16, 16, 16, 16, 16, -16, 16],
    [16, -16, 16, 16, 16, 16, 16, -16],
    [16, 16, -16, -16, -16, -16, -16, -16],
    [16, 16, -16, -16, -16, -16, 16, 16],
    [16, 16, -16, -16, -16, 16, -16, 16],
    [16, 16, -16, -16, -16, 16, 16, -16],
    [16, 16, -16, -16, 16, -16, -16, 16],
    [16, 16, -16, -16, 16, -16, 16, -16],
    [16, 16, -16, -16, 16, 16, -16, -16],
    [16, 16, -16, -16, 16, 16, 16, 16],
    [16, 16, -16, 16, -16, -16, -16, 16],
    [16, 16, -16, 16, -16, -16, 16, -16],
    [16, 16, -16, 16, -16, 16, -16, -16],
    [16, 16, -16, 16, -16, 16, 16, 16],
    [16, 16, -16, 16, 16, -16, -16, -16],
    [16, 16, -16, 16, 16, -16, 16, 16],
    [16, 16, -16, 16, 16, 16, -16, 16],
    [16, 16, -16, 16, 16, 16, 16, -16],
    [16, 16, 16, -16, -16, -16, -16, 16],
    [16, 16, 16, -16, -16, -16, 16, -16],
    [16, 16, 16, -16, -16, 16, -16, -16],
    [16, 16, 16, -16, -16, 16, 16, 16],
    [16, 16, 16, -16, 16, -16, -16, -16],
    [16, 16, 16, -16, 16, -16, 16, 16],
    [16, 16, 16, -16, 16, 16, -16, 16],
    [16, 16, 16, -16, 16, 16, 16, -16],
    [16, 16, 16, 16, -16, -16, -16, -16],
    [16, 16, 16, 16, -16, -16, 16, 16],
    [16, 16, 16, 16, -16, 16, -16, 16],
    [16, 16, 16, 16, -16, 16, 16, -16],
    [16, 16, 16, 16, 16, -16, -16, 16],
    [16, 16, 16, 16, 16, -16, 16, -16],
    [16, 16, 16, 16, 16, 16, -16, -16],
    [16, 16, 16, 16, 16, 16, 16, 16],
    [32, -32, 0, 0, 0, 0, 0, 0],
    [32, 0, -32, 0, 0, 0, 0, 0],
    [32, 0, 0, -32, 0, 0, 0, 0],
    [32, 0, 0, 0, -32, 0, 0, 0],
    [32, 0, 0, 0, 0, -32, 0, 0],
    [32, 0, 0, 0, 0, 0, -32, 0],
    [32, 0, 0, 0, 0, 0, 0, -32],
    [32, 0, 0, 0, 0, 0, 0, 32],
    [32, 0, 0, 0, 0, 0, 32, 0],
    [32, 0, 0, 0, 0, 32, 0, 0],
    [32, 0, 0, 0, 32, 0, 0, 0],
    [32, 0, 0, 32, 0, 0, 0, 0],
    [32, 0, 32, 0, 0, 0, 0, 0],
    [32, 32, 0, 0, 0, 0, 0, 0],
    [-32, -32, -32, -32, 0, 0, 0, 0],
    [-32, -32, -32, 0, -32, 0, 0, 0],
    [-32, -32, -32, 0, 0, -32, 0, 0],
    [-32, -32, -32, 0, 0, 0, -32, 0],
    [-32, -32, -32, 0, 0, 0, 0, -32],
    [-32, -32, -32, 0, 0, 0, 0, 32],
    [-32, -32, -32, 0, 0, 0, 32, 0],
    [-32, -32, -32, 0, 0, 32, 0, 0],
    [-32, -32, -32, 0, 32, 0, 0, 0],
    [-32, -32, -32, 32, 0, 0, 0, 0],
    [-32, -32, 0, -32, -32, 0, 0, 0],
    [-32, -32, 0, -32, 0, -32, 0, 0],
    [-32, -32, 0, -32, 0, 0, -32, 0],
    [-32, -32, 0, -32, 0, 0, 0, -32],
    [-32, -32, 0, -32, 0, 0, 0, 32],
];

/// The E8 codebook over a declared alphabet.
///
/// `None` when the declared bound does not admit every codeword --- the same
/// O(1) non-existence check every other borrowed table gets, decided before any
/// arithmetic (C6).
pub fn e8_table<Bd: Bound>() -> Option<[[Alphabet<i8, Bd>; BLOCK]; ENTRIES]> {
    let mut out = [[Alphabet::ZERO; BLOCK]; ENTRIES];
    for (entry, row) in E8_I8.iter().zip(out.iter_mut()) {
        for (v, slot) in entry.iter().zip(row.iter_mut()) {
            if v.magnitude() > Bd::VALUE {
                return None;
            }
            *slot = Alphabet::new(*v).ok()?;
        }
    }
    Some(out)
}

/// Borrow a prepared table as a codec.
pub fn e8_codec<Bd: Bound>(
    table: &[[Alphabet<i8, Bd>; BLOCK]; ENTRIES],
) -> Book<'_, i8, Bd, ENTRIES, BLOCK> {
    Book::new(table)
}

/// The same codebook, at the byte width.
///
/// 256 entries fit a byte, so this spelling stores its index stream at one
/// byte per codeword --- 0.125 bytes per decoded weight, the residency
/// `CG-03` records --- against the `u16` spelling's two. The same decode, the
/// same enumeration, and a gather that reads the operand's own bytes: `CK-15`
/// pins all three byte for byte.
pub fn e8_codec_u8<Bd: Bound>(
    table: &[[Alphabet<i8, Bd>; BLOCK]; ENTRIES],
) -> Book<'_, i8, Bd, ENTRIES, BLOCK, u8> {
    Book::new(table)
}

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

    /// The table is the lattice it claims to be.
    ///
    /// Re-derives E8's defining property rather than comparing against a
    /// second copy: every codeword is either all-integer with an even
    /// coordinate sum, or all-half-integer. A table that had drifted would fail
    /// this even if it still had 256 rows of the right shape.
    #[test]
    fn e8_is_the_lattice_it_claims() {
        assert_eq!(E8_I8.len(), ENTRIES);
        // Code 0 is the origin, so a zeroed code stream decodes to zero.
        assert_eq!(E8_I8[0], [0i8; BLOCK]);

        for (i, w) in E8_I8.iter().enumerate().skip(1) {
            let all_32 = w.iter().all(|v| v % 32 == 0);
            let all_16 = w.iter().all(|v| v % 32 == 16 || v % 32 == -16);
            assert!(all_32 || all_16, "codeword {i} is in neither coset: {w:?}");
            if all_32 {
                let sum: i32 = w.iter().map(|&v| v as i32 / 32).sum();
                assert_eq!(sum % 2, 0, "codeword {i} has an odd coordinate sum");
            }
        }

        // Every entry is distinct: a codebook with a repeat wastes a code.
        // Sorted in place on the stack, so the check needs no heap even here.
        let mut seen = E8_I8;
        seen.sort_unstable();
        for pair in seen.windows(2) {
            assert_ne!(
                pair[0], pair[1],
                "the codebook has a duplicate: {:?}",
                pair[0]
            );
        }

        // 2048 bytes, which is the figure the model records.
        assert_eq!(core::mem::size_of_val(&E8_I8), 2048);
    }

    /// A bound that cannot hold the codewords has no such table, and that is
    /// reported at construction rather than at decode.
    #[test]
    fn a_bound_too_narrow_has_no_table() {
        assert!(e8_table::<Full<i8>>().is_some());
        assert!(e8_table::<uor_matmul_core::Bnd<8>>().is_none());
    }
}