calyx_opt/passes_experimental/fsm_builder.rs
1use crate::analysis::{IncompleteTransition, StaticSchedule};
2use crate::traversal::{Action, ConstructVisitor, Named, Visitor};
3use calyx_ir::{self as ir, GetAttributes};
4use calyx_utils::CalyxResult;
5const ACYCLIC: ir::Attribute =
6 ir::Attribute::Internal(ir::InternalAttr::ACYCLIC);
7const UNROLL: ir::Attribute = ir::Attribute::Internal(ir::InternalAttr::UNROLL);
8const OFFLOAD: ir::Attribute =
9 ir::Attribute::Internal(ir::InternalAttr::OFFLOAD);
10const INLINE: ir::Attribute = ir::Attribute::Internal(ir::InternalAttr::INLINE);
11
12pub struct FSMBuilder {}
13
14impl Named for FSMBuilder {
15 fn name() -> &'static str {
16 "fsm-builder"
17 }
18 fn description() -> &'static str {
19 "translate control into structure using medium-sized explicit FSMs"
20 }
21}
22
23impl ConstructVisitor for FSMBuilder {
24 fn from(_ctx: &ir::Context) -> CalyxResult<Self> {
25 Ok(FSMBuilder {})
26 }
27 fn clear_data(&mut self) {}
28}
29
30pub struct Component {
31 non_promoted_static_component: Option<bool>,
32 static_control_component: bool,
33 // In the future we'll want to incorporate dynamic components.
34}
35
36// Helper functions to get attributes for each part of the control.
37/// Gets the `@ACYCLIC` attribute
38fn is_acyclic<T: GetAttributes>(control: &T) -> bool {
39 matches!(control.get_attributes().get(ACYCLIC), Some(1))
40}
41/// Gets the `@UNROLL` attribute
42fn is_unroll<T: GetAttributes>(control: &T) -> bool {
43 matches!(control.get_attributes().get(UNROLL), Some(1))
44}
45/// Gets the `@OFFLOAD` attribute
46fn is_offload<T: GetAttributes>(control: &T) -> bool {
47 matches!(control.get_attributes().get(OFFLOAD), Some(1))
48}
49/// Gets the `@INLINE` attribute
50fn is_inline<T: GetAttributes>(control: &T) -> bool {
51 matches!(control.get_attributes().get(INLINE), Some(1))
52}
53
54// A `StaticSchedule` is an abstract representation of fsms and maps out transitions, states, and assignments.
55// This implmentation includes functions to build up static schedules and transform them to `ir::RRC::FSM`s.
56impl StaticSchedule<'_, '_> {
57 /// Provided a static control node, calling the `build_abstract` method on an empty `StaticSchedule`
58 /// `sch` will build out the `latency` and `state2assigns` fields of `sch`, in
59 /// preparation to replace the `StaticControl` node with an instance of `ir::FSM`.
60 /// Every static assignment collected into `state2assigns` will have its existing guard
61 /// "anded" with `guard`. The `looped_once_guard` is used to encode the "doneness" of a FSM.
62 fn build_abstract(
63 &mut self,
64 scon: &ir::StaticControl,
65 guard: ir::Guard<ir::Nothing>,
66 mut transitions_to_curr: Vec<IncompleteTransition>,
67 looped_once_guard: Option<ir::Guard<ir::Nothing>>,
68 ) -> (Vec<IncompleteTransition>, Option<ir::Guard<ir::Nothing>>) {
69 match scon {
70 ir::StaticControl::Empty(_) => (transitions_to_curr, None),
71 ir::StaticControl::Enable(sen) => {
72 if is_acyclic(sen) {
73 // The `@ACYCLIC` attribute requires that one state is allocated per cycle in a static enable.
74 // For all parts of the FSM that want to transition to this enable,
75 // register their transitions in self.state2trans.
76 self.register_transitions(
77 self.state,
78 &mut transitions_to_curr,
79 guard.clone(),
80 );
81
82 sen.group.borrow().assignments.iter().for_each(|sassign| {
83 sassign
84 .guard
85 .compute_live_states(sen.group.borrow().latency)
86 .into_iter()
87 .for_each(|offset| {
88 // convert the static assignment to a normal one
89 let mut assign: ir::Assignment<ir::Nothing> =
90 ir::Assignment::from(sassign.clone());
91 // "and" the assignment's guard with argument guard
92 assign.and_guard(guard.clone());
93 // add this assignment to the list of assignments
94 // that are supposed to be valid at this state
95 self.state2assigns
96 .entry(self.state + offset)
97 .and_modify(|other_assigns| {
98 other_assigns.push(assign.clone())
99 })
100 .or_insert(vec![assign]);
101 })
102 });
103 // On an acyclic annotated node, we allocate N states to make N cycles elapse.
104 self.state += sen.group.borrow().latency;
105 // Don't know where to transition next; let the parent that called
106 // `build_abstract` deal with registering the transition
107 // from the state(s) we just built.
108 (
109 vec![IncompleteTransition::new(
110 self.state - 1,
111 ir::Guard::True,
112 )],
113 looped_once_guard,
114 )
115 } else {
116 // In the absence of `@ACYCLIC`, the node must contain cycles,
117 // or have children that contain cycles; We'll run this placeholder code that
118 // creates one state for now.
119 self.register_transitions(
120 self.state,
121 &mut transitions_to_curr,
122 guard.clone(),
123 );
124
125 let final_state_guard =
126 self.leave_one_state_condition(guard, sen);
127
128 self.state += 1;
129 (
130 vec![IncompleteTransition::new(
131 self.state - 1,
132 final_state_guard,
133 )],
134 None,
135 )
136 }
137 }
138 ir::StaticControl::Seq(sseq) => (
139 // For now, we'll automatically inline static seqs.
140 // Added support for more annotations will be coming in the future.
141 sseq.stmts.iter().fold(
142 transitions_to_curr,
143 |transitions_to_this_stmt, stmt| {
144 self.build_abstract(
145 stmt,
146 guard.clone(),
147 transitions_to_this_stmt,
148 looped_once_guard.clone(),
149 )
150 .0
151 },
152 ),
153 None,
154 ),
155 ir::StaticControl::Repeat(srep) => {
156 // Matching for the `@ACYCLIC` attribute coming soon.
157 if is_unroll(srep) {
158 // In the encounter of a `@UNROLL` attribute, we'll want to create a state for each child.
159 (
160 (0..srep.num_repeats).fold(
161 transitions_to_curr,
162 |transitions_to_this_body, _| {
163 self.build_abstract(
164 &srep.body,
165 guard.clone(),
166 transitions_to_this_body,
167 looped_once_guard.clone(),
168 )
169 .0
170 },
171 ),
172 None,
173 )
174 } else if is_offload(srep) {
175 todo!()
176 } else if is_inline(srep) {
177 // In the case of inline, we'll want to assign as many states as children of this loop,
178 // but create a register to count the amount of times looped, and assign
179 // backward edges from the end to the beginning.
180 todo!()
181 } else {
182 todo!()
183 }
184 }
185 ir::StaticControl::If(_sif) => {
186 todo!()
187 }
188 ir::StaticControl::Par(_spar) => {
189 todo!()
190 }
191 ir::StaticControl::Invoke(_) => {
192 unreachable!(
193 "`build_abstract` encountered a `static_invoke` node. \
194 Should have been compiled away."
195 )
196 }
197 }
198 }
199 /// Returns the FSM implementing the given control node, as well as the builder
200 /// object from which it was built.
201 fn fsm_build(
202 &mut self,
203 control: &ir::StaticControl,
204 build_component_type: Component, // need to get better type name. Some(True) means non-promoted-static-component. False means promoted/static island. Otherwise it's a
205 ) -> ir::RRC<ir::FSM> {
206 let signal_on = self.builder.add_constant(1, 1);
207
208 // Declare the FSM
209 let fsm = self.builder.add_fsm("fsm");
210
211 let (mut remaining_assignments, additional_looped_once_guard) =
212 self.build_abstract(control, ir::Guard::True, vec![], None);
213
214 // add loopback transitions to first state
215 self.register_transitions(
216 0,
217 &mut remaining_assignments,
218 ir::Guard::True,
219 );
220
221 let (mut assignments, transitions, state2wires) =
222 self.build_fsm_pieces(ir::RRC::clone(&fsm));
223
224 // We'll build the fsm different based off of what kind of component this node is.
225 match build_component_type {
226 Component {
227 non_promoted_static_component: Some(true),
228 static_control_component: true,
229 } => {
230 // If the component is static by design, there will be exactly one
231 // FSM allocated to it. We will get rid of the FSMEnable node from the
232 // control in this case, so we need to manually add fsm[start] = comp[go]
233 // because wire-inliner will not get to it.
234
235 // (We get rid of the FSMEnable node because the FSM will not have a
236 // DONE state, and hence no way to terminate the control. )
237 let assign_fsm_start = self.builder.build_assignment(
238 fsm.borrow().get("start"),
239 self.builder
240 .component
241 .signature
242 .borrow()
243 .find_unique_with_attr(ir::NumAttr::Go)
244 .unwrap()
245 .unwrap(),
246 ir::Guard::True,
247 );
248 self.builder
249 .add_continuous_assignments(vec![assign_fsm_start]);
250 }
251 Component {
252 non_promoted_static_component: Some(false),
253 static_control_component: true,
254 } => {
255 // In this case, the component is either a promoted static component
256 // or the control is a static island that needs to handshake with its
257 // surrounding dynamic context. In either event, we want to assign
258 // fsm[done] to maintain the dynamic interface. We'll do this in state 0:
259
260 // register to store whether the FSM has been run exactly one time when
261 // we return to state 0
262 let looped_once: ir::RRC<ir::Cell> =
263 self.builder.add_primitive("looped_once", "std_reg", &[1]);
264
265 looped_once
266 .borrow_mut()
267 .add_attribute(ir::BoolAttr::FSMControl, 1);
268
269 let (assign_looped_once, assign_looped_once_we, fsm_done) = (
270 self.builder.build_assignment(
271 looped_once.borrow().get("in"),
272 signal_on.borrow().get("out"),
273 match additional_looped_once_guard {
274 None => ir::guard!(fsm["start"]),
275 Some(g) => ir::guard!(fsm["start"]).and(g),
276 },
277 ),
278 self.builder.build_assignment(
279 looped_once.borrow().get("write_en"),
280 signal_on.borrow().get("out"),
281 ir::Guard::True,
282 ),
283 self.builder.build_assignment(
284 fsm.borrow().get("done"),
285 looped_once.borrow().get("out"),
286 ir::Guard::True,
287 ),
288 );
289
290 assignments.first_mut().unwrap().extend(vec![
291 assign_looped_once,
292 assign_looped_once_we,
293 fsm_done,
294 ]);
295 }
296 Component {
297 non_promoted_static_component: None,
298 static_control_component: true,
299 } => {
300 // Do nothing because we want to build a subset of static control component.
301 // Think ifs, repeats, pars, which don't rely on doneness.
302 }
303 Component {
304 non_promoted_static_component: _, // This branch doesn't matter in a dynamic component.
305 static_control_component: false,
306 } => {
307 todo!("Dynamic component!")
308 }
309 }
310
311 // Build up the fsm here and return.
312
313 // For test cases, we want to maintain a reliable order!
314 let mut state_assigns: Vec<_> = self.state2assigns.drain().collect();
315 state_assigns.sort_by_key(|(state, _)| *state);
316
317 // Build up the fsm here and return.
318 self.builder.add_continuous_assignments(
319 state_assigns
320 .into_iter()
321 .flat_map(|(state, mut assigns)| {
322 assigns.iter_mut().for_each(|assign| {
323 assign.and_guard(ir::Guard::port(
324 state2wires
325 .get(state as usize)
326 .unwrap()
327 .borrow()
328 .get("out"),
329 ));
330 });
331 assigns
332 })
333 .collect(),
334 );
335
336 // Instantiate the FSM with the assignments and transitions we built.
337 fsm.borrow_mut().extend_fsm(assignments, transitions);
338 fsm
339 }
340}
341
342impl Visitor for FSMBuilder {
343 /// `finish_static_control` is called once, at the very end of traversing the control tree,
344 /// when all child nodes have been traversed. We traverse the static control node from parent to
345 /// child, and recurse inward to inline children.
346 fn finish_static_control(
347 &mut self,
348 scon: &mut calyx_ir::StaticControl,
349 comp: &mut calyx_ir::Component,
350 sigs: &calyx_ir::LibrarySignatures,
351 _comps: &[calyx_ir::Component],
352 ) -> crate::traversal::VisResult {
353 let non_promoted_static_component = comp.is_static()
354 && !(comp
355 .attributes
356 .has(ir::Attribute::Bool(ir::BoolAttr::Promoted)));
357 // Implementation for single static enable components and static seqs for now.
358 let mut builder = ir::Builder::new(comp, sigs);
359
360 let mut ssch = StaticSchedule::from(&mut builder);
361
362 Ok(Action::change(ir::Control::fsm_enable(ssch.fsm_build(
363 scon,
364 Component {
365 non_promoted_static_component: Some(
366 non_promoted_static_component,
367 ),
368 static_control_component: true,
369 },
370 ))))
371 }
372}