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#[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
26pub struct StaticSchedule<'b, 'a: 'b> {
29 pub builder: &'b mut ir::Builder<'a>,
31 pub state: u64,
33 pub state2assigns: HashMap<u64, Vec<ir::Assignment<ir::Nothing>>>,
35 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 let width = bits_needed_for(group_latency);
63 let counter =
64 self.builder
65 .add_primitive("group_counter", "std_reg", &[width]);
66
67 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 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 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 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 let state_wire: ir::RRC<ir::Cell> = self.builder.add_primitive(
182 format!("{}_{state}", fsm.borrow().name()),
183 "std_wire",
184 &[1],
185 );
186 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 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 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 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 ir::Transition::Conditional(vec![
243 (guard!(fsm["start"]), 1 % self.state),
244 (ir::Guard::True, 0),
245 ])
246 } else {
247 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}