1use calyx_ir::{self as ir};
2use calyx_ir::{Nothing, build_assignments};
3use calyx_ir::{guard, structure};
4use calyx_utils::math::bits_needed_for;
5use std::collections::HashMap;
6use std::rc::Rc;
7
8#[derive(Debug, Clone, Copy, Default)]
9pub enum FSMEncoding {
11 #[default]
12 Binary,
13 OneHot,
14}
15
16#[derive(Debug)]
17pub struct StaticFSM {
19 fsm_cell: ir::RRC<ir::Cell>,
21 encoding: FSMEncoding,
23 bitwidth: u64,
26 queries: HashMap<(u64, u64), ir::RRC<ir::Port>>,
28}
29impl StaticFSM {
30 pub fn from_basic_info(
32 num_states: u64,
33 encoding: FSMEncoding,
34 builder: &mut ir::Builder,
35 ) -> Self {
36 let fsm_size = match encoding {
38 FSMEncoding::Binary => bits_needed_for(num_states + 1),
40 FSMEncoding::OneHot => num_states,
41 };
42 let register = match encoding {
44 FSMEncoding::Binary => {
45 builder.add_primitive("fsm", "std_reg", &[fsm_size])
46 }
47 FSMEncoding::OneHot => {
48 builder.add_primitive("fsm", "init_one_reg", &[fsm_size])
49 }
50 };
51
52 StaticFSM {
53 encoding,
54 fsm_cell: register,
55 bitwidth: fsm_size,
56 queries: HashMap::new(),
57 }
58 }
59
60 pub fn build_incrementer(
65 &self,
66 builder: &mut ir::Builder,
67 ) -> (Vec<ir::Assignment<Nothing>>, ir::RRC<ir::Cell>) {
68 let fsm_cell = Rc::clone(&self.fsm_cell);
69 let adder = match self.encoding {
73 FSMEncoding::Binary => {
74 builder.add_primitive("adder", "std_add", &[self.bitwidth])
75 }
76 FSMEncoding::OneHot => {
77 builder.add_primitive("lsh", "std_lsh", &[self.bitwidth])
78 }
79 };
80 let const_one = builder.add_constant(1, self.bitwidth);
81 let incr_assigns = build_assignments!(
82 builder;
83 adder["left"] = ? fsm_cell["out"];
85 adder["right"] = ? const_one["out"];
86 )
87 .to_vec();
88 (incr_assigns, adder)
89 }
90
91 pub fn conditional_increment(
98 &self,
99 guard: ir::Guard<Nothing>,
100 adder: ir::RRC<ir::Cell>,
101 builder: &mut ir::Builder,
102 ) -> Vec<ir::Assignment<Nothing>> {
103 let fsm_cell = Rc::clone(&self.fsm_cell);
104 let signal_on = builder.add_constant(1, 1);
105 let my_assigns = build_assignments!(
106 builder;
107 fsm_cell["in"] = guard ? adder["out"];
109 fsm_cell["write_en"] = guard ? signal_on["out"];
110 );
111 my_assigns.to_vec()
112 }
113
114 pub fn conditional_reset(
121 &self,
122 guard: ir::Guard<Nothing>,
123 builder: &mut ir::Builder,
124 ) -> Vec<ir::Assignment<Nothing>> {
125 let fsm_cell = Rc::clone(&self.fsm_cell);
126 let signal_on = builder.add_constant(1, 1);
127 let const_0 = match self.encoding {
128 FSMEncoding::Binary => builder.add_constant(0, self.bitwidth),
129 FSMEncoding::OneHot => builder.add_constant(1, self.bitwidth),
130 };
131 let assigns = build_assignments!(
132 builder;
133 fsm_cell["in"] = guard ? const_0["out"];
134 fsm_cell["write_en"] = guard ? signal_on["out"];
135 );
136 assigns.to_vec()
137 }
138
139 pub fn query_between(
142 &mut self,
143 builder: &mut ir::Builder,
144 query: (u64, u64),
145 ) -> Box<ir::Guard<Nothing>> {
146 let (beg, end) = query;
147 let fsm_cell = Rc::clone(&self.fsm_cell);
149 if matches!(self.encoding, FSMEncoding::OneHot) {
150 let g = self.get_one_hot_query(fsm_cell, (beg, end), builder);
151 return Box::new(g);
152 }
153
154 if beg + 1 == end {
155 let interval_const = builder.add_constant(beg, self.bitwidth);
157 let g = guard!(fsm_cell["out"] == interval_const["out"]);
158 Box::new(g)
159 } else if beg == 0 {
160 let end_const = builder.add_constant(end, self.bitwidth);
162 let lt: ir::Guard<Nothing> =
163 guard!(fsm_cell["out"] < end_const["out"]);
164 Box::new(lt)
165 } else {
166 let beg_const = builder.add_constant(beg, self.bitwidth);
168 let end_const = builder.add_constant(end, self.bitwidth);
169 let beg_guard: ir::Guard<Nothing> =
170 guard!(fsm_cell["out"] >= beg_const["out"]);
171 let end_guard: ir::Guard<Nothing> =
172 guard!(fsm_cell["out"] < end_const["out"]);
173 Box::new(ir::Guard::And(Box::new(beg_guard), Box::new(end_guard)))
174 }
175 }
176
177 fn get_one_hot_query(
182 &mut self,
183 fsm_cell: ir::RRC<ir::Cell>,
184 (lb, ub): (u64, u64),
185 builder: &mut ir::Builder,
186 ) -> ir::Guard<Nothing> {
187 match self.queries.get(&(lb, ub)) {
188 None => {
189 let port = Self::build_one_hot_query(
190 Rc::clone(&fsm_cell),
191 self.bitwidth,
192 (lb, ub),
193 builder,
194 );
195 self.queries.insert((lb, ub), Rc::clone(&port));
196 ir::Guard::port(port)
197 }
198 Some(port) => ir::Guard::port(Rc::clone(port)),
199 }
200 }
201
202 fn build_one_hot_query(
207 fsm_cell: ir::RRC<ir::Cell>,
208 fsm_bitwidth: u64,
209 (lb, ub): (u64, u64),
210 builder: &mut ir::Builder,
211 ) -> ir::RRC<ir::Port> {
212 let formatted_name = format!("bw_{lb}_{ub}");
214 let wire: ir::RRC<ir::Cell> =
215 builder.add_primitive(formatted_name, "std_wire", &[1]);
216 let wire_out = wire.borrow().get("out");
217
218 let assigns = {
220 let in_width = fsm_bitwidth;
221 let start_index = lb;
223 let end_index = ub - 1;
225 let out_width = ub - lb; structure!(builder;
227 let slicer = prim std_bit_slice(in_width, start_index, end_index, out_width);
228 let const_slice_0 = constant(0, out_width);
229 let signal_on = constant(1,1);
230 );
231 let slicer_neq_0 = guard!(slicer["out"] != const_slice_0["out"]);
232 let my_assigns = build_assignments!(builder;
234 slicer["in"] = ? fsm_cell["out"];
235 wire["in"] = slicer_neq_0 ? signal_on["out"];
236 );
237 my_assigns.to_vec()
238 };
239 builder.add_continuous_assignments(assigns);
240 wire_out
241 }
242
243 pub fn get_unique_id(&self) -> ir::Id {
246 self.fsm_cell.borrow().name()
247 }
248
249 pub fn get_bitwidth(&self) -> u64 {
251 self.bitwidth
252 }
253}