[mlir][Standard] Add a canonicalization to simplify cond_br when the successors are identical

[rriddle]

Summary

This revision adds support for canonicalizing the following:

cond_br %cond, ^bb1(A, ..., N), ^bb1(A, ..., N)

br ^bb1(A, ..., N)

If the operands to the successor are different and the cond_br is the only predecessor, we emit selects for the branch operands.

cond_br %cond, ^bb1(A), ^bb1(B)

%select = select %cond, A, B
br ^bb1(%select)

Depends On D78681

llvm/include/llvm/ADT/STLExtras.h

public:
   iterator begin() const { return iterator(base, 0); }
   iterator end() const { return iterator(base, count); }
   ReferenceT operator[](unsigned index) const {
     assert(index < size() && "invalid index for value range");
     return DerivedT::dereference_iterator(base, index);
   }
 
   /// Compare this range with another.
-  template <typename OtherT> bool operator==(const OtherT &other) {
+  template <typename OtherT> bool operator==(const OtherT &other) const {
     return size() == std::distance(other.begin(), other.end()) &&
            std::equal(begin(), end(), other.begin());
   }
+  template <typename OtherT> bool operator!=(const OtherT &other) const {
+    return !(*this == other);
+  }
 
   /// Return the size of this range.
   size_t size() const { return count; }
 
   /// Return if the range is empty.
   bool empty() const { return size() == 0; }
 
   /// Drop the first N elements, and keep M elements.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let description = [{
 
     // Vector selection is element-wise
     %vx = "std.select"(%vcond, %vtrue, %vfalse)
         : (vector<42xi1>, vector<42xf32>, vector<42xf32>) -> vector<42xf32>
     ```
   }];
 
   let arguments = (ins BoolLike:$condition,
-                       SignlessIntegerOrFloatLike:$true_value,
-                       SignlessIntegerOrFloatLike:$false_value);
-  let results = (outs SignlessIntegerOrFloatLike:$result);
+                       AnyType:$true_value,
+                       AnyType:$false_value);
+  let results = (outs AnyType:$result);
   let verifier = ?;
 
   let builders = [OpBuilder<
     "Builder *builder, OperationState &result, Value condition,"
     "Value trueValue, Value falseValue", [{
       result.addOperands({condition, trueValue, falseValue});
       result.addTypes(trueValue.getType());
   }]>];

mlir/include/mlir/IR/Block.h

public:
   /// If this block has exactly one predecessor, return it.  Otherwise, return
   /// null.
   ///
   /// Note that if a block has duplicate predecessors from a single block (e.g.
   /// if you have a conditional branch with the same block as the true/false
   /// destinations) is not considered to be a single predecessor.
   Block *getSinglePredecessor();
 
+  /// If this block has a unique predecessor, i.e., all incoming edges originate
+  /// from one block, return it. Otherwise, return null.
+  Block *getUniquePredecessor();
+
   // Indexed successor access.
   unsigned getNumSuccessors();
   Block *getSuccessor(unsigned i);
 
   // Successor iteration.
   using succ_iterator = SuccessorRange::iterator;
   succ_iterator succ_begin() { return getSuccessors().begin(); }
   succ_iterator succ_end() { return getSuccessors().end(); }

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

void CallIndirectOp::getCanonicalizationPatterns(
   results.insert<SimplifyIndirectCallWithKnownCallee>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // General helpers for comparison ops
 //===----------------------------------------------------------------------===//
 
 // Return the type of the same shape (scalar, vector or tensor) containing i1.
-static Type getCheckedI1SameShape(Type type) {
+static Type getI1SameShape(Type type) {
   auto i1Type = IntegerType::get(1, type.getContext());
-  if (type.isSignlessIntOrIndexOrFloat())
-    return i1Type;
   if (auto tensorType = type.dyn_cast<RankedTensorType>())
     return RankedTensorType::get(tensorType.getShape(), i1Type);
   if (type.isa<UnrankedTensorType>())
     return UnrankedTensorType::get(i1Type);
   if (auto vectorType = type.dyn_cast<VectorType>())
     return VectorType::get(vectorType.getShape(), i1Type);
-  return Type();
-}
-
-static Type getI1SameShape(Type type) {
-  Type res = getCheckedI1SameShape(type);
-  assert(res && "expected type with valid i1 shape");
-  return res;
+  return i1Type;
 }
 
 //===----------------------------------------------------------------------===//
 // CmpIOp
 //===----------------------------------------------------------------------===//
 
 static void buildCmpIOp(Builder *build, OperationState &result,
                         CmpIPredicate predicate, Value lhs, Value rhs) {
OpFoldResult CmpFOp::fold(ArrayRef<Attribute> operands) {
   return IntegerAttr::get(IntegerType::get(1, getContext()), APInt(1, val));
 }
 
 //===----------------------------------------------------------------------===//
 // CondBranchOp
 //===----------------------------------------------------------------------===//
 
 namespace {
-/// cond_br true, ^bb1, ^bb2 -> br ^bb1
-/// cond_br false, ^bb1, ^bb2 -> br ^bb2
+/// cond_br true, ^bb1, ^bb2
+///  -> br ^bb1
+/// cond_br false, ^bb1, ^bb2
+///  -> br ^bb2
 ///
 struct SimplifyConstCondBranchPred : public OpRewritePattern<CondBranchOp> {
   using OpRewritePattern<CondBranchOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(CondBranchOp condbr,
                                 PatternRewriter &rewriter) const override {
     if (matchPattern(condbr.getCondition(), m_NonZero())) {
       // True branch taken.

 };
 
 ///   cond_br %cond, ^bb1, ^bb2
 /// ^bb1
 ///   br ^bbN(...)
 /// ^bb2
 ///   br ^bbK(...)
 ///
-///   cond_br %cond, ^bbN(...), ^bbK(...)
+///  -> cond_br %cond, ^bbN(...), ^bbK(...)
 ///
 struct SimplifyPassThroughCondBranch : public OpRewritePattern<CondBranchOp> {
   using OpRewritePattern<CondBranchOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(CondBranchOp condbr,
                                 PatternRewriter &rewriter) const override {
     Block *trueDest = condbr.trueDest(), *falseDest = condbr.falseDest();
     ValueRange trueDestOperands = condbr.getTrueOperands();
for (Value operand : operands) {
       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)
+///
+struct SimplifyCondBranchIdenticalSuccessors
+    : public OpRewritePattern<CondBranchOp> {
+  using OpRewritePattern<CondBranchOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CondBranchOp condbr,
+                                PatternRewriter &rewriter) const override {
+    // Check that the true and false destinations are the same and have the same
+    // operands.
+    Block *trueDest = condbr.trueDest();
+    if (trueDest != condbr.falseDest())
+      return failure();
+
+    // If all of the operands match, no selects need to be generated.
+    OperandRange trueOperands = condbr.getTrueOperands();
+    OperandRange falseOperands = condbr.getFalseOperands();
+    if (trueOperands == falseOperands) {
+      rewriter.replaceOpWithNewOp<BranchOp>(condbr, trueDest, trueOperands);
+      return success();
+    }
+
+    // Otherwise, if the current block is the only predecessor insert selects
+    // for any mismatched branch operands.
+    if (trueDest->getUniquePredecessor() != condbr.getOperation()->getBlock())
+      return failure();
+
+    // TODO: ATM Tensor/Vector SelectOp requires that the condition has the same
+    // shape as the operands. We should relax that to allow an i1 to signify
+    // that everything is selected.
+    auto doesntSupportsScalarI1 = [](Type type) {
+      return type.isa<TensorType>() || type.isa<VectorType>();
+    };
+    if (llvm::any_of(trueOperands.getTypes(), doesntSupportsScalarI1))
+      return failure();
+
+    // Generate a select for any operands that differ between the two.
+    SmallVector<Value, 8> mergedOperands;
+    mergedOperands.reserve(trueOperands.size());
+    Value condition = condbr.getCondition();
+    for (auto it : llvm::zip(trueOperands, falseOperands)) {
+      if (std::get<0>(it) == std::get<1>(it))
+        mergedOperands.push_back(std::get<0>(it));
+      else
+        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();
+  }
+};
 } // end anonymous namespace
 
 void CondBranchOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
-  results.insert<SimplifyConstCondBranchPred, SimplifyPassThroughCondBranch>(
-      context);
+  results.insert<SimplifyConstCondBranchPred, SimplifyPassThroughCondBranch,
+                 SimplifyCondBranchIdenticalSuccessors>(context);
 }
 
 Optional<OperandRange> CondBranchOp::getSuccessorOperands(unsigned index) {
   assert(index < getNumSuccessors() && "invalid successor index");
   return index == trueIndex ? getTrueOperands() : getFalseOperands();
 }
 
 bool CondBranchOp::canEraseSuccessorOperand() { return true; }

mlir/lib/IR/Block.cpp

Block *Block::getSinglePredecessor() {
   auto it = pred_begin();
   if (it == pred_end())
     return nullptr;
   auto *firstPred = *it;
   ++it;
   return it == pred_end() ? firstPred : nullptr;
 }
 
+/// If this block has a unique predecessor, i.e., all incoming edges originate
+/// from one block, return it. Otherwise, return null.
+Block *Block::getUniquePredecessor() {
+  auto it = pred_begin(), e = pred_end();
+  if (it == e)
+    return nullptr;
+
+  // Check for any conflicting predecessors.
+  auto *firstPred = *it;
+  for (++it; it != e; ++it)
+    if (*it != firstPred)
+      return nullptr;
+  return firstPred;
+}
+
 //===----------------------------------------------------------------------===//
 // Other
 //===----------------------------------------------------------------------===//
 
 /// Split the block into two blocks before the specified operation or
 /// iterator.
 ///
 /// Note that all operations BEFORE the specified iterator stay as part of

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.
+/// 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 the folding of CondBranchOp with a constant condition.
+/// 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: cond_br %{{.*}}, ^bb1, ^bb1
+  // CHECK-NEXT: return
 
   %false_cond = constant 0 : i1
   %true_cond = constant 1 : i1
   cond_br %cond, ^bb1, ^bb2(%a : i32)
 
 ^bb1:
   cond_br %true_cond, ^bb3, ^bb2(%a : i32)
 
 ^bb2(%x : i32):
   cond_br %false_cond, ^bb2(%x : i32), ^bb3
 
 ^bb3:
-  // CHECK: ^bb1:
+  return
+}
+
+/// Test the folding of CondBranchOp when the successors are identical.
+
+// CHECK-LABEL: func @cond_br_same_successor(
+func @cond_br_same_successor(%cond : i1, %a : i32) {
   // CHECK-NEXT: return
 
+  cond_br %cond, ^bb1(%a : i32), ^bb1(%a : i32)
+
+^bb1(%result : i32):
   return
 }
 
-// Test the compound folding of BranchOp and CondBranchOp.
+/// Test the folding of CondBranchOp when the successors are identical, but the
+/// arguments are different.
+
+// CHECK-LABEL: func @cond_br_same_successor_insert_select(
+// CHECK-SAME: %[[COND:.*]]: i1, %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32
+func @cond_br_same_successor_insert_select(%cond : i1, %a : i32, %b : i32) -> i32 {
+  // CHECK: %[[RES:.*]] = select %[[COND]], %[[ARG0]], %[[ARG1]]
+  // CHECK: return %[[RES]]
+
+  cond_br %cond, ^bb1(%a : i32), ^bb1(%b : i32)
+
+^bb1(%result : i32):
+  return %result : i32
+}
+
+/// Check that we don't generate a select if the type requires a splat.
+/// TODO: SelectOp should allow for matching a vector/tensor with i1.
+
+// CHECK-LABEL: func @cond_br_same_successor_no_select_tensor(
+func @cond_br_same_successor_no_select_tensor(%cond : i1, %a : tensor<2xi32>,
+                                              %b : tensor<2xi32>) -> tensor<2xi32>{
+  // CHECK: cond_br
+
+  cond_br %cond, ^bb1(%a : tensor<2xi32>), ^bb1(%b : tensor<2xi32>)
+
+^bb1(%result : tensor<2xi32>):
+  return %result : tensor<2xi32>
+}
+
+// CHECK-LABEL: func @cond_br_same_successor_no_select_vector(
+func @cond_br_same_successor_no_select_vector(%cond : i1, %a : vector<2xi32>,
+                                              %b : vector<2xi32>) -> vector<2xi32> {
+  // CHECK: cond_br
+
+  cond_br %cond, ^bb1(%a : vector<2xi32>), ^bb1(%b : vector<2xi32>)
+
+^bb1(%result : vector<2xi32>):
+  return %result : vector<2xi32>
+}
+
+/// Test the compound folding of BranchOp and CondBranchOp.
 
 // CHECK-LABEL: func @cond_br_and_br_folding(
 func @cond_br_and_br_folding(%a : i32) {
   // CHECK-NEXT: return
 
   %false_cond = constant 0 : i1
   %true_cond = constant 1 : i1
   cond_br %true_cond, ^bb2, ^bb1(%a : i32)
 
 ^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-SAME: %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32
+// 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) {
-  // CHECK: cond_br %{{.*}}, ^bb1(%[[ARG0]], %[[ARG1]] : i32, i32), ^bb1(%[[ARG2]], %[[ARG2]] : 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)
 
 ^bb2(%arg4: i32, %arg5: i32):
-  // CHECK: ^bb1(%[[RET0:.*]]: i32, %[[RET1:.*]]: i32):
-  // CHECK-NEXT: return %[[RET0]], %[[RET1]]
-
   return %arg4, %arg5 : i32, i32
 }
 
 /// Test the failure modes of collapsing CondBranchOp pass-throughs successors.
 
 // CHECK-LABEL: func @cond_br_pass_through_fail(
 func @cond_br_pass_through_fail(%cond : i1) {
   // CHECK: cond_br %{{.*}}, ^bb1, ^bb2

mlir/test/IR/invalid-ops.mlir

 func @func_with_ops(i1, tensor<42xi32>, tensor<?xi32>) {
 ^bb0(%cond : i1, %t : tensor<42xi32>, %f : tensor<?xi32>):
   // expected-error@+1 {{ op requires the same shape for all operands and results}}
   %r = "std.select"(%cond, %t, %f) : (i1, tensor<42xi32>, tensor<?xi32>) -> tensor<42xi32>
 }
 
 // -----
 
-func @invalid_select_shape(%cond : i1, %idx : () -> ()) {
-  // expected-error@+1 {{'result' must be signless-integer-like or floating-point-like, but got '() -> ()'}}
-  %sel = select %cond, %idx, %idx : () -> ()
-
-// -----
-
 func @invalid_cmp_shape(%idx : () -> ()) {
   // expected-error@+1 {{'lhs' must be signless-integer-like, but got '() -> ()'}}
   %cmp = cmpi "eq", %idx, %idx : () -> ()
 
 // -----
 
 func @dma_no_src_memref(%m : f32, %tag : f32, %c0 : index) {
   // expected-error@+1 {{expected source to be of memref type}}