From 805c3024114bffd4769391a5bec9a3752904673c Mon Sep 17 00:00:00 2001
From: Ethan Mahintorabi <ethanmoon@google.com>
Date: Wed, 29 Apr 2026 20:48:14 +0000
Subject: [PATCH] simplemap: Moves $pmux mapping from techmap.v to simple map

This Fixes the slow downs I observed in techmap.v, which we
attempted to fix via the simplify ast.h route originally. This
turned out to be rather complex though.

By moving $pmux to simplemap we can just avoid that code. My
test case now runs in 310s which is 40s faster than the baseline
change.

B:507898959
Signed-off-by: Ethan Mahintorabi <ethanmoon@google.com>
---
 passes/techmap/simplemap.cc | 45 ++++++++++++++++++++++++++++++++++++-
 techlibs/common/techmap.v   | 44 +-----------------------------------
 2 files changed, 45 insertions(+), 44 deletions(-)

diff --git a/passes/techmap/simplemap.cc b/passes/techmap/simplemap.cc
index 0c7d1930e..93001df35 100644
--- a/passes/techmap/simplemap.cc
+++ b/passes/techmap/simplemap.cc
@@ -438,6 +438,48 @@ void simplemap_ff(RTLIL::Module *, RTLIL::Cell *cell)
 	}
 }
 
+void simplemap_pmux(RTLIL::Module *module, RTLIL::Cell *cell)
+{
+	RTLIL::SigSpec sig_a = cell->getPort(ID::A);
+	RTLIL::SigSpec sig_b = cell->getPort(ID::B);
+	RTLIL::SigSpec sig_s = cell->getPort(ID::S);
+	RTLIL::SigSpec sig_y = cell->getPort(ID::Y);
+
+	int width = GetSize(sig_a);
+	int s_width = GetSize(sig_s);
+
+	// Implement: |S
+	RTLIL::SigSpec any_s = sig_s;
+	logic_reduce(module, any_s, cell);
+
+	for (int i = 0; i < width; i++) {
+		RTLIL::SigSpec b_and_bits;
+
+		// Implement: B_AND_BITS = B_AND_S[WIDTH*j+i]
+		for (int j = 0; j < s_width; j++) {
+			RTLIL::Cell *and_gate = module->addCell(NEW_ID, ID($_AND_));
+			transfer_src(and_gate, cell);
+			and_gate->setPort(ID::A, sig_b[j * width + i]);
+			and_gate->setPort(ID::B, sig_s[j]);
+
+			RTLIL::SigSpec and_y = module->addWire(NEW_ID, 1);
+			and_gate->setPort(ID::Y, and_y);
+			b_and_bits.append(and_y);
+		}
+
+		// Implement: Y_B[i] = |B_AND_BITS
+		logic_reduce(module, b_and_bits, cell);
+
+		// Implement: Y[i] = |S ? Y_B[i] : A[i]
+		RTLIL::Cell *mux_gate = module->addCell(NEW_ID, ID($_MUX_));
+		transfer_src(mux_gate, cell);
+		mux_gate->setPort(ID::A, sig_a[i]);
+		mux_gate->setPort(ID::B, b_and_bits);
+		mux_gate->setPort(ID::S, any_s);
+		mux_gate->setPort(ID::Y, sig_y[i]);
+	}
+}
+
 void simplemap_get_mappers(dict<IdString, void(*)(RTLIL::Module*, RTLIL::Cell*)> &mappers)
 {
 	mappers[ID($not)]         = simplemap_not;
@@ -461,6 +503,7 @@ void simplemap_get_mappers(dict<IdString, void(*)(RTLIL::Module*, RTLIL::Cell*)>
 	mappers[ID($ne)]          = simplemap_eqne;
 	mappers[ID($nex)]         = simplemap_eqne;
 	mappers[ID($mux)]         = simplemap_mux;
+	mappers[ID($pmux)]        = simplemap_pmux;
 	mappers[ID($bwmux)]       = simplemap_bwmux;
 	mappers[ID($tribuf)]      = simplemap_tribuf;
 	mappers[ID($bmux)]        = simplemap_bmux;
@@ -515,7 +558,7 @@ struct SimplemapPass : public Pass {
 		log("\n");
 		log("  $not, $pos, $and, $or, $xor, $xnor\n");
 		log("  $reduce_and, $reduce_or, $reduce_xor, $reduce_xnor, $reduce_bool\n");
-		log("  $logic_not, $logic_and, $logic_or, $mux, $tribuf\n");
+		log("  $logic_not, $logic_and, $logic_or, $mux, $pmux, $tribuf\n");
 		log("  $sr, $ff, $dff, $dffe, $dffsr, $dffsre, $adff, $adffe, $aldff, $aldffe, $sdff,\n");
 		log("  $sdffe, $sdffce, $dlatch, $adlatch, $dlatchsr\n");
 		log("\n");
diff --git a/techlibs/common/techmap.v b/techlibs/common/techmap.v
index c3364e628..8061b2e81 100644
--- a/techlibs/common/techmap.v
+++ b/techlibs/common/techmap.v
@@ -59,7 +59,7 @@ module _90_simplemap_compare_ops;
 endmodule
 
 (* techmap_simplemap *)
-(* techmap_celltype = "$buf $pos $slice $concat $mux $tribuf $bmux $bwmux $bweqx" *)
+(* techmap_celltype = "$buf $pos $slice $concat $mux $pmux $tribuf $bmux $bwmux $bweqx" *)
 module _90_simplemap_various;
 endmodule
 
@@ -563,48 +563,6 @@ module _90_pow (A, B, Y);
 	wire _TECHMAP_FAIL_ = 1;
 endmodule
 
-
-// --------------------------------------------------------
-// Parallel Multiplexers
-// --------------------------------------------------------
-
-(* techmap_celltype = "$pmux" *)
-module _90_pmux (A, B, S, Y);
-	parameter WIDTH = 1;
-	parameter S_WIDTH = 1;
-
-	(* force_downto *)
-	input [WIDTH-1:0] A;
-	(* force_downto *)
-	input [WIDTH*S_WIDTH-1:0] B;
-	(* force_downto *)
-	input [S_WIDTH-1:0] S;
-	(* force_downto *)
-	output [WIDTH-1:0] Y;
-
-	(* force_downto *)
-	wire [WIDTH-1:0] Y_B;
-
-	genvar i, j;
-	generate
-		(* force_downto *)
-		wire [WIDTH*S_WIDTH-1:0] B_AND_S;
-		for (i = 0; i < S_WIDTH; i = i + 1) begin:B_AND
-			assign B_AND_S[WIDTH*(i+1)-1:WIDTH*i] = B[WIDTH*(i+1)-1:WIDTH*i] & {WIDTH{S[i]}};
-		end:B_AND
-		for (i = 0; i < WIDTH; i = i + 1) begin:B_OR
-			(* force_downto *)
-			wire [S_WIDTH-1:0] B_AND_BITS;
-			for (j = 0; j < S_WIDTH; j = j + 1) begin:B_AND_BITS_COLLECT
-				assign B_AND_BITS[j] = B_AND_S[WIDTH*j+i];
-			end:B_AND_BITS_COLLECT
-			assign Y_B[i] = |B_AND_BITS;
-		end:B_OR
-	endgenerate
-
-	assign Y = |S ? Y_B : A;
-endmodule
-
 // --------------------------------------------------------
 // Demultiplexers
 // --------------------------------------------------------