weirflow 0.1.0

GPU-first dataflow analysis primitives for Vyre and Santh compiler pipelines.
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
121
122
123
124

//! `range_check`  -  interval-VSA bound check. Composes the per-node
//! interval table with a comparison against a known bound.
//!
//! Per node `n`, write 1 to `out[n]` iff `interval(n).hi < bound[n]`.
//! Dataflow consumers use this for CWE-787 / CWE-190 to assert "size is
//! bounded by destination buffer length."

use std::sync::Arc;
use vyre::ir::Program;
use vyre_foundation::ir::model::expr::Ident;
use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr, Node};

pub(crate) const OP_ID: &str = "weir::range_check";

/// Build a range-check Program. Inputs:
/// - `interval_hi`: per-node u32 buffer of interval upper bounds.
/// - `bound`:       per-node u32 buffer of allowed maxima.
/// - `out`:         per-node bitset; bit `n` set iff `interval_hi[n] < bound[n]`.
#[must_use]
pub fn range_check(node_count: u32, interval_hi: &str, bound: &str, out: &str) -> Program {
    let domain = crate::graph_layout::LinearDomain::new(node_count);
    let t = Expr::InvocationId { axis: 0 };
    let body = vec![Node::store(
        out,
        t.clone(),
        Expr::select(
            Expr::lt(
                Expr::load(interval_hi, t.clone()),
                Expr::load(bound, t.clone()),
            ),
            Expr::u32(1),
            Expr::u32(0),
        ),
    )];
    Program::wrapped(
        vec![
            BufferDecl::storage(interval_hi, 0, BufferAccess::ReadOnly, DataType::U32)
                .with_count(domain.element_count()),
            BufferDecl::storage(bound, 1, BufferAccess::ReadOnly, DataType::U32)
                .with_count(domain.element_count()),
            BufferDecl::storage(out, 2, BufferAccess::ReadWrite, DataType::U32)
                .with_count(domain.element_count()),
        ],
        [256, 1, 1],
        vec![Node::Region {
            generator: Ident::from(OP_ID),
            source_region: None,
            body: Arc::new(vec![Node::if_then(
                Expr::lt(t.clone(), Expr::u32(node_count)),
                body,
            )]),
        }],
    )
}

/// reference oracle.
#[must_use]
#[cfg(any(test, feature = "cpu-parity"))]
#[deprecated(
    note = "reference oracle only; production code must dispatch the Weir Program on a concrete GPU backend or use weir::oracle for parity evidence"
)]
pub(crate) fn cpu_ref(interval_hi: &[u32], bound: &[u32]) -> Vec<u32> {
    let n = interval_hi.len().min(bound.len());
    let mut out = crate::staging_reserve::reserved_vec(n, "range-check CPU oracle output")
        .unwrap_or_else(|error| {
            panic!("range-check CPU oracle output reservation failed: {error}")
        });
    out.extend((0..n).map(|i| u32::from(interval_hi[i] < bound[i])));
    out
}

/// Soundness marker for [`range_check`].
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RangeCheck;
impl super::soundness::SoundnessTagged for RangeCheck {
    fn soundness(&self) -> super::soundness::Soundness {
        super::soundness::Soundness::Exact
    }
}

inventory::submit! {
    vyre_harness::OpEntry::new(
        OP_ID,
        || range_check(4, "hi", "bound", "out"),
        Some(|| {
            let u32s = crate::dispatch_decode::pack_u32;
            vec![vec![
                u32s(&[5, 10, 11, 0]),
                u32s(&[6, 10, 9, 1]),
                u32s(&[0; 4]),
            ]]
        }),
        Some(|| {
            let u32s = crate::dispatch_decode::pack_u32;
            vec![vec![u32s(&[1, 0, 0, 1])]]
        }),
    )
}

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

    #[test]
    fn within_bound_returns_one() {
        assert_eq!(cpu_ref(&[5, 10, 15], &[10, 20, 30]), vec![1, 1, 1]);
    }

    #[test]
    fn at_bound_returns_zero() {
        assert_eq!(cpu_ref(&[10, 20], &[10, 20]), vec![0, 0]);
    }

    #[test]
    fn over_bound_returns_zero() {
        assert_eq!(cpu_ref(&[100], &[10]), vec![0]);
    }

    #[test]
    fn mixed_bounds() {
        assert_eq!(cpu_ref(&[1, 100, 5], &[10, 10, 100]), vec![1, 0, 1]);
    }
}