From 83b095ab6ca294cb8ccd8504e5dcdae54841199a Mon Sep 17 00:00:00 2001
From: Anhijkt <mihail.ershov.zp@gmail.com>
Date: Sun, 30 Mar 2025 15:43:41 +0300
Subject: [PATCH] opt_expr: optimize pow of 2 cells

---
 passes/opt/opt_expr.cc | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/passes/opt/opt_expr.cc b/passes/opt/opt_expr.cc
index 62a0ffc48..5089959ae 100644
--- a/passes/opt/opt_expr.cc
+++ b/passes/opt/opt_expr.cc
@@ -1690,7 +1690,43 @@ skip_identity:
 			else if (inA == inB)
 				ACTION_DO(ID::Y, cell->getPort(ID::A));
 		}
+		if (cell->type == ID($pow) && cell->getPort(ID::A).is_fully_const() && !cell->parameters[ID::B_SIGNED].as_bool()) {
+			SigSpec sig_a = assign_map(cell->getPort(ID::A));
+			SigSpec sig_y = assign_map(cell->getPort(ID::Y));
+			int y_size = GetSize(sig_y);
 
+			unsigned int bits = unsigned(sig_a.as_int());
+			int bit_count = 0;
+			for (; bits; bits >>= 1)
+				bit_count += (bits & 1);
+
+			if (bit_count == 1) {
+				if (sig_a.as_int() == 2) {
+					log_debug("Replacing pow cell `%s' in module `%s' with left-shift\n", 
+							cell->name.c_str(), module->name.c_str());
+					cell->type = ID($shl);
+					cell->parameters[ID::A_WIDTH] = 1;
+					cell->setPort(ID::A, Const(1, 1));
+				}
+				else {
+					log_debug("Replacing pow cell `%s' in module `%s' with multiply and left-shift\n", 
+							cell->name.c_str(), module->name.c_str());
+					cell->type = ID($mul);
+					cell->parameters[ID::A_SIGNED] = 0;
+
+					int left_shift;
+					sig_a.is_onehot(&left_shift);
+					cell->setPort(ID::A, Const(left_shift, cell->parameters[ID::A_WIDTH].as_int()));
+					
+					SigSpec y_wire = module->addWire(NEW_ID, y_size);
+					cell->setPort(ID::Y, y_wire);
+					
+					module->addShl(NEW_ID, Const(1, 1), y_wire, sig_y);
+				}
+				did_something = true;
+				goto next_cell;
+			}
+		}
 		if (!keepdc && cell->type == ID($mul))
 		{
 			bool a_signed = cell->parameters[ID::A_SIGNED].as_bool();