From 2c9c6e693fddd66df83888ead9a5f43bc3f08dbd Mon Sep 17 00:00:00 2001
From: Akash Levy <akashlevy@gmail.com>
Date: Fri, 30 Aug 2024 04:45:01 -0700
Subject: [PATCH] Add muxadd peepopt

---
 Makefile                        |  1 +
 passes/pmgen/Makefile.inc       |  1 +
 passes/pmgen/peepopt.cc         |  3 +++
 passes/pmgen/peepopt_muxadd.pmg | 42 +++++++++++++++++++++++++++++++++
 4 files changed, 47 insertions(+)
 create mode 100644 passes/pmgen/peepopt_muxadd.pmg

diff --git a/Makefile b/Makefile
index 2846be497..917369b80 100644
--- a/Makefile
+++ b/Makefile
@@ -716,6 +716,7 @@ OBJS += passes/cmds/clean_zerowidth.o
 OBJS += passes/cmds/splitfanout.o
 
 include $(YOSYS_SRC)/passes/memory/Makefile.inc
+include $(YOSYS_SRC)/passes/pmgen/Makefile.inc
 
 include $(YOSYS_SRC)/passes/proc/Makefile.inc
 include $(YOSYS_SRC)/passes/opt/Makefile.inc
diff --git a/passes/pmgen/Makefile.inc b/passes/pmgen/Makefile.inc
index ed22d3209..44d2225c3 100644
--- a/passes/pmgen/Makefile.inc
+++ b/passes/pmgen/Makefile.inc
@@ -57,6 +57,7 @@ PEEPOPT_PATTERN  = passes/pmgen/peepopt_shiftmul_right.pmg
 PEEPOPT_PATTERN += passes/pmgen/peepopt_shiftmul_left.pmg
 PEEPOPT_PATTERN += passes/pmgen/peepopt_shiftadd.pmg
 PEEPOPT_PATTERN += passes/pmgen/peepopt_muldiv.pmg
+PEEPOPT_PATTERN += passes/pmgen/peepopt_muxadd.pmg
 
 passes/pmgen/peepopt_pm.h: passes/pmgen/pmgen.py $(PEEPOPT_PATTERN)
 	$(P) mkdir -p passes/pmgen && $(PYTHON_EXECUTABLE) $< -o $@ -p peepopt $(filter-out $<,$^)
diff --git a/passes/pmgen/peepopt.cc b/passes/pmgen/peepopt.cc
index 5638ec3c2..7de4e9fba 100644
--- a/passes/pmgen/peepopt.cc
+++ b/passes/pmgen/peepopt.cc
@@ -42,6 +42,8 @@ struct PeepoptPass : public Pass {
 		log("\n");
 		log("This pass employs the following rules:\n");
 		log("\n");
+		log("   * muxadd - Replace S?(A+B):A with A+(S?B:0)\n");
+		log("\n");
 		log("   * muldiv - Replace (A*B)/B with A\n");
 		log("\n");
 		log("   * shiftmul - Replace A>>(B*C) with A'>>(B<<K) where C and K are constants\n");
@@ -90,6 +92,7 @@ struct PeepoptPass : public Pass {
 				pm.run_shiftmul_right();
 				pm.run_shiftmul_left();
 				pm.run_muldiv();
+				pm.run_muxadd();
 			}
 		}
 	}
diff --git a/passes/pmgen/peepopt_muxadd.pmg b/passes/pmgen/peepopt_muxadd.pmg
new file mode 100644
index 000000000..b211fd38c
--- /dev/null
+++ b/passes/pmgen/peepopt_muxadd.pmg
@@ -0,0 +1,42 @@
+pattern muxadd
+//
+// Transforms add->mux into mux->add:
+// y = s ? (a + b) : a   ===>   y = a + (s ? b : 0)
+//
+
+state <SigSpec> add_y add_a add_b
+
+match add
+	select add->type == $add
+endmatch
+
+code add_y add_a add_b
+	add_y = port(add, \Y);
+	add_a = port(add, \A);
+	add_b = port(add, \B);
+	branch;
+	std::swap(add_a, add_b);
+endcode
+
+match mux
+	select mux->type == $mux
+	index <SigSpec> port(mux, \A) === SigSpec({Const(State::S0, GetSize(add_y)-GetSize(add_a)), add_a})
+	index <SigSpec> port(mux, \B) === add_y
+endmatch
+
+code
+	SigSpec mux_y = port(mux, \Y);
+	SigSpec mid = module->addWire(NEW_ID, GetSize(add_b));
+
+	mux->setPort(\A, Const(State::S0, GetSize(add_b)));
+	mux->setPort(\B, add_b);
+	mux->setPort(\Y, mid);
+	add->setPort(\B, mid);
+	add->setPort(\Y, mux_y);
+
+	log("muxadd pattern in %s: mux=%s, add=%s\n", log_id(module), log_id(mux), log_id(add));
+
+	mux->fixup_parameters();
+	accept;
+endcode
+