1use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};
2
3pub fn loop_carry_smoke() -> Program {
4 Program::wrapped(
5 vec![
6 BufferDecl::storage("haystack", 0, BufferAccess::ReadOnly, DataType::U32).with_count(8),
7 BufferDecl::storage("out_count", 1, BufferAccess::ReadWrite, DataType::U32)
8 .with_count(1),
9 ],
10 [1, 1, 1],
11 vec![Node::if_then(
12 Expr::eq(Expr::InvocationId { axis: 0 }, Expr::u32(0)),
13 vec![
14 Node::let_bind("cursor", Expr::u32(0)),
15 Node::let_bind("tok_idx", Expr::u32(0)),
16 Node::loop_for(
17 "iter",
18 Expr::u32(0),
19 Expr::u32(8),
20 vec![Node::if_then(
21 Expr::lt(Expr::var("cursor"), Expr::u32(8)),
22 vec![
23 Node::let_bind("byte", Expr::load("haystack", Expr::var("cursor"))),
24 Node::let_bind("emit", Expr::u32(0)),
25 Node::if_then(
26 Expr::and(
27 Expr::eq(Expr::var("emit"), Expr::u32(0)),
28 Expr::ne(Expr::var("byte"), Expr::u32(0)),
29 ),
30 vec![Node::assign("emit", Expr::u32(1))],
31 ),
32 Node::if_then(
33 Expr::eq(Expr::var("emit"), Expr::u32(1)),
34 vec![Node::assign(
35 "tok_idx",
36 Expr::add(Expr::var("tok_idx"), Expr::u32(1)),
37 )],
38 ),
39 Node::assign("cursor", Expr::add(Expr::var("cursor"), Expr::u32(1))),
40 ],
41 )],
42 ),
43 Node::store("out_count", Expr::u32(0), Expr::var("tok_idx")),
44 ],
45 )],
46 )
47}