[HWLegalizeModules] transform mux of array to muxes of elements (#9230)

* When disallowPackedArrays is set, transform muxes of arrays to muxes of elements. See the added tests for examples.
* Some minor refactoring.

Author: ollef

lib/Dialect/SV/Transforms/HWLegalizeModules.cpp

@@ -14,6 +14,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "circt/Dialect/Comb/CombOps.h"
 #include "circt/Dialect/HW/HWOps.h"
 #include "circt/Dialect/HW/HWTypes.h"
 #include "circt/Dialect/SV/SVOps.h"
@@ -42,6 +43,10 @@ struct HWLegalizeModulesPass
 private:
   void processPostOrder(Block &block);
   bool tryLoweringPackedArrayOp(Operation &op);
+  template <typename ElementType>
+  SmallVector<std::pair<Value, Value>>
+  createIndexValuePairs(OpBuilder &builder, LocationAttr loc, hw::ArrayType ty,
+                        Value array);
   Value lowerLookupToCasez(Operation &op, Value input, Value index,
                            mlir::Type elementType,
                            SmallVector<Value> caseValues);
@@ -89,43 +94,20 @@ bool HWLegalizeModulesPass::tryLoweringPackedArrayOp(Operation &op) {
       })
       .Case<hw::ArrayConcatOp>([&](hw::ArrayConcatOp concatOp) {
         // Redirect individual element uses (if any) to the input arguments.
-        SmallVector<std::pair<Value, uint64_t>> arrays;
+        SmallVector<Value> values;
+        OpBuilder builder(concatOp);
         for (auto array : llvm::reverse(concatOp.getInputs())) {
           auto ty = hw::type_cast<hw::ArrayType>(array.getType());
-          arrays.emplace_back(array, ty.getNumElements());
-        }
-        for (auto *user :
-             llvm::make_early_inc_range(concatOp.getResult().getUsers())) {
-          if (TypeSwitch<Operation *, bool>(user)
-                  .Case<hw::ArrayGetOp>([&](hw::ArrayGetOp getOp) {
-                    if (auto indexAndBitWidth =
-                            tryExtractIndexAndBitWidth(getOp.getIndex())) {
-                      auto [indexValue, bitWidth] = *indexAndBitWidth;
-                      // FIXME: More efficient search
-                      for (const auto &[array, size] : arrays) {
-                        if (indexValue >= size) {
-                          indexValue -= size;
-                          continue;
-                        }
-                        OpBuilder builder(getOp);
-                        getOp.getInputMutable().set(array);
-                        getOp.getIndexMutable().set(
-                            builder.createOrFold<hw::ConstantOp>(
-                                getOp.getLoc(), APInt(bitWidth, indexValue)));
-                        return true;
-                      }
-                    }
-
-                    return false;
-                  })
-                  .Default([](auto op) { return false; }))
-            continue;
-
-          op.emitError("unsupported packed array expression");
-          signalPassFailure();
+          const auto indexValues = createIndexValuePairs<hw::ArrayGetOp>(
+              builder, concatOp.getLoc(), ty, array);
+          for (const auto &[_, value] : indexValues) {
+            values.push_back(value);
+          }
         }
+        if (!processUsers(op, concatOp.getResult(), values))
+          return false;
 
-        // Remove the original op.
+        // Remove original op.
         return true;
       })
       .Case<hw::ArrayCreateOp>([&](hw::ArrayCreateOp createOp) {
@@ -147,15 +129,12 @@ bool HWLegalizeModulesPass::tryLoweringPackedArrayOp(Operation &op) {
         // Generate case value element lookups.
         auto ty = hw::type_cast<hw::ArrayType>(getOp.getInput().getType());
         OpBuilder builder(getOp);
+        auto loc = op.getLoc();
+        const auto indexValues = createIndexValuePairs<hw::ArrayGetOp>(
+            builder, loc, ty, getOp.getInput());
         SmallVector<Value> caseValues;
-        for (size_t i = 0, e = ty.getNumElements(); i < e; i++) {
-          auto loc = op.getLoc();
-          auto index = builder.createOrFold<hw::ConstantOp>(
-              loc, APInt(llvm::Log2_64_Ceil(e), i));
-          auto element =
-              hw::ArrayGetOp::create(builder, loc, getOp.getInput(), index);
-          caseValues.push_back(element);
-        }
+        for (const auto &[_, value] : indexValues)
+          caseValues.push_back(value);
 
         // Transform array index op into casez statement.
         auto theWire = lowerLookupToCasez(op, getOp.getInput(), index,
@@ -183,13 +162,11 @@ bool HWLegalizeModulesPass::tryLoweringPackedArrayOp(Operation &op) {
         // Generate case value element lookups.
         auto ty = hw::type_cast<hw::ArrayType>(inout.getElementType());
         OpBuilder builder(&op);
+        auto loc = op.getLoc();
+        const auto indexValues = createIndexValuePairs<sv::ArrayIndexInOutOp>(
+            builder, loc, ty, indexOp.getInput());
         SmallVector<Value> caseValues;
-        for (size_t i = 0, e = ty.getNumElements(); i < e; i++) {
-          auto loc = op.getLoc();
-          auto index = builder.createOrFold<hw::ConstantOp>(
-              loc, APInt(llvm::Log2_64_Ceil(e), i));
-          auto element = sv::ArrayIndexInOutOp::create(
-              builder, loc, indexOp.getInput(), index);
+        for (const auto &[_, element] : indexValues) {
           auto readElement = sv::ReadInOutOp::create(builder, loc, element);
           caseValues.push_back(readElement);
         }
@@ -211,14 +188,12 @@ bool HWLegalizeModulesPass::tryLoweringPackedArrayOp(Operation &op) {
           return false;
 
         OpBuilder builder(assignOp);
-        for (size_t i = 0, e = ty.getNumElements(); i < e; i++) {
-          auto loc = op.getLoc();
-          auto index = builder.createOrFold<hw::ConstantOp>(
-              loc, APInt(llvm::Log2_64_Ceil(e), i));
+        auto loc = op.getLoc();
+        const auto indexValues = createIndexValuePairs<hw::ArrayGetOp>(
+            builder, loc, ty, assignOp.getSrc());
+        for (const auto &[index, srcElement] : indexValues) {
           auto dstElement = sv::ArrayIndexInOutOp::create(
               builder, loc, assignOp.getDest(), index);
-          auto srcElement =
-              hw::ArrayGetOp::create(builder, loc, assignOp.getSrc(), index);
           sv::PAssignOp::create(builder, loc, dstElement, srcElement);
         }
 
@@ -251,9 +226,53 @@ bool HWLegalizeModulesPass::tryLoweringPackedArrayOp(Operation &op) {
         // Remove original reg.
         return true;
       })
+      .Case<comb::MuxOp>([&](comb::MuxOp muxOp) {
+        // Transform array mux into individual element muxes.
+        auto ty = hw::type_dyn_cast<hw::ArrayType>(muxOp.getType());
+        if (!ty)
+          return false;
+
+        OpBuilder builder(muxOp);
+
+        auto trueValues = createIndexValuePairs<hw::ArrayGetOp>(
+            builder, muxOp.getLoc(), ty, muxOp.getTrueValue());
+        auto falseValues = createIndexValuePairs<hw::ArrayGetOp>(
+            builder, muxOp.getLoc(), ty, muxOp.getFalseValue());
+
+        SmallVector<Value> muxedValues;
+
+        for (size_t i = 0, e = trueValues.size(); i < e; i++) {
+          const auto &[trueIndex, trueValue] = trueValues[i];
+          const auto &[falseIndex, falseValue] = falseValues[i];
+          muxedValues.push_back(
+              comb::MuxOp::create(builder, muxOp.getLoc(), muxOp.getCond(),
+                                  trueValue, falseValue, muxOp.getTwoState()));
+        }
+
+        if (!processUsers(op, muxOp.getResult(), muxedValues))
+          return false;
+
+        // Remove original mux.
+        return true;
+      })
       .Default([&](auto op) { return false; });
 }
 
+template <typename ElementType>
+SmallVector<std::pair<Value, Value>>
+HWLegalizeModulesPass::createIndexValuePairs(OpBuilder &builder,
+                                             LocationAttr loc, hw::ArrayType ty,
+                                             Value array) {
+  SmallVector<std::pair<Value, Value>> result;
+  for (size_t i = 0, e = ty.getNumElements(); i < e; i++) {
+    auto index = builder.createOrFold<hw::ConstantOp>(
+        loc, APInt(llvm::Log2_64_Ceil(e), i));
+    auto element = ElementType::create(builder, loc, array, index);
+    result.emplace_back(index, element);
+  }
+  return result;
+}
+
 Value HWLegalizeModulesPass::lowerLookupToCasez(Operation &op, Value input,
                                                 Value index,
                                                 mlir::Type elementType,

test/Dialect/SV/hw-legalize-modules-packed-arrays.mlir

@@ -81,6 +81,46 @@ hw.module @array_create_get_default(in %arg0: i8, in %arg1: i8, in %arg2: i8, in
   }
 }
 
+// CHECK-LABEL: hw.module @array_muxed_create_get_default
+hw.module @array_muxed_create_get_default(in %arg0: i8, in %arg1: i8, in %arg2: i8, in %arg3: i8, in %arg4: i8, in %arg5: i8,
+                            in %array_sel: i1, in %index_sel: i2) {
+  // CHECK: sv.initial  {
+  sv.initial {
+    %three_array1 = hw.array_create %arg2, %arg1, %arg0 : i8
+    %three_array2 = hw.array_create %arg5, %arg4, %arg3 : i8
+
+    // CHECK: %0 = comb.mux %array_sel, %arg0, %arg3 : i8
+    // CHECK: %1 = comb.mux %array_sel, %arg1, %arg4 : i8
+    // CHECK: %2 = comb.mux %array_sel, %arg2, %arg5 : i8
+    %muxed = comb.mux %array_sel, %three_array1, %three_array2 : !hw.array<3xi8>
+
+    // CHECK: %x_i8 = sv.constantX : i8
+    // CHECK: sv.case casez %index_sel : i2
+    // CHECK: case b00: {
+    // CHECK:   sv.bpassign %casez_tmp, %0 : i8
+    // CHECK: }
+    // CHECK: case b01: {
+    // CHECK:   sv.bpassign %casez_tmp, %1 : i8
+    // CHECK: }
+    // CHECK: case b10: {
+    // CHECK:   sv.bpassign %casez_tmp, %2 : i8
+    // CHECK: }
+    // CHECK: default: {
+    // CHECK:   sv.bpassign %casez_tmp, %x_i8 : i8
+    // CHECK: }
+
+    // CHECK: %3 = sv.read_inout %casez_tmp : !hw.inout<i8>
+    %2 = hw.array_get %muxed[%index_sel] : !hw.array<3xi8>, i2
+
+    // CHECK: %4 = comb.icmp eq %3, %arg2 : i8
+    // CHECK: sv.if %4  {
+    %cond = comb.icmp eq %2, %arg2 : i8
+    sv.if %cond {
+      sv.fatal 1
+    }
+  }
+}
+
 // CHECK-LABEL: hw.module @array_create_concat_get_default
 hw.module @array_create_concat_get_default(in %arg0: i8, in %arg1: i8, in %arg2: i8, in %arg3: i8,
                             in %sel: i2) {
@@ -143,6 +183,38 @@ hw.module @array_constant_get_comb(in %sel: i2, out a: i8) {
   hw.output %1 : i8
 }
 
+// CHECK-LABEL: hw.module @array_muxed_constant_get_comb
+hw.module @array_muxed_constant_get_comb(in %array_sel: i1, in %index_sel: i2, out a: i8) {
+  // CHECK: %0 = comb.mux %array_sel, %c3_i8, %c7_i8 : i8
+  // CHECK: %1 = comb.mux %array_sel, %c2_i8, %c6_i8 : i8
+  // CHECK: %2 = comb.mux %array_sel, %c1_i8, %c5_i8 : i8
+  // CHECK: %3 = comb.mux %array_sel, %c0_i8, %c4_i8 : i8
+  // CHECK: %casez_tmp = sv.reg : !hw.inout<i8>
+  // CHECK: sv.alwayscomb {
+  // CHECK:  sv.case casez %index_sel : i2
+  // CHECK:  case b00: {
+  // CHECK:    sv.bpassign %casez_tmp, %0 : i8
+  // CHECK:   }
+  // CHECK:   case b01: {
+  // CHECK:     sv.bpassign %casez_tmp, %1 : i8
+  // CHECK:   }
+  // CHECK:   case b10: {
+  // CHECK:     sv.bpassign %casez_tmp, %2 : i8
+  // CHECK:   }
+  // CHECK:   default: {
+  // CHECK:     sv.bpassign %casez_tmp, %3 : i8
+  // CHECK:   }
+  // CHECK: }
+  %0 = hw.aggregate_constant [0 : i8, 1 : i8, 2 : i8, 3 : i8] : !hw.array<4xi8>
+  %1 = hw.aggregate_constant [4 : i8, 5 : i8, 6 : i8, 7 : i8] : !hw.array<4xi8>
+  %muxed = comb.mux %array_sel, %0, %1 : !hw.array<4xi8>
+  // CHECK: %4 = sv.read_inout %casez_tmp : !hw.inout<i8>
+  %3 = hw.array_get %muxed[%index_sel] : !hw.array<4xi8>, i2
+
+  // CHECK: hw.output %4 : i8
+  hw.output %3 : i8
+}
+
 // CHECK-LABEL: hw.module @array_reg_mux_2
 hw.module @array_reg_mux_2(in %clock: i1, in %arg0: i8, in %arg1: i8, in %sel: i1, out a: i8) {
   // CHECK: %reg = sv.reg : !hw.inout<i8>