wai-quantum 0.6.0

//! The gate on a compute tile's vector engine: the programs the tile's first
//! RISC-V core runs to drive it, built with [`rv32`] and [`tensix`].
//!
//! **One pass.** The host writes up to [`PAIRS`] pairs' four columns to the
//! tile's memory at [`IN`]. The core writes the unpacker's and the packer's
//! configuration, then hands the tile its instructions inline: the unpacker
//! copies the columns, as 32-bit integers, into the destination registers; the
//! vector engine forms each pair's four rows, 32 lanes at a time; the packer
//! copies the rows back to [`OUT`]; and the core writes [`DONE`] once the
//! packer has finished.
//!
//! **Exact arithmetic.** The vector engine multiplies 23-bit numbers
//! (`SFPMUL24`: the low or the high 23 bits of the 46-bit product). Each
//! operand is split into 23 and 9 bits; the four partial products give the
//! unsigned product's high word, corrected for sign to the signed product's
//! `P_hi`, and the bits of the low word that the rounding reaches. With the
//! rounding bit `r` and the shift `s`, `fx = (P + 2^r) >> s` is
//! `P_hi · 2^(32−s) + t`, `t = ((P_lo >> r) + 1) >> (s − r)`, held as 30 low
//! bits and the rest; a row adds or subtracts its four terms in that form,
//! carrying after each, and is stored as its 64-bit value's two words.

use crate::quantum_fabric::rv32::{A0, Asm, S1, T0, T1, T2, T3, ZERO};
use crate::quantum_fabric::tensix::word;

/// Pairs per pass, per tile: half a 32 × 32 tile of each operand.
pub const PAIRS: usize = 512;
/// Where the columns go in the tile's memory: four of [`PAIRS`] 32-bit values.
pub const IN: u32 = 0x4_0000;
/// Where the rows come back: for each row, [`PAIRS`] low words then [`PAIRS`]
/// high words.
pub const OUT: u32 = 0x5_0000;
/// The word the core sets to 1 when the pass is done.
pub const DONE: u32 = 0x3_F000;

/// The kernel's constants: the rounding bit `r` (rounding constant `2^r`),
/// the shift, and whether row 0's last real term is added.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Constants {
    pub round_bit: u32,
    pub shift: u32,
    pub last_add: bool,
}

const INT32: u32 = 4;
/// The vector register holding `2^30 − 1`.
const M30: u32 = 12;
/// Where the operands sit in the destination registers (rows).
const INPUTS: u32 = 256;

