weir/loop_sum/
program.rs

1use std::sync::Arc;
2
3use super::OP_ID;
4use vyre::ir::{BufferAccess, BufferDecl, DataType, Expr, Ident, Node, Program};
5
6/// One widening step. `cfg_in` holds the previous-iteration
7/// intervals `[prev_lo_v, prev_hi_v, …]`; `ranges_in` holds the
8/// new iteration's intervals; `summary_out` receives the widened
9/// result per variable.
10#[must_use]
11pub fn loop_summarize(cfg_in: &str, ranges_in: &str, summary_out: &str) -> Program {
12    loop_summarize_with_slots(cfg_in, ranges_in, summary_out, 4, 8)
13}
14
15/// Checked version of [`loop_summarize_with_count`] for production GPU wrappers.
16///
17/// # Errors
18///
19/// Returns an actionable construction error when the variable domain cannot
20/// fit the `lo/hi` slot layout used by the GPU program.
21pub fn try_loop_summarize_with_count(
22    cfg_in: &str,
23    ranges_in: &str,
24    summary_out: &str,
25    var_count: u32,
26) -> Result<Program, String> {
27    let slots = var_count.checked_mul(2).ok_or_else(|| {
28        format!(
29            "loop_summarize var_count={var_count} overflows u32 buffer slots. Fix: shard loop summarization before building the GPU Program."
30        )
31    })?.max(1);
32    Ok(loop_summarize_with_slots(
33        cfg_in,
34        ranges_in,
35        summary_out,
36        var_count,
37        slots,
38    ))
39}
40
41/// Version that takes `var_count` explicitly.
42#[must_use]
43#[cfg(any(test, feature = "legacy-infallible"))]
44pub fn loop_summarize_with_count(
45    cfg_in: &str,
46    ranges_in: &str,
47    summary_out: &str,
48    var_count: u32,
49) -> Program {
50    let slots = var_count.checked_mul(2).unwrap_or_else(|| {
51        panic!(
52            "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."
53        )
54    }).max(1);
55    loop_summarize_with_slots(cfg_in, ranges_in, summary_out, var_count, slots)
56}
57
58fn loop_summarize_with_slots(
59    cfg_in: &str,
60    ranges_in: &str,
61    summary_out: &str,
62    var_count: u32,
63    slots: u32,
64) -> Program {
65    let v = Expr::InvocationId { axis: 0 };
66    const NEG_INF: u32 = 0;
67    const POS_INF: u32 = u32::MAX;
68
69    let body = vec![
70        Node::let_bind("lo_idx", Expr::mul(v.clone(), Expr::u32(2))),
71        Node::let_bind("hi_idx", Expr::add(Expr::var("lo_idx"), Expr::u32(1))),
72        Node::let_bind("prev_lo", Expr::load(cfg_in, Expr::var("lo_idx"))),
73        Node::let_bind("prev_hi", Expr::load(cfg_in, Expr::var("hi_idx"))),
74        Node::let_bind("new_lo", Expr::load(ranges_in, Expr::var("lo_idx"))),
75        Node::let_bind("new_hi", Expr::load(ranges_in, Expr::var("hi_idx"))),
76        // widen lo: if new_lo < prev_lo (lower bound decreasing),
77        // jump to -∞ (NEG_INF); else keep prev_lo.
78        Node::let_bind(
79            "wide_lo",
80            Expr::select(
81                Expr::lt(Expr::var("new_lo"), Expr::var("prev_lo")),
82                Expr::u32(NEG_INF),
83                Expr::var("prev_lo"),
84            ),
85        ),
86        // widen hi: if new_hi > prev_hi (upper bound increasing),
87        // jump to +∞ (POS_INF); else keep prev_hi.
88        Node::let_bind(
89            "wide_hi",
90            Expr::select(
91                Expr::lt(Expr::var("prev_hi"), Expr::var("new_hi")),
92                Expr::u32(POS_INF),
93                Expr::var("prev_hi"),
94            ),
95        ),
96        Node::store(summary_out, Expr::var("lo_idx"), Expr::var("wide_lo")),
97        Node::store(summary_out, Expr::var("hi_idx"), Expr::var("wide_hi")),
98    ];
99
100    let buffers = vec![
101        BufferDecl::storage(cfg_in, 0, BufferAccess::ReadOnly, DataType::U32).with_count(slots),
102        BufferDecl::storage(ranges_in, 1, BufferAccess::ReadOnly, DataType::U32).with_count(slots),
103        BufferDecl::storage(summary_out, 2, BufferAccess::ReadWrite, DataType::U32)
104            .with_count(slots),
105    ];
106
107    Program::wrapped(
108        buffers,
109        [256, 1, 1],
110        vec![Node::Region {
111            generator: Ident::from(OP_ID),
112            source_region: None,
113            body: Arc::new(vec![Node::if_then(
114                Expr::lt(v.clone(), Expr::u32(var_count)),
115                body,
116            )]),
117        }],
118    )
119}
weir/loop_sum/program.rs

weir/loop_sum/
program.rs
