vyre-libs 0.7.2

//! Surface-presence tests for vyre-libs items that graduated from the
//! gap ledger.
//! Each `#[test]` asserts the named function exists; the test fails at
//! compile time if a migrated item disappears.

#![allow(deprecated)]
#![cfg(all(
    feature = "math-linalg",
    feature = "nn-norm",
    feature = "nn-attention",
    feature = "matching-substring",
    feature = "matching-dfa",
    feature = "crypto-blake3",
))]

// Gaps L-1..L-6 and L-8 CLOSED  -  the Cat-A compositions ship. These
// tests now assert their presence as structural smoke tests. Byte-level
// oracle checks live in `tests/cat_a_conform.rs`.

/// Split a program's buffers into the caller-visible bindings and the
/// workgroup-local scratch the kernel allocates for itself.
///
/// A raw `buffers().len()` conflates the two, and several of these contracts
/// were written before their builder reduced through workgroup memory. When
/// `layer_norm` gained three scratch tiles, a count of 2 became a count of 5
/// and the assertion read as "the surface changed" when the surface had not
/// moved at all. Splitting the two makes the contract say what it means: the
/// caller still binds exactly these buffers, and the scratch is an
/// implementation detail whose names are pinned separately.
fn split_buffers(program: &vyre::ir::Program) -> (Vec<String>, Vec<String>) {
    let mut bindings = Vec::new();
    let mut scratch = Vec::new();
    for buffer in program.buffers() {
        if buffer.kind() == vyre::ir::MemoryKind::Shared {
            scratch.push(buffer.name().to_string());
        } else {
            bindings.push(buffer.name().to_string());
        }
    }
    (bindings, scratch)
}

#[test]
fn contract_nn_softmax_exists() {
    use vyre_libs::nn::attention::softmax;
    let p = softmax("x", "y", 64);
    let (bindings, scratch) = split_buffers(&p);
    assert_eq!(bindings, vec!["x", "y"]);
    assert_eq!(scratch, vec!["softmax_scratch", "softmax_max"]);
}

#[test]
fn contract_nn_layer_norm_exists() {
    use vyre_libs::nn::norm::layer_norm;
    let p = layer_norm("x", "out", 64, 1e-5);
    // Two caller bindings, unchanged. The three scratch tiles are the
    // workgroup reduction that computes mean and variance in one pass.
    let (bindings, scratch) = split_buffers(&p);
    assert_eq!(bindings, vec!["x", "out"]);
    assert_eq!(
        scratch,
        vec!["ln_sum_scratch", "ln_sq_scratch", "ln_stats"],
        "layer_norm reduces through workgroup memory; these are its scratch tiles"
    );
}

#[test]
fn contract_nn_attention_exists() {
    use vyre_libs::nn::attention::attention;
    // s = 64 is above the direct-unroll threshold, so this builds the tiled
    // kernel. See `contract_nn_attention_small_shapes_unroll_directly` for the
    // other side of that branch.
    let p = attention("q", "k", "v", "out", 64, 8);
    let (bindings, scratch) = split_buffers(&p);
    assert_eq!(bindings, vec!["q", "k", "v", "out"]);
    assert_eq!(scratch, vec!["attention_scratch"]);
}

/// Small attention shapes take the fully unrolled path and allocate no scratch.
///
/// `attention` branches on shape: at `s <= 8 && d <= 16` it emits straight-line
/// code with no workgroup reduction, so it has four buffers and a `[1, 1, 1]`
/// workgroup. That branch was invisible to these contracts, which is how
/// `attention_default_is_query_row_parallel` came to assert the tiled
/// workgroup size against a shape that never reaches the tiled kernel.
#[test]
fn contract_nn_attention_small_shapes_unroll_directly() {
    use vyre_libs::nn::attention::attention;
    let p = attention("q", "k", "v", "out", 8, 4);
    let (bindings, scratch) = split_buffers(&p);
    assert_eq!(bindings, vec!["q", "k", "v", "out"]);
    assert!(
        scratch.is_empty(),
        "the unrolled path performs no workgroup reduction, so it needs no scratch: {scratch:?}"
    );
    assert_eq!(
        p.workgroup_size(),
        [1, 1, 1],
        "straight-line code runs one invocation, not a cooperative tile"
    );
}

#[test]
fn contract_matching_aho_corasick_exists() {
    use vyre_libs::scan::aho_corasick;
    let p = aho_corasick("haystack", "transitions", "accept", "matches", 16, 8);
    assert_eq!(p.buffers().len(), 4);
}

