weir/range/
program.rs

1use std::sync::Arc;
2
3use super::OP_ID;
4use vyre::ir::{BufferAccess, BufferDecl, DataType, Expr, Ident, Node, Program};
5
6/// Build one forward interval-propagation step.
7#[must_use]
8pub fn range_propagate(defs_in: &str, edges_in: &str, ranges_out: &str) -> Program {
9    range_propagate_with_slots(defs_in, edges_in, ranges_out, 4, 8)
10}
11
12/// Checked version of [`range_propagate`] for callers that want uniform
13/// fallible construction.
14pub fn try_range_propagate(
15    defs_in: &str,
16    edges_in: &str,
17    ranges_out: &str,
18) -> Result<Program, String> {
19    try_range_propagate_with_count(defs_in, edges_in, ranges_out, 4)
20}
21
22/// Checked version that takes `var_count` explicitly.
23pub fn try_range_propagate_with_count(
24    defs_in: &str,
25    edges_in: &str,
26    ranges_out: &str,
27    var_count: u32,
28) -> Result<Program, String> {
29    let domain = crate::graph_layout::LinearDomain::new(var_count);
30    let slots = domain.scalar_slots("range_propagate", 2)?;
31    Ok(range_propagate_with_slots(
32        defs_in, edges_in, ranges_out, var_count, slots,
33    ))
34}
35
36/// Version that takes `var_count` explicitly.
37#[must_use]
38#[cfg(any(test, feature = "legacy-infallible"))]
39pub fn range_propagate_with_count(
40    defs_in: &str,
41    edges_in: &str,
42    ranges_out: &str,
43    var_count: u32,
44) -> Program {
45    try_range_propagate_with_count(defs_in, edges_in, ranges_out, var_count)
46        .expect("range_propagate var_count overflows scalar slots in legacy infallible caller")
47}
48
49fn range_propagate_with_slots(
50    defs_in: &str,
51    edges_in: &str,
52    ranges_out: &str,
53    var_count: u32,
54    slots: u32,
55) -> Program {
56    let v = Expr::InvocationId { axis: 0 };
57
58    let body = vec![
59        Node::let_bind("lo_idx", Expr::mul(v.clone(), Expr::u32(2))),
60        Node::let_bind("hi_idx", Expr::add(Expr::var("lo_idx"), Expr::u32(1))),
61        Node::let_bind("lo", Expr::load(defs_in, Expr::var("lo_idx"))),
62        Node::let_bind("hi", Expr::load(defs_in, Expr::var("hi_idx"))),
63        Node::let_bind("t_lo", Expr::load(edges_in, Expr::var("lo_idx"))),
64        Node::let_bind("t_hi", Expr::load(edges_in, Expr::var("hi_idx"))),
65        Node::store(
66            ranges_out,
67            Expr::var("lo_idx"),
68            Expr::add(Expr::var("lo"), Expr::var("t_lo")),
69        ),
70        Node::store(
71            ranges_out,
72            Expr::var("hi_idx"),
73            Expr::add(Expr::var("hi"), Expr::var("t_hi")),
74        ),
75    ];
76
77    let buffers = vec![
78        BufferDecl::storage(defs_in, 0, BufferAccess::ReadOnly, DataType::U32).with_count(slots),
79        BufferDecl::storage(edges_in, 1, BufferAccess::ReadOnly, DataType::U32).with_count(slots),
80        BufferDecl::storage(ranges_out, 2, BufferAccess::ReadWrite, DataType::U32)
81            .with_count(slots),
82    ];
83
84    Program::wrapped(
85        buffers,
86        [256, 1, 1],
87        vec![Node::Region {
88            generator: Ident::from(OP_ID),
89            source_region: None,
90            body: Arc::new(vec![Node::if_then(
91                Expr::lt(v.clone(), Expr::u32(var_count)),
92                body,
93            )]),
94        }],
95    )
96}
weir/range/program.rs

weir/range/
program.rs
