From 69a0a742c1f86a27f2ef15b6b1aa22d4a0b630a4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Povi=C5=A1er?= <povik@cutebit.org>
Date: Mon, 7 Aug 2023 13:35:13 +0200
Subject: [PATCH] peepopt: Add 'bmux' and 'wshift' modes for testing

---
 passes/pmgen/peepopt.cc                 | 12 +++++
 passes/pmgen/peepopt_shiftmul_right.pmg | 69 +++++++++++++++++++++++++
 2 files changed, 81 insertions(+)

diff --git a/passes/pmgen/peepopt.cc b/passes/pmgen/peepopt.cc
index 2b225623dba..b87fbbe50de 100644
--- a/passes/pmgen/peepopt.cc
+++ b/passes/pmgen/peepopt.cc
@@ -45,6 +45,7 @@ struct PeepoptPass : public Pass {
 
 		log_header(design, "Executing PEEPOPT pass (run peephole optimizers).\n");
 
+		bool bmux_mode = false, wshift_mode = false;
 		size_t argidx;
 		for (argidx = 1; argidx < args.size(); argidx++)
 		{
@@ -52,6 +53,14 @@ struct PeepoptPass : public Pass {
 				genmode = args[++argidx];
 				continue;
 			}
+			if (args[argidx] == "-bmux") {
+				bmux_mode = true;
+				continue;
+			}
+			if (args[argidx] == "-wshift") {
+				wshift_mode = true;
+				continue;
+			}
 			break;
 		}
 		extra_args(args, argidx, design);
@@ -79,6 +88,9 @@ struct PeepoptPass : public Pass {
 
 				pm.setup(module->selected_cells());
 
+				pm.ud_shiftmul_right.bmux_mode = bmux_mode;
+				pm.ud_shiftmul_right.wshift_mode = wshift_mode;
+
 				pm.run_shiftmul_right();
 				pm.run_shiftmul_left();
 				pm.run_muldiv();
diff --git a/passes/pmgen/peepopt_shiftmul_right.pmg b/passes/pmgen/peepopt_shiftmul_right.pmg
index 71e0980234a..ff30a3b17b2 100644
--- a/passes/pmgen/peepopt_shiftmul_right.pmg
+++ b/passes/pmgen/peepopt_shiftmul_right.pmg
@@ -3,6 +3,8 @@ pattern shiftmul_right
 // Optimize mul+shift pairs that result from expressions such as foo[s*W+:W]
 //
 
+udata <bool> bmux_mode wshift_mode
+
 match shift
 	select shift->type.in($shift, $shiftx, $shr)
 	filter !port(shift, \B).empty()
@@ -68,6 +70,7 @@ code
 	// make sure there's no overlap in the signal
 	// selections by the shiftmul pattern
 	if (GetSize(port(shift, \Y)) > mul_const.as_int())
+	if (!bmux_mode && !wshift_mode)
 		reject;
 
 	int factor_bits = ceil_log2(mul_const.as_int());
@@ -78,6 +81,72 @@ code
 	did_something = true;
 	log("right shiftmul pattern in %s: shift=%s, mul=%s\n", log_id(module), log_id(shift), log_id(mul));
 
+	if (bmux_mode) {
+		SigSpec out = port(shift, \Y);
+		int shift_stride = mul_const.as_int();
+		SigSpec shift_vector = port(shift, \A);
+		SigBit padbit = (shift->type != $shiftx) ? (param(shift, \A_SIGNED)).as_bool() ?
+						port(shift, \A).bits().back() : State::S0 : State::Sx;
+
+		// Construct the input vector for the new bmux cell
+		SigSpec bmux_vector;
+		for (int i = 0; i * shift_stride < GetSize(shift_vector); i++) {
+			int validbits = std::min(GetSize(out),
+									 GetSize(shift_vector) - i * shift_stride);
+			bmux_vector.append(shift_vector.extract(i * shift_stride, validbits));
+		}
+
+		// Prep the select signal
+		SigSpec bmux_sel = mul_din;
+		bmux_sel.extend_u0(ceil_log2(GetSize(bmux_vector)));
+
+		// Pad the vector
+		int vector_targetlen = GetSize(out) << GetSize(bmux_sel);
+		if (vector_targetlen > GetSize(bmux_vector))
+			bmux_vector.append(SigSpec(padbit, vector_targetlen - GetSize(bmux_vector)));
+		else
+			bmux_vector.remove(vector_targetlen, GetSize(bmux_vector) - vector_targetlen);
+
+		// Build the bmux, add an extra mux to handle top bits on the old select signal
+		SigSpec bmux_out, out_of_range;
+		std::string src = shift->get_src_attribute();
+		bmux_out = module->Bmux(NEW_ID, bmux_vector, bmux_sel, src);
+		if (GetSize(mul_din) > GetSize(bmux_sel)) {
+			out_of_range = module->ReduceBool(NEW_ID, mul_din.extract_end(GetSize(bmux_sel)), false, src);
+			module->addMux(NEW_ID, bmux_out, SigSpec(padbit, GetSize(out)), out_of_range, out, src);
+		} else {
+			module->connect(out, bmux_out);
+		}
+		autoremove(shift);
+		accept;
+	} else if (wshift_mode) {
+		SigSpec out = port(shift, \Y);
+		SigSpec shift_vector = port(shift, \A);
+		int shift_stride = mul_const.as_int();
+		for (int j = 0; j < shift_stride; j++) {
+			SigSpec slice_vector;
+			SigSpec slice_out;
+			for (int off = j; off < GetSize(shift_vector); off += shift_stride)
+				slice_vector.append(shift_vector.extract(off, 1));
+			if (param(shift, \A_SIGNED).as_bool())
+				slice_vector.append(shift_vector.bits().back());
+			for (int off = j; off < GetSize(out); off += shift_stride)
+				slice_out.append(out.extract(off, 1));
+
+			if (slice_out.empty()) continue;
+			Cell *slice = module->addCell(NEW_ID, shift);
+			slice->setParam(\A_WIDTH, GetSize(slice_vector));
+			slice->setPort(\A, slice_vector);
+			slice->setParam(\B_WIDTH, GetSize(mul_din));
+			slice->setParam(\B_SIGNED, Const(1, 0));
+			slice->setPort(\B, mul_din);
+			slice->setParam(\Y_WIDTH, GetSize(slice_out));
+			slice->setPort(\Y, slice_out);
+		}
+		autoremove(shift);
+		accept;
+	}
+
 	int const_factor = mul_const.as_int();
 	int new_const_factor = 1 << factor_bits;
 	SigSpec padding(State::Sx, new_const_factor-const_factor);