From a99bdeebc538e0e0811322b80e593c9365f0537e Mon Sep 17 00:00:00 2001
From: Parth Sarkar <93054509+parthsarkar17@users.noreply.github.com>
Date: Tue, 17 Dec 2024 20:59:00 -0500
Subject: [PATCH] regrouped assignments outside of always_comb case block
 (#2379)

---
 calyx-backend/src/verilog.rs               | 22 ++++++++++++----------
 calyx-ir/src/structure.rs                  | 16 ++++++++++++++--
 calyx-opt/src/passes/dyn_fsm_allocation.rs | 15 +++------------
 3 files changed, 29 insertions(+), 24 deletions(-)

diff --git a/calyx-backend/src/verilog.rs b/calyx-backend/src/verilog.rs
index 9b6afdf8a..d21ad4af6 100644
--- a/calyx-backend/src/verilog.rs
+++ b/calyx-backend/src/verilog.rs
@@ -468,7 +468,7 @@ fn emit_fsm<F: io::Write>(fsm: &RRC<ir::FSM>, f: &mut F) -> io::Result<()> {
 
     // dump assignments to enable in this state
     emit_fsm_dependent_assignments(
-        &fsm.borrow().assignments,
+        fsm.borrow().merge_assignments(),
         &state_reg,
         reg_bitwidth,
         f,
@@ -527,14 +527,15 @@ fn emit_fsm_assignments<F: io::Write>(
 }
 
 fn emit_fsm_dependent_assignments<F: io::Write>(
-    assignments: &Vec<Vec<ir::Assignment<Nothing>>>,
+    assignments: Vec<Vec<(usize, 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;
+    for collection in assignments.iter() {
+        let dst_ref = &collection.first().unwrap().1.dst;
+        writeln!(f, "assign {} =", VerilogPortRef(dst_ref))?;
+        for (i, (case, assign)) in collection.iter().enumerate() {
             let case_guard =
                 format!("{} == {}'d{}", fsm_out, reg_bitwidth, case);
 
@@ -557,15 +558,16 @@ fn emit_fsm_dependent_assignments<F: io::Write>(
                 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
+                VerilogPortRef(&assign.src)
             )?;
+
+            if i + 1 == collection.len() {
+                writeln!(f, " {guard_unmet_value};")?;
+            }
         }
     }
     io::Result::Ok(())
diff --git a/calyx-ir/src/structure.rs b/calyx-ir/src/structure.rs
index 86d5c27a9..0d6a57862 100644
--- a/calyx-ir/src/structure.rs
+++ b/calyx-ir/src/structure.rs
@@ -897,13 +897,25 @@ impl FSM {
 
     pub fn merge_assignments(&self) -> Vec<Vec<(usize, Assignment<Nothing>)>> {
         let mut gathered_assigns: HashMap<
-            Id,
+            String,
             Vec<(usize, Assignment<Nothing>)>,
         > = HashMap::new();
         for (case, assigns_at_state) in self.assignments.iter().enumerate() {
             for assign in assigns_at_state.iter() {
+                let port = assign.dst.borrow();
+                let dest = match &port.parent {
+                    PortParent::Cell(cell) => {
+                        format!(
+                            "{}_{}",
+                            cell.upgrade().borrow().name,
+                            port.name
+                        )
+                    }
+                    _ => unreachable!(),
+                };
+
                 gathered_assigns
-                    .entry(assign.dst.borrow().name)
+                    .entry(dest)
                     .and_modify(|gathered| {
                         gathered.push((case, assign.clone()));
                     })
diff --git a/calyx-opt/src/passes/dyn_fsm_allocation.rs b/calyx-opt/src/passes/dyn_fsm_allocation.rs
index 46fb4847a..9a6de4608 100644
--- a/calyx-opt/src/passes/dyn_fsm_allocation.rs
+++ b/calyx-opt/src/passes/dyn_fsm_allocation.rs
@@ -427,16 +427,11 @@ impl<'b, 'a> Schedule<'b, 'a> {
             .drain()
             .sorted_by(|(s1, _), (s2, _)| s1.cmp(s2))
             .map(|(state, mut cond_dsts)| {
-                let mut assigns = match self.fsm_enables.get(&(state - 1)) {
+                let assigns = match self.fsm_enables.get(&(state - 1)) {
                     None => vec![],
                     Some(assigns) => assigns.clone(),
                 };
-                let signal_off = self.builder.add_constant(0, 1);
-                let true_guard = ir::Guard::True;
-                let not_done = build_assignments!(self.builder;
-                    fsm["done"] = true_guard ? signal_off["out"];
-                );
-                assigns.extend(not_done);
+
                 // self-loop if all other guards are not met;
                 // should be at the end of the conditional destinations vec!
                 cond_dsts.push((ir::Guard::True, state));
@@ -447,11 +442,7 @@ impl<'b, 'a> Schedule<'b, 'a> {
 
         // insert transition condition from 0 to 1
         let true_guard = ir::Guard::True;
-        let signal_off = self.builder.add_constant(0, 1);
-        let not_done = build_assignments!(self.builder;
-            fsm["done"] = true_guard ? signal_off["out"];
-        );
-        assignments.push_front(not_done.to_vec());
+        assignments.push_front(vec![]);
         transitions.push_front(ir::Transition::Conditional(vec![
             (guard!(fsm["start"]), 1),
             (true_guard.clone(), 0),