calyx_opt/passes/
compile_repeat.rs

1use crate::traversal::{Action, Named, VisResult, Visitor};
2use calyx_ir::structure;
3use calyx_ir::{self as ir, LibrarySignatures};
4use calyx_utils::math::bits_needed_for;
5
6use ir::{build_assignments, guard};
7/// Compiles [`ir::Invoke`](calyx_ir::Invoke) statements into an [`ir::Enable`](calyx_ir::Enable)
8/// that runs the invoked component.
9#[derive(Default)]
10pub struct CompileRepeat;
11
12impl Named for CompileRepeat {
13    fn name() -> &'static str {
14        "compile-repeat"
15    }
16
17    fn description() -> &'static str {
18        "Rewrites repeat statements into while loops"
19    }
20}
21
22impl Visitor for CompileRepeat {
23    fn finish_repeat(
24        &mut self,
25        s: &mut ir::Repeat,
26        comp: &mut ir::Component,
27        ctx: &LibrarySignatures,
28        _comps: &[ir::Component],
29    ) -> VisResult {
30        let num_repeats = s.num_repeats;
31        if num_repeats == 0 {
32            // 0 repeats is the same thing as an empty control statement.
33            Ok(Action::change(ir::Control::empty()))
34        } else if num_repeats == 1 {
35            // 1 repeat means we can just replace the repeat stmt with the body.
36            Ok(Action::change(s.body.take_control()))
37        } else {
38            // Otherwise we should build a while loop.
39            let mut builder = ir::Builder::new(comp, ctx);
40            let idx_size = bits_needed_for(num_repeats + 1);
41            structure!( builder;
42                // holds the idx of the iteration
43                let idx = prim std_reg(idx_size);
44                // cond_reg.out will be condition port for the while loop
45                let cond_reg = prim std_reg(1);
46                let adder = prim std_add(idx_size);
47                let lt = prim std_lt(idx_size);
48                let const_zero = constant(0, idx_size);
49                let const_one = constant(1, idx_size);
50                let num_repeats = constant(num_repeats, idx_size);
51                let signal_on = constant(1,1);
52            );
53            // regs_done is `cond_reg.done & idx.done`
54            let regs_done: ir::Guard<ir::Nothing> =
55                guard!(cond_reg["done"] & idx["done"]);
56            // init_group sets cond_reg to 1 and idx to 0
57            let init_group = builder.add_group("init_repeat");
58            let init_assigns = build_assignments!(
59              builder;
60              // initial state for idx and cond_reg;
61              idx["write_en"] = ? signal_on["out"];
62              idx["in"] = ? const_zero["out"];
63              cond_reg["write_en"] = ? signal_on["out"];
64              cond_reg["in"] = ? signal_on["out"];
65              init_group["done"] = regs_done ? signal_on["out"];
66            )
67            .to_vec();
68            init_group.borrow_mut().assignments = init_assigns;
69            init_group
70                .borrow_mut()
71                .attributes
72                .insert(ir::NumAttr::Promotable, 1);
73            // incr_group:
74            // 1) writes results of idx + 1 into idx (i.e., increments idx)
75            // 2) writes the result of (idx + 1 < num_repeats) into cond_reg,
76            let incr_group = builder.add_group("incr_repeat");
77            let idx_incr_assigns = build_assignments!(
78              builder;
79              adder["left"] = ? idx["out"];
80              adder["right"] = ? const_one["out"];
81              lt["left"] = ? adder["out"];
82              lt["right"] = ? num_repeats["out"];
83              cond_reg["write_en"] = ? signal_on["out"];
84              cond_reg["in"] = ? lt["out"];
85              idx["write_en"] = ? signal_on["out"];
86              idx["in"] = ? adder["out"];
87              incr_group["done"] = regs_done ? signal_on["out"];
88            )
89            .to_vec();
90            incr_group.borrow_mut().assignments = idx_incr_assigns;
91            incr_group
92                .borrow_mut()
93                .attributes
94                .insert(ir::NumAttr::Promotable, 1);
95            // create control:
96            // init_group; while cond_reg.out {repeat_body; incr_group;}
97            let while_body = ir::Control::seq(vec![
98                s.body.take_control(),
99                ir::Control::enable(incr_group),
100            ]);
101            let while_loop = ir::Control::while_(
102                cond_reg.borrow().get("out"),
103                None,
104                Box::new(while_body),
105            );
106            let while_seq = ir::Control::Seq(ir::Seq {
107                stmts: vec![ir::Control::enable(init_group), while_loop],
108                attributes: std::mem::take(&mut s.attributes),
109            });
110            Ok(Action::change(while_seq))
111        }
112    }
113}
calyx_opt/passes/compile_repeat.rs

calyx_opt/passes/
compile_repeat.rs
