calyx_opt/analysis/
fsm_construction.rs

1use calyx_ir::{self as ir, build_assignments, guard};
2use calyx_utils::math::bits_needed_for;
3use core::ops::Not;
4use itertools::Itertools;
5use std::collections::HashMap;
6
7type FSMPieces = (
8    Vec<Vec<ir::Assignment<ir::Nothing>>>,
9    Vec<ir::Transition>,
10    Vec<ir::RRC<ir::Cell>>,
11);
12
13/// Represents an FSM transition that doesn't yet have a destination state.
14#[derive(Clone)]
15pub struct IncompleteTransition {
16    source: u64,
17    guard: ir::Guard<ir::Nothing>,
18}
19
20impl IncompleteTransition {
21    pub fn new(source: u64, guard: ir::Guard<ir::Nothing>) -> Self {
22        Self { source, guard }
23    }
24}
25
26/// An instance of `StaticSchedule` is constrainted to live at least as long as
27/// the component in which the static island that it represents lives.
28pub struct StaticSchedule<'b, 'a: 'b> {
29    /// Builder construct to add hardware to the component it's built from
30    pub builder: &'b mut ir::Builder<'a>,
31    /// Number of cycles to which the static schedule should count up
32    pub state: u64,
33    /// Maps every FSM state to assignments that should be active in that state
34    pub state2assigns: HashMap<u64, Vec<ir::Assignment<ir::Nothing>>>,
35    /// Parital map from FSM state to transitions out of that state.
36    /// If a state has no mapping, assume it's an unconditional transition to
37    /// state + 1.
38    pub state2trans: HashMap<u64, ir::Transition>,
39}
40
41impl<'b, 'a> From<&'b mut ir::Builder<'a>> for StaticSchedule<'b, 'a> {
42    fn from(builder: &'b mut ir::Builder<'a>) -> Self {
43        StaticSchedule {
44            builder,
45            state: 0,
46            state2assigns: HashMap::new(),
47            state2trans: HashMap::new(),
48        }
49    }
50}
51
52impl StaticSchedule<'_, '_> {
53    pub fn leave_one_state_condition(
54        &mut self,
55        guard: ir::Guard<ir::Nothing>,
56        sen: &ir::StaticEnable,
57    ) -> ir::Guard<ir::Nothing> {
58        let signal_on = self.builder.add_constant(1, 1);
59        let group_latency = sen.group.borrow().get_latency();
60
61        // instantiate a local counter register
62        let width = bits_needed_for(group_latency);
63        let counter =
64            self.builder
65                .add_primitive("group_counter", "std_reg", &[width]);
66
67        // transform all assignments in the static group to read
68        // from the local counter
69        let mut assigns = sen
70            .group
71            .borrow_mut()
72            .assignments
73            .clone()
74            .drain(..)
75            .map(|mut sassign| {
76                sassign.guard.replace_static_timing(
77                    self.builder,
78                    &counter,
79                    &width,
80                    &group_latency,
81                );
82                let mut assign = ir::Assignment::from(sassign);
83                assign.and_guard(guard.clone());
84                assign
85            })
86            .collect_vec();
87
88        // guard reprsenting if counter is in final state
89        let final_state_const =
90            self.builder.add_constant(group_latency - 1, width);
91        let final_state_wire: ir::RRC<ir::Cell> = self.builder.add_primitive(
92            format!("const{}_{}_", group_latency - 1, width),
93            "std_wire",
94            &[width],
95        );
96        let final_state_guard = ir::Guard::CompOp(
97            ir::PortComp::Eq,
98            counter.borrow().get("out"),
99            final_state_wire.borrow().get("out"),
100        );
101        let not_final_state_guard = final_state_guard.clone().not();
102
103        // build assignments to increment / reset the counter
104        let adder = self.builder.add_primitive("adder", "std_add", &[width]);
105        let const_one = self.builder.add_constant(1, width);
106        let const_zero = self.builder.add_constant(0, width);
107        let incr_counter_assigns = build_assignments!(self.builder;
108            final_state_wire["in"] = ? final_state_const["out"];
109            adder["left"] = ? counter["out"];
110            adder["right"] = ? const_one["out"];
111            counter["write_en"] = ? signal_on["out"];
112            counter["in"] = final_state_guard ? const_zero["out"];
113            counter["in"] = not_final_state_guard ? adder["out"];
114        );
115
116        assigns.extend(incr_counter_assigns.to_vec());
117
118        // push these assignments into the one state allocated for this
119        // enable
120        self.state2assigns
121            .entry(self.state)
122            .and_modify(|other_assigns| {
123                other_assigns.extend(assigns.clone());
124            })
125            .or_insert(assigns);
126
127        final_state_guard
128    }
129
130    pub fn register_transitions(
131        &mut self,
132        curr_state: u64,
133        transitions_to_curr: &mut Vec<IncompleteTransition>,
134        and_guard: ir::Guard<ir::Nothing>,
135    ) {
136        transitions_to_curr.drain(..).for_each(
137            |IncompleteTransition { source, guard }| {
138                let complete_transition =
139                    match (guard, &and_guard) {
140                        (ir::Guard::True, ir::Guard::True) => {
141                            ir::Transition::Unconditional(curr_state)
142                        }
143                        (ir::Guard::True, _) => ir::Transition::Conditional(
144                            vec![(and_guard.clone(), curr_state)],
145                        ),
146                        (guard, ir::Guard::True) => {
147                            ir::Transition::Conditional(vec![(
148                                guard, curr_state,
149                            )])
150                        }
151                        (guard, and_guard) => ir::Transition::Conditional(
152                            vec![(guard.and(and_guard.clone()), curr_state)],
153                        ),
154                    };
155
156                self.state2trans
157                    .entry(source)
158                    .and_modify(|existing_transition| {
159                        match (existing_transition, complete_transition.clone())
160                        {
161                            (ir::Transition::Unconditional(_), _)
162                            | (_, ir::Transition::Unconditional(_)) => (),
163                            (
164                                ir::Transition::Conditional(existing_conds),
165                                ir::Transition::Conditional(new_conds),
166                            ) => {
167                                existing_conds.extend(new_conds);
168                            }
169                        };
170                    })
171                    .or_insert(complete_transition);
172            },
173        );
174    }
175
176    pub fn build_fsm_pieces(&mut self, fsm: ir::RRC<ir::FSM>) -> FSMPieces {
177        let signal_on = self.builder.add_constant(1, 1);
178        (0..self.state)
179            .map(|state| {
180                // construct a wire to represent this state
181                let state_wire: ir::RRC<ir::Cell> = self.builder.add_primitive(
182                    format!("{}_{state}", fsm.borrow().name()),
183                    "std_wire",
184                    &[1],
185                );
186                // build assignment to indicate that we're in this state
187                let mut state_assign: ir::Assignment<ir::Nothing> =
188                    self.builder.build_assignment(
189                        state_wire.borrow().get("in"),
190                        signal_on.borrow().get("out"),
191                        ir::Guard::True,
192                    );
193
194                // merge `idle` and first `calc` state
195                if state == 0 {
196                    state_assign.and_guard(ir::guard!(fsm["start"]));
197                }
198
199                let transition_from_state = match self
200                    .state2trans
201                    .remove(&state)
202                {
203                    Some(mut transition) => {
204                        // if in first state, transition conditioned on fsm[start]
205                        let transition_mut_ref = &mut transition;
206                        if state == 0 {
207                            match transition_mut_ref {
208                                ir::Transition::Unconditional(next_state) => {
209                                    *transition_mut_ref =
210                                        ir::Transition::Conditional(vec![
211                                            (guard!(fsm["start"]), *next_state),
212                                            (ir::Guard::True, 0),
213                                        ]);
214                                }
215                                ir::Transition::Conditional(conditions) => {
216                                    conditions.iter_mut().for_each(
217                                        |(condition, _)| {
218                                            condition.update(|g| {
219                                                g.and(guard!(fsm["start"]))
220                                            });
221                                        },
222                                    );
223                                }
224                            }
225                        }
226
227                        // add a default self-loop for every conditional transition
228                        // if it doesn't already have it
229                        if let ir::Transition::Conditional(trans) =
230                            &mut transition
231                        {
232                            if !(trans.last_mut().unwrap().0.is_true()) {
233                                trans.push((ir::Guard::True, state))
234                            }
235                        }
236                        transition
237                    }
238                    None => {
239                        if state == 0 {
240                            // set transition out of first state, which is
241                            // conditional on reading fsm[start]
242                            ir::Transition::Conditional(vec![
243                                (guard!(fsm["start"]), 1 % self.state),
244                                (ir::Guard::True, 0),
245                            ])
246                        } else {
247                            // loopback to start at final state, and increment
248                            // state otherwise
249                            ir::Transition::Unconditional(
250                                if state + 1 == self.state {
251                                    0
252                                } else {
253                                    state + 1
254                                },
255                            )
256                        }
257                    }
258                };
259
260                (vec![state_assign], transition_from_state, state_wire)
261            })
262            .multiunzip()
263    }
264}