Add bitwise `$bweqx` and `$bwmux` cells

The new bitwise case equality (`$bweqx`) and bitwise mux (`$bwmux`)
cells enable compact encoding and decoding of 3-valued logic signals
using multiple 2-valued signals.

5 files changed, 112 insertions, 4 deletions

diff --git a/kernel/calc.cc b/kernel/calc.cc
index 32e8a49fe..9b02a6e30 100644
--- a/kernel/calc.cc
+++ b/kernel/calc.cc
@@ -690,5 +690,28 @@ RTLIL::Const RTLIL::const_demux(const RTLIL::Const &arg1, const RTLIL::Const &ar
 	return res;
 }
 
+RTLIL::Const RTLIL::const_bweqx(const RTLIL::Const &arg1, const RTLIL::Const &arg2)
+{
+	log_assert(arg2.size() == arg1.size());
+	RTLIL::Const result(RTLIL::State::S0, arg1.size());
+	for (int i = 0; i < arg1.size(); i++)
+		result[i] = arg1[i] == arg2[i] ? State::S1 : State::S0;
+
+	return result;
+}
+
+RTLIL::Const RTLIL::const_bwmux(const RTLIL::Const &arg1, const RTLIL::Const &arg2, const RTLIL::Const &arg3)
+{
+	log_assert(arg2.size() == arg1.size());
+	log_assert(arg3.size() == arg1.size());
+	RTLIL::Const result(RTLIL::State::Sx, arg1.size());
+	for (int i = 0; i < arg1.size(); i++) {
+		if (arg3[i] != State::Sx || arg1[i] == arg2[i])
+			result[i] = arg3[i] == State::S1 ? arg2[i] : arg1[i];
+	}
+
+	return result;
+}
+
 YOSYS_NAMESPACE_END
 
diff --git a/kernel/celltypes.h b/kernel/celltypes.h
index e617b4603..63e7408c1 100644
--- a/kernel/celltypes.h
+++ b/kernel/celltypes.h
@@ -116,7 +116,8 @@ struct CellTypes
 			ID($shl), ID($shr), ID($sshl), ID($sshr), ID($shift), ID($shiftx),
 			ID($lt), ID($le), ID($eq), ID($ne), ID($eqx), ID($nex), ID($ge), ID($gt),
 			ID($add), ID($sub), ID($mul), ID($div), ID($mod), ID($divfloor), ID($modfloor), ID($pow),
-			ID($logic_and), ID($logic_or), ID($concat), ID($macc)
+			ID($logic_and), ID($logic_or), ID($concat), ID($macc),
+			ID($bweqx)
 		};
 
 		for (auto type : unary_ops)
@@ -125,7 +126,7 @@ struct CellTypes
 		for (auto type : binary_ops)
 			setup_type(type, {ID::A, ID::B}, {ID::Y}, true);
 
-		for (auto type : std::vector<RTLIL::IdString>({ID($mux), ID($pmux)}))
+		for (auto type : std::vector<RTLIL::IdString>({ID($mux), ID($pmux), ID($bwmux)}))
 			setup_type(type, {ID::A, ID::B, ID::S}, {ID::Y}, true);
 
 		for (auto type : std::vector<RTLIL::IdString>({ID($bmux), ID($demux)}))
@@ -430,6 +431,11 @@ struct CellTypes
 			return const_demux(arg1, arg2);
 		}
 
