1 files changed, 55 insertions, 10 deletions

diff --git a/passes/memory/memory_dff.cc b/passes/memory/memory_dff.cc
index 5584f27c3..2eec0207b 100644
--- a/passes/memory/memory_dff.cc
+++ b/passes/memory/memory_dff.cc
@@ -31,6 +31,8 @@ struct MemoryDffWorker
 	vector<Cell*> dff_cells;
 	dict<SigBit, SigBit> invbits;
 	dict<SigBit, int> sigbit_users_count;
+	dict<SigSpec, Cell*> mux_cells_a, mux_cells_b;
+	pool<Cell*> forward_merged_dffs, candidate_dffs;
 
 	MemoryDffWorker(Module *module) : module(module), sigmap(module) { }
 
@@ -45,6 +47,9 @@ struct MemoryDffWorker
 
 			for (auto cell : dff_cells)
 			{
+				if (after && forward_merged_dffs.count(cell))
+					continue;
+
 				SigSpec this_clk = cell->getPort("\\CLK");
 				bool this_clk_polarity = cell->parameters["\\CLK_POLARITY"].as_bool();
 
@@ -70,6 +75,7 @@ struct MemoryDffWorker
 				bit = d;
 				clk = this_clk;
 				clk_polarity = this_clk_polarity;
+				candidate_dffs.insert(cell);
 				goto replaced_this_bit;
 			}
 
@@ -86,6 +92,7 @@ struct MemoryDffWorker
 
 		RTLIL::SigSpec clk = RTLIL::SigSpec(RTLIL::State::Sx);
 		bool clk_polarity = 0;
+		candidate_dffs.clear();
 
 		RTLIL::SigSpec sig_addr = cell->getPort("\\ADDR");
 		if (!find_sig_before_dff(sig_addr, clk, clk_polarity)) {
@@ -105,13 +112,18 @@ struct MemoryDffWorker
 			return;
 		}
 
-		if (clk != RTLIL::SigSpec(RTLIL::State::Sx)) {
+		if (clk != RTLIL::SigSpec(RTLIL::State::Sx))
+		{
+			for (auto cell : candidate_dffs)
+				forward_merged_dffs.insert(cell);
+
 			cell->setPort("\\CLK", clk);
 			cell->setPort("\\ADDR", sig_addr);
 			cell->setPort("\\DATA", sig_data);
 			cell->setPort("\\EN", sig_en);
 			cell->parameters["\\CLK_ENABLE"] = RTLIL::Const(1);
 			cell->parameters["\\CLK_POLARITY"] = RTLIL::Const(clk_polarity);
+
 			log("merged $dff to cell.\n");
 			return;
 		}
@@ -150,16 +162,44 @@ struct MemoryDffWorker
 			if (sigbit_users_count[bit] > 1)
 				goto skip_ff_after_read_merging;
 
-		if (find_sig_before_dff(sig_data, clk_data, clk_polarity, true) && clk_data != RTLIL::SigSpec(RTLIL::State::Sx))
+		if (mux_cells_a.count(sig_data) || mux_cells_b.count(sig_data))
 		{
-			disconnect_dff(sig_data);
-			cell->setPort("\\CLK", clk_data);
-			cell->setPort("\\DATA", sig_data);
-			cell->parameters["\\CLK_ENABLE"] = RTLIL::Const(1);
-			cell->parameters["\\CLK_POLARITY"] = RTLIL::Const(clk_polarity);
-			cell->parameters["\\TRANSPARENT"] = RTLIL::Const(0);
-			log("merged data $dff to cell.\n");
-			return;
+			bool enable_invert = mux_cells_a.count(sig_data);
+			Cell *mux = enable_invert ? mux_cells_a.at(sig_data) : mux_cells_b.at(sig_data);
+			SigSpec check_q = sigmap(mux->getPort(enable_invert ? "\\B" : "\\A"));
+
+			sig_data = sigmap(mux->getPort("\\Y"));
+			for (auto bit : sig_data)
+				if (sigbit_users_count[bit] > 1)
+					goto skip_ff_after_read_merging;
+
+			if (find_sig_before_dff(sig_data, clk_data, clk_polarity, true) && clk_data != RTLIL::SigSpec(RTLIL::State::Sx) && sig_data == check_q)
+			{
+				disconnect_dff(sig_data);
+				cell->setPort("\\CLK", clk_data);
+				cell->setPort("\\EN", enable_invert ? module->LogicNot(NEW_ID, mux->getPort("\\S")) : mux->getPort("\\S"));
+				cell->setPort("\\DATA", sig_data);
+				cell->parameters["\\CLK_ENABLE"] = RTLIL::Const(1);
+				cell->parameters["\\CLK_POLARITY"] = RTLIL::Const(clk_polarity);
+				cell->parameters["\\TRANSPARENT"] = RTLIL::Const(0);
+				log("merged data $dff with rd enable to cell.\n");
+				return;
+			}
+		}
+		else
+		{
+			if (find_sig_before_dff(sig_data, clk_data, clk_polarity, true) && clk_data != RTLIL::SigSpec(RTLIL::State::Sx))
+			{
+				disconnect_dff(sig_data);
+				cell->setPort("\\CLK", clk_data);
+				cell->setPort("\\EN", State::S1);
+				cell->setPort("\\DATA", sig_data);
+				cell->parameters["\\CLK_ENABLE"] = RTLIL::Const(1);
+				cell->parameters["\\CLK_POLARITY"] = RTLIL::Const(clk_polarity);
+				cell->parameters["\\TRANSPARENT"] = RTLIL::Const(0);
+				log("merged data $dff to cell.\n");
+				return;
+			}
 		}
 
 	skip_ff_after_read_merging:;
@@ -169,6 +209,7 @@ struct MemoryDffWorker
 				clk_addr != RTLIL::SigSpec(RTLIL::State::Sx))
 		{
 			cell->setPort("\\CLK", clk_addr);
+			cell->setPort("\\EN", State::S1);
 			cell->setPort("\\ADDR", sig_addr);
 			cell->parameters["\\CLK_ENABLE"] = RTLIL::Const(1);
 			cell->parameters["\\CLK_POLARITY"] = RTLIL::Const(clk_polarity);
@@ -191,6 +232,10 @@ struct MemoryDffWorker
 		for (auto cell : module->cells()) {
 			if (cell->type == "$dff")
 				dff_cells.push_back(cell);
+			if (cell->type == "$mux") {
+				mux_cells_a[sigmap(cell->getPort("\\A"))] = cell;
+				mux_cells_b[sigmap(cell->getPort("\\B"))] = cell;
+			}
 			if (cell->type == "$not" || cell->type == "$_NOT_" || (cell->type == "$logic_not" && GetSize(cell->getPort("\\A")) == 1)) {
 				SigSpec sig_a = cell->getPort("\\A");
 				SigSpec sig_y = cell->getPort("\\Y");