[mlir][Standard] Add support for canonicalizing branches to passthrough blocks

[rriddle]

Summary

This revision adds support for canonicalizing the following:

  br ^bb1
^bb1
  br ^bbN(...)

br ^bbN(...)

Depends On D78682

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

return op.emitOpError() << "types mismatch between yield op: " << resultType
                             << " and its parent: " << parentType;
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // BranchOp
 //===----------------------------------------------------------------------===//
 
+/// Given a successor, try to collapse it to a new destination if it only
+/// contains a passthrough unconditional branch. If the successor is
+/// collapsable, `successor` and `successorOperands` are updated to reference
+/// the new destination and values. `argStorage` is an optional storage to use
+/// if operands to the collapsed successor need to be remapped.
+static LogicalResult collapseBranch(Block *&successor,
+                                    ValueRange &successorOperands,
+                                    SmallVectorImpl<Value> &argStorage) {
+  // Check that the successor only contains a unconditional branch.
+  if (std::next(successor->begin()) != successor->end())
+    return failure();
+  // Check that the terminator is an unconditional branch.
+  BranchOp successorBranch = dyn_cast<BranchOp>(successor->getTerminator());
+  if (!successorBranch)
+    return failure();
+  // Check that the arguments are only used within the terminator.
+  for (BlockArgument arg : successor->getArguments()) {
+    for (Operation *user : arg.getUsers())
+      if (user != successorBranch)
+        return failure();
+  }
+  // Don't try to collapse branches to infinite loops.
+  Block *successorDest = successorBranch.getDest();
+  if (successorDest == successor)
+    return failure();
+
+  // Update the operands to the successor. If the branch parent has no
+  // arguments, we can use the branch operands directly.
+  OperandRange operands = successorBranch.getOperands();
+  if (successor->args_empty()) {
+    successor = successorDest;
+    successorOperands = operands;
+    return success();
+  }
+
+  // Otherwise, we need to remap any argument operands.
+  for (Value operand : operands) {
+    BlockArgument argOperand = operand.dyn_cast<BlockArgument>();
+    if (argOperand && argOperand.getOwner() == successor)
+      argStorage.push_back(successorOperands[argOperand.getArgNumber()]);
+    else
+      argStorage.push_back(operand);
+  }
+  successor = successorDest;
+  successorOperands = argStorage;
+  return success();
+}
+
 namespace {
 /// Simplify a branch to a block that has a single predecessor. This effectively
 /// merges the two blocks.
 struct SimplifyBrToBlockWithSinglePred : public OpRewritePattern<BranchOp> {
   using OpRewritePattern<BranchOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(BranchOp op,
                                 PatternRewriter &rewriter) const override {
     // Check that the successor block has a single predecessor.
     Block *succ = op.getDest();
     Block *opParent = op.getOperation()->getBlock();
     if (succ == opParent || !llvm::hasSingleElement(succ->getPredecessors()))
       return failure();
 
     // Merge the successor into the current block and erase the branch.
     rewriter.mergeBlocks(succ, opParent, op.getOperands());
     rewriter.eraseOp(op);
     return success();
   }
 };
+
+///   br ^bb1
+/// ^bb1
+///   br ^bbN(...)
+///
+///  -> br ^bbN(...)
+///
+struct SimplifyPassThroughBr : public OpRewritePattern<BranchOp> {
+  using OpRewritePattern<BranchOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(BranchOp op,
+                                PatternRewriter &rewriter) const override {
+    Block *dest = op.getDest();
+    ValueRange destOperands = op.getOperands();
+    SmallVector<Value, 4> destOperandStorage;
+
+    // Try to collapse the successor if it points somewhere other than this
+    // block.
+    if (dest == op.getOperation()->getBlock() ||
+        failed(collapseBranch(dest, destOperands, destOperandStorage)))
+      return failure();
+
+    // Create a new branch with the collapsed successor.
+    rewriter.replaceOpWithNewOp<BranchOp>(op, dest, destOperands);
+    return success();
+  }
+};
 } // end anonymous namespace.
 
 Block *BranchOp::getDest() { return getSuccessor(); }
 
 void BranchOp::setDest(Block *block) { return setSuccessor(block); }
 
 void BranchOp::eraseOperand(unsigned index) {
   getOperation()->eraseOperand(index);
 }
 
 void BranchOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                            MLIRContext *context) {
-  results.insert<SimplifyBrToBlockWithSinglePred>(context);
+  results.insert<SimplifyBrToBlockWithSinglePred, SimplifyPassThroughBr>(
+      context);
 }
 
 Optional<OperandRange> BranchOp::getSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
   return getOperands();
 }
 
 bool BranchOp::canEraseSuccessorOperand() { return true; }
if (failed(collapsedTrue) && failed(collapsedFalse))
       return failure();
 
     // Create a new branch with the collapsed successors.
     rewriter.replaceOpWithNewOp<CondBranchOp>(condbr, condbr.getCondition(),
                                               trueDest, trueDestOperands,
                                               falseDest, falseDestOperands);
     return success();
   }
