open-standards/wai/wai-quantum/examples/stablehlo_gate.rs

//! Writes the gate as a portable tensor program, a batch of inputs, and the
//! exact outputs a compiler's build of it must give:
//!
//! ```text
//! cargo run --example stablehlo_gate --features quantum_fabric -- <dir>
//! ```
//!
//! In `<dir>`: `gate.mlir` (conforming) and `gate_m1.mlir` … `gate_m3.mlir`
//! (mutants, each of which must fail); for `rows`, the inputs `e0.npy` …
//! `e7.npy`, `a0r.npy` … `a1i.npy` and the expected `r0.bin` … `r3.bin`; for
//! `probe`, `pa.npy`, `pb.npy` and the expected `prod.bin`, `fx.bin`. Each
//! expected file is the raw little-endian 64-bit values, `4096xi64`, so a
//! runtime compares them under the module's own (signless) element type.
use wai_quantum::quantum_fabric::lanes;
use wai_quantum::quantum_fabric::stablehlo::{CONFORMING, Constants, limb_module, module, npy_i32, npy_scalar_i64, reference_coefficients, reference_rows};
use wai_quantum::quantum_fabric::tensor_engine;

/// Pairs per batch.
const N: usize = 4096;

fn raw(v: &[i64]) -> Vec<u8> {
    v.iter().flat_map(|x| x.to_le_bytes()).collect()
}

fn write(dir: &std::path::Path, name: &str, bytes: &[u8]) {
    std::fs::write(dir.join(name), bytes).unwrap_or_else(|e| panic!("{name}: {e}"));
}

fn main() {
    let dir = std::path::PathBuf::from(std::env::args().nth(1).expect("an output directory"));
    std::fs::create_dir_all(&dir).expect("the directory");
    write(&dir, "gate.mlir", module(N, CONFORMING).as_bytes());
    for (name, c) in [
        ("gate_m1.mlir", Constants { round: 1 << 28, ..CONFORMING }),
        ("gate_m2.mlir", Constants { shift: 29, ..CONFORMING }),
        ("gate_m3.mlir", Constants { last_add: true, ..CONFORMING }),
    ] {
        write(&dir, name, module(N, c).as_bytes());
    }
    // A gate whose entries reach the 32-bit edges, and states that do too.
    let m: [i64; 8] = [i64::from(i32::MAX), i64::from(i32::MIN), 1 << 30, -(1 << 30), 0x1234_5678, -0x0765_4321, 1, -1];
    let mut s = 0x5eed_u64;
    let mut next = move || {
        s = s.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mut z = s;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        (z ^ (z >> 31)) as u32 as i32
    };
    let edges = [0, 1, -1, i32::MAX, i32::MIN, 1 << 29, -(1 << 29) - 1];
    let cols: [Vec<i32>; 4] = std::array::from_fn(|k| (0..N).map(|i| if (i + k) % 9 == 0 { edges[(i / 9 + k) % edges.len()] } else { next() }).collect());
    for (k, e) in m.iter().enumerate() {
        write(&dir, &format!("e{k}.npy"), &npy_scalar_i64(*e));
    }
    for (name, c) in ["a0r", "a0i", "a1r", "a1i"].iter().zip(&cols) {
        write(&dir, &format!("{name}.npy"), &npy_i32(c));
    }
    for (r, v) in reference_rows(&m, &cols, CONFORMING).iter().enumerate() {
        write(&dir, &format!("r{r}.bin"), &raw(v));
    }
    // The probe: every pair of the edge set, then random pairs, to N.
    let mut pairs = lanes::probe_pairs(N, 0x0cda);
    pairs.truncate(N);
    let (pa, pb): (Vec<i32>, Vec<i32>) = pairs.iter().copied().unzip();
    write(&dir, "pa.npy", &npy_i32(&pa));
    write(&dir, "pb.npy", &npy_i32(&pb));
    write(&dir, "prod.bin", &raw(&pairs.iter().map(|&(a, b)| i64::from(a) * i64::from(b)).collect::<Vec<i64>>()));
    write(&dir, "fx.bin", &raw(&pairs.iter().map(|&(a, b)| (i64::from(a) * i64::from(b) + (1 << 29)) >> 30).collect::<Vec<i64>>()));
    // The products on a matrix unit: the gate's limb weights against the
    // batch's limbs, and the exact coefficients. Then nine-bit limbs, which
    // bf16 cannot hold: their exact coefficients are not what any bf16 unit
    // gives, so a run that matches them is not comparing.
    write(&dir, "limbs.mlir", limb_module(N).as_bytes());
    let cases: Vec<[i32; 4]> = (0..N).map(|q| [cols[0][q], cols[1][q], cols[2][q], cols[3][q]]).collect();
    let m32: [i32; 8] = m.map(|e| e as i32);
    for (bits, suffix) in [(8, ""), (9, "9")] {
        let int_limbs = |x: i32| tensor_engine::limbs(x, bits);
        let mut w = vec![0i32; tensor_engine::K * tensor_engine::M];
        for (p, &(k, c)) in tensor_engine::TERMS.iter().enumerate() {
            for j in 0..4 {
                for (i, &l) in int_limbs(m32[k]).iter().enumerate() {
                    w[(c * 4 + j) * tensor_engine::M + p * tensor_engine::COEFFICIENTS + i + j] = l;
                }
            }
        }
        let mut a = vec![0i32; tensor_engine::K * N];
        for (q, case) in cases.iter().enumerate() {
            for (c, &x) in case.iter().enumerate() {
                for (j, &l) in int_limbs(x).iter().enumerate() {
                    a[(c * 4 + j) * N + q] = l;
                }
            }
        }
        let bf = |v: &[i32]| v.iter().flat_map(|&x| tensor_engine::bf16(x).to_le_bytes()).collect::<Vec<u8>>();
        write(&dir, &format!("w{suffix}.bin"), &bf(&w));
        write(&dir, &format!("a{suffix}.bin"), &bf(&a));
        write(&dir, &format!("c{suffix}.bin"), &reference_coefficients(&w, &a, N).iter().flat_map(|x| x.to_le_bytes()).collect::<Vec<u8>>());
    }
    println!("wrote the module, its mutants, the limb module, and a batch of {N} pairs to {}", dir.display());
}