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

use std::sync::Arc;

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

/// Build one forward interval-propagation step.
#[must_use]
pub fn range_propagate(defs_in: &str, edges_in: &str, ranges_out: &str) -> Program {
    range_propagate_with_slots(defs_in, edges_in, ranges_out, 4, 8)
}

/// Checked version of [`range_propagate`] for callers that want uniform
/// fallible construction.
pub fn try_range_propagate(
    defs_in: &str,
    edges_in: &str,
    ranges_out: &str,
) -> Result<Program, String> {
    try_range_propagate_with_count(defs_in, edges_in, ranges_out, 4)
}

/// Checked version that takes `var_count` explicitly.
pub fn try_range_propagate_with_count(
    defs_in: &str,
    edges_in: &str,
    ranges_out: &str,
    var_count: u32,
) -> Result<Program, String> {
    let domain = crate::graph_layout::LinearDomain::new(var_count);
    let slots = domain.scalar_slots("range_propagate", 2)?;
    Ok(range_propagate_with_slots(
        defs_in, edges_in, ranges_out, var_count, slots,
    ))
}

/// Version that takes `var_count` explicitly.
#[must_use]
#[cfg(any(test, feature = "legacy-infallible"))]
pub fn range_propagate_with_count(
    defs_in: &str,
    edges_in: &str,
    ranges_out: &str,
    var_count: u32,
) -> Program {
    try_range_propagate_with_count(defs_in, edges_in, ranges_out, var_count)
        .expect("range_propagate var_count overflows scalar slots in legacy infallible caller")
}

fn range_propagate_with_slots(
    defs_in: &str,
    edges_in: &str,
    ranges_out: &str,
    var_count: u32,
    slots: u32,
) -> Program {
    let v = Expr::InvocationId { axis: 0 };

    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("lo", Expr::load(defs_in, Expr::var("lo_idx"))),
        Node::let_bind("hi", Expr::load(defs_in, Expr::var("hi_idx"))),
        Node::let_bind("t_lo", Expr::load(edges_in, Expr::var("lo_idx"))),
        Node::let_bind("t_hi", Expr::load(edges_in, Expr::var("hi_idx"))),
        Node::store(
            ranges_out,
            Expr::var("lo_idx"),
            Expr::add(Expr::var("lo"), Expr::var("t_lo")),
        ),
        Node::store(
            ranges_out,
            Expr::var("hi_idx"),
            Expr::add(Expr::var("hi"), Expr::var("t_hi")),
        ),
    ];

    let buffers = vec![
        BufferDecl::storage(defs_in, 0, BufferAccess::ReadOnly, DataType::U32).with_count(slots),
        BufferDecl::storage(edges_in, 1, BufferAccess::ReadOnly, DataType::U32).with_count(slots),
        BufferDecl::storage(ranges_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,
            )]),
        }],
    )
}