calyx_opt/passes_experimental/
simplify_if_comb.rs

1use crate::traversal::{Action, ConstructVisitor, Named, VisResult, Visitor};
2use calyx_ir::{self as ir, LibrarySignatures, Rewriter, rewriter::RewriteMap};
3use calyx_utils::CalyxResult;
4use itertools::Itertools;
5
6/// Transforms `if`s with `comb` groups into `if`s with the condition being computed
7/// via continuous assignments.
8///
9/// The cell used for the condition (and any other cells on the LHS of an assignment
10/// in the comb group) will be cloned. Therefore, it is important that dead-cell-removal
11/// and dead-group-removal is ran after this pass.
12///
13/// # Example
14/// ```
15/// comb cond_comb {
16///     lt.left = x.out;
17///     lt.right = 32'd10;
18/// }
19/// ...
20/// if lt.out with cond_comb {
21///     then_group;
22/// }
23/// ```
24/// into
25///
26/// ```
27/// // continuous assignments
28/// lt0.left = x.out;
29/// lt0.right = 32'd10;
30/// ...
31/// if lt0.out {
32///     then_group;
33/// }
34/// ```
35pub struct SimplifyIfComb {}
36
37impl Named for SimplifyIfComb {
38    fn name() -> &'static str {
39        "simplify-if-comb"
40    }
41
42    fn description() -> &'static str {
43        "Transform `if` with comb groups into `if` with continuous assignments when there is only one enable in the `then` block and there is no `else` block."
44    }
45
46    fn opts() -> Vec<crate::traversal::PassOpt> {
47        vec![]
48    }
49}
50
51impl ConstructVisitor for SimplifyIfComb {
52    fn from(_ctx: &ir::Context) -> CalyxResult<Self>
53    where
54        Self: Sized + Named,
55    {
56        Ok(SimplifyIfComb {})
57    }
58
59    fn clear_data(&mut self) {}
60}
61
62impl Visitor for SimplifyIfComb {
63    fn finish_if(
64        &mut self,
65        s: &mut calyx_ir::If,
66        comp: &mut calyx_ir::Component,
67        sigs: &LibrarySignatures,
68        _comps: &[calyx_ir::Component],
69    ) -> VisResult {
70        let mut builder = ir::Builder::new(comp, sigs);
71        let mut rewrite_map = RewriteMap::new();
72        let mut new_continuous_assignments = Vec::new();
73
74        let Some(cond_group_ref) = &s.cond else {
75            return Ok(Action::Continue);
76        };
77        // create new cell for all cells in LHS of cond group assignments
78        for cond_group_asgn in &cond_group_ref.borrow().assignments {
79            if let calyx_ir::PortParent::Cell(c) =
80                &cond_group_asgn.dst.borrow().parent
81            {
82                let c_ref = c.upgrade();
83                let c_name = c_ref.borrow().name();
84                if !rewrite_map.contains_key(&c_name)
85                    && let ir::CellType::Primitive {
86                        name,
87                        param_binding,
88                        ..
89                    } = &c_ref.borrow().prototype
90                {
91                    let new_cell = builder.add_primitive(
92                        c_name,
93                        *name,
94                        &param_binding.iter().map(|(_, v)| *v).collect_vec(),
95                    );
96                    rewrite_map.insert(c_name, new_cell);
97                }
98            }
99        }
100
101        // move all assignments in cond group to continuous and rewrite cells
102        for cond_group_asgn in &cond_group_ref.borrow().assignments {
103            let new_asgn = cond_group_asgn.clone();
104            new_continuous_assignments.push(new_asgn);
105        }
106
107        let rewrite = Rewriter {
108            cell_map: rewrite_map,
109            ..Default::default()
110        };
111        for asgn in new_continuous_assignments.iter_mut() {
112            rewrite.rewrite_assign(asgn);
113            comp.continuous_assignments.push(asgn.clone());
114        }
115
116        // NOTE: Technically no assignments within `s.tbranch` and `s.fbranch` should reference
117        // any cells that are being set in the original `s.comb` so we probably don't need
118        // to rewrite the whole `if`, but doing so in case someone refers to cells set in `s.comb`
119        let mut new_if = calyx_ir::Control::if_(
120            s.port.clone(),
121            None,
122            Box::new(s.tbranch.take_control()),
123            Box::new(s.fbranch.take_control()),
124        );
125        rewrite.rewrite_control(&mut new_if);
126        Ok(Action::Change(Box::new(new_if)))
127    }
128}
calyx_opt/passes_experimental/simplify_if_comb.rs

calyx_opt/passes_experimental/
simplify_if_comb.rs
