1 files changed, 125 insertions, 23 deletions

diff --git a/passes/memory/memory_dff.cc b/passes/memory/memory_dff.cc
index 726a5c1ff..4adcb462e 100644
--- a/passes/memory/memory_dff.cc
+++ b/passes/memory/memory_dff.cc
@@ -20,6 +20,7 @@
 #include <algorithm>
 #include "kernel/yosys.h"
 #include "kernel/sigtools.h"
+#include "kernel/ffinit.h"
 
 USING_YOSYS_NAMESPACE
 PRIVATE_NAMESPACE_BEGIN
@@ -34,66 +35,158 @@ 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;
+	FfInitVals initvals;
 
 	MemoryDffWorker(Module *module) : module(module), sigmap(module)
 	{
-		for (auto wire : module->wires()) {
-			if (wire->attributes.count(ID::init) == 0)
-				continue;
-			SigSpec sig = sigmap(wire);
-			Const initval = wire->attributes.at(ID::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]);
-		}
+		initvals.set(&sigmap, module);
 	}
 
-	bool find_sig_before_dff(RTLIL::SigSpec &sig, RTLIL::SigSpec &clk, bool &clk_polarity, bool after = false)
+	bool find_sig_before_dff(RTLIL::SigSpec &sig, RTLIL::SigSpec &clk, bool &clk_polarity)
 	{
 		sigmap.apply(sig);
 
+		dict<SigBit, SigBit> cache;
+
 		for (auto &bit : sig)
 		{
+			if (cache.count(bit)) {
+				bit = cache[bit];
+				continue;
+			}
+
 			if (bit.wire == NULL)
 				continue;
 
-			if (!after && init_bits.count(sigmap(bit)))
+			if (initvals(bit) != State::Sx)
 				return false;
 
 			for (auto cell : dff_cells)
 			{
-				if (after && forward_merged_dffs.count(cell))
+				SigSpec this_clk = cell->getPort(ID::CLK);
+				bool this_clk_polarity = cell->parameters[ID::CLK_POLARITY].as_bool();
+
+				if (invbits.count(this_clk)) {
+					this_clk = invbits.at(this_clk);
+					this_clk_polarity = !this_clk_polarity;
+				}
+
+				if (clk != RTLIL::SigSpec(RTLIL::State::Sx)) {
+					if (this_clk != clk)
+						continue;
+					if (this_clk_polarity != clk_polarity)
+						continue;
+				}
+
+				RTLIL::SigSpec q_norm = cell->getPort(ID::Q);
+				sigmap.apply(q_norm);
+
+				RTLIL::SigSpec d = q_norm.extract(bit, &cell->getPort(ID::D));
+				if (d.size() != 1)
+					continue;
+
+				if (cell->type == ID($sdffce)) {
+					SigSpec rval = cell->parameters[ID::SRST_VALUE];
+					SigSpec rbit = q_norm.extract(bit, &rval);
+					if (cell->parameters[ID::SRST_POLARITY].as_bool())
+						d = module->Mux(NEW_ID, d, rbit, cell->getPort(ID::SRST));
+					else
+						d = module->Mux(NEW_ID, rbit, d, cell->getPort(ID::SRST));
+				}
+
+				if (cell->type.in(ID($dffe), ID($sdffe), ID($sdffce))) {
+					if (cell->parameters[ID::EN_POLARITY].as_bool())
+						d = module->Mux(NEW_ID, bit, d, cell->getPort(ID::EN));
+					else
+						d = module->Mux(NEW_ID, d, bit, cell->getPort(ID::EN));
+				}
+
+				if (cell->type.in(ID($sdff), ID($sdffe))) {
+					SigSpec rval = cell->parameters[ID::SRST_VALUE];
+					SigSpec rbit = q_norm.extract(bit, &rval);
+					if (cell->parameters[ID::SRST_POLARITY].as_bool())
+						d = module->Mux(NEW_ID, d, rbit, cell->getPort(ID::SRST));
+					else
+						d = module->Mux(NEW_ID, rbit, d, cell->getPort(ID::SRST));
+				}
+
+				cache[bit] = d;
+				bit = d;
+				clk = this_clk;
+				clk_polarity = this_clk_polarity;
+				candidate_dffs.insert(cell);
+				goto replaced_this_bit;
+			}
+
+			return false;
+		replaced_this_bit:;
+		}
+
+		return true;
+	}
+
+	bool find_sig_after_dffe(RTLIL::SigSpec &sig, RTLIL::SigSpec &clk, bool &clk_polarity, RTLIL::SigSpec &en, bool &en_polarity)
+	{
+		sigmap.apply(sig);
+
+		for (auto &bit : sig)
+		{
+			if (bit.wire == NULL)
+				continue;
+
+			for (auto cell : dff_cells)
+			{
+				if (forward_merged_dffs.count(cell))
+					continue;
+				if (!cell->type.in(ID($dff), ID($dffe)))
 					continue;
 
 				SigSpec this_clk = cell->getPort(ID::CLK);
 				bool this_clk_polarity = cell->parameters[ID::CLK_POLARITY].as_bool();
+				SigSpec this_en = State::S1;
+				bool this_en_polarity = true;
+
+				if (cell->type == ID($dffe)) {
+					this_en = cell->getPort(ID::EN);
+					this_en_polarity = cell->parameters[ID::EN_POLARITY].as_bool();
+				}
 
 				if (invbits.count(this_clk)) {
 					this_clk = invbits.at(this_clk);
 					this_clk_polarity = !this_clk_polarity;
 				}
 
+				if (invbits.count(this_en)) {
+					this_en = invbits.at(this_en);
+					this_en_polarity = !this_en_polarity;
+				}
+
 				if (clk != RTLIL::SigSpec(RTLIL::State::Sx)) {
 					if (this_clk != clk)
 						continue;
 					if (this_clk_polarity != clk_polarity)
 						continue;
+					if (this_en != en)
+						continue;
+					if (this_en_polarity != en_polarity)
+						continue;
 				}
 
-				RTLIL::SigSpec q_norm = cell->getPort(after ? ID::D : ID::Q);
+				RTLIL::SigSpec q_norm = cell->getPort(ID::D);
 				sigmap.apply(q_norm);
 
-				RTLIL::SigSpec d = q_norm.extract(bit, &cell->getPort(after ? ID::Q : ID::D));
+				RTLIL::SigSpec d = q_norm.extract(bit, &cell->getPort(ID::Q));
 				if (d.size() != 1)
 					continue;
 
-				if (after && init_bits.count(d))
+				if (initvals(d) != State::Sx)
 					return false;
 
 				bit = d;
 				clk = this_clk;
 				clk_polarity = this_clk_polarity;
+				en = this_en;
+				en_polarity = this_en_polarity;
 				candidate_dffs.insert(cell);
 				goto replaced_this_bit;
 			}
@@ -161,7 +254,7 @@ struct MemoryDffWorker
 		RTLIL::SigSpec new_sig = module->addWire(sstr.str(), sig.size());
 
 		for (auto cell : module->cells())
-			if (cell->type == ID($dff)) {
+			if (cell->type.in(ID($dff), ID($dffe))) {
 				RTLIL::SigSpec new_q = cell->getPort(ID::Q);
 				new_q.replace(sig, new_sig);
 				cell->setPort(ID::Q, new_q);
@@ -173,8 +266,10 @@ struct MemoryDffWorker
 		log("Checking cell `%s' in module `%s': ", cell->name.c_str(), module->name.c_str());
 
 		bool clk_polarity = 0;
+		bool en_polarity = 0;
 
 		RTLIL::SigSpec clk_data = RTLIL::SigSpec(RTLIL::State::Sx);
+		RTLIL::SigSpec en_data;
 		RTLIL::SigSpec sig_data = cell->getPort(ID::DATA);
 
 		for (auto bit : sigmap(sig_data))
@@ -198,9 +293,14 @@ 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 (find_sig_after_dffe(sig_data, clk_data, clk_polarity, en_data, en_polarity) && clk_data != RTLIL::SigSpec(RTLIL::State::Sx) &&
 					std::all_of(check_q.begin(), check_q.end(), [&](const SigSpec &cq) {return cq == sig_data; }))
 			{
+				if (en_data != State::S1 || !en_polarity) {
+					if (!en_polarity)
+						en_data = module->LogicNot(NEW_ID, en_data);
+					en.append(en_data);
+				}
 				disconnect_dff(sig_data);
 				cell->setPort(ID::CLK, clk_data);
 				cell->setPort(ID::EN, en.size() > 1 ? module->ReduceAnd(NEW_ID, en) : en);
@@ -214,11 +314,13 @@ struct MemoryDffWorker
 		}
 		else
 		{
-			if (find_sig_before_dff(sig_data, clk_data, clk_polarity, true) && clk_data != RTLIL::SigSpec(RTLIL::State::Sx))
+			if (find_sig_after_dffe(sig_data, clk_data, clk_polarity, en_data, en_polarity) && clk_data != RTLIL::SigSpec(RTLIL::State::Sx))
 			{
+				if (!en_polarity)
+					en_data = module->LogicNot(NEW_ID, en_data);
 				disconnect_dff(sig_data);
 				cell->setPort(ID::CLK, clk_data);
-				cell->setPort(ID::EN, State::S1);
+				cell->setPort(ID::EN, en_data);
 				cell->setPort(ID::DATA, sig_data);
 				cell->parameters[ID::CLK_ENABLE] = RTLIL::Const(1);
 				cell->parameters[ID::CLK_POLARITY] = RTLIL::Const(clk_polarity);
@@ -256,7 +358,7 @@ struct MemoryDffWorker
 		}
 
 		for (auto cell : module->cells()) {
-			if (cell->type == ID($dff))
+			if (cell->type.in(ID($dff), ID($dffe), ID($sdff), ID($sdffe), ID($sdffce)))
 				dff_cells.push_back(cell);
 			if (cell->type == ID($mux)) {
 				mux_cells_a[sigmap(cell->getPort(ID::A))] = cell;
@@ -291,7 +393,7 @@ struct MemoryDffWorker
 
 struct MemoryDffPass : public Pass {
 	MemoryDffPass() : Pass("memory_dff", "merge input/output DFFs into memories") { }
-	void help() YS_OVERRIDE
+	void help() override
 	{
 		//   |---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|
 		log("\n");
@@ -305,7 +407,7 @@ struct MemoryDffPass : public Pass {
 		log("        do not merge registers on read ports\n");
 		log("\n");
 	}
-	void execute(std::vector<std::string> args, RTLIL::Design *design) YS_OVERRIDE
+	void execute(std::vector<std::string> args, RTLIL::Design *design) override
 	{
 		bool flag_wr_only = false;