fn ld(l: u32, addr: u32) -> u32 {
    word("SFPLOAD", &[("lreg_ind", l), ("instr_mod0", INT32), ("dest_reg_addr", addr)])
}
fn st(l: u32, addr: u32) -> u32 {
    word("SFPSTORE", &[("lreg_ind", l), ("instr_mod0", INT32), ("dest_reg_addr", addr)])
}
fn imm12(v: i32) -> u32 {
    (v as u32) & 0xfff
}
/// `d = c << n` (n ≥ 0) or `d = c >>u −n`.
fn shift(d: u32, c: u32, n: i32) -> u32 {
    word("SFPSHFT", &[("lreg_dest", d), ("lreg_c", c), ("imm12_math", imm12(n)), ("instr_mod1", 5)])
}
/// `d = c >>a n`.
fn shra(d: u32, c: u32, n: i32) -> u32 {
    word("SFPSHFT", &[("lreg_dest", d), ("lreg_c", c), ("imm12_math", imm12(-n)), ("instr_mod1", 7)])
}
/// `d = c + d` (the condition codes untouched).
fn add(d: u32, c: u32) -> u32 {
    word("SFPIADD", &[("lreg_dest", d), ("lreg_c", c), ("instr_mod1", 4)])
}
/// `d = c − d`.
fn rsub(d: u32, c: u32) -> u32 {
    word("SFPIADD", &[("lreg_dest", d), ("lreg_c", c), ("instr_mod1", 6)])
}
/// `d = c + v`.
fn addi(d: u32, c: u32, v: i32) -> u32 {
    word("SFPIADD", &[("lreg_dest", d), ("lreg_c", c), ("imm12_math", imm12(v)), ("instr_mod1", 5)])
}
/// `d = L[l] & c`.
fn and_with(d: u32, l: u32, c: u32) -> u32 {
    word("SFPAND", &[("lreg_dest", d), ("lreg_c", c), ("imm12_math", l), ("instr_mod1", 1)])
}
/// `d = d & c`.
fn and(d: u32, c: u32) -> u32 {
    word("SFPAND", &[("lreg_dest", d), ("lreg_c", c)])
}
/// `d = d | c`.
fn or(d: u32, c: u32) -> u32 {
    word("SFPOR", &[("lreg_dest", d), ("lreg_c", c)])
}
/// `d = lo23(a · b)` or `hi23` (`high`), then a gap: silicon does not wait for
/// the multiplier's result before the next instruction reads it.
fn mul(w: &mut Vec<(u32, String)>, d: u32, a: u32, b: u32, high: bool) {
    let m = word("SFPMUL24", &[("lreg_src_a", a), ("lreg_src_b", b), ("lreg_src_c", 9), ("lreg_dest", d), ("instr_mod1", u32::from(high))]);
    w.push((m, format!("SFPMUL24 L{d} = {}(L{a} * L{b})", if high { "hi23" } else { "lo23" })));
    w.push((word("SFPNOP", &[]), "SFPNOP".into()));
}
fn push(w: &mut Vec<(u32, String)>, word: u32, what: &str) {
    w.push((word, what.into()));
}

/// The product of the operands in L0 and L1: `P_hi` into L3, the 9 bits of
/// `P_lo` above bit 22 (`d1 & 511`, at the bottom of L2's word), with L0 and
/// L1 kept.
fn product(w: &mut Vec<(u32, String)>) {
    push(w, shift(2, 0, -23), "L2 = a1 = a >>u 23");
    push(w, shift(3, 1, -23), "L3 = e1 = e >>u 23");
    mul(w, 4, 2, 3, false); // a1·e1, exact
    mul(w, 5, 2, 1, true);
    push(w, add(4, 5), "d2 += hi(a1, e0)");
    mul(w, 5, 0, 3, true);
    push(w, add(4, 5), "d2 += hi(a0, e1)");
    mul(w, 2, 2, 1, false);
    mul(w, 3, 0, 3, false);
    push(w, add(2, 3), "d1 = lo(a1, e0) + lo(a0, e1)");
    mul(w, 3, 0, 1, true);
    push(w, add(2, 3), "d1 += hi(a0, e0)");
    push(w, shra(3, 0, 31), "L3 = a < 0 ? -1 : 0");
    push(w, and(3, 1), "L3 = a < 0 ? e : 0");
    push(w, shra(5, 1, 31), "L5 = e < 0 ? -1 : 0");
    push(w, and(5, 0), "L5 = e < 0 ? a : 0");
    push(w, add(3, 5), "the sign correction");
    push(w, shift(4, 4, 14), "d2 << 14");
    push(w, shift(5, 2, -9), "d1 >>u 9");
    push(w, add(4, 5), "U_hi");
    push(w, rsub(3, 4), "L3 = P_hi = U_hi - correction");
}