+		if (cell->type == ID($bweqx))
+		{
+			return const_bweqx(arg1, arg2);
+		}
+
 		if (cell->type == ID($lut))
 		{
 			int width = cell->parameters.at(ID::WIDTH).as_int();
@@ -490,6 +496,8 @@ struct CellTypes
 	{
 		if (cell->type.in(ID($mux), ID($_MUX_)))
 			return const_mux(arg1, arg2, arg3);
+		if (cell->type == ID($bwmux))
+			return const_bwmux(arg1, arg2, arg3);
 		if (cell->type == ID($pmux))
 			return const_pmux(arg1, arg2, arg3);
 		if (cell->type == ID($_AOI3_))
diff --git a/kernel/rtlil.cc b/kernel/rtlil.cc
index 5211c3b3f..eee014c54 100644
--- a/kernel/rtlil.cc
+++ b/kernel/rtlil.cc
@@ -1613,6 +1613,23 @@ namespace {
 				return;
 			}
 
+			if (cell->type == ID($bweqx)) {
+				port(ID::A, param(ID::WIDTH));
+				port(ID::B, param(ID::WIDTH));
+				port(ID::Y, param(ID::WIDTH));
+				check_expected();
+				return;
+			}
+
+			if (cell->type == ID($bwmux)) {
+				port(ID::A, param(ID::WIDTH));
+				port(ID::B, param(ID::WIDTH));
+				port(ID::S, param(ID::WIDTH));
+				port(ID::Y, param(ID::WIDTH));
+				check_expected();
+				return;
+			}
+
 			if (cell->type.in(ID($assert), ID($assume), ID($live), ID($fair), ID($cover))) {
 				port(ID::A, 1);
 				port(ID::EN, 1);
@@ -2466,6 +2483,7 @@ DEF_METHOD(Sshr,     sig_a.size(), ID($sshr))
 		return sig_y;                                             \
 	}
 DEF_METHOD(Mux,      ID($mux),        0)
+DEF_METHOD(Bwmux,    ID($bwmux),      0)
 DEF_METHOD(Pmux,     ID($pmux),       1)
 #undef DEF_METHOD
 
@@ -2489,6 +2507,24 @@ DEF_METHOD(Bmux,     ID($bmux),       0)
 DEF_METHOD(Demux,    ID($demux),      1)
 #undef DEF_METHOD
 
+#define DEF_METHOD(_func, _type) \
+	RTLIL::Cell* RTLIL::Module::add ## _func(RTLIL::IdString name, const RTLIL::SigSpec &sig_a, const RTLIL::SigSpec &sig_b, const RTLIL::SigSpec &sig_y, const std::string &src) { \
+		RTLIL::Cell *cell = addCell(name, _type);                 \
+		cell->parameters[ID::WIDTH] = sig_a.size();               \
+		cell->setPort(ID::A, sig_a);                              \
+		cell->setPort(ID::B, sig_b);                              \
+		cell->setPort(ID::Y, sig_y);                              \
+		cell->set_src_attribute(src);                             \
+		return cell;                                              \
+	} \
+	RTLIL::SigSpec RTLIL::Module::_func(RTLIL::IdString name, const RTLIL::SigSpec &sig_a, const RTLIL::SigSpec &sig_s, const std::string &src) { \
+		RTLIL::SigSpec sig_y = addWire(NEW_ID, sig_a.size());     \
+		add ## _func(name, sig_a, sig_s, sig_y, src);             \
+		return sig_y;                                             \
+	}
+DEF_METHOD(Bweqx,    ID($bweqx))
+#undef DEF_METHOD
+
 #define DEF_METHOD_2(_func, _type, _P1, _P2) \
 	RTLIL::Cell* RTLIL::Module::add ## _func(RTLIL::IdString name, const RTLIL::SigBit &sig1, const RTLIL::SigBit &sig2, const std::string &src) { \
 		RTLIL::Cell *cell = addCell(name, _type);         \
diff --git a/kernel/rtlil.h b/kernel/rtlil.h
index 755abf534..42bb66da8 100644
--- a/kernel/rtlil.h
+++ b/kernel/rtlil.h
@@ -505,6 +505,9 @@ namespace RTLIL
 	RTLIL::Const const_bmux        (const RTLIL::Const &arg1, const RTLIL::Const &arg2);
 	RTLIL::Const const_demux       (const RTLIL::Const &arg1, const RTLIL::Const &arg2);
 
+	RTLIL::Const const_bweqx       (const RTLIL::Const &arg1, const RTLIL::Const &arg2);
+	RTLIL::Const const_bwmux       (const RTLIL::Const &arg1, const RTLIL::Const &arg2, const RTLIL::Const &arg3);
+
 
 	// This iterator-range-pair is used for Design::modules(), Module::wires() and Module::cells().
 	// It maintains a reference counter that is used to make sure that the container is not modified while being iterated over.
@@ -1303,6 +1306,9 @@ public:
 	RTLIL::Cell* addBmux (RTLIL::IdString name, const RTLIL::SigSpec &sig_a, const RTLIL::SigSpec &sig_s, const RTLIL::SigSpec &sig_y, const std::string &src = "");
 	RTLIL::Cell* addDemux (RTLIL::IdString name, const RTLIL::SigSpec &sig_a, const RTLIL::SigSpec &sig_s, const RTLIL::SigSpec &sig_y, const std::string &src = "");
 
+	RTLIL::Cell* addBweqx  (RTLIL::IdString name, const RTLIL::SigSpec &sig_a, const RTLIL::SigSpec &sig_b, const RTLIL::SigSpec &sig_y, const std::string &src = "");
+	RTLIL::Cell* addBwmux  (RTLIL::IdString name, const RTLIL::SigSpec &sig_a, const RTLIL::SigSpec &sig_b, const RTLIL::SigSpec &sig_s, const RTLIL::SigSpec &sig_y, const std::string &src = "");
+
 	RTLIL::Cell* addSlice  (RTLIL::IdString name, const RTLIL::SigSpec &sig_a, const RTLIL::SigSpec &sig_y, RTLIL::Const offset, const std::string &src = "");
 	RTLIL::Cell* addConcat (RTLIL::IdString name, const RTLIL::SigSpec &sig_a, const RTLIL::SigSpec &sig_b, const RTLIL::SigSpec &sig_y, const std::string &src = "");
 	RTLIL::Cell* addLut    (RTLIL::IdString name, const RTLIL::SigSpec &sig_a, const RTLIL::SigSpec &sig_y, RTLIL::Const lut, const std::string &src = "");
@@ -1432,6 +1438,9 @@ public:
 	RTLIL::SigSpec Bmux     (RTLIL::IdString name, const RTLIL::SigSpec &sig_a, const RTLIL::SigSpec &sig_s, const std::string &src = "");
 	RTLIL::SigSpec Demux     (RTLIL::IdString name, const RTLIL::SigSpec &sig_a, const RTLIL::SigSpec &sig_s, const std::string &src = "");
 
+	RTLIL::SigSpec Bweqx      (RTLIL::IdString name, const RTLIL::SigSpec &sig_a, const RTLIL::SigSpec &sig_b, const std::string &src = "");
+	RTLIL::SigSpec Bwmux      (RTLIL::IdString name, const RTLIL::SigSpec &sig_a, const RTLIL::SigSpec &sig_b, const RTLIL::SigSpec &sig_s, const std::string &src = "");
+
 	RTLIL::SigBit BufGate    (RTLIL::IdString name, const RTLIL::SigBit &sig_a, const std::string &src = "");
 	RTLIL::SigBit NotGate    (RTLIL::IdString name, const RTLIL::SigBit &sig_a, const std::string &src = "");
 	RTLIL::SigBit AndGate    (RTLIL::IdString name, const RTLIL::SigBit &sig_a, const RTLIL::SigBit &sig_b, const std::string &src = "");
diff --git a/kernel/satgen.cc b/kernel/satgen.cc
index 3c21a8f99..3a2fa4735 100644
--- a/kernel/satgen.cc
+++ b/kernel/satgen.cc
@@ -223,7 +223,33 @@ bool SatGen::importCell(RTLIL::Cell *cell, int timestep)
 		return true;
 	}
 
