From 2dedfb48ec9152896fa6e6cdc388db64e25ac10c Mon Sep 17 00:00:00 2001
From: Parth Sarkar <ps679@cornell.edu>
Date: Tue, 17 Dec 2024 15:27:37 -0500
Subject: [PATCH] Added method to group all assignments that write to the same
 wire, just at different cases

---
 calyx-backend/src/verilog.rs | 67 +++++++++++++++++++++++++++++-------
 calyx-ir/src/structure.rs    | 19 ++++++++++
 2 files changed, 73 insertions(+), 13 deletions(-)

diff --git a/calyx-backend/src/verilog.rs b/calyx-backend/src/verilog.rs
index 3c88e1cd2..9b6afdf8a 100644
--- a/calyx-backend/src/verilog.rs
+++ b/calyx-backend/src/verilog.rs
@@ -454,6 +454,8 @@ fn cell_instance(cell: &ir::Cell) -> Option<v::Instance> {
 /// block to transition the FSM and drive assignments that read from the FSM
 /// register
 fn emit_fsm<F: io::Write>(fsm: &RRC<ir::FSM>, f: &mut F) -> io::Result<()> {
+    let reg_bitwidth =
+        get_bit_width_from(fsm.borrow().assignments.len() as u64);
     // generate fsm logic variables
     let (state_reg, state_next) = vec!["out", "in"]
         .drain(..)
@@ -462,10 +464,13 @@ fn emit_fsm<F: io::Write>(fsm: &RRC<ir::FSM>, f: &mut F) -> io::Result<()> {
         .unwrap();
 
     // emit inlined register
-    emit_fsm_inlined_reg(
+    emit_fsm_inlined_reg(&state_reg, &state_next, reg_bitwidth, f)?;
+
+    // dump assignments to enable in this state
+    emit_fsm_dependent_assignments(
+        &fsm.borrow().assignments,
         &state_reg,
-        &state_next,
-        get_bit_width_from(fsm.borrow().assignments.len() as u64),
+        reg_bitwidth,
         f,
     )?;
 
@@ -473,18 +478,9 @@ fn emit_fsm<F: io::Write>(fsm: &RRC<ir::FSM>, f: &mut F) -> io::Result<()> {
     writeln!(f, "always @(*) begin")?;
     writeln!(f, "  case ({})", state_reg)?;
 
-    for (case, (assigns, trans)) in fsm
-        .borrow()
-        .assignments
-        .iter()
-        .zip(fsm.borrow().transitions.iter())
-        .enumerate()
-    {
+    for (case, trans) in fsm.borrow().transitions.iter().enumerate() {
         writeln!(f, "{}{}: begin", " ".repeat(4), case)?;
 
-        // dump assignments to enable in this state
-        emit_fsm_assignments(assigns, f)?;
-
         // write transitions
         emit_fsm_transition(trans, &state_next, f)?;
 
@@ -530,6 +526,51 @@ fn emit_fsm_assignments<F: io::Write>(
     io::Result::Ok(())
 }
 
+fn emit_fsm_dependent_assignments<F: io::Write>(
+    assignments: &Vec<Vec<ir::Assignment<Nothing>>>,
+    fsm_out: &String,
+    reg_bitwidth: u64,
+    f: &mut F,
+) -> io::Result<()> {
+    for (case, assigns) in assignments.iter().enumerate() {
+        for assign in assigns.iter() {
+            let dst_ref = &assign.dst;
+            let case_guard =
+                format!("{} == {}'d{}", fsm_out, reg_bitwidth, case);
+
+            // string representing the new guard on the assignment
+            let case_guarded_assign_guard = if assign.guard.is_true() {
+                case_guard
+            } else {
+                format!(
+                    "(({}) & ({}))",
+                    case_guard,
+                    unflattened_guard(&assign.guard)
+                )
+            };
+
+            // value for the wire to take if either fsm is not in relevant state
+            // or if the assignment's original condition is not met
+            let guard_unmet_value = if is_data_port(dst_ref) {
+                format!("'x")
+            } else {
+                format!("{}'d0", dst_ref.borrow().width)
+            };
+
+            // emit assignment dependent on both case and the assignment's original guard
+            writeln!(
+                f,
+                "assign {} = {} ? {} : {};",
+                VerilogPortRef(dst_ref),
+                case_guarded_assign_guard,
+                VerilogPortRef(&assign.src),
+                guard_unmet_value
+            )?;
+        }
+    }
+    io::Result::Ok(())
+}
+
 fn emit_fsm_inlined_reg<F: io::Write>(
     state_reg_logic_var: &String,
     state_next_logic_var: &String,
diff --git a/calyx-ir/src/structure.rs b/calyx-ir/src/structure.rs
index b897bb3b2..86d5c27a9 100644
--- a/calyx-ir/src/structure.rs
+++ b/calyx-ir/src/structure.rs
@@ -10,6 +10,7 @@ use calyx_frontend::{Attribute, BoolAttr};
 use calyx_utils::{CalyxResult, Error, GetName};
 use itertools::Itertools;
 use smallvec::{smallvec, SmallVec};
+use std::collections::HashMap;
 use std::hash::Hash;
 use std::rc::Rc;
 
@@ -893,6 +894,24 @@ impl FSM {
             })
             .collect()
     }
+
+    pub fn merge_assignments(&self) -> Vec<Vec<(usize, Assignment<Nothing>)>> {
+        let mut gathered_assigns: HashMap<
+            Id,
+            Vec<(usize, Assignment<Nothing>)>,
+        > = HashMap::new();
+        for (case, assigns_at_state) in self.assignments.iter().enumerate() {
+            for assign in assigns_at_state.iter() {
+                gathered_assigns
+                    .entry(assign.dst.borrow().name)
+                    .and_modify(|gathered| {
+                        gathered.push((case, assign.clone()));
+                    })
+                    .or_insert(vec![(case, assign.clone())]);
+            }
+        }
+        gathered_assigns.drain().map(|(_, v)| v).collect()
+    }
 }
 
 impl GetName for Cell {