Recursively generate one-hot multiplexers for aggregates (#1557)

Test all cases in ParameterizedOneHotTesters

Co-authored-by: Jack Koenig <koenig@sifive.com>

2 files changed, 23 insertions, 19 deletions

diff --git a/core/src/main/scala/chisel3/SeqUtils.scala b/core/src/main/scala/chisel3/SeqUtils.scala
index 9f068898..97bdc505 100644
--- a/core/src/main/scala/chisel3/SeqUtils.scala
+++ b/core/src/main/scala/chisel3/SeqUtils.scala
@@ -113,10 +113,17 @@ private[chisel3] object SeqUtils {
             buildAndOrMultiplexor(sels.zip(inWidthMatched))
           }
 
-        case _: Aggregate =>
+        case agg: Aggregate =>
           val allDefineWidth = in.forall { case (_, element) => element.widthOption.isDefined }
-          if(allDefineWidth) {
-            buildAndOrMultiplexor(in)
+          if (allDefineWidth) {
+            val out = Wire(agg)
+            val (sel, inData) = in.unzip
+            val inElts = inData.map(_.asInstanceOf[Aggregate].getElements)
+            // We want to iterate on the columns of inElts, so we transpose
+            out.getElements.zip(inElts.transpose).foreach { case (outElt, elts) =>
+              outElt := oneHotMux(sel.zip(elts))
+            }
+            out.asInstanceOf[T]
           }
           else {
             throwException(s"Cannot Mux1H with aggregates with inferred widths")
diff --git a/src/test/scala/chiselTests/OneHotMuxSpec.scala b/src/test/scala/chiselTests/OneHotMuxSpec.scala
index cc359e8e..fb6edb2c 100644
--- a/src/test/scala/chiselTests/OneHotMuxSpec.scala
+++ b/src/test/scala/chiselTests/OneHotMuxSpec.scala
@@ -26,8 +26,7 @@ class OneHotMuxSpec extends AnyFreeSpec with Matchers with ChiselRunners {
     assertTesterPasses(new AllSameFixedPointOneHotTester)
   }
   "simple one hot mux with all same parameterized sint values should work" in {
-    val values: Seq[SInt] = Seq((-3).S, (-5).S, (-7).S, (-11).S)
-    assertTesterPasses(new ParameterizedOneHotTester(values, SInt(8.W), -5.S(8.W)))
+    assertTesterPasses(new ParameterizedOneHotTester)
   }
   "simple one hot mux with all same parameterized aggregates containing fixed values should work" in {
     assertTesterPasses(new ParameterizedAggregateOneHotTester)
@@ -121,14 +120,14 @@ class AllSameFixedPointOneHotTester extends BasicTester {
   stop()
 }
 
-class ParameterizedOneHotTester[T <: Data](values: Seq[T], outGen: T, expected: T) extends BasicTester {
-  val dut = Module(new ParameterizedOneHot(values, outGen))
-  dut.io.selectors(0) := false.B
-  dut.io.selectors(1) := true.B
-  dut.io.selectors(2) := false.B
-  dut.io.selectors(3) := false.B
+class ParameterizedOneHotTester extends BasicTester {
+  val values: Seq[Int] = Seq(-3, -5, -7, -11)
+  for ((v, i) <- values.zipWithIndex) {
+    val dut = Module(new ParameterizedOneHot(values.map(_.S), SInt(8.W)))
+    dut.io.selectors := (1 << i).U(4.W).asBools
 
-  assert(dut.io.out.asUInt() === expected.asUInt())
+    assert(dut.io.out.asUInt() === v.S(8.W).asUInt())
+  }
 
   stop()
 }
@@ -178,14 +177,12 @@ object Agg2 extends HasMakeLit[Agg2] {
 
 class ParameterizedAggregateOneHotTester extends BasicTester {
   val values = (0 until 4).map { n => Agg1.makeLit(n) }
+  for ((v, i) <- values.zipWithIndex) {
+    val dut = Module(new ParameterizedAggregateOneHot(Agg1, new Agg1))
+    dut.io.selectors := (1 << i).U(4.W).asBools
 
-  val dut = Module(new ParameterizedAggregateOneHot(Agg1, new Agg1))
-  dut.io.selectors(0) := false.B
-  dut.io.selectors(1) := true.B
-  dut.io.selectors(2) := false.B
-  dut.io.selectors(3) := false.B
-
-  assert(dut.io.out.asUInt() === values(1).asUInt())
+    assert(dut.io.out.asUInt() === values(i).asUInt())
+  }
 
   stop()
 }