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

use std::sync::Arc;

use super::OP_ID;
use vyre::ir::{BufferAccess, BufferDecl, DataType, Expr, Ident, Node, Program};

/// One widening step. `cfg_in` holds the previous-iteration
/// intervals `[prev_lo_v, prev_hi_v, …]`; `ranges_in` holds the
/// new iteration's intervals; `summary_out` receives the widened
/// result per variable.
#[must_use]
pub fn loop_summarize(cfg_in: &str, ranges_in: &str, summary_out: &str) -> Program {
    loop_summarize_with_slots(cfg_in, ranges_in, summary_out, 4, 8)
}

/// Checked version of [`loop_summarize_with_count`] for production GPU wrappers.
///
/// # Errors
///
/// Returns an actionable construction error when the variable domain cannot
/// fit the `lo/hi` slot layout used by the GPU program.
pub fn try_loop_summarize_with_count(
    cfg_in: &str,
    ranges_in: &str,
    summary_out: &str,
    var_count: u32,
) -> Result<Program, String> {
    let slots = var_count.checked_mul(2).ok_or_else(|| {
        format!(
            "loop_summarize var_count={var_count} overflows u32 buffer slots. Fix: shard loop summarization before building the GPU Program."
        )
    })?.max(1);
    Ok(loop_summarize_with_slots(
        cfg_in,
        ranges_in,
        summary_out,
        var_count,
        slots,
    ))
}

/// Version that takes `var_count` explicitly.
#[must_use]
#[cfg(any(test, feature = "legacy-infallible"))]
pub fn loop_summarize_with_count(
    cfg_in: &str,
    ranges_in: &str,
    summary_out: &str,
    var_count: u32,
) -> Program {
    let slots = var_count.checked_mul(2).unwrap_or_else(|| {
        panic!(
            "loop_summarize var_count={var_count} overflows u32 buffer slots. Fix: call try_loop_summarize_with_count and shard loop summarization before building the GPU Program."
        )
    }).max(1);
    loop_summarize_with_slots(cfg_in, ranges_in, summary_out, var_count, slots)
}

fn loop_summarize_with_slots(
    cfg_in: &str,
    ranges_in: &str,
    summary_out: &str,
    var_count: u32,
    slots: u32,
) -> Program {
    let v = Expr::InvocationId { axis: 0 };
    const NEG_INF: u32 = 0;
    const POS_INF: u32 = u32::MAX;

    let body = vec![
        Node::let_bind("lo_idx", Expr::mul(v.clone(), Expr::u32(2))),
        Node::let_bind("hi_idx", Expr::add(Expr::var("lo_idx"), Expr::u32(1))),
        Node::let_bind("prev_lo", Expr::load(cfg_in, Expr::var("lo_idx"))),
        Node::let_bind("prev_hi", Expr::load(cfg_in, Expr::var("hi_idx"))),
        Node::let_bind("new_lo", Expr::load(ranges_in, Expr::var("lo_idx"))),
        Node::let_bind("new_hi", Expr::load(ranges_in, Expr::var("hi_idx"))),
        // widen lo: if new_lo < prev_lo (lower bound decreasing),
        // jump to -∞ (NEG_INF); else keep prev_lo.
        Node::let_bind(
            "wide_lo",
            Expr::select(
                Expr::lt(Expr::var("new_lo"), Expr::var("prev_lo")),
                Expr::u32(NEG_INF),
                Expr::var("prev_lo"),
            ),
        ),
        // widen hi: if new_hi > prev_hi (upper bound increasing),
        // jump to +∞ (POS_INF); else keep prev_hi.
        Node::let_bind(
            "wide_hi",
            Expr::select(
                Expr::lt(Expr::var("prev_hi"), Expr::var("new_hi")),
                Expr::u32(POS_INF),
                Expr::var("prev_hi"),
            ),
        ),
        Node::store(summary_out, Expr::var("lo_idx"), Expr::var("wide_lo")),
        Node::store(summary_out, Expr::var("hi_idx"), Expr::var("wide_hi")),
    ];

    let buffers = vec![
        BufferDecl::storage(cfg_in, 0, BufferAccess::ReadOnly, DataType::U32).with_count(slots),
        BufferDecl::storage(ranges_in, 1, BufferAccess::ReadOnly, DataType::U32).with_count(slots),
        BufferDecl::storage(summary_out, 2, BufferAccess::ReadWrite, DataType::U32)
            .with_count(slots),
    ];

    Program::wrapped(
        buffers,
        [256, 1, 1],
        vec![Node::Region {
            generator: Ident::from(OP_ID),
            source_region: None,
            body: Arc::new(vec![Node::if_then(
                Expr::lt(v.clone(), Expr::u32(var_count)),
                body,
            )]),
        }],
    )
}