/// From [`product`]: `fx` as 30 low bits in L4 and the rest in L3.
fn fx_from_product(w: &mut Vec<(u32, String)>, c: Constants) {
    let k = 32 - c.shift as i32;
    push(w, shift(2, 2, 23), "the bits of P_lo the rounding reaches, at the top");
    push(w, shift(2, 2, -(c.round_bit as i32)), "P_lo >>u r");
    push(w, addi(2, 2, 1), "+ 1");
    if c.shift > c.round_bit {
        push(w, shift(2, 2, -((c.shift - c.round_bit) as i32)), "t = ((P_lo >> r) + 1) >> (s - r)");
    }
    push(w, shift(4, 3, k), "P_hi << (32 - s)");
    push(w, add(4, 2), "+ t");
    push(w, and_with(4, M30, 4), "fx low 30 bits");
    push(w, shift(2, 2, -k), "t >> (32 - s)");
    push(w, add(3, 2), "P_hi + (t >> (32 - s))");
    push(w, shra(3, 3, 30 - k), "fx >> 30");
}

/// The row accumulator (L6 low 30 bits, L7 the rest) plus or minus `fx`, carried.
fn accumulate(w: &mut Vec<(u32, String)>, sign: i32, first: bool) {
    if first {
        push(w, addi(6, 4, 0), "R_lo = fx_lo");
        push(w, addi(7, 3, 0), "R_hi = fx_hi");
        return;
    }
    if sign > 0 {
        push(w, add(6, 4), "R_lo += fx_lo");
        push(w, add(7, 3), "R_hi += fx_hi");
    } else {
        push(w, rsub(4, 6), "R_lo - fx_lo");
        push(w, addi(6, 4, 0), "R_lo -= fx_lo");
        push(w, rsub(3, 7), "R_hi - fx_hi");
        push(w, addi(7, 3, 0), "R_hi -= fx_hi");
    }
    push(w, shra(5, 6, 30), "the carry");
    push(w, add(7, 5), "R_hi += carry");
    push(w, and_with(6, M30, 6), "R_lo &= 2^30 - 1");
}

/// One iteration: the four rows of 32 lanes. `m` the gate's entries.
fn rows_iteration(m: &[i32; 8], c: Constants) -> Vec<(u32, String)> {
    let last = if c.last_add { 1 } else { -1 };
    let rows: [[(usize, u32, i32); 4]; 4] = [
        [(0, 0, 1), (1, 1, -1), (2, 2, 1), (3, 3, last)],
        [(0, 1, 1), (1, 0, 1), (2, 3, 1), (3, 2, 1)],
        [(4, 0, 1), (5, 1, -1), (6, 2, 1), (7, 3, -1)],
        [(4, 1, 1), (5, 0, 1), (6, 3, 1), (7, 2, 1)],
    ];
    let mut w = Vec::new();
    for (r, terms) in rows.iter().enumerate() {
        for (i, &(e, op, sign)) in terms.iter().enumerate() {
            push(&mut w, ld(0, INPUTS + 32 * op), &format!("L0 = column {op}"));
            let v = m[e] as u32;
            push(&mut w, word("SFPLOADI", &[("lreg_ind", 1), ("instr_mod0", 10), ("imm16", v & 0xffff)]), &format!("L1 = m{e} (low half)"));
            push(&mut w, word("SFPLOADI", &[("lreg_ind", 1), ("instr_mod0", 8), ("imm16", v >> 16)]), &format!("L1 = m{e} (high half)"));
            product(&mut w);
            fx_from_product(&mut w, c);
            accumulate(&mut w, sign, i == 0);
        }
        push(&mut w, shift(4, 7, 30), "R_hi << 30");
        push(&mut w, or(4, 6), "the row's low word");
        push(&mut w, shra(7, 7, 2), "the row's high word");
        push(&mut w, st(4, 64 * r as u32), &format!("row {r} low"));
        push(&mut w, st(7, 64 * r as u32 + 32), &format!("row {r} high"));
    }
    push(&mut w, word("INCRWC", &[("rwc_d", 2)]), "the next 32 lanes");
    w
}

