vyre-libs 0.6.3

vyre Category A library ecosystem - pure-IR compositions over vyre-ops hardware primitives
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
114
115
116
117
118
119
120

//! Overflow-guard regression tests for Cat-A composition builders.
//!
//! When a caller passes shape parameters whose product overflows u32,
//! the builders MUST panic with an actionable message rather than
//! silently produce an under-sized buffer + OOB memory access at
//! runtime. The guards use `checked_mul`; these tests prove every
//! exit path fires at the correct boundary.

#![cfg(all(feature = "math-linalg", feature = "nn-attention"))]

use vyre_libs::math::{matmul, matmul_tiled, Matmul, MatmulTiled};
use vyre_libs::nn::{attention, Attention};
use vyre_libs::tensor_ref::TensorRef;

/// Near-overflow bound: any `a*b > u32::MAX` must panic.
const PIVOT: u32 = 1u32 << 16; // 65_536; 65_536 * 65_536 = u32::MAX + 1

#[test]
fn matmul_m_times_k_overflow_returns_error() {
    let error = Matmul::new(
        TensorRef::u32_2d("a", PIVOT, PIVOT),
        TensorRef::u32_2d("b", PIVOT, 1),
        TensorRef::u32_2d("out", PIVOT, 1),
    )
    .build()
    .expect_err("Fix: m*k overflow must return an error contract");
    assert!(error.to_string().contains("element-count overflows u32"));
}

#[test]
fn matmul_k_times_n_overflow_returns_error() {
    // m=1 keeps m*k safe; k*n is the overflow.
    let error = Matmul::new(
        TensorRef::u32_2d("a", 1, PIVOT),
        TensorRef::u32_2d("b", PIVOT, PIVOT),
        TensorRef::u32_2d("out", 1, PIVOT),
    )
    .build()
    .expect_err("Fix: k*n overflow must return an error contract");
    assert!(error.to_string().contains("element-count overflows u32"));
}

#[test]
fn matmul_m_times_n_overflow_returns_error() {
    // m*k safe (m=1), k*n safe (k=1), m*n overflows.
    let error = Matmul::new(
        TensorRef::u32_2d("a", PIVOT, 1),
        TensorRef::u32_2d("b", 1, PIVOT),
        TensorRef::u32_2d("out", PIVOT, PIVOT),
    )
    .build()
    .expect_err("Fix: m*n overflow must return an error contract");
    assert!(error.to_string().contains("element-count overflows u32"));
}

#[test]
fn matmul_tiled_m_times_k_overflow_returns_error() {
    let error = MatmulTiled::new(
        TensorRef::u32_2d("a", PIVOT, PIVOT),
        TensorRef::u32_2d("b", PIVOT, 1),
        TensorRef::u32_2d("out", PIVOT, 1),
        16,
    )
    .build()
    .expect_err("Fix: tiled m*k overflow must return an error contract");
    assert!(error.to_string().contains("element-count overflows u32"));
}

#[test]
fn matmul_tiled_k_times_n_overflow_returns_error() {
    let error = MatmulTiled::new(
        TensorRef::u32_2d("a", 1, PIVOT),
        TensorRef::u32_2d("b", PIVOT, PIVOT),
        TensorRef::u32_2d("out", 1, PIVOT),
        16,
    )
    .build()
    .expect_err("Fix: tiled k*n overflow must return an error contract");
    assert!(error.to_string().contains("element-count overflows u32"));
}

#[test]
fn matmul_tiled_m_times_n_overflow_returns_error() {
    let error = MatmulTiled::new(
        TensorRef::u32_2d("a", PIVOT, 1),
        TensorRef::u32_2d("b", 1, PIVOT),
        TensorRef::u32_2d("out", PIVOT, PIVOT),
        16,
    )
    .build()
    .expect_err("Fix: tiled m*n overflow must return an error contract");
    assert!(error.to_string().contains("element-count overflows u32"));
}

#[test]
fn attention_s_times_d_overflow_returns_error() {
    let error = Attention::new(
        TensorRef::f32_2d("q", PIVOT, PIVOT),
        TensorRef::f32_2d("k", PIVOT, PIVOT),
        TensorRef::f32_2d("v", PIVOT, PIVOT),
        TensorRef::f32_2d("out", PIVOT, PIVOT),
    )
    .build()
    .expect_err("Fix: attention s*d overflow must return an error contract");
    assert!(error.to_string().contains("element-count overflows u32"));
}

#[test]
fn reasonable_shapes_build_without_panic() {
    // Shapes under the pivot must continue to work end-to-end.
    let _ = matmul("a", "b", "out", 64, 64, 64);
    let _ = matmul_tiled("a", "b", "out", 64, 64, 64, 16);
    let _ = attention("q", "k", "v", "out", 128, 64);
}

#[test]
fn just_under_pivot_succeeds_for_one_operand() {
    // 65_535 * 65_535 = 4_294_836_225 < u32::MAX. Still safe.
    let _ = matmul("a", "b", "out", 65_535, 1, 65_535);
}