1 files changed, 36 insertions, 10 deletions

diff --git a/passes/memory/memory_dff.cc b/passes/memory/memory_dff.cc
index 40691d160..be4b3c100 100644
--- a/passes/memory/memory_dff.cc
+++ b/passes/memory/memory_dff.cc
@@ -17,6 +17,7 @@
  *
  */
 
+#include <algorithm>
 #include "kernel/yosys.h"
 #include "kernel/sigtools.h"
 
@@ -33,8 +34,20 @@ struct MemoryDffWorker
 	dict<SigBit, int> sigbit_users_count;
 	dict<SigSpec, Cell*> mux_cells_a, mux_cells_b;
 	pool<Cell*> forward_merged_dffs, candidate_dffs;
+	pool<SigBit> init_bits;
 
-	MemoryDffWorker(Module *module) : module(module), sigmap(module) { }
+	MemoryDffWorker(Module *module) : module(module), sigmap(module)
+	{
+		for (auto wire : module->wires()) {
+			if (wire->attributes.count("\\init") == 0)
+				continue;
+			SigSpec sig = sigmap(wire);
+			Const initval = wire->attributes.at("\\init");
+			for (int i = 0; i < GetSize(sig) && i < GetSize(initval); i++)
+				if (initval[i] == State::S0 || initval[i] == State::S1)
+					init_bits.insert(sig[i]);
+		}
+	}
 
 	bool find_sig_before_dff(RTLIL::SigSpec &sig, RTLIL::SigSpec &clk, bool &clk_polarity, bool after = false)
 	{
@@ -45,6 +58,9 @@ struct MemoryDffWorker
 			if (bit.wire == NULL)
 				continue;
 
+			if (!after && init_bits.count(sigmap(bit)))
+				return false;
+
 			for (auto cell : dff_cells)
 			{
 				if (after && forward_merged_dffs.count(cell))
@@ -72,6 +88,9 @@ struct MemoryDffWorker
 				if (d.size() != 1)
 					continue;
 
+				if (after && init_bits.count(d))
+					return false;
+
 				bit = d;
 				clk = this_clk;
 				clk_polarity = this_clk_polarity;
@@ -164,20 +183,27 @@ struct MemoryDffWorker
 
 		if (mux_cells_a.count(sig_data) || mux_cells_b.count(sig_data))
 		{
-			bool enable_invert = mux_cells_a.count(sig_data) != 0;
-			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"));
+			RTLIL::SigSpec en;
+			std::vector<RTLIL::SigSpec> check_q;
+
+			do {
+				bool enable_invert = mux_cells_a.count(sig_data) != 0;
+				Cell *mux = enable_invert ? mux_cells_a.at(sig_data) : mux_cells_b.at(sig_data);
+				check_q.push_back(sigmap(mux->getPort(enable_invert ? "\\B" : "\\A")));
+				sig_data = sigmap(mux->getPort("\\Y"));
+				en.append(enable_invert ? module->LogicNot(NEW_ID, mux->getPort("\\S")) : mux->getPort("\\S"));
+			} while (mux_cells_a.count(sig_data) || mux_cells_b.count(sig_data));
 
-			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)
+			if (find_sig_before_dff(sig_data, clk_data, clk_polarity, true) && clk_data != RTLIL::SigSpec(RTLIL::State::Sx) &&
+					std::all_of(check_q.begin(), check_q.end(), [&](const SigSpec &cq) {return cq == sig_data; }))
 			{
 				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("\\EN", en.size() > 1 ? module->ReduceAnd(NEW_ID, en) : en);
 				cell->setPort("\\DATA", sig_data);
 				cell->parameters["\\CLK_ENABLE"] = RTLIL::Const(1);
 				cell->parameters["\\CLK_POLARITY"] = RTLIL::Const(clk_polarity);
@@ -236,7 +262,7 @@ struct MemoryDffWorker
 				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)) {
+			if (cell->type.in("$not", "$_NOT_") || (cell->type == "$logic_not" && GetSize(cell->getPort("\\A")) == 1)) {
 				SigSpec sig_a = cell->getPort("\\A");
 				SigSpec sig_y = cell->getPort("\\Y");
 				if (cell->type == "$not")
@@ -265,7 +291,7 @@ struct MemoryDffWorker
 
 struct MemoryDffPass : public Pass {
 	MemoryDffPass() : Pass("memory_dff", "merge input/output DFFs into memories") { }
-	virtual void help()
+	void help() YS_OVERRIDE
 	{
 		//   |---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|
 		log("\n");
@@ -279,7 +305,7 @@ struct MemoryDffPass : public Pass {
 		log("        do not merge registers on read ports\n");
 		log("\n");
 	}
-	virtual void execute(std::vector<std::string> args, RTLIL::Design *design)
+	void execute(std::vector<std::string> args, RTLIL::Design *design) YS_OVERRIDE
 	{
 		bool flag_wr_only = false;