Index: mlir/lib/Transforms/Utils/RegionUtils.cpp =================================================================== --- mlir/lib/Transforms/Utils/RegionUtils.cpp +++ mlir/lib/Transforms/Utils/RegionUtils.cpp @@ -518,6 +518,23 @@ // 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; } Index: mlir/test/Transforms/canonicalize-block-merge.mlir =================================================================== --- mlir/test/Transforms/canonicalize-block-merge.mlir +++ mlir/test/Transforms/canonicalize-block-merge.mlir @@ -251,3 +251,27 @@ ^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 +} Index: mlir/test/lib/Dialect/Test/TestDialect.cpp =================================================================== --- mlir/test/lib/Dialect/Test/TestDialect.cpp +++ mlir/test/lib/Dialect/Test/TestDialect.cpp @@ -341,6 +341,19 @@ return getTargetOperandsMutable(); } +//===----------------------------------------------------------------------===// +// TestProducingBranchOp +//===----------------------------------------------------------------------===// + +Optional +TestProducingBranchOp::getMutableSuccessorOperands(unsigned index) { + assert(index <= 1 && "invalid successor index"); + if (index == 1) { + return getFirstOperandsMutable(); + } + return getSecondOperandsMutable(); +} + //===----------------------------------------------------------------------===// // TestDialectCanonicalizerOp //===----------------------------------------------------------------------===// Index: mlir/test/lib/Dialect/Test/TestOps.td =================================================================== --- mlir/test/lib/Dialect/Test/TestOps.td +++ mlir/test/lib/Dialect/Test/TestOps.td @@ -617,6 +617,15 @@ let successors = (successor AnySuccessor:$target); } +def TestProducingBranchOp : TEST_Op<"producing_br", + [DeclareOpInterfaceMethods, Terminator, + AttrSizedOperandSegments]> { + let arguments = (ins Variadic:$firstOperands, + Variadic:$secondOperands); + let results = (outs I32:$dummy); + let successors = (successor AnySuccessor:$first,AnySuccessor:$second); +} + def AttrSizedOperandOp : TEST_Op<"attr_sized_operands", [AttrSizedOperandSegments]> { let arguments = (ins