From 721b5044799d891620f734081a31221617fbec84 Mon Sep 17 00:00:00 2001
From: Natalia <nataliakokoromyti@gmail.com>
Date: Thu, 18 Dec 2025 13:06:22 -0800
Subject: [PATCH] lut2mux: add -word option and test

---
 passes/techmap/lut2mux.cc | 38 +++++++++++++++++++++++------------
 tests/techmap/lut2mux.ys  | 42 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 13 deletions(-)
 create mode 100644 tests/techmap/lut2mux.ys

diff --git a/passes/techmap/lut2mux.cc b/passes/techmap/lut2mux.cc
index ef76e0deb..28f466874 100644
--- a/passes/techmap/lut2mux.cc
+++ b/passes/techmap/lut2mux.cc
@@ -23,7 +23,7 @@
 USING_YOSYS_NAMESPACE
 PRIVATE_NAMESPACE_BEGIN
 
-int lut2mux(Cell *cell)
+int lut2mux(Cell *cell, bool word_mode)
 {
 	SigSpec sig_a = cell->getPort(ID::A);
 	SigSpec sig_y = cell->getPort(ID::Y);
@@ -32,7 +32,10 @@ int lut2mux(Cell *cell)
 
 	if (GetSize(sig_a) == 1)
 	{
-		cell->module->addMuxGate(NEW_ID, lut.extract(0)[0], lut.extract(1)[0], sig_a, sig_y);
+		if (!word_mode)
+			cell->module->addMuxGate(NEW_ID, lut.extract(0)[0], lut.extract(1)[0], sig_a, sig_y); 
+		else
+		    cell->module->addMux(NEW_ID, lut.extract(0)[0], lut.extract(1)[0], sig_a, sig_y);
 	}
 	else
 	{
@@ -44,10 +47,13 @@ int lut2mux(Cell *cell)
 		Const lut1 = lut.extract(0, GetSize(lut)/2);
 		Const lut2 = lut.extract(GetSize(lut)/2, GetSize(lut)/2);
 
-		count += lut2mux(cell->module->addLut(NEW_ID, sig_a_lo, sig_y1, lut1));
-		count += lut2mux(cell->module->addLut(NEW_ID, sig_a_lo, sig_y2, lut2));
+		count += lut2mux(cell->module->addLut(NEW_ID, sig_a_lo, sig_y1, lut1), word_mode); 
+		count += lut2mux(cell->module->addLut(NEW_ID, sig_a_lo, sig_y2, lut2), word_mode);
 
-		cell->module->addMuxGate(NEW_ID, sig_y1, sig_y2, sig_a_hi, sig_y);
+		if (!word_mode)
+			cell->module->addMuxGate(NEW_ID, sig_y1, sig_y2, sig_a_hi, sig_y); 
+		else
+			cell->module->addMux(NEW_ID, sig_y1, sig_y2, sig_a_hi, sig_y);
 	}
 
 	cell->module->remove(cell);
@@ -55,35 +61,41 @@ int lut2mux(Cell *cell)
 }
 
 struct Lut2muxPass : public Pass {
-	Lut2muxPass() : Pass("lut2mux", "convert $lut to $_MUX_") { }
+	Lut2muxPass() : Pass("lut2mux", "convert $lut to $mux/$_MUX_") { }
 	void help() override
 	{
 		//   |---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|
 		log("\n");
 		log("    lut2mux [options] [selection]\n");
 		log("\n");
-		log("This pass converts $lut cells to $_MUX_ gates.\n");
+		log("This pass converts $lut cells to $mux/$_MUX_ gates.\n");
+		log("\n");
+		log("    -word\n");
+		log("        Convert $lut cells with a single input to word-level $mux gates.\n");
+		log("   	   The default is to convert them to bit-level $_MUX_ gates.\n");
 		log("\n");
 	}
 	void execute(std::vector<std::string> args, RTLIL::Design *design) override
 	{
-		log_header(design, "Executing LUT2MUX pass (convert $lut to $_MUX_).\n");
+		log_header(design, "Executing LUT2MUX pass (convert $lut to $mux/$_MUX_).\n");
+			log("ARGS:"); for (auto &a: args) log(" [%s]", a.c_str()); log("\n");
 
 		size_t argidx;
+		bool word_mode = false;
 		for (argidx = 1; argidx < args.size(); argidx++)
 		{
-			// if (args[argidx] == "-v") {
-			// 	continue;
-			// }
+			if (args[argidx] == "-word") {
+				word_mode = true;
+				continue;
+			}
 			break;
 		}
-		extra_args(args, argidx, design);
 
 		for (auto module : design->selected_modules())
 		for (auto cell : module->selected_cells()) {
 			if (cell->type == ID($lut)) {
 				IdString cell_name = cell->name;
-				int count = lut2mux(cell);
+				int count = lut2mux(cell, word_mode);
 				log("Converted %s.%s to %d MUX cells.\n", log_id(module), log_id(cell_name), count);
 			}
 		}
diff --git a/tests/techmap/lut2mux.ys b/tests/techmap/lut2mux.ys
new file mode 100644
index 000000000..212003756
--- /dev/null
+++ b/tests/techmap/lut2mux.ys
@@ -0,0 +1,42 @@
+# Test lut2mux pass using a directly constructed $lut (avoids frontend/synth differences in test-verific)
+
+read_rtlil << EOT
+module \top
+  wire width 2 input 1 \a
+  wire width 1 output 2 \y
+  cell $lut \u_lut
+    parameter \WIDTH 2
+    parameter \LUT 4'0110
+    connect \A \a
+    connect \Y \y
+  end
+end
+EOT
+
+select -assert-count 1 t:$lut
+
+# default mode -> gate-level $_MUX_
+design -save gold
+lut2mux
+rename \top \gate
+select -assert-count 3 gate/t:$_MUX_
+select -assert-count 0 gate/t:$mux
+select -assert-count 0 gate/t:$lut
+
+# -word mode -> word-level $mux
+design -copy-from gold -as top \top
+select -none
+select top
+lut2mux -word
+select -clear
+rename \top \word
+select -assert-count 3 word/t:$mux
+select -assert-count 0 word/t:$_MUX_
+select -assert-count 0 gate/t:$lut
+
+# equivalence
+equiv_make \gate \word equiv
+hierarchy -top equiv
+equiv_simple
+equiv_induct
+equiv_status -assert