/// One iteration of the probe: each pair's product and `fx`, both 64-bit.
fn probe_iteration(c: Constants) -> Vec<(u32, String)> {
    let mut w = Vec::new();
    push(&mut w, ld(0, INPUTS), "L0 = a");
    push(&mut w, ld(1, INPUTS + 32), "L1 = b");
    product(&mut w);
    mul(&mut w, 6, 0, 1, false);
    push(&mut w, shift(5, 2, 23), "(d1 & 511) << 23");
    push(&mut w, or(6, 5), "P_lo");
    push(&mut w, st(6, 0), "the product's low word");
    push(&mut w, st(3, 32), "the product's high word");
    fx_from_product(&mut w, c);
    push(&mut w, shift(5, 3, 30), "fx_hi << 30");
    push(&mut w, or(5, 4), "fx low word");
    push(&mut w, shra(3, 3, 2), "fx high word");
    push(&mut w, st(5, 64), "fx low");
    push(&mut w, st(3, 96), "fx high");
    push(&mut w, word("INCRWC", &[("rwc_d", 2)]), "the next 32 lanes");
    w
}

/// The configuration word `idx` of the tile, through the core's view of it.
fn cfg(p: &mut Asm, idx: u32, value: u32) {
    p.li(T1, value).sw(T1, (4 * idx) as i32, T0);
}

/// Lanes one iteration covers.
pub const LANES: usize = 32;

/// The whole program around one iteration's `body`, run `iterations` times
/// (1 to 16, each [`LANES`] more pairs): configure, unpack `faces` faces of
/// columns, iterate, pack `rows` destination rows, signal.
fn program(body: &[(u32, String)], iterations: u32, faces: u32, rows: u32) -> (Vec<u32>, String) {
    assert!((1..=(PAIRS / LANES) as u32).contains(&iterations));
    let mut p = Asm::default();
    p.li(T0, 0xffef_0000);
    // Unpacker 0: 32-bit integers into the destination registers from row 256.
    for (idx, v) in [(64, 0x18), (65, 0x0020_0001), (72, 0x808), (73, 0x11), (76, IN / 16 - 1), (84, 0x1040_1040), (86, 0x0100_0100), (57, 0x400)] {
        cfg(&mut p, idx, v);
    }
    // Packer 0: 32-bit integers from the destination registers from row 0.
    for (idx, v) in [(1, 0x1000_0000), (12, 0x0040_0000), (13, 0x1000_0400), (18, 1), (24, 0xffff), (28, 0x1000), (69, OUT / 16 - 1), (70, 0x881)] {
        cfg(&mut p, idx, v);
    }
    // Read a configuration word back before the tile acts on them.
    p.lw(T1, 4 * 76, T0).add(T1, T1, ZERO);
    p.li(T2, 0xffe0_0000).sw(ZERO, 0xfc, T2).lw(T1, 0xfc, T2);
    p.tt(word("SETC16", &[("setc16_reg", 37), ("setc16_value", 4)]), "SETC16: the packer steps 4 rows");
    p.tt(word("SETADCZW", &[("bit_mask", 0xf), ("cnt_set_mask", 1)]), "SETADCZW unpacker 0");
    p.tt(word("SETADCXY", &[("bit_mask", 0xf), ("cnt_set_mask", 1)]), "SETADCXY unpacker 0");
    p.tt(word("SETADCXX", &[("x_end2", 255), ("cnt_set_mask", 1)]), "SETADCXX unpacker 0: a face");
    for f in 0..faces {
        p.tt(word("UNPACR", &[("last", 1), ("ovrd_thread_id", 1), ("addr_mode", 0b0001_0001)]), &format!("UNPACR face {f}"));
    }
    p.tt(word("STALLWAIT", &[("stall_res", 1 << 8), ("wait_res", 1 << 1)]), "STALLWAIT: the vector engine waits for the unpacker");
    p.tt(word("SFPLOADI", &[("lreg_ind", 0), ("instr_mod0", 10), ("imm16", 0xffff)]), "L0 = 2^30 - 1 (low half)");
    p.tt(word("SFPLOADI", &[("lreg_ind", 0), ("instr_mod0", 8), ("imm16", 0x3fff)]), "L0 = 2^30 - 1 (high half)");
    p.tt(word("SFPCONFIG", &[("config_dest", M30)]), "L12 = L0");
    p.tt(word("SETRWC", &[("bit_mask", 4)]), "SETRWC: lanes from row 0");
    p.li(S1, iterations);
    p.label("iteration");
    for (w, what) in body {
        p.tt(*w, what);
    }
    p.addi(S1, S1, -1).bne(S1, ZERO, "iteration");
    p.tt(word("STALLWAIT", &[("stall_res", 1 << 2), ("wait_res", 1 << 11)]), "STALLWAIT: the packer waits for the vector engine");
    p.tt(word("SETADCXY", &[("bit_mask", 0xb), ("cnt_set_mask", 4)]), "SETADCXY packer");
    p.tt(word("SETADCZW", &[("bit_mask", 0xf), ("cnt_set_mask", 4)]), "SETADCZW packer");
    p.tt(word("SETADCXX", &[("x_end2", 15), ("cnt_set_mask", 4)]), "SETADCXX packer: a row");
    for k in 0..rows / 4 {
        p.tt(word("PACR", &[("last", u32::from(k == rows / 4 - 1))]), &format!("PACR rows {}..{}", 4 * k, 4 * k + 3));
    }
    p.tt(word("STALLWAIT", &[("stall_res", 1 << 5), ("wait_res", 1 << 3)]), "STALLWAIT: until the packer is done");
    p.tt(word("SETGPR", &[("reg_index_16b", 126), ("payload_sig_sel", 1)]), "SETGPR: register 63 = 1");
    p.label("wait").lw(T1, 0xfc, T2).beq(T1, ZERO, "wait");
    p.li(T3, DONE).addi(A0, ZERO, 1).sw(A0, 0, T3);
    p.label("idle").j("idle");
    p.finish()
}

