Improve Static FSM Compilation (#2215)

Implements the ideas from #2207. 

Apologies for the gigantic PR, there are a couple reasons why it is so
big:
1) It represents a big change in the compiler 
2) I didn't know if it was going to even be worth it to implement these
changes to the compiler, so I implemented some improvements to the
compilation process that complicated the code (but
[improved](https://github.com/orgs/calyxir/discussions/2202#discussioncomment-10014608)
results).

There are some minor changes to `static_inline.rs` (in particular,
inlining `static par` blocks is more complicated now because we can't
merge always just merge all threads of a `static par` into the same
group). There are some changes to `compile_static.rs`, but the main
contribution of this PR is the file `static_tree.rs`.

I'm still going to write some tests to make sure I'm getting all edge
cases for this new tree-looking FSM compilation.

calyx-frontend/src/attribute.rs

@@ -69,6 +69,10 @@ pub enum BoolAttr {
     #[strum(serialize = "promoted")]
     /// denotes a static component or control promoted from dynamic
     Promoted,
+    #[strum(serialize = "par")]
+    /// Denotes a group that was generated from a `staticpar` during static
+    /// inlining.
+    ParCtrl,
     #[strum(serialize = "fast")]
     /// https://github.com/calyxir/calyx/issues/1828
     Fast,
@@ -202,8 +206,8 @@ impl FromStr for Attribute {
 #[derive(Default, Debug, Clone, PartialEq, Eq)]
 /// Inline storage for boolean attributes.
 pub(super) struct InlineAttributes {
-    /// Boolean attributes stored in a 16-bit number.
-    attrs: u16,
+    /// Boolean attributes stored in a 32-bit number.
+    attrs: u32,
 }
 
 impl InlineAttributes {

calyx-opt/src/analysis/graph_coloring.rs

@@ -169,6 +169,17 @@ where
             .collect()
     }
 
+    // Reverses a coloring by mapping color C -> vec of nodes colored C.
+    pub fn reverse_coloring(coloring: &HashMap<T, T>) -> HashMap<T, Vec<T>> {
+        let mut rev_coloring: HashMap<T, Vec<T>> = HashMap::new();
+        for (node, color) in coloring {
+            rev_coloring
+                .entry(color.clone())
+                .or_default()
+                .push(node.clone());
+        }
+        rev_coloring
+    }
     pub fn welsh_powell_coloring(&self) -> HashMap<T, T> {
         let mut coloring: HashMap<T, T> = HashMap::new();
 

calyx-opt/src/analysis/mod.rs

@@ -20,8 +20,9 @@ pub mod reaching_defns;
 mod read_write_set;
 mod schedule_conflicts;
 mod share_set;
+mod static_fsm;
 mod static_par_timing;
-mod static_schedule;
+mod static_tree;
 mod variable_detection;
 
 pub use compaction_analysis::CompactionAnalysis;
@@ -42,6 +43,7 @@ pub use promotion_analysis::PromotionAnalysis;
 pub use read_write_set::{AssignmentAnalysis, ReadWriteSet};
 pub use schedule_conflicts::ScheduleConflicts;
 pub use share_set::ShareSet;
+pub use static_fsm::{FSMEncoding, StaticFSM};
 pub use static_par_timing::StaticParTiming;
-pub use static_schedule::{StaticFSM, StaticSchedule};
+pub use static_tree::{Node, ParNodes, SingleNode, StateType};
 pub use variable_detection::VariableDetection;

calyx-opt/src/analysis/static_fsm.rs

@@ -0,0 +1,253 @@
+use crate::passes::math_utilities::get_bit_width_from;
+use calyx_ir::{self as ir};
+use calyx_ir::{build_assignments, Nothing};
+use calyx_ir::{guard, structure};
+use std::collections::HashMap;
+use std::rc::Rc;
+
+#[derive(Debug, Clone, Copy, Default)]
+// Define an FSMEncoding Enum
+pub enum FSMEncoding {
+    #[default]
+    Binary,
+    OneHot,
+}
+
+#[derive(Debug)]
+/// Represents a static FSM (i.e., the actual register in hardware that counts)
+pub struct StaticFSM {
+    /// The actual register cell
+    fsm_cell: ir::RRC<ir::Cell>,
+    /// Type of encoding (binary or one-hot)
+    encoding: FSMEncoding,
+    /// The fsm's bitwidth (this redundant information bc  we have `cell`)
+    /// but makes it easier if we easily have access to this.
+    bitwidth: u64,
+    /// Mapping of queries: (u64, u64) -> Port
+    queries: HashMap<(u64, u64), ir::RRC<ir::Port>>,
+}
+impl StaticFSM {
+    // Builds a static_fsm from: num_states and encoding type.
+    pub fn from_basic_info(
+        num_states: u64,
+        encoding: FSMEncoding,
+        builder: &mut ir::Builder,
+    ) -> Self {
+        // Determine number of bits needed in the register.
+        let fsm_size = match encoding {
+            /* represent 0..latency */
+            FSMEncoding::Binary => get_bit_width_from(num_states + 1),
+            FSMEncoding::OneHot => num_states,
+        };
+        // OHE needs an initial value of 1.
+        let register = match encoding {
+            FSMEncoding::Binary => {
+                builder.add_primitive("fsm", "std_reg", &[fsm_size])
+            }
+            FSMEncoding::OneHot => {
+                builder.add_primitive("fsm", "init_one_reg", &[fsm_size])
+            }
+        };
+
+        StaticFSM {
+            encoding,
+            fsm_cell: register,
+            bitwidth: fsm_size,
+            queries: HashMap::new(),
+        }
+    }
+
+    // Builds an incrementer, and returns the assignments and incrementer cell itself.
+    // assignments are:
+    // adder.left = fsm.out; adder.right = 1;
+    // Returns tuple: (assignments, adder)
+    pub fn build_incrementer(
+        &self,
+        builder: &mut ir::Builder,
+    ) -> (Vec<ir::Assignment<Nothing>>, ir::RRC<ir::Cell>) {
+        let fsm_cell = Rc::clone(&self.fsm_cell);
+        // For OHE, the "adder" can just be a shifter.
+        // For OHE the first_state = 1 rather than 0.
+        // Final state is encoded differently for OHE vs. Binary
+        let adder = match self.encoding {
+            FSMEncoding::Binary => {
+                builder.add_primitive("adder", "std_add", &[self.bitwidth])
+            }
+            FSMEncoding::OneHot => {
+                builder.add_primitive("lsh", "std_lsh", &[self.bitwidth])
+            }
+        };
+        let const_one = builder.add_constant(1, self.bitwidth);
+        let incr_assigns = build_assignments!(
+          builder;
+          // increments the fsm
+          adder["left"] = ? fsm_cell["out"];
+          adder["right"] = ? const_one["out"];
+        )
+        .to_vec();
+        (incr_assigns, adder)
+    }
+
+    // Returns the assignments that conditionally increment the fsm,
+    // based on guard.
+    // The assignments are:
+    // fsm.in = guard ? adder.out;
+    // fsm.write_en = guard ? 1'd1;
+    // Returns a vec of these assignments.
+    pub fn conditional_increment(
+        &self,
+        guard: ir::Guard<Nothing>,
+        adder: ir::RRC<ir::Cell>,
+        builder: &mut ir::Builder,
+    ) -> Vec<ir::Assignment<Nothing>> {
+        let fsm_cell = Rc::clone(&self.fsm_cell);
+        let signal_on = builder.add_constant(1, 1);
+        let my_assigns = build_assignments!(
+          builder;
+          // increments the fsm
+          fsm_cell["in"] = guard ? adder["out"];
+          fsm_cell["write_en"] = guard ? signal_on["out"];
+        );
+        my_assigns.to_vec()
+    }
+
+    // Returns the assignments that conditionally resets the fsm to 0,
+    // but only if guard is true.
+    // The assignments are:
+    // fsm.in = guard ? 0;
+    // fsm.write_en = guard ? 1'd1;
+    // Returns a vec of these assignments.
+    pub fn conditional_reset(
+        &self,
+        guard: ir::Guard<Nothing>,
+        builder: &mut ir::Builder,
+    ) -> Vec<ir::Assignment<Nothing>> {
+        let fsm_cell = Rc::clone(&self.fsm_cell);
+        let signal_on = builder.add_constant(1, 1);
+        let const_0 = match self.encoding {
+            FSMEncoding::Binary => builder.add_constant(0, self.bitwidth),
+            FSMEncoding::OneHot => builder.add_constant(1, self.bitwidth),
+        };
+        let assigns = build_assignments!(
+          builder;
+          fsm_cell["in"] = guard ? const_0["out"];
+          fsm_cell["write_en"] = guard ? signal_on["out"];
+        );
+        assigns.to_vec()
+    }
+
+    // Returns a guard that takes a (beg, end) `query`, and returns the equivalent
+    // guard to `beg <= fsm.out < end`.
+    pub fn query_between(
+        &mut self,
+        builder: &mut ir::Builder,
+        query: (u64, u64),
+    ) -> Box<ir::Guard<Nothing>> {
+        let (beg, end) = query;
+        // Querying OHE is easy, since we already have `self.get_one_hot_query()`
+        let fsm_cell = Rc::clone(&self.fsm_cell);
+        if matches!(self.encoding, FSMEncoding::OneHot) {
+            let g = self.get_one_hot_query(fsm_cell, (beg, end), builder);
+            return Box::new(g);
+        }
+
+        if beg + 1 == end {
+            // if beg + 1 == end then we only need to check if fsm == beg
+            let interval_const = builder.add_constant(beg, self.bitwidth);
+            let g = guard!(fsm_cell["out"] == interval_const["out"]);
+            Box::new(g)
+        } else if beg == 0 {
+            // if beg == 0, then we only need to check if fsm < end
+            let end_const = builder.add_constant(end, self.bitwidth);
+            let lt: ir::Guard<Nothing> =
+                guard!(fsm_cell["out"] < end_const["out"]);
+            Box::new(lt)
+        } else {
+            // otherwise, check if fsm >= beg & fsm < end
+            let beg_const = builder.add_constant(beg, self.bitwidth);
+            let end_const = builder.add_constant(end, self.bitwidth);
+            let beg_guard: ir::Guard<Nothing> =
+                guard!(fsm_cell["out"] >= beg_const["out"]);
+            let end_guard: ir::Guard<Nothing> =
+                guard!(fsm_cell["out"] < end_const["out"]);
+            Box::new(ir::Guard::And(Box::new(beg_guard), Box::new(end_guard)))
+        }
+    }
+
+    // Given a one-hot query, it will return a guard corresponding to that query.
+    // If it has already built the query (i.e., added the wires/continuous assigments),
+    // it just uses the same port.
+    // Otherwise it will build the query.
+    fn get_one_hot_query(
+        &mut self,
+        fsm_cell: ir::RRC<ir::Cell>,
+        (lb, ub): (u64, u64),
+        builder: &mut ir::Builder,
+    ) -> ir::Guard<Nothing> {
+        match self.queries.get(&(lb, ub)) {
+            None => {
+                let port = Self::build_one_hot_query(
+                    Rc::clone(&fsm_cell),
+                    self.bitwidth,
+                    (lb, ub),
+                    builder,
+                );
+                self.queries.insert((lb, ub), Rc::clone(&port));
+                ir::Guard::port(port)
+            }
+            Some(port) => ir::Guard::port(Rc::clone(port)),
+        }
+    }
+
+    // Given a (lb, ub) query, and an fsm (and for convenience, a bitwidth),
+    // Returns a `port`: port is a `wire.out`, where `wire` holds
+    // whether or not the query is true, i.e., whether the FSM really is
+    // between [lb, ub).
+    fn build_one_hot_query(
+        fsm_cell: ir::RRC<ir::Cell>,
+        fsm_bitwidth: u64,
+        (lb, ub): (u64, u64),
+        builder: &mut ir::Builder,
+    ) -> ir::RRC<ir::Port> {
+        // The wire that holds the query
+        let formatted_name = format!("bw_{}_{}", lb, ub);
+        let wire: ir::RRC<ir::Cell> =
+            builder.add_primitive(formatted_name, "std_wire", &[1]);
+        let wire_out = wire.borrow().get("out");
+
+        // Continuous assignments to check the FSM
+        let assigns = {
+            let in_width = fsm_bitwidth;
+            // Since 00...00 is the initial state, we need to check lb-1.
+            let start_index = lb;
+            // Since verilog slices are inclusive.
+            let end_index = ub - 1;
+            let out_width = ub - lb; // == (end_index - start_index + 1)
+            structure!(builder;
+                let slicer = prim std_bit_slice(in_width, start_index, end_index, out_width);
+                let const_slice_0 = constant(0, out_width);
+                let signal_on = constant(1,1);
+            );
+            let slicer_neq_0 = guard!(slicer["out"] != const_slice_0["out"]);
+            // Extend the continuous assignmments to include this particular query for FSM state;
+            let my_assigns = build_assignments!(builder;
+                slicer["in"] = ? fsm_cell["out"];
+                wire["in"] = slicer_neq_0 ? signal_on["out"];
+            );
+            my_assigns.to_vec()
+        };
+        builder.add_continuous_assignments(assigns);
+        wire_out
+    }
+
+    // Return a unique id (i.e., get_unique_id for each FSM in the same component
+    // will be different).
+    pub fn get_unique_id(&self) -> ir::Id {
+        self.fsm_cell.borrow().name()
+    }
+
+    // Return the bitwidth of an FSM object
+    pub fn get_bitwidth(&self) -> u64 {
+        self.bitwidth
+    }
+}