[mlir] Fix block merging with the result of a terminator

[zero9178]

Summary

When the current implementation merges two blocks that have operands defined outside of their block respectively, it will merge these by adding a block argument in the resulting merged block and adding successor arguments to the predecessors.
There is a special case where this is incorrect however: If one of predecessors terminator produce the operand, inserting the block argument and updating the predecessor would lead to the terminator using its own result as successor argument.
IR Example:

  %0 = "test.producing_br"()[^bb1, ^bb2] {
        operand_segment_sizes = dense<0> : vector<2 x i32>
	} : () -> i32

^bb1:
  "test.br"(%0)[^bb4] : (i32) -> ()

where ^bb1 is then merged with another block would lead to:

%0 = "test.producing_br"(%0)[^bb1, ^bb2]

This patch fixes that issue during clustering by making sure that if the operand is from an outside block, that it is not produced by the terminator of a predecessor.

(Note: The tests use "test.br" instead of cf.br to avoid other canonicalizations changing the IR and not causing the block merging)

[rriddle]

mlir/lib/Transforms/Utils/RegionUtils.cpp

536

Can we remove the mlir:: here?

mlir/test/lib/Dialect/Test/TestDialect.cpp

351

Drop the trivial braces.

[zero9178]

Thanks for noticing! I corrected the code-style issues in https://github.com/llvm/llvm-project/commit/a7865228b30dc019c0d13d0c81247d90fcad325b

mlir/lib/Transforms/Utils/RegionUtils.cpp

for (int operand : llvm::seq<int>(0, lhsIt->getNumOperands())) {
       // Check that these uses are both external, or both internal.
       bool lhsIsInBlock = lhsOperand.getParentBlock() == leaderBlock;
       bool rhsIsInBlock = rhsOperand.getParentBlock() == mergeBlock;
       if (lhsIsInBlock != rhsIsInBlock)
         return failure();
       // Let the operands differ if they are defined in a different block. These
       // will become new arguments if the blocks get merged.
       if (!lhsIsInBlock) {
+
+        // Check whether the operands aren't the result of an immediate
+        // predecessors terminator. In that case we are not able to use it as a
+        // successor operand when branching to the merged block as it does not
+        // dominate its producing operation.
+        auto isValidSuccessorArg = [](Block *block, Value operand) {
+          if (operand.getDefiningOp() !=
+              operand.getParentBlock()->getTerminator())
+            return true;
+          return !llvm::is_contained(block->getPredecessors(),
+                                     operand.getParentBlock());
+        };
+
+        if (!isValidSuccessorArg(leaderBlock, lhsOperand) ||
+            !isValidSuccessorArg(mergeBlock, rhsOperand))
+          return mlir::failure();
+
         mismatchedOperands.emplace_back(opI, operand);
         continue;
       }
 
       // Otherwise, these operands must have the same logical order within the
       // parent block.
       if (leaderData.getOrderOf(lhsOperand) != blockData.getOrderOf(rhsOperand))
         return failure();

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

^bb4(%j: i32):  // 2 preds: ^bb2, ^bb7
   %jcmp = arith.cmpi slt, %j, %arg0 : i32
 // CHECK-NOT:  call @print(%[[arg1:.+]], %[[arg1]])
   call @print(%j, %ip1) : (i32, i32) -> ()
   cf.cond_br %jcmp, ^bb7, ^bb3
 
 ^bb3:  // pred: ^bb1
   return
 }
+
+
+// CHECK-LABEL: func @mismatch_dominance(
+func @mismatch_dominance() -> i32 {
+  // CHECK: %[[RES:.*]] = "test.producing_br"()
+  %0 = "test.producing_br"()[^bb1, ^bb2] {
+        operand_segment_sizes = dense<0> : vector<2 x i32>
+  } : () -> i32
+
+^bb1:
+  // CHECK: "test.br"(%[[RES]])[^[[MERGE_BLOCK:.*]]]
+  "test.br"(%0)[^bb4] : (i32) -> ()
+
+^bb2:
+  %1 = "foo.def"() : () -> i32
+  "test.br"()[^bb3] : () -> ()
+
+^bb3:
+  // CHECK: "test.br"(%{{.*}})[^[[MERGE_BLOCK]]]
+  "test.br"(%1)[^bb4] : (i32) -> ()
+
+^bb4(%3: i32):
+  return %3 : i32
+}

mlir/test/lib/Dialect/Test/TestDialect.cpp

 
 Optional<MutableOperandRange>
 TestBranchOp::getMutableSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
   return getTargetOperandsMutable();
 }
 
 //===----------------------------------------------------------------------===//
+// TestProducingBranchOp
+//===----------------------------------------------------------------------===//
+
+Optional<MutableOperandRange>
+TestProducingBranchOp::getMutableSuccessorOperands(unsigned index) {
+  assert(index <= 1 && "invalid successor index");
+  if (index == 1) {
+    return getFirstOperandsMutable();
+  }
+  return getSecondOperandsMutable();
+}
+
+//===----------------------------------------------------------------------===//
 // TestDialectCanonicalizerOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult
 dialectCanonicalizationPattern(TestDialectCanonicalizerOp op,
                                PatternRewriter &rewriter) {
   rewriter.replaceOpWithNewOp<arith::ConstantOp>(
       op, rewriter.getI32IntegerAttr(42));

mlir/test/lib/Dialect/Test/TestOps.td

def UpdateAttr : Pat<(I32ElementsAttrOp $attr),
                      [(IsNotScalar $attr)]>;
 
 def TestBranchOp : TEST_Op<"br",
     [DeclareOpInterfaceMethods<BranchOpInterface>, Terminator]> {
   let arguments = (ins Variadic<AnyType>:$targetOperands);
   let successors = (successor AnySuccessor:$target);
 }
 
+def TestProducingBranchOp : TEST_Op<"producing_br",
+    [DeclareOpInterfaceMethods<BranchOpInterface>, Terminator,
+     AttrSizedOperandSegments]> {
+  let arguments = (ins Variadic<AnyType>:$firstOperands,
+                       Variadic<AnyType>:$secondOperands);
+  let results = (outs I32:$dummy);
+  let successors = (successor AnySuccessor:$first,AnySuccessor:$second);
+}
+
 def AttrSizedOperandOp : TEST_Op<"attr_sized_operands",
                                  [AttrSizedOperandSegments]> {
   let arguments = (ins
     Variadic<I32>:$a,
     Variadic<I32>:$b,
     I32:$c,
     Variadic<I32>:$d,
     I32ElementsAttr:$operand_segment_sizes