/// The rows program for one gate over `iterations` × [`LANES`] pairs a tile:
/// the gate's entries are in the program.
pub fn rows_program(m: &[i32; 8], c: Constants, iterations: u32) -> (Vec<u32>, String) {
    program(&rows_iteration(m, c), iterations, 8, 256)
}

/// The probe program: each pair's `a·b` and `fx(a, b)` (a in column 0, b in
/// column 1); the results are four blocks of [`PAIRS`] words (the product's
/// low and high words, then `fx`'s).
pub fn probe_program(c: Constants, iterations: u32) -> (Vec<u32>, String) {
    program(&probe_iteration(c), iterations, 4, 128)
}

/// Clocks a pass of `iterations` takes, about: under nine hundred instructions
/// an iteration, and the unpacking and packing around them.
pub fn pass_clocks(iterations: u32) -> u64 {
    2_000 + 900 * u64::from(iterations)
}

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

    #[test]
    fn the_programs_fit_and_their_mutants_differ() {
        let c = Constants { round_bit: 29, shift: 30, last_add: false };
        let m = [1, -2, 3, -4, 5, -6, 7, -8];
        let (w, text) = rows_program(&m, c, 16);
        assert!(4 * w.len() < DONE as usize && text.contains("SFPMUL24"));
        for other in [Constants { round_bit: 28, ..c }, Constants { shift: 29, ..c }, Constants { last_add: true, ..c }] {
            assert_ne!(rows_program(&m, other, 16).1, text);
        }
        assert_ne!(rows_program(&[0; 8], c, 16).0, w, "the entries are in the program");
        assert_ne!(rows_program(&m, c, 1).0, w, "so is the pass's length");
        let (pw, _) = probe_program(c, 16);
        assert!(4 * pw.len() < DONE as usize);
        assert!(IN + 4 * 4 * PAIRS as u32 <= OUT && OUT + 8 * 4 * PAIRS as u32 <= 0x8_0000);
    }
}