-	if (cell->type.in(ID($_MUX_), ID($mux), ID($_NMUX_)))
+	if (cell->type == ID($bweqx))
+	{
+		std::vector<int> a = importDefSigSpec(cell->getPort(ID::A), timestep);
+		std::vector<int> b = importDefSigSpec(cell->getPort(ID::B), timestep);
+		std::vector<int> y = importDefSigSpec(cell->getPort(ID::Y), timestep);
+
+		std::vector<int> bweqx = ez->vec_not(ez->vec_xor(a, b));
+
+		if (model_undef)
+		{
+			std::vector<int> undef_a = importUndefSigSpec(cell->getPort(ID::A), timestep);
+			std::vector<int> undef_b = importUndefSigSpec(cell->getPort(ID::B), timestep);
+			std::vector<int> undef_y = importUndefSigSpec(cell->getPort(ID::Y), timestep);
+
+			std::vector<int> both_undef = ez->vec_and(undef_a, undef_b);
+			std::vector<int> both_def = ez->vec_and(ez->vec_not(undef_a), ez->vec_not(undef_b));
+
+			bweqx = ez->vec_or(both_undef, ez->vec_and(both_def, bweqx));
+
+			for (int yx : undef_y)
+				ez->assume(ez->NOT(yx));
+		}
+		ez->assume(ez->vec_eq(bweqx, y));
+		return true;
+	}
+
+	if (cell->type.in(ID($_MUX_), ID($mux), ID($_NMUX_), ID($bwmux)))
 	{
 		std::vector<int> a = importDefSigSpec(cell->getPort(ID::A), timestep);
 		std::vector<int> b = importDefSigSpec(cell->getPort(ID::B), timestep);
@@ -233,6 +259,8 @@ bool SatGen::importCell(RTLIL::Cell *cell, int timestep)
 		std::vector<int> yy = model_undef ? ez->vec_var(y.size()) : y;
 		if (cell->type == ID($_NMUX_))
 			ez->assume(ez->vec_eq(ez->vec_not(ez->vec_ite(s.at(0), b, a)), yy));
+		else if (cell->type == ID($bwmux))
+			ez->assume(ez->vec_eq(ez->vec_ite(s, b, a), yy));
 		else
 			ez->assume(ez->vec_eq(ez->vec_ite(s.at(0), b, a), yy));
 
@@ -245,7 +273,11 @@ bool SatGen::importCell(RTLIL::Cell *cell, int timestep)
 
 			std::vector<int> unequal_ab = ez->vec_not(ez->vec_iff(a, b));
 			std::vector<int> undef_ab = ez->vec_or(unequal_ab, ez->vec_or(undef_a, undef_b));
-			std::vector<int> yX = ez->vec_ite(undef_s.at(0), undef_ab, ez->vec_ite(s.at(0), undef_b, undef_a));
+			std::vector<int> yX;
+			if (cell->type == ID($bwmux))
+				yX = ez->vec_ite(undef_s, undef_ab, ez->vec_ite(s, undef_b, undef_a));
+			else
+				yX = ez->vec_ite(undef_s.at(0), undef_ab, ez->vec_ite(s.at(0), undef_b, undef_a));
 			ez->assume(ez->vec_eq(yX, undef_y));
 			undefGating(y, yy, undef_y);
 		}