calyx_opt/passes_experimental/
fsm_annotator.rs

1use crate::traversal::{
2    Action, ConstructVisitor, Named, ParseVal, PassOpt, VisResult, Visitor,
3};
4use calyx_ir::{self as ir};
5
6/// From the perspective of a parent control node, the annotation `@NUM_STATES(n)`
7/// on one of its child nodes means that the parent should allocate `n` states
8/// in order to implement the child's schedule.
9const NUM_STATES: ir::Attribute =
10    ir::Attribute::Internal(ir::InternalAttr::NUM_STATES);
11
12/// From the perspective of a parent control node, if the attribute `@ACYCLIC` is
13/// present on one of its children, then the following implication is guaranteed:
14///
15/// If the child's FSM has an @UNROLL or @INLINE attribute, then
16/// the states implementing the child's schedule should have the property of
17/// having one state for every cycle.
18///
19/// If the attribute `@ACYCLIC` is not present on a child node, then this simply
20/// means that the states implementing the child's control has a backedge.
21///
22/// For now, it is undefined behavior for one child to have both `@ACYCLIC` and
23/// `@OFFLOAD` attributes attached.
24const ACYCLIC: ir::Attribute =
25    ir::Attribute::Internal(ir::InternalAttr::ACYCLIC);
26
27fn is_acyclic(ctrl: &ir::StaticControl) -> bool {
28    ctrl.get_attribute(ACYCLIC).is_some()
29}
30
31fn get_num_states(ctrl: &ir::Control) -> u64 {
32    ctrl.get_attribute(NUM_STATES).unwrap()
33}
34
35fn get_num_states_static(ctrl: &ir::StaticControl) -> u64 {
36    ctrl.get_attribute(NUM_STATES).unwrap()
37}
38
39/// Given the attributes field of a control node, insert the annotations derived
40/// in this pass into the `control {...}` section of a Calyx program.
41fn transfer_attributes(
42    attributes_map: &mut ir::Attributes,
43    FSMImplementation {
44        num_states,
45        acyclic,
46        attr,
47    }: &FSMImplementation,
48) {
49    attributes_map.insert(NUM_STATES, *num_states);
50    let ir_attribute = match attr {
51        ControlNodeAnnotation::Inline(inline_impl) => match inline_impl {
52            InlineImplementation::StandardInline => ir::InternalAttr::INLINE,
53            InlineImplementation::Unroll => ir::InternalAttr::UNROLL,
54        },
55        ControlNodeAnnotation::Offload => ir::InternalAttr::OFFLOAD,
56    };
57    attributes_map.insert(ir::Attribute::Internal(ir_attribute), 1);
58    if *acyclic {
59        attributes_map.insert(ACYCLIC, 1);
60    }
61}
62
63/// Encodes the possible FSM implementations of a control node. From the perspective
64/// of a parent FSM, the states implementing each of its children can
65/// either be brought into the parent FSM (i.e. inlined) or can be outsourced
66/// to another FSM (offloaded).
67enum ControlNodeAnnotation {
68    Inline(InlineImplementation),
69    Offload,
70}
71
72/// Encodes the possible ways for a control node's FSM to be inlined into that of its parent.
73/// Here, `Unroll` can only be used on ir::Repeat or ir::StaticRepeat nodes. In the event
74/// that `StandardInline` is used on one of these nodes, then, within the parent
75/// FSM, there will be a backedge from the bottom to the top of the repeat.
76///
77/// The annotation `StandardInline` is used to indicate inlining for every other
78/// control node.
79enum InlineImplementation {
80    Unroll,
81    StandardInline,
82}
83
84/// An instance of this struct is computed for every control node. It will be used
85/// by its parent control node.
86struct FSMImplementation {
87    num_states: u64,
88    acyclic: bool,
89    attr: ControlNodeAnnotation,
90}
91
92trait FSMPolicy {
93    /// Given a control node, returns a number of states allocated for the FSM,
94    /// whether backedges exist in the FSM, and an attribute for this node.
95    fn policy(ctrl: &mut Self, child_fsm_cutoff: u64) -> FSMImplementation;
96}
97
98impl FSMPolicy for ir::FSMEnable {
99    fn policy(_: &mut Self, _: u64) -> FSMImplementation {
100        FSMImplementation {
101            num_states: 1,
102            acyclic: false,
103            attr: ControlNodeAnnotation::Offload,
104        }
105    }
106}
107
108impl FSMPolicy for ir::StaticEnable {
109    fn policy(
110        ctrl: &mut ir::StaticEnable,
111        child_fsm_cutoff: u64,
112    ) -> FSMImplementation {
113        let (num_states, acyclic) = {
114            let latency = ctrl.group.borrow().get_latency();
115            if latency < child_fsm_cutoff {
116                (latency, true)
117            } else {
118                (1, false)
119            }
120        };
121
122        FSMImplementation {
123            num_states,
124            acyclic,
125            attr: ControlNodeAnnotation::Inline(
126                InlineImplementation::StandardInline,
127            ),
128        }
129    }
130}
131
132impl FSMPolicy for ir::StaticSeq {
133    fn policy(ctrl: &mut ir::StaticSeq, _: u64) -> FSMImplementation {
134        let (num_states, acyclic) =
135            ctrl.stmts
136                .iter()
137                .fold((0, true), |(num_states, acyclic), stmt| {
138                    let stmt_latency = get_num_states_static(stmt);
139                    let stmt_acyclic = is_acyclic(stmt);
140                    (num_states + stmt_latency, acyclic && stmt_acyclic)
141                });
142
143        FSMImplementation {
144            num_states,
145            acyclic,
146            attr: ControlNodeAnnotation::Inline(
147                InlineImplementation::StandardInline,
148            ),
149        }
150    }
151}
152
153impl FSMPolicy for ir::StaticPar {
154    fn policy(ctrl: &mut ir::StaticPar, _: u64) -> FSMImplementation {
155        let (num_states, acyclic, attr) = if ctrl.stmts.iter().all(is_acyclic) {
156            (
157                ctrl.latency,
158                true,
159                ControlNodeAnnotation::Inline(
160                    InlineImplementation::StandardInline,
161                ),
162            )
163        } else {
164            (1, false, ControlNodeAnnotation::Offload)
165        };
166
167        FSMImplementation {
168            num_states,
169            acyclic,
170            attr,
171        }
172    }
173}
174
175impl FSMPolicy for ir::StaticRepeat {
176    fn policy(
177        ctrl: &mut ir::StaticRepeat,
178        child_fsm_cutoff: u64,
179    ) -> FSMImplementation {
180        let (body_num_states, body_is_acyclic) =
181            (get_num_states_static(&ctrl.body), is_acyclic(&ctrl.body));
182        let unrolled_num_states = ctrl.num_repeats * body_num_states;
183        let (num_states, acyclic, attr) =
184            if body_is_acyclic && (unrolled_num_states < child_fsm_cutoff) {
185                (
186                    unrolled_num_states,
187                    true,
188                    ControlNodeAnnotation::Inline(InlineImplementation::Unroll),
189                )
190            } else if body_num_states < child_fsm_cutoff {
191                (
192                    body_num_states,
193                    false,
194                    ControlNodeAnnotation::Inline(
195                        InlineImplementation::StandardInline,
196                    ),
197                )
198            } else {
199                (1, false, ControlNodeAnnotation::Offload)
200            };
201
202        FSMImplementation {
203            num_states,
204            acyclic,
205            attr,
206        }
207    }
208}
209
210impl FSMPolicy for ir::StaticIf {
211    fn policy(ctrl: &mut ir::StaticIf, _: u64) -> FSMImplementation {
212        let (num_states, acyclic, attr) =
213            if is_acyclic(&ctrl.tbranch) && is_acyclic(&ctrl.fbranch) {
214                (
215                    ctrl.latency,
216                    true,
217                    ControlNodeAnnotation::Inline(
218                        InlineImplementation::StandardInline,
219                    ),
220                )
221            } else {
222                (1, false, ControlNodeAnnotation::Offload)
223            };
224
225        FSMImplementation {
226            num_states,
227            acyclic,
228            attr,
229        }
230    }
231}
232
233impl FSMPolicy for ir::Enable {
234    fn policy(_ctrl: &mut ir::Enable, _: u64) -> FSMImplementation {
235        FSMImplementation {
236            num_states: 1,
237            acyclic: false,
238            attr: ControlNodeAnnotation::Inline(
239                InlineImplementation::StandardInline,
240            ),
241        }
242    }
243}
244
245impl FSMPolicy for ir::Seq {
246    fn policy(ctrl: &mut ir::Seq, _: u64) -> FSMImplementation {
247        FSMImplementation {
248            num_states: ctrl.stmts.iter().map(get_num_states).sum(),
249            acyclic: false,
250            attr: ControlNodeAnnotation::Inline(
251                InlineImplementation::StandardInline,
252            ),
253        }
254    }
255}
256
257impl FSMPolicy for ir::Par {
258    fn policy(_ctrl: &mut ir::Par, _: u64) -> FSMImplementation {
259        FSMImplementation {
260            num_states: 1,
261            acyclic: false,
262            attr: ControlNodeAnnotation::Offload,
263        }
264    }
265}
266
267impl FSMPolicy for ir::If {
268    fn policy(_ctrl: &mut ir::If, _: u64) -> FSMImplementation {
269        FSMImplementation {
270            num_states: 1,
271            acyclic: false,
272            attr: ControlNodeAnnotation::Offload,
273        }
274    }
275}
276
277impl FSMPolicy for ir::While {
278    fn policy(ctrl: &mut ir::While, _: u64) -> FSMImplementation {
279        FSMImplementation {
280            num_states: ctrl.body.get_attribute(NUM_STATES).unwrap(),
281            acyclic: false,
282            attr: ControlNodeAnnotation::Inline(
283                InlineImplementation::StandardInline,
284            ),
285        }
286    }
287}
288
289impl FSMPolicy for ir::Repeat {
290    fn policy(ctrl: &mut ir::Repeat, _: u64) -> FSMImplementation {
291        FSMImplementation {
292            num_states: ctrl.body.get_attribute(NUM_STATES).unwrap(),
293            acyclic: false,
294            attr: ControlNodeAnnotation::Inline(
295                InlineImplementation::StandardInline,
296            ),
297        }
298    }
299}
300
301pub struct FSMAnnotator {
302    child_fsm_cutoff: u64,
303}
304
305impl Named for FSMAnnotator {
306    fn name() -> &'static str {
307        "fsm-annotator"
308    }
309
310    fn description() -> &'static str {
311        "annotate a control program, determining how FSMs should be allocated"
312    }
313
314    fn opts() -> Vec<PassOpt> {
315        vec![PassOpt::new(
316            "child-fsm-cutoff",
317            "The maximum number of states a child FSM can have, before it is offloaded",
318            ParseVal::Num(300),
319            PassOpt::parse_num,
320        )]
321    }
322}
323impl ConstructVisitor for FSMAnnotator {
324    fn from(ctx: &ir::Context) -> calyx_utils::CalyxResult<Self> {
325        let opts = Self::get_opts(ctx);
326        Ok(FSMAnnotator {
327            child_fsm_cutoff: opts[&"child-fsm-cutoff"]
328                .pos_num()
329                .expect("requires non-negative OHE cutoff parameter"),
330        })
331    }
332    fn clear_data(&mut self) {}
333}
334
335impl Visitor for FSMAnnotator {
336    fn empty(
337        &mut self,
338        s: &mut ir::Empty,
339        _comp: &mut ir::Component,
340        _sigs: &ir::LibrarySignatures,
341        _comps: &[ir::Component],
342    ) -> VisResult {
343        s.attributes.insert(NUM_STATES, 0);
344        Ok(Action::Continue)
345    }
346    fn enable(
347        &mut self,
348        s: &mut ir::Enable,
349        _comp: &mut ir::Component,
350        _sigs: &ir::LibrarySignatures,
351        _comps: &[ir::Component],
352    ) -> VisResult {
353        let fsm_impl = ir::Enable::policy(s, self.child_fsm_cutoff);
354        transfer_attributes(&mut s.attributes, &fsm_impl);
355        Ok(Action::Continue)
356    }
357
358    fn static_enable(
359        &mut self,
360        s: &mut ir::StaticEnable,
361        _comp: &mut ir::Component,
362        _sigs: &ir::LibrarySignatures,
363        _comps: &[ir::Component],
364    ) -> VisResult {
365        let fsm_impl = ir::StaticEnable::policy(s, self.child_fsm_cutoff);
366        transfer_attributes(&mut s.attributes, &fsm_impl);
367        Ok(Action::Continue)
368    }
369
370    fn fsm_enable(
371        &mut self,
372        s: &mut ir::FSMEnable,
373        _comp: &mut ir::Component,
374        _sigs: &ir::LibrarySignatures,
375        _comps: &[ir::Component],
376    ) -> VisResult {
377        let fsm_impl = ir::FSMEnable::policy(s, self.child_fsm_cutoff);
378        transfer_attributes(&mut s.attributes, &fsm_impl);
379        Ok(Action::Continue)
380    }
381
382    fn finish_static_seq(
383        &mut self,
384        s: &mut ir::StaticSeq,
385        _comp: &mut ir::Component,
386        _sigs: &ir::LibrarySignatures,
387        _comps: &[ir::Component],
388    ) -> VisResult {
389        let fsm_impl = ir::StaticSeq::policy(s, self.child_fsm_cutoff);
390        transfer_attributes(&mut s.attributes, &fsm_impl);
391        Ok(Action::Continue)
392    }
393
394    fn finish_static_par(
395        &mut self,
396        s: &mut ir::StaticPar,
397        _comp: &mut ir::Component,
398        _sigs: &ir::LibrarySignatures,
399        _comps: &[ir::Component],
400    ) -> VisResult {
401        let fsm_impl = ir::StaticPar::policy(s, self.child_fsm_cutoff);
402        transfer_attributes(&mut s.attributes, &fsm_impl);
403        Ok(Action::Continue)
404    }
405
406    fn finish_static_if(
407        &mut self,
408        s: &mut ir::StaticIf,
409        _comp: &mut ir::Component,
410        _sigs: &ir::LibrarySignatures,
411        _comps: &[ir::Component],
412    ) -> VisResult {
413        let fsm_impl = ir::StaticIf::policy(s, self.child_fsm_cutoff);
414        transfer_attributes(&mut s.attributes, &fsm_impl);
415        Ok(Action::Continue)
416    }
417
418    fn finish_static_repeat(
419        &mut self,
420        s: &mut ir::StaticRepeat,
421        _comp: &mut ir::Component,
422        _sigs: &ir::LibrarySignatures,
423        _comps: &[ir::Component],
424    ) -> VisResult {
425        let fsm_impl = ir::StaticRepeat::policy(s, self.child_fsm_cutoff);
426        transfer_attributes(&mut s.attributes, &fsm_impl);
427        Ok(Action::Continue)
428    }
429
430    fn finish_seq(
431        &mut self,
432        s: &mut ir::Seq,
433        _comp: &mut ir::Component,
434        _sigs: &ir::LibrarySignatures,
435        _comps: &[ir::Component],
436    ) -> VisResult {
437        let fsm_impl = ir::Seq::policy(s, self.child_fsm_cutoff);
438        transfer_attributes(&mut s.attributes, &fsm_impl);
439        Ok(Action::Continue)
440    }
441
442    fn finish_par(
443        &mut self,
444        s: &mut ir::Par,
445        _comp: &mut ir::Component,
446        _sigs: &ir::LibrarySignatures,
447        _comps: &[ir::Component],
448    ) -> VisResult {
449        let fsm_impl = ir::Par::policy(s, self.child_fsm_cutoff);
450        transfer_attributes(&mut s.attributes, &fsm_impl);
451        Ok(Action::Continue)
452    }
453
454    fn finish_if(
455        &mut self,
456        s: &mut ir::If,
457        _comp: &mut ir::Component,
458        _sigs: &ir::LibrarySignatures,
459        _comps: &[ir::Component],
460    ) -> VisResult {
461        let fsm_impl = ir::If::policy(s, self.child_fsm_cutoff);
462        transfer_attributes(&mut s.attributes, &fsm_impl);
463        Ok(Action::Continue)
464    }
465
466    fn finish_while(
467        &mut self,
468        s: &mut ir::While,
469        _comp: &mut ir::Component,
470        _sigs: &ir::LibrarySignatures,
471        _comps: &[ir::Component],
472    ) -> VisResult {
473        let fsm_impl = ir::While::policy(s, self.child_fsm_cutoff);
474        transfer_attributes(&mut s.attributes, &fsm_impl);
475        Ok(Action::Continue)
476    }
477
478    fn finish_repeat(
479        &mut self,
480        s: &mut calyx_ir::Repeat,
481        _comp: &mut calyx_ir::Component,
482        _sigs: &calyx_ir::LibrarySignatures,
483        _comps: &[calyx_ir::Component],
484    ) -> VisResult {
485        let fsm_impl = ir::Repeat::policy(s, self.child_fsm_cutoff);
486        transfer_attributes(&mut s.attributes, &fsm_impl);
487        Ok(Action::Continue)
488    }
489}
calyx_opt/passes_experimental/fsm_annotator.rs

calyx_opt/passes_experimental/
fsm_annotator.rs
