calyx_opt/passes/
wire_inliner.rs

1use crate::traversal::{Action, Named, VisResult, Visitor};
2use calyx_ir as ir;
3use ir::{LibrarySignatures, build_assignments, guard, structure};
4use ir::{Nothing, RRC};
5use itertools::Itertools;
6use std::{collections::HashMap, rc::Rc};
7
8#[derive(Default)]
9/// Alternate hole inliner that removes groups and group holes by instantiating
10/// wires that hold the value for each signal.
11pub struct WireInliner;
12
13type HoleMapping = HashMap<ir::Id, (RRC<ir::Cell>, RRC<ir::Cell>)>;
14
15impl Named for WireInliner {
16    fn name() -> &'static str {
17        "wire-inliner"
18    }
19
20    fn description() -> &'static str {
21        "inlines holes using wires"
22    }
23}
24
25fn rewrite(map: &HoleMapping, port: &RRC<ir::Port>) -> Option<RRC<ir::Cell>> {
26    let port_cell = &port.borrow().parent;
27    if let ir::PortParent::Group(g) = port_cell {
28        let (go, done) = &map[&g.upgrade().borrow().name()];
29        let cell = if port.borrow().name == "go" { go } else { done };
30        Some(Rc::clone(cell))
31    } else if let ir::PortParent::FSM(f) = port_cell {
32        let fsm_name = &f.upgrade().borrow().name();
33        let (start, done) = &map[fsm_name];
34        let cell = if port.borrow().name == "start" {
35            start
36        } else {
37            done
38        };
39        Some(Rc::clone(cell))
40    } else {
41        None
42    }
43}
44
45fn rewrite_assign(map: &HoleMapping, assign: &mut ir::Assignment<Nothing>) {
46    if let Some(cell) = rewrite(map, &assign.dst) {
47        assign.dst = cell.borrow().get("in");
48    }
49    if let Some(cell) = rewrite(map, &assign.src) {
50        assign.src = cell.borrow().get("out");
51    }
52    assign.guard.for_each(&mut |port| {
53        rewrite(map, &port)
54            .map(|cell| ir::Guard::port(cell.borrow().get("out")))
55    });
56}
57
58fn rewrite_guard(map: &HoleMapping, guard: &mut ir::Guard<Nothing>) {
59    match guard {
60        ir::Guard::True | ir::Guard::Info(_) => (),
61        // update the port of a port guard to read from appropriate
62        // group wire, if it is dependent on a group's port in the first place
63        ir::Guard::Port(p) => {
64            if let Some(cell) = rewrite(map, p) {
65                let _ = std::mem::replace(p, cell.borrow().get("out"));
66            }
67        }
68        // update the ports of a port-comparison guard to read from appropriate
69        // group wire, if these are dependent on groups' ports in the first place
70        ir::Guard::CompOp(_, p1, p2) => {
71            if let Some(cell) = rewrite(map, p1) {
72                let _ = std::mem::replace(p1, cell.borrow().get("out"));
73            }
74            if let Some(cell) = rewrite(map, p2) {
75                let _ = std::mem::replace(p2, cell.borrow().get("out"));
76            }
77        }
78        ir::Guard::Not(b) => {
79            rewrite_guard(map, &mut *b);
80        }
81        ir::Guard::And(b1, b2) | ir::Guard::Or(b1, b2) => {
82            rewrite_guard(map, &mut *b1);
83            rewrite_guard(map, &mut *b2);
84        }
85    }
86}
87
88impl Visitor for WireInliner {
89    fn start(
90        &mut self,
91        comp: &mut ir::Component,
92        sigs: &LibrarySignatures,
93        _comps: &[ir::Component],
94    ) -> VisResult {
95        let control_ref = Rc::clone(&comp.control);
96        let control = control_ref.borrow();
97
98        match &*control {
99            ir::Control::Enable(..) | ir::Control::FSMEnable(..) => {
100                let this = Rc::clone(&comp.signature);
101                let mut builder = ir::Builder::new(comp, sigs);
102
103                let this_go_port = this
104                    .borrow()
105                    .find_unique_with_attr(ir::NumAttr::Go)?
106                    .unwrap();
107
108                structure!(builder;
109                    let one = constant(1, 1);
110                );
111                let assigns = match &*control {
112                    ir::Control::Enable(en) => {
113                        let group = &en.group;
114                        let this_done_port = this
115                            .borrow()
116                            .find_unique_with_attr(ir::NumAttr::Done)?
117                            .unwrap();
118
119                        let group_done =
120                            guard!(group[this_done_port.borrow().name]);
121                        build_assignments!(builder;
122                            group["go"] = ? this[this_go_port.borrow().name];
123                            this["done"] = group_done ? one["out"];
124                        )
125                    }
126                    ir::Control::FSMEnable(fsm_en) => {
127                        let fsm = &fsm_en.fsm;
128                        let fsm_done = guard!(fsm["done"]);
129                        build_assignments!(builder;
130                            fsm["start"] = ? this[this_go_port.borrow().name];
131                            this["done"] = fsm_done ? one["out"];
132                        )
133                    }
134                    _ => unreachable!(),
135                };
136                comp.continuous_assignments.extend(assigns);
137            }
138            ir::Control::Empty(_) => {}
139            _ => {
140                return Err(calyx_utils::Error::malformed_control(format!(
141                    "{}: Structure has more than one group",
142                    Self::name()
143                )));
144            }
145        }
146
147        // assume static groups is empty
148        let groups = comp.get_groups_mut().drain().collect_vec();
149        let mut fsms = comp.get_fsms_mut().drain().collect_vec();
150        let mut builder = ir::Builder::new(comp, sigs);
151        // for each group, instantiate wires to hold its `go` and `done` signals.
152        let mut hole_map: HoleMapping = groups
153            .iter()
154            .map(|gr| {
155                let name = gr.borrow().name();
156                let go = builder.add_primitive(
157                    format!("{name}_go"),
158                    "std_wire",
159                    &[1],
160                );
161                let done = builder.add_primitive(
162                    format!("{name}_done"),
163                    "std_wire",
164                    &[1],
165                );
166                (name, (go, done))
167            })
168            .collect();
169
170        // add wires to hold `start` and `done` signals for each FSM, which we
171        // compile differently
172        let fsm_interface = fsms
173            .iter()
174            .map(|fsm| {
175                let name = fsm.borrow().name();
176                let start = builder.add_primitive(
177                    format!("{name}_start"),
178                    "std_wire",
179                    &[1],
180                );
181                let done = builder.add_primitive(
182                    format!("{name}_done"),
183                    "std_wire",
184                    &[1],
185                );
186                (name, (start, done))
187            })
188            .collect_vec();
189        hole_map.extend(fsm_interface);
190
191        // Rewrite all assignments first
192        groups.iter().for_each(|gr| {
193            // Detach assignment from the group first because the rewrite
194            // method will try to borrow the underlying group.
195            // This means that the group's borrow_mut cannot be active when
196            // rewrite_assign is called.
197            let mut assigns =
198                gr.borrow_mut().assignments.drain(..).collect_vec();
199            assigns
200                .iter_mut()
201                .for_each(|asgn| rewrite_assign(&hole_map, asgn));
202            gr.borrow_mut().assignments = assigns;
203        });
204
205        // rewrite the contents of each fsm to use instantiated std_wire prims
206        fsms.iter().for_each(|fsm| {
207            // rewrite all assignments at each state within every fsm
208            let mut assigns =
209                fsm.borrow_mut().assignments.drain(..).collect_vec();
210            for assigns_at_state in assigns.iter_mut() {
211                for asgn in assigns_at_state.iter_mut() {
212                    rewrite_assign(&hole_map, asgn);
213                }
214            }
215            fsm.borrow_mut().assignments = assigns;
216
217            // rewrite all guards in transitions that depend on group port values
218            let mut transitions =
219                fsm.borrow_mut().transitions.drain(..).collect_vec();
220            for trans_at_state in transitions.iter_mut() {
221                if let ir::Transition::Conditional(cond_trans_at_state) =
222                    trans_at_state
223                {
224                    for (cond_trans, _) in cond_trans_at_state {
225                        rewrite_guard(&hole_map, cond_trans)
226                    }
227                }
228            }
229            fsm.borrow_mut().transitions = transitions;
230        });
231
232        comp.continuous_assignments
233            .iter_mut()
234            .for_each(|assign| rewrite_assign(&hole_map, assign));
235
236        let mut group_assigns = groups
237            .into_iter()
238            .flat_map(|g| g.borrow_mut().assignments.drain(..).collect_vec())
239            .collect_vec();
240
241        comp.continuous_assignments.append(&mut group_assigns);
242
243        // reinstate fsms into the component
244        let comp_fsms_collection = comp.get_fsms_mut();
245        fsms.drain(..).for_each(|fsm| {
246            comp_fsms_collection.add(fsm);
247        });
248
249        // remove group from control
250        Ok(Action::change(ir::Control::empty()))
251    }
252}