[mlir] Add canonicalization for cond_br that feed into a cond_br on the same condition

[rriddle]

Summary

  ...
  cond_br %cond, ^bb1(...), ^bb2(...)
...
^bb1: // has single predecessor
  ...
  cond_br %cond, ^bb3(...), ^bb4(...)

->

  ...
  cond_br %cond, ^bb1(...), ^bb2(...)
...
^bb1: // has single predecessor
  ...
  br ^bb3(...)

[jpienaar]

Nice, micro nit: in description add ... as when I initially read it I thought there was requirement for empty block bb1

[mehdi_amini]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1105

Trivial braces?

[bondhugula]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1068

Consider adding a line on the simplification - the example below doesn't capture the fact that ^bb1 has only one predessor.

1077

Nit: please add `// has single predecessor```

1098

Nit: a -> an

[rriddle]

Resolve comments

[rriddle]

Thanks for the review!

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

for (auto it : llvm::zip(trueOperands, falseOperands)) {
         mergedOperands.push_back(rewriter.create<SelectOp>(
             condbr.getLoc(), condition, std::get<0>(it), std::get<1>(it)));
     }
 
     rewriter.replaceOpWithNewOp<BranchOp>(condbr, trueDest, mergedOperands);
     return success();
   }
 };
+
+///   cond_br %cond, ^bb1(...), ^bb2(...)
+/// ...
+/// ^bb1:
+///   cond_br %cond, ^bb3(...), ^bb4(...)
+///
+/// ->
+///
+///   cond_br %cond, ^bb1(...), ^bb2(...)
+/// ...
+/// ^bb1:
+///   br ^bb3(...)
+///
+struct SimplifyCondBranchFromCondBranchOnSameCondition
+    : public OpRewritePattern<CondBranchOp> {
+  using OpRewritePattern<CondBranchOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CondBranchOp condbr,
+                                PatternRewriter &rewriter) const override {
+    // Check that we have a single distinct predecessor.
+    Block *currentBlock = condbr.getOperation()->getBlock();
+    Block *predecessor = currentBlock->getSinglePredecessor();
+    if (!predecessor)
+      return failure();
+
+    // Check that the predecessor terminates with a conditional branch to this
+    // block and that it branches on the same condition.
+    auto predBranch = dyn_cast<CondBranchOp>(predecessor->getTerminator());
+    if (!predBranch || condbr.getCondition() != predBranch.getCondition())
+      return failure();
+
+    // Fold this branch to a unconditional branch.
+    if (currentBlock == predBranch.trueDest()) {
+      rewriter.replaceOpWithNewOp<BranchOp>(condbr, condbr.trueDest(),
+                                            condbr.trueDestOperands());
+    } else {
+      rewriter.replaceOpWithNewOp<BranchOp>(condbr, condbr.falseDest(),
+                                            condbr.falseDestOperands());
+    }
+    return success();
+  }
+};
 } // end anonymous namespace
 
 void CondBranchOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<SimplifyConstCondBranchPred, SimplifyPassThroughCondBranch,
-                 SimplifyCondBranchIdenticalSuccessors>(context);
+                 SimplifyCondBranchIdenticalSuccessors,
+                 SimplifyCondBranchFromCondBranchOnSameCondition>(context);
 }
 
 Optional<MutableOperandRange>
 CondBranchOp::getMutableSuccessorOperands(unsigned index) {
   assert(index < getNumSuccessors() && "invalid successor index");
   return index == trueIndex ? trueDestOperandsMutable()
                             : falseDestOperandsMutable();
 }

mlir/test/Dialect/Standard/canonicalize-cf.mlir

^bb1:
   // Successors can't be collapsed if they contain other operations.
   "foo.op"() : () -> ()
   br ^bb2
 
 ^bb2:
   return
 }
 
+/// Test folding conditional branches that are successors of conditional
+/// branches with the same condition.
+
+// CHECK-LABEL: func @cond_br_from_cond_br_with_same_condition
+func @cond_br_from_cond_br_with_same_condition(%cond : i1) {
+  // CHECK:   cond_br %{{.*}}, ^bb1, ^bb2
+  // CHECK: ^bb1:
+  // CHECK:   return
+
+  cond_br %cond, ^bb1, ^bb2
+
+^bb1:
+  cond_br %cond, ^bb3, ^bb2
+
+^bb2:
+  "foo.terminator"() : () -> ()
+
+^bb3:
+  return
+}
+
 // -----
 
 // Erase assertion if condition is known to be true at compile time.
 // CHECK-LABEL: @assert_true
 func @assert_true() {
   // CHECK-NOT: assert
   %true = constant true
   assert %true, "Computer says no"

mlir/test/Transforms/canonicalize-block-merge.mlir

^bb2:
   }
   return
 }
 
 // Check that properly handles back edges and the case where a value from one
 // block is used in another.
 
 // CHECK-LABEL: func @mismatch_loop(
-// CHECK-SAME: %[[ARG:.*]]: i1
-func @mismatch_loop(%cond : i1) {
-  // CHECK: cond_br %{{.*}}, ^bb1(%[[ARG]] : i1), ^bb2
+// CHECK-SAME: %[[ARG:.*]]: i1, %[[ARG2:.*]]: i1
+func @mismatch_loop(%cond : i1, %cond2 : i1) {
+  // CHECK: cond_br %{{.*}}, ^bb1(%[[ARG2]] : i1), ^bb2
 
   cond_br %cond, ^bb2, ^bb3
 
 ^bb1:
-  // CHECK: ^bb1(%[[ARG2:.*]]: i1):
+  // CHECK: ^bb1(%[[ARG3:.*]]: i1):
   // CHECK-NEXT: %[[LOOP_CARRY:.*]] = "foo.op"
-  // CHECK-NEXT: cond_br %[[ARG2]], ^bb1(%[[LOOP_CARRY]] : i1), ^bb2
+  // CHECK-NEXT: cond_br %[[ARG3]], ^bb1(%[[LOOP_CARRY]] : i1), ^bb2
 
   %ignored = "foo.op"() : () -> (i1)
-  cond_br %cond2, ^bb1, ^bb3
+  cond_br %cond3, ^bb1, ^bb3
 
 ^bb2:
-  %cond2 = "foo.op"() : () -> (i1)
-  cond_br %cond, ^bb1, ^bb3
+  %cond3 = "foo.op"() : () -> (i1)
+  cond_br %cond2, ^bb1, ^bb3
 
 ^bb3:
   // CHECK: ^bb2:
   // CHECK-NEXT: return
 
   return
 }