calyx_opt/analysis/
static_fsm.rs

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)]
9// Define an FSMEncoding Enum
10pub enum FSMEncoding {
11    #[default]
12    Binary,
13    OneHot,
14}
15
16#[derive(Debug)]
17/// Represents a static FSM (i.e., the actual register in hardware that counts)
18pub struct StaticFSM {
19    /// The actual register cell
20    fsm_cell: ir::RRC<ir::Cell>,
21    /// Type of encoding (binary or one-hot)
22    encoding: FSMEncoding,
23    /// The fsm's bitwidth (this redundant information bc  we have `cell`)
24    /// but makes it easier if we easily have access to this.
25    bitwidth: u64,
26    /// Mapping of queries: (u64, u64) -> Port
27    queries: HashMap<(u64, u64), ir::RRC<ir::Port>>,
28}
29impl StaticFSM {
30    // Builds a static_fsm from: num_states and encoding type.
31    pub fn from_basic_info(
32        num_states: u64,
33        encoding: FSMEncoding,
34        builder: &mut ir::Builder,
35    ) -> Self {
36        // Determine number of bits needed in the register.
37        let fsm_size = match encoding {
38            /* represent 0..latency */
39            FSMEncoding::Binary => bits_needed_for(num_states + 1),
40            FSMEncoding::OneHot => num_states,
41        };
42        // OHE needs an initial value of 1.
43        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    // Builds an incrementer, and returns the assignments and incrementer cell itself.
61    // assignments are:
62    // adder.left = fsm.out; adder.right = 1;
63    // Returns tuple: (assignments, adder)
64    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        // For OHE, the "adder" can just be a shifter.
70        // For OHE the first_state = 1 rather than 0.
71        // Final state is encoded differently for OHE vs. Binary
72        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          // increments the fsm
84          adder["left"] = ? fsm_cell["out"];
85          adder["right"] = ? const_one["out"];
86        )
87        .to_vec();
88        (incr_assigns, adder)
89    }
90
91    // Returns the assignments that conditionally increment the fsm,
92    // based on guard.
93    // The assignments are:
94    // fsm.in = guard ? adder.out;
95    // fsm.write_en = guard ? 1'd1;
96    // Returns a vec of these assignments.
97    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          // increments the fsm
108          fsm_cell["in"] = guard ? adder["out"];
109          fsm_cell["write_en"] = guard ? signal_on["out"];
110        );
111        my_assigns.to_vec()
112    }
113
114    // Returns the assignments that conditionally resets the fsm to 0,
115    // but only if guard is true.
116    // The assignments are:
117    // fsm.in = guard ? 0;
118    // fsm.write_en = guard ? 1'd1;
119    // Returns a vec of these assignments.
120    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    // Returns a guard that takes a (beg, end) `query`, and returns the equivalent
140    // guard to `beg <= fsm.out < end`.
141    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        // Querying OHE is easy, since we already have `self.get_one_hot_query()`
148        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            // if beg + 1 == end then we only need to check if fsm == beg
156            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            // if beg == 0, then we only need to check if fsm < end
161            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            // otherwise, check if fsm >= beg & fsm < end
167            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    // Given a one-hot query, it will return a guard corresponding to that query.
178    // If it has already built the query (i.e., added the wires/continuous assigments),
179    // it just uses the same port.
180    // Otherwise it will build the query.
181    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    // Given a (lb, ub) query, and an fsm (and for convenience, a bitwidth),
203    // Returns a `port`: port is a `wire.out`, where `wire` holds
204    // whether or not the query is true, i.e., whether the FSM really is
205    // between [lb, ub).
206    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        // The wire that holds the query
213        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        // Continuous assignments to check the FSM
219        let assigns = {
220            let in_width = fsm_bitwidth;
221            // Since 00...00 is the initial state, we need to check lb-1.
222            let start_index = lb;
223            // Since verilog slices are inclusive.
224            let end_index = ub - 1;
225            let out_width = ub - lb; // == (end_index - start_index + 1)
226            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            // Extend the continuous assignmments to include this particular query for FSM state;
233            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    // Return a unique id (i.e., get_unique_id for each FSM in the same component
244    // will be different).
245    pub fn get_unique_id(&self) -> ir::Id {
246        self.fsm_cell.borrow().name()
247    }
248
249    // Return the bitwidth of an FSM object
250    pub fn get_bitwidth(&self) -> u64 {
251        self.bitwidth
252    }
253}