#[test]
fn contract_crypto_blake3_exists() {
    use vyre_libs::hash::blake3_compress;
    let p = blake3_compress("chaining_in", "message", "params", "chaining_out");
    assert_eq!(p.buffers().len(), 4);
}

// Substring search: substring_search performs a real byte-by-byte
// equality via `load(haystack, i+k) == load(needle, k)`. A regressed
// constant-true predicate would make every byte position match; this
// test combines (a) a structural assertion that the Load/Load Eq
// pair exists in the inner k-loop with (b) an end-to-end execution
// through the reference interpreter to prove the bitmap is correct
// on a canonical input. Region execution is wired into the
// interpreter so the full path is exercised.
#[test]
fn contract_substring_real_byte_compare() {
    use vyre::ir::{Expr, Node};
    use vyre_libs::scan::substring_search;

    let program = substring_search("haystack", "needle", "matches", 5, 2);

    fn contains_load_load_eq(nodes: &[Node]) -> bool {
        nodes.iter().any(node_contains)
    }
    fn node_contains(node: &Node) -> bool {
        match node {
            Node::Block(children) | Node::Loop { body: children, .. } => {
                contains_load_load_eq(children)
            }
            Node::If {
                then,
                otherwise,
                cond,
            } => {
                expr_contains(cond)
                    || contains_load_load_eq(then)
                    || contains_load_load_eq(otherwise)
            }
            Node::Let { value, .. } | Node::Assign { value, .. } => expr_contains(value),
            Node::Region { body, .. } => contains_load_load_eq(body),
            _ => false,
        }
    }
    fn expr_contains(expr: &Expr) -> bool {
        use Expr::*;
        match expr {
            BinOp { op, left, right } => {
                matches!(op, vyre::ir::BinOp::Eq)
                    && matches!(left.as_ref(), Load { .. })
                    && matches!(right.as_ref(), Load { .. })
                    || expr_contains(left)
                    || expr_contains(right)
            }
            Select {
                cond,
                true_val,
                false_val,
            } => expr_contains(cond) || expr_contains(true_val) || expr_contains(false_val),
            _ => false,
        }
    }

    assert!(
        contains_load_load_eq(program.entry()),
        "substring_search must contain a Load-vs-Load equality inside its k-loop; \
         if this regresses, the inner compare has been replaced with a constant predicate (LAW 1)"
    );

    // End-to-end execution: "hello" contains "lo" at byte offset 3.
    use vyre_reference::value::Value;
    let haystack_bytes: Vec<u8> = "hello"
        .bytes()
        .flat_map(|b| u32::from(b).to_le_bytes())
        .chain(std::iter::repeat_n(0u8, 12))
        .collect();
    let needle_bytes: Vec<u8> = "lo"
        .bytes()
        .flat_map(|b| u32::from(b).to_le_bytes())
        .collect();
    let matches_bytes = vec![0u8; 5 * 4];
    let inputs = [
        Value::from(haystack_bytes),
        Value::from(needle_bytes),
        Value::from(matches_bytes),
    ];
    let outputs =
        vyre_reference::reference_eval(&program, &inputs).expect("execute substring_search");
    assert_eq!(outputs.len(), 1, "only matches buffer is ReadWrite");
    let raw = outputs[0].to_bytes();
    let words: Vec<u32> = raw
        .chunks_exact(4)
        .map(|c| u32::from_le_bytes(c.try_into().unwrap()))
        .collect();
    assert_eq!(words.len(), 5, "matches buffer has 5 u32 slots");
    assert_eq!(words[3], 1, "match at byte offset 3");
    for (i, w) in words.iter().enumerate() {
        if i == 3 {
            continue;
        }
        assert_eq!(*w, 0, "no match at byte offset {i}, got {w}");
    }
}

#[test]
fn contract_math_matmul_tiled_exists() {
    use vyre_libs::math::linalg::matmul_tiled;
    let p = matmul_tiled("a", "b", "c", 64, 64, 64, 16);
    // Three caller bindings, unchanged, plus the two cooperative tiles the
    // kernel stages `a` and `b` through. The old count of 3 predated tiling.
    let (bindings, scratch) = split_buffers(&p);
    assert_eq!(bindings, vec!["a", "b", "c"]);
    assert_eq!(
        scratch.len(),
        2,
        "matmul_tiled stages one workgroup tile per input: {scratch:?}"
    );
    // matmul_tiled flattens the 2D workgroup into a 1D dispatch so that
    // InvocationId { axis: 0 } is a unique linear index.
    assert_eq!(p.workgroup_size(), [256, 1, 1]);
}