calyx_opt/passes/
wire_inliner.rs1use 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)]
9pub 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 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 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 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 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 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 groups.iter().for_each(|gr| {
193 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 fsms.iter().for_each(|fsm| {
207 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 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 let comp_fsms_collection = comp.get_fsms_mut();
245 fsms.drain(..).for_each(|fsm| {
246 comp_fsms_collection.add(fsm);
247 });
248
249 Ok(Action::change(ir::Control::empty()))
251 }
252}