[Comb] Fix another recursive mux canonicalizer (#9144)

Author: TaoBi22

lib/Dialect/Comb/CombFolds.cpp

@@ -2552,7 +2552,8 @@ LogicalResult MuxRewriter::matchAndRewrite(MuxOp op,
   if (auto falseMux = op.getFalseValue().getDefiningOp<MuxOp>();
       falseMux && falseMux != op) {
     // mux(selector, x, mux(selector, y, z) = mux(selector, x, z)
-    if (op.getCond() == falseMux.getCond()) {
+    if (op.getCond() == falseMux.getCond() &&
+        falseMux.getFalseValue() != falseMux) {
       replaceOpWithNewOpAndCopyNamehint<MuxOp>(
           rewriter, op, op.getCond(), op.getTrueValue(),
           falseMux.getFalseValue(), op.getTwoStateAttr());

test/Dialect/Comb/canonicalization.mlir

@@ -1284,7 +1284,7 @@ hw.module @muxCommon(in %cond: i1, in %cond2: i1, in %cond3: i1,
                      in %arg0 : i32, in %arg1 : i32, in %arg2: i32, in %arg3: i32,
   out o1: i32, out o2: i32, out o3: i32, out o4: i32, out o5: i32,
   out orResult: i32, out o6: i32, out o7: i32, out o8 : i1, out o9: i32,
-  out o10: i32, out o11: i1, out o12: i32) {
+  out o10: i32, out o11: i1, out o12: i32, out o13: i32) {
   // CHECK: [[TRUE:%.+]] = hw.constant true
   // CHECK: [[FALSE:%.+]] = hw.constant false
   %true = hw.constant true
@@ -1352,10 +1352,13 @@ hw.module @muxCommon(in %cond: i1, in %cond2: i1, in %cond3: i1,
   %3 = comb.mux %cond2, %arg0, %2 : i32
   %o12 = comb.mux %cond3, %arg0, %3 : i32
 
+  // CHECK: [[O13:%.+]] = comb.mux %cond, %arg0, [[LOOPMUX]] : i32
+  %o13 = comb.mux %cond, %arg0, %2 : i32
+
   // CHECK: hw.output [[O1]], [[O2]], [[O3]], [[O4]], [[O5]], [[ORRESULT]],
   // CHECK: [[O6]], [[O7]], [[O8]], [[O9]], [[O10]]
-  hw.output %o1, %o2, %o3, %o4, %o5, %orResult, %o6, %o7, %o8, %o9, %o10, %o11, %o12
-    : i32, i32, i32, i32, i32, i32, i32, i32, i1, i32, i32, i1, i32
+  hw.output %o1, %o2, %o3, %o4, %o5, %orResult, %o6, %o7, %o8, %o9, %o10, %o11, %o12, %o13
+    : i32, i32, i32, i32, i32, i32, i32, i32, i1, i32, i32, i1, i32, i32
 }
 
 // CHECK-LABEL: @flatten_multi_use_and