-
-  /// Given a successor, try to collapse it to a new destination if it only
-  /// contains a passthrough unconditional branch. If the successor is
-  /// collapsable, `successor` and `successorOperands` are updated to reference
-  /// the new destination and values. `argStorage` is an optional storage to use
-  /// if operands to the collapsed successor need to be remapped.
-  LogicalResult collapseBranch(Block *&successor, ValueRange &successorOperands,
-                               SmallVectorImpl<Value> &argStorage) const {
-    // Check that the successor only contains a unconditional branch.
-    if (std::next(successor->begin()) != successor->end())
-      return failure();
-    // Check that the terminator is an unconditional branch.
-    BranchOp successorBranch = dyn_cast<BranchOp>(successor->getTerminator());
-    if (!successorBranch)
-      return failure();
-    // Check that the arguments are only used within the terminator.
-    for (BlockArgument arg : successor->getArguments()) {
-      for (Operation *user : arg.getUsers())
-        if (user != successorBranch)
-          return failure();
-    }
-    // Don't try to collapse branches to infinite loops.
-    Block *successorDest = successorBranch.getDest();
-    if (successorDest == successor)
-      return failure();
-
-    // Update the operands to the successor. If the branch parent has no
-    // arguments, we can use the branch operands directly.
-    OperandRange operands = successorBranch.getOperands();
-    if (successor->args_empty()) {
-      successor = successorBranch.getDest();
-      successorOperands = operands;
-      return success();
-    }
-
-    // Otherwise, we need to remap any argument operands.
-    for (Value operand : operands) {
-      BlockArgument argOperand = operand.dyn_cast<BlockArgument>();
-      if (argOperand && argOperand.getOwner() == successor)
-        argStorage.push_back(successorOperands[argOperand.getArgNumber()]);
-      else
-        argStorage.push_back(operand);
-    }
-    successor = successorBranch.getDest();
-    successorOperands = argStorage;
-    return success();
-  }
 };
 
 /// cond_br %cond, ^bb1(A, ..., N), ^bb1(A, ..., N)
 ///  -> br ^bb1(A, ..., N)
 ///
 /// cond_br %cond, ^bb1(A), ^bb1(B)
 ///  -> %select = select %cond, A, B
 ///     br ^bb1(%select)

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

 // RUN: mlir-opt %s -allow-unregistered-dialect -pass-pipeline='func(canonicalize)' -split-input-file | FileCheck %s
 
 /// Test the folding of BranchOp.
 
 // CHECK-LABEL: func @br_folding(
 func @br_folding() -> i32 {
   // CHECK-NEXT: %[[CST:.*]] = constant 0 : i32
   // CHECK-NEXT: return %[[CST]] : i32
   %c0_i32 = constant 0 : i32
   br ^bb1(%c0_i32 : i32)
 ^bb1(%x : i32):
   return %x : i32
 }
 
+/// Test that pass-through successors of BranchOp get folded.
+
+// CHECK-LABEL: func @br_passthrough(
+// CHECK-SAME: %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32
+func @br_passthrough(%arg0 : i32, %arg1 : i32) -> (i32, i32) {
+  "foo.switch"() [^bb1, ^bb2, ^bb3] : () -> ()
+
+^bb1:
+  // CHECK: ^bb1:
+  // CHECK-NEXT: br ^bb3(%[[ARG0]], %[[ARG1]] : i32, i32)
+
+  br ^bb2(%arg0 : i32)
+
+^bb2(%arg2 : i32):
+  br ^bb3(%arg2, %arg1 : i32, i32)
+
+^bb3(%arg4 : i32, %arg5 : i32):
+  return %arg4, %arg5 : i32, i32
+}
+
 /// Test the folding of CondBranchOp with a constant condition.
 
 // CHECK-LABEL: func @cond_br_folding(
 func @cond_br_folding(%cond : i1, %a : i32) {
   // CHECK-NEXT: return
 
   %false_cond = constant 0 : i1
   %true_cond = constant 1 : i1
^bb1(%x : i32):
   cond_br %false_cond, ^bb1(%x : i32), ^bb2
 
 ^bb2:
   return
 }
 
 /// Test that pass-through successors of CondBranchOp get folded.
 
-// CHECK-LABEL: func @cond_br_pass_through(
+// CHECK-LABEL: func @cond_br_passthrough(
 // CHECK-SAME: %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32, %[[COND:.*]]: i1
-func @cond_br_pass_through(%arg0 : i32, %arg1 : i32, %arg2 : i32, %cond : i1) -> (i32, i32) {
+func @cond_br_passthrough(%arg0 : i32, %arg1 : i32, %arg2 : i32, %cond : i1) -> (i32, i32) {
   // CHECK: %[[RES:.*]] = select %[[COND]], %[[ARG0]], %[[ARG2]]
   // CHECK: %[[RES2:.*]] = select %[[COND]], %[[ARG1]], %[[ARG2]]
   // CHECK: return %[[RES]], %[[RES2]]
 
   cond_br %cond, ^bb1(%arg0 : i32), ^bb2(%arg2, %arg2 : i32, i32)
 
 ^bb1(%arg3: i32):
   br ^bb2(%arg3, %arg1 : i32, i32)