vyre-libs 0.7.2

vyre Category A library ecosystem - pure-IR compositions over foundation IR and primitive-owned kernels
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

//! Verifies that INT4 Cat-A wrappers retain their Tier-2.5 primitive provenance.

#![cfg(feature = "nn-linear-4bit")]

use vyre::ir::{Node, Program};
use vyre_libs::nn::quant::{
    int4_batched_matmul_f32_scaled, int4_batched_matmul_top1_f32_scaled,
    int4_batched_matvec_f32_scaled, int4_dot_f32_scaled, int4_dot_i32, int4_matvec_f32_scaled,
};

fn assert_composition_chain(program: &Program, parent_id: &str, primitive_id: &str) {
    let [Node::Region {
        generator,
        source_region: None,
        body,
    }] = program.entry()
    else {
        panic!("expected one Cat-A parent region");
    };
    assert_eq!(generator.as_ref(), parent_id);
    let [Node::Region {
        generator,
        source_region: Some(parent),
        ..
    }] = body.as_slice()
    else {
        panic!("expected one Tier-2.5 child region");
    };
    assert_eq!(generator.as_ref(), primitive_id);
    assert_eq!(parent.name, parent_id);
}

/// This test prevents the two public INT4 dot wrappers from presenting primitive IR as parent-owned implementation work.
#[test]
fn dot_wrappers_record_cat_a_parent_and_primitive_child() {
    let dot = int4_dot_i32("lhs", "rhs", "out", 8);
    assert_composition_chain(
        &dot,
        "vyre-libs::quant::int4_dot_i32",
        "vyre-primitives::math::quantized::i4x8_dot_i32",
    );

    let scaled = int4_dot_f32_scaled("lhs", "rhs", "lhs_scale", "rhs_scale", "out", 8);
    assert_composition_chain(
        &scaled,
        "vyre-libs::quant::int4_dot_f32_scaled",
        "vyre-primitives::math::quantized::i4x8_dot_f32_scaled",
    );
}

/// This test locks every matrix-vector and matrix-matrix wrapper to its canonical Tier-2.5 implementation instead of a copied body.
#[test]
fn matrix_wrappers_record_their_exact_primitive_children() {
    let matvec = int4_matvec_f32_scaled("weights", "x", "row_scales", "out", 3, 9);
    assert_composition_chain(
        &matvec,
        "vyre-libs::quant::int4_matvec_f32_scaled",
        "vyre-primitives::math::quantized::i4x8_matvec_f32_scaled",
    );

    let batched_matvec =
        int4_batched_matvec_f32_scaled("weights", "x", "row_scales", "out", 2, 3, 9);
    assert_composition_chain(
        &batched_matvec,
        "vyre-libs::quant::int4_batched_matvec_f32_scaled",
        "vyre-primitives::math::quantized::i4x8_batched_matvec_f32_scaled",
    );

    let batched_matmul = int4_batched_matmul_f32_scaled(
        "weights",
        "activations",
        "row_scales",
        "batch_scales",
        "out",
        2,
        3,
        9,
    );
    assert_composition_chain(
        &batched_matmul,
        "vyre-libs::quant::int4_batched_matmul_f32_scaled",
        "vyre-primitives::math::quantized::i4x8_batched_matmul_f32_scaled",
    );

    let top1 = int4_batched_matmul_top1_f32_scaled(
        "weights",
        "activations",
        "row_scales",
        "batch_scales",
        "out",
        2,
        3,
        9,
    );
    assert_composition_chain(
        &top1,
        "vyre-libs::quant::int4_batched_matmul_top1_f32_scaled",
        "vyre-primitives::math::quantized::i4x8_batched_matmul_top1_f32_scaled",
    );
}

/// This boundary test proves invalid zero-lane requests keep their primitive diagnostic body beneath the Cat-A wrapper rather than bypassing provenance.
#[test]
fn zero_lane_dot_keeps_composition_chain_and_trap() {
    let program = int4_dot_i32("lhs", "rhs", "out", 0);

    assert_composition_chain(
        &program,
        "vyre-libs::quant::int4_dot_i32",
        "vyre-primitives::math::quantized::i4x8_dot_i32",
    );
    assert!(program.stats().trap());
}