From ccdf989025e57da7dfd5ab609676ebe3cfb2c2d6 Mon Sep 17 00:00:00 2001
From: Eddie Hung <eddie@fpgeh.com>
Date: Thu, 6 Jun 2019 13:51:22 -0700
Subject: [PATCH] Support cascading $pmux.A with $mux.A and $mux.B

---
 passes/opt/muxpack.cc    | 42 ++++++++++++++++++++++++----------------
 tests/various/muxpack.v  | 25 ++++++++++++++++++++++++
 tests/various/muxpack.ys | 15 ++++++++++++++
 3 files changed, 65 insertions(+), 17 deletions(-)

diff --git a/passes/opt/muxpack.cc b/passes/opt/muxpack.cc
index ae4b67db2..f9e5c8f09 100644
--- a/passes/opt/muxpack.cc
+++ b/passes/opt/muxpack.cc
@@ -51,10 +51,12 @@ struct MuxpackWorker
 
 		for (auto cell : module->cells())
 		{
-			if (cell->type.in("$mux") && !cell->get_bool_attribute("\\keep"))
+			if (cell->type.in("$mux", "$pmux") && !cell->get_bool_attribute("\\keep"))
 			{
 				SigSpec a_sig = sigmap(cell->getPort("\\A"));
-				SigSpec b_sig = sigmap(cell->getPort("\\B"));
+				SigSpec b_sig;
+				if (cell->type == "$mux")
+					b_sig = sigmap(cell->getPort("\\B"));
 				SigSpec y_sig = sigmap(cell->getPort("\\Y"));
    
 				if (sig_chain_next.count(a_sig))
@@ -65,12 +67,14 @@ struct MuxpackWorker
 					candidate_cells.insert(cell);
 				}
 
-				if (sig_chain_next.count(b_sig))
-					for (auto b_bit : b_sig.bits())
-						sigbit_with_non_chain_users.insert(b_bit);
-				else {
-					sig_chain_next[b_sig] = cell;
-					candidate_cells.insert(cell);
+				if (!b_sig.empty()) {
+					if (sig_chain_next.count(b_sig))
+						for (auto b_bit : b_sig.bits())
+							sigbit_with_non_chain_users.insert(b_bit);
+					else {
+						sig_chain_next[b_sig] = cell;
+						candidate_cells.insert(cell);
+					}
 				}
 
 				sig_chain_prev[y_sig] = cell;
@@ -88,10 +92,16 @@ struct MuxpackWorker
 	{
 		for (auto cell : candidate_cells)
 		{
+			log_debug("Considering %s (%s)\n", log_id(cell), log_id(cell->type));
+
 			SigSpec next_sig = cell->getPort("\\A");
 			if (sig_chain_prev.count(next_sig) == 0) {
-				next_sig = cell->getPort("\\B");
-				if (sig_chain_prev.count(next_sig) == 0)
+				if (cell->type == "$mux") {
+					next_sig = cell->getPort("\\B");
+					if (sig_chain_prev.count(next_sig) == 0)
+						goto start_cell;
+				}
+				else
 					goto start_cell;
 			}
 
@@ -103,10 +113,7 @@ struct MuxpackWorker
 				Cell *c1 = sig_chain_prev.at(next_sig);
 				Cell *c2 = cell;
 
-				if (c1->type != c2->type)
-					goto start_cell;
-
-				if (c1->parameters != c2->parameters)
+				if (c1->getParam("\\WIDTH") != c2->getParam("\\WIDTH"))
 					goto start_cell;
 			}
 
@@ -220,15 +227,16 @@ struct MuxpackWorker
 };
 
 struct MuxpackPass : public Pass {
-	MuxpackPass() : Pass("muxpack", "$mux cell cascades to $pmux") { }
+	MuxpackPass() : Pass("muxpack", "$mux/$pmux cascades to $pmux") { }
 	void help() YS_OVERRIDE
 	{
 		//   |---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|
 		log("\n");
 		log("    muxpack [selection]\n");
 		log("\n");
-		log("This pass converts cascaded chains of $mux cells (e.g. those created by if-else\n");
-		log("constructs) into $pmux cells.\n");
+		log("This pass converts cascaded chains of $pmux cells (e.g. those create from case\n");
+		log("constructs) and $mux cells (e.g. those created by if-else constructs) into \n");
+		log("into $pmux cells.\n");
 		log("\n");
 	}
 	void execute(std::vector<std::string> args, RTLIL::Design *design) YS_OVERRIDE
diff --git a/tests/various/muxpack.v b/tests/various/muxpack.v
index e847fef27..fe0150532 100644
--- a/tests/various/muxpack.v
+++ b/tests/various/muxpack.v
@@ -85,3 +85,28 @@ always @* begin
 	if (s == 0) o <= i[2*W+:W];
 end
 endmodule
+
+module mux_case_unbal_7_7#(parameter N=7, parameter W=7) (input [N*W-1:0] i, input [$clog2(N)-1:0] s, output reg [W-1:0] o);
+always @* begin
+    o <= {W{1'bx}};
+    case (s)
+    0: o <= i[0*W+:W];
+    default:
+        case (s)
+        1: o <= i[1*W+:W];
+        2: o <= i[2*W+:W];
+        default:
+            case (s)
+            3: o <= i[3*W+:W];
+            4: o <= i[4*W+:W];
+            5: o <= i[5*W+:W];
+            default:
+                case (s)
+                    6: o <= i[6*W+:W];
+                    default: o <= i[7*W+:W];
+                endcase
+            endcase
+        endcase
+    endcase
+end
+endmodule
diff --git a/tests/various/muxpack.ys b/tests/various/muxpack.ys
index 178860b88..4dcb9ed89 100644
--- a/tests/various/muxpack.ys
+++ b/tests/various/muxpack.ys
@@ -118,3 +118,18 @@ design -import gold -as gold
 design -import gate -as gate
 miter -equiv -flatten -make_assert -make_outputs gold gate miter
 sat -verify -prove-asserts -show-ports miter
+
+design -load read
+hierarchy -top mux_case_unbal_7_7
+prep
+design -save gold
+muxpack
+opt
+stat
+select -assert-count 0 t:$mux
+select -assert-count 1 t:$pmux
+design -stash gate
+design -import gold -as gold
+design -import gate -as gate
+miter -equiv -flatten -make_assert -make_outputs gold gate miter
+sat -verify -prove-asserts -show-ports miter
-- 
2.30.2