add pmux option to bmuxmap for better fsm detection with verific frontend

2 files changed, 75 insertions, 6 deletions

diff --git a/passes/techmap/bmuxmap.cc b/passes/techmap/bmuxmap.cc
index 03673c278..15b149239 100644
--- a/passes/techmap/bmuxmap.cc
+++ b/passes/techmap/bmuxmap.cc
@@ -36,10 +36,16 @@ struct BmuxmapPass : public Pass {
 	}
 	void execute(std::vector<std::string> args, RTLIL::Design *design) override
 	{
+		bool pmux_mode = false;
+
 		log_header(design, "Executing BMUXMAP pass.\n");
 
 		size_t argidx;
 		for (argidx = 1; argidx < args.size(); argidx++) {
+			if (args[argidx] == "-pmux") {
+				pmux_mode = true;
+				continue;
+			}
 			break;
 		}
 		extra_args(args, argidx, design);
@@ -53,18 +59,36 @@ struct BmuxmapPass : public Pass {
 			SigSpec sel = cell->getPort(ID::S);
 			SigSpec data = cell->getPort(ID::A);
 			int width = GetSize(cell->getPort(ID::Y));
+			int s_width = GetSize(cell->getPort(ID::S));
 
-			for (int idx = 0; idx < GetSize(sel); idx++) {
-				SigSpec new_data = module->addWire(NEW_ID, GetSize(data)/2);
-				for (int i = 0; i < GetSize(new_data); i += width) {
-					RTLIL::Cell *mux = module->addMux(NEW_ID,
+			if(pmux_mode)
+			{
+				int num_cases = 1 << s_width;
+				SigSpec new_a = SigSpec(State::Sx, width);
+				SigSpec new_s = module->addWire(NEW_ID, num_cases);
+				SigSpec new_data = module->addWire(NEW_ID, width);
+				for (int val = 0; val < num_cases; val++)
+				{
+					module->addEq(NEW_ID, sel, SigSpec(val, GetSize(sel)), new_s[val]);
+				}
+				RTLIL::Cell *pmux = module->addPmux(NEW_ID, new_a, data, new_s, new_data);
+				pmux->add_strpool_attribute(ID::src, cell->get_strpool_attribute(ID::src));
+				data = new_data;
+			}
+			else
+			{
+				for (int idx = 0; idx < GetSize(sel); idx++) {
+					SigSpec new_data = module->addWire(NEW_ID, GetSize(data)/2);
+					for (int i = 0; i < GetSize(new_data); i += width) {
+						RTLIL::Cell *mux = module->addMux(NEW_ID,
 							data.extract(i*2, width),
 							data.extract(i*2+width, width),
 							sel[idx],
 							new_data.extract(i, width));
-					mux->add_strpool_attribute(ID::src, cell->get_strpool_attribute(ID::src));
+						mux->add_strpool_attribute(ID::src, cell->get_strpool_attribute(ID::src));
+					}
+					data = new_data;
 				}
-				data = new_data;
 			}
 
 			module->connect(cell->getPort(ID::Y), data);
diff --git a/tests/techmap/bmuxmap_pmux.ys b/tests/techmap/bmuxmap_pmux.ys
new file mode 100644
index 000000000..c75d981e7
--- /dev/null
+++ b/tests/techmap/bmuxmap_pmux.ys
@@ -0,0 +1,45 @@
+read_ilang << EOT
+
+module \top
+  wire width 4 input 0 \S
+  wire width 5 output 1 \Y
+
+  cell $bmux $0
+    parameter \WIDTH 5
+    parameter \S_WIDTH 4
+    connect \A 80'10110100011101110001110010001110101010111000110011111111111110100000110100111000
+    connect \S \S
+    connect \Y \Y
+  end
+end
+
+EOT
+
+hierarchy -auto-top
+equiv_opt -assert bmuxmap -pmux
+
+###
+design -reset
+
+read_ilang << EOT
+
+module \top
+  wire width 10 input 0 \A
+  wire input 1 \S
+  wire width 5 output 2 \Y
+
+  cell $bmux $0
+    parameter \WIDTH 5
+    parameter \S_WIDTH 1
+    connect \A \A
+    connect \S \S
+    connect \Y \Y
+  end
+end
+
+EOT
+
+hierarchy -auto-top
+equiv_opt -assert bmuxmap -pmux
+
+