opt: Remove O(n²) opt routines across the codebase for pmux

$pmux have bad runtime behavior with large cases and output
bit widths. On one Google test case synthesis takes 57 hours
with these changes it takes 350s.

Signed-off-by: Ethan Mahintorabi <ethanmoon@google.com>

passes/opt/opt_expr.cc

@@ -1491,9 +1491,12 @@ skip_identity:
 				replace_cell(assign_map, module, cell, "mux_undef", ID::Y, cell->getPort(ID::A));
 				goto next_cell;
 			}
+
+			std::vector<RTLIL::SigBit> cell_b_bits = cell->getPort(ID::B).to_sigbit_vector();
 			for (int i = 0; i < cell->getPort(ID::S).size(); i++) {
-				RTLIL::SigSpec old_b = cell->getPort(ID::B).extract(i*width, width);
-				RTLIL::SigSpec old_s = cell->getPort(ID::S).extract(i, 1);
+				auto start = cell_b_bits.begin() + i * width;
+				RTLIL::SigSpec old_b(std::vector<RTLIL::SigBit>{start, start + width});
+				RTLIL::SigSpec old_s{cell->getPort(ID::S)[i]};
 				if (old_b.is_fully_undef() || old_s.is_fully_undef())
 					continue;
 				new_b.append(old_b);

passes/opt/opt_merge.cc

@@ -112,8 +112,12 @@ struct OptMergeThreadWorker
 		int width = GetSize(sig_b) / s_width;
 
 		hashlib::commutative_hash comm;
-		for (int i = 0; i < s_width; i++)
-			comm.eat(hash_pair(sig_s[i], sig_b.extract(i*width, width)));
+		std::vector<RTLIL::SigBit> sig_b_bits = sig_b.to_sigbit_vector();
+		for (int i = 0; i < s_width; i++){
+			auto start = sig_b_bits.begin() + i*width;
+			SigSpec sig_b_i(std::vector<RTLIL::SigBit>{start, start + width});
+			comm.eat(hash_pair(SigSpec{sig_s[i]}, sig_b_i));
+		}
 
 		return comm.hash_into(h);
 	}

passes/opt/opt_muxtree.cc

@@ -21,11 +21,9 @@
 #include "kernel/sigtools.h"
 #include "kernel/log.h"
 #include "kernel/celltypes.h"
+#include <cstddef>
 #include <stdlib.h>
 #include <stdio.h>
-#include <unordered_map>
-#include <unordered_set>
-#include <set>
 
 USING_YOSYS_NAMESPACE
 PRIVATE_NAMESPACE_BEGIN
@@ -126,9 +124,13 @@ struct OptMuxtreeWorker
 
 		// Analyze port B
 		// In case of $pmux, port B is multiple slices, concatenated, one per bit of port S
+		std::vector<RTLIL::SigBit> sig_b_bits = sig_b.to_sigbit_vector();
 		for (int i = 0; i < GetSize(sig_s); i++) {
-			RTLIL::SigSpec sig = sig_b.extract(i*GetSize(sig_a), GetSize(sig_a));
-			RTLIL::SigSpec ctrl_sig = assign_map(sig_s.extract(i, 1));
+			auto start = sig_b_bits.begin() + i * GetSize(sig_a);
+			RTLIL::SigSpec sig(std::vector<RTLIL::SigBit>{
+         		start, start + GetSize(sig_a)
+      		});
+			RTLIL::SigSpec ctrl_sig = assign_map(SigSpec{sig_s[i]});
 			portinfo_t portinfo = used_port_bit(sig, this_mux_idx);
 			portinfo.ctrl_sig = sig2bits(ctrl_sig, false).front();
 			portinfo.const_activated = ctrl_sig.is_fully_const() && ctrl_sig.as_bool();
@@ -342,14 +344,21 @@ struct OptMuxtreeWorker
 		// The payload is a reference counter used to manage the list
 		// When it is non-zero, the signal in known to be inactive
 		// When it reaches zero, the map element is removed
-		std::unordered_map<int, int> known_inactive;
+		std::vector<int> known_inactive;
 
 		// database of known active signals
-		std::unordered_map<int, int> known_active;
+		std::vector<int> known_active;
 
 		// this is just used to keep track of visited muxes in order to prohibit
 		// endless recursion in mux loops
-		std::unordered_set<int> visited_muxes;
+		std::vector<bool> visited_muxes;
+
+		// Initialize with the maximum possible sizes
+		knowledge_t(int num_bits, int num_muxes) {
+			known_inactive.assign(num_bits, 0);
+			known_active.assign(num_bits, 0);
+			visited_muxes.assign(num_muxes, false);
+		}
 	};
 
 	static void activate_port(knowledge_t &knowledge, int port_idx, const muxinfo_t &muxinfo) {
@@ -366,11 +375,10 @@ struct OptMuxtreeWorker
 	}
 
 	static void deactivate_port(knowledge_t &knowledge, int port_idx, const muxinfo_t &muxinfo) {
-		auto unlearn = [](std::unordered_map<int, int>& knowns, int i) {
-			auto it = knowns.find(i);
-			if (it != knowns.end())
-				if (--it->second == 0)
-					knowns.erase(it);
+		auto unlearn = [](std::vector<int>& knowns, int bit_idx) {
+			if (knowns[bit_idx] > 0) {
+				--knowns[bit_idx];
+			}
 		};
 
 		if (port_idx < GetSize(muxinfo.ports)-1 && !muxinfo.ports[port_idx].const_activated)
@@ -418,9 +426,9 @@ struct OptMuxtreeWorker
 
 		vector<int> input_mux_queue;
 		for (int m : muxinfo.ports[port_idx].input_muxes) {
-			if (knowledge.visited_muxes.count(m))
+			if (knowledge.visited_muxes[m])
 				continue;
-			knowledge.visited_muxes.insert(m);
+			knowledge.visited_muxes[m] = true;
 			input_mux_queue.push_back(m);
 		}
 		for (int m : input_mux_queue) {
@@ -453,7 +461,7 @@ struct OptMuxtreeWorker
 		// Allow revisiting input muxes, since evaluating other ports should
 		// revisit these input muxes with different activation assumptions
 		for (int m : input_mux_queue)
-			knowledge.visited_muxes.erase(m);
+			knowledge.visited_muxes[m] = false;
 
 		// Undo our assumptions that the port is active
 		deactivate_port(knowledge, port_idx, muxinfo);
@@ -475,11 +483,11 @@ struct OptMuxtreeWorker
 		vector<int> bits = sig2bits(sig, false);
 		for (int i = 0; i < GetSize(bits); i++) {
 			if (bits[i] >= 0) {
-				if (knowledge.known_inactive.count(bits[i]) > 0) {
+				if (knowledge.known_inactive[bits[i]] > 0) {
 					sig[i] = State::S0;
 					did_something = true;
 				} else
-				if (knowledge.known_active.count(bits[i]) > 0) {
+				if (knowledge.known_active[bits[i]] > 0) {
 					sig[i] = State::S1;
 					did_something = true;
 				}
@@ -552,7 +560,7 @@ struct OptMuxtreeWorker
 			portinfo_t &portinfo = muxinfo.ports[port_idx];
 			if (portinfo.const_deactivated)
 				continue;
-			if (knowledge.known_active.count(portinfo.ctrl_sig) > 0) {
+			if (knowledge.known_active[portinfo.ctrl_sig] > 0) {
 				eval_mux_port(knowledge, mux_idx, port_idx, limits);
 				return;
 			}
@@ -566,7 +574,7 @@ struct OptMuxtreeWorker
 			if (portinfo.const_deactivated)
 				continue;
 			if (port_idx < GetSize(muxinfo.ports)-1)
-				if (knowledge.known_inactive.count(portinfo.ctrl_sig) > 0)
+				if (knowledge.known_inactive[portinfo.ctrl_sig] > 0)
 					continue;
 			eval_mux_port(knowledge, mux_idx, port_idx, limits);
 
@@ -578,8 +586,8 @@ struct OptMuxtreeWorker
 	void eval_root_mux(int mux_idx)
 	{
 		log_assert(glob_evals_left > 0);
-		knowledge_t knowledge;
-		knowledge.visited_muxes.insert(mux_idx);
+		knowledge_t knowledge(GetSize(bit2info), GetSize(mux2info));
+		knowledge.visited_muxes[mux_idx] = true;
 		limits_t limits = {};
 		limits.do_mark_ports_observable = root_enable_muxes.at(mux_idx);
 		eval_mux(knowledge, mux_idx, limits);

passes/opt/opt_reduce.cc

@@ -110,22 +110,21 @@ struct OptReduceWorker
 		RTLIL::SigSpec sig_s = assign_map(cell->getPort(ID::S));
 
 		RTLIL::SigSpec new_sig_b, new_sig_s;
-		pool<RTLIL::SigSpec> handled_sig;
+		dict<RTLIL::SigSpec, std::vector<RTLIL::SigBit>> grouped_b_to_s;
 
-		handled_sig.insert(sig_a);
-		for (int i = 0; i < sig_s.size(); i++)
-		{
-			RTLIL::SigSpec this_b = sig_b.extract(i*sig_a.size(), sig_a.size());
-			if (handled_sig.count(this_b) > 0)
-				continue;
-
-			RTLIL::SigSpec this_s = sig_s.extract(i, 1);
-			for (int j = i+1; j < sig_s.size(); j++) {
-				RTLIL::SigSpec that_b = sig_b.extract(j*sig_a.size(), sig_a.size());
-				if (this_b == that_b)
-					this_s.append(sig_s.extract(j, 1));
+		std::vector<SigBit> sig_b_bits = sig_b.to_sigbit_vector();
+		for (int i = 0; i < sig_s.size(); i++) {
+			auto start = sig_b_bits.begin() + i*sig_a.size();
+			RTLIL::SigSpec this_b(std::vector<RTLIL::SigBit>{start, start + sig_a.size()});
+			if (grouped_b_to_s.count(this_b)) {
+				grouped_b_to_s[this_b].push_back(sig_s[i]);
+			} else {
+				grouped_b_to_s[this_b] = {sig_s[i]};
 			}
+		}
 
+		for (auto &[this_b, this_s_bit] : grouped_b_to_s) {
+			RTLIL::SigSpec this_s{this_s_bit};
 			if (this_s.size() > 1)
 			{
 				RTLIL::Cell *reduce_or_cell = module->addCell(NEW_ID, ID($reduce_or));
@@ -141,7 +140,6 @@ struct OptReduceWorker
 
 			new_sig_b.append(this_b);
 			new_sig_s.append(this_s);
-			handled_sig.insert(this_b);
 		}
 
 		if (new_sig_s.size() == 0)

passes/opt/opt_share.cc

@@ -423,12 +423,15 @@ struct OptSharePass : public Pass {
 				std::vector<std::set<OpMuxConn>> mux_port_conns(mux_port_num);
 				int found = 0;
 
+				std::vector<RTLIL::SigBit> mux_b_bits = mux->getPort(ID::B).to_sigbit_vector();
 				for (int mux_port_id = 0; mux_port_id < mux_port_num; mux_port_id++) {
 					SigSpec mux_insig;
 					if (mux_port_id == mux_port_num - 1) {
 						mux_insig = sigmap(mux->getPort(ID::A));
 					} else {
-						mux_insig = sigmap(mux->getPort(ID::B).extract(mux_port_id * mux_port_size, mux_port_size));
+						auto mux_b_bit_slice = mux_b_bits.begin() + mux_port_id * mux_port_size;
+						SigSpec b_extracted_sigspec(std::vector<RTLIL::SigBit>{mux_b_bit_slice, mux_b_bit_slice + mux_port_size});
+						mux_insig = sigmap(b_extracted_sigspec);
 					}
 
 					for (int mux_port_offset = 0; mux_port_offset < mux_port_size; ++mux_port_offset) {