[mlir] Region/BranchOpInterface: Allow implicit type conversions along control-flow edges

[Mogball]

Summary

RegionBranchOpInterface and BranchOpInterface are allowed to make implicit type conversions along control-flow edges. In effect, this adds an interface method, areTypesCompatible, to both interfaces, which should return whether the types of corresponding successor operands and block arguments are compatible. Users of the interfaces, here on forth, must be aware that types may mismatch, although current users (in MLIR core), are not affected by this change. By default, type equality is used.

async.execute already has unequal types along control-flow edges (!async.value<f32> vs. f32), but it opted out of calling RegionBranchOpInterface::verifyTypes in its verifier. That method has now been removed and RegionBranchOpInterface will verify types along control edges by default in its verifier.

[rriddle]

I would change the title to be more specific in that this is just about the interfaces (you could already have implicit type changes if you wanted).

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

53

Can you sink this down to where it is actually used?

mlir/test/Transforms/control-flow-sink.mlir

168

Why was this changed?

mlir/unittests/Interfaces/ControlFlowInterfacesTest.cpp

57

You should be able to just do index == 0.

[Mogball]

review comments

[Mogball]

mlir/test/Transforms/control-flow-sink.mlir

168

I was incorrectly using the test op in this test. It turns out that the control-flow semantics of the op according to the interface implementation is that all its operands are passed as block arguments to the then and else regions, and the operands to the yield of those regions are passed as block arguments to the join region (with no indication of which of the then and else regions are actually selected...)

mlir/include/mlir/Dialect/Async/IR/AsyncOps.td

 
 // Base class for the operation in this dialect
 class Async_Op<string mnemonic, list<Trait> traits = []> :
     Op<AsyncDialect, mnemonic, traits>;
 
 def Async_ExecuteOp :
   Async_Op<"execute", [SingleBlockImplicitTerminator<"YieldOp">,
                        DeclareOpInterfaceMethods<RegionBranchOpInterface,
-                                                 ["getSuccessorEntryOperands"]>,
+                                                 ["getSuccessorEntryOperands",
+                                                  "areTypesCompatible"]>,
                        AttrSizedOperandSegments,
                        AutomaticAllocationScope]> {
   let summary = "Asynchronous execute operation";
   let description = [{
     The `body` region attached to the `async.execute` operation semantically
     can be executed concurrently with the successor operation. In the followup
     example "compute0" can be executed concurrently with "compute1".
 

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let description = [{
 
     If `memref.alloca_scope` returns no value, the `memref.alloca_scope.return ` can
     be left out, and will be inserted implicitly.
   }];
 
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$bodyRegion);
   let hasCustomAssemblyFormat = 1;
-  let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // AllocaScopeReturnOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_AllocaScopeReturnOp : MemRef_Op<"alloca_scope.return",
       [HasParent<"AllocaScopeOp">,

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

def Shape_AssumingOp : Shape_Op<"assuming", [
 
   let builders = [
     OpBuilder<(ins "Value":$witness,
         CArg<"function_ref<SmallVector<Value, 2>(OpBuilder &, Location)>">)>
   ];
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
-  let hasVerifier = 1;
 }
 
 def Shape_AssumingYieldOp : Shape_Op<"assuming_yield",
        [NoSideEffect, ReturnLike, Terminator, HasParent<"AssumingOp">]> {
   let summary = "Yield operation";
   let description = [{
     This yield operation represents a return operation within the
     `shape.assuming` operation region. The operation takes variable number of

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

 //===----------------------------------------------------------------------===//
 // BranchOpInterface
 //===----------------------------------------------------------------------===//
 
 def BranchOpInterface : OpInterface<"BranchOpInterface"> {
   let description = [{
     This interface provides information for branching terminator operations,
     i.e. terminator operations with successors.
+
+    This interface is meant to model well-defined cases of control-flow of
+    value propagation, where what occurs along control-flow edges is assumed to
+    be side-effect free. For example, corresponding successor operands and
+    successor block arguments may have different types. In such cases,
+    `areTypesCompatible` can be implemented to compare types along control-flow
+    edges. By default, type equality is used.
   }];
   let cppNamespace = "::mlir";
 
   let methods = [
     InterfaceMethod<[{
         Returns a mutable range of operands that correspond to the arguments of
         successor at the given index. Returns None if the operands to the
         successor are non-materialized values, i.e. they are internal to the
let methods = [
     >,
     InterfaceMethod<[{
         Returns the successor that would be chosen with the given constant
         operands. Returns nullptr if a single successor could not be chosen.
       }],
       "::mlir::Block *", "getSuccessorForOperands",
       (ins "::mlir::ArrayRef<::mlir::Attribute>":$operands), [{}],
       /*defaultImplementation=*/[{ return nullptr; }]
-    >
+    >,
+    InterfaceMethod<[{
+        This method is called to compare types along control-flow edges. By
+        default, the types are checked as equal.
+      }],
+      "bool", "areTypesCompatible",
+      (ins "::mlir::Type":$lhs, "::mlir::Type":$rhs), [{}],
+       [{ return lhs == rhs; }]
+    >,
   ];
 
   let verify = [{
     auto concreteOp = ::mlir::cast<ConcreteOp>($_op);
     for (unsigned i = 0, e = $_op->getNumSuccessors(); i != e; ++i) {
       ::mlir::Optional<OperandRange> operands = concreteOp.getSuccessorOperands(i);
       if (::mlir::failed(::mlir::detail::verifyBranchSuccessorOperands($_op, i, operands)))
         return ::mlir::failure();
     }
     return ::mlir::success();
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // RegionBranchOpInterface
 //===----------------------------------------------------------------------===//
 
 def RegionBranchOpInterface : OpInterface<"RegionBranchOpInterface"> {
   let description = [{
     This interface provides information for region operations that contain
     branching behavior between held regions, i.e. this interface allows for
     expressing control flow information for region holding operations.
+
+    This interface is meant to model well-defined cases of control-flow of
+    value propagation, where what occurs along control-flow edges is assumed to
+    be side-effect free. For example, corresponding successor operands and
+    successor block arguments may have different types. In such cases,
+    `areTypesCompatible` can be implemented to compare types along control-flow
+    edges. By default, type equality is used.
   }];
   let cppNamespace = "::mlir";
 
   let methods = [
     InterfaceMethod<[{
         Returns the operands of this operation used as the entry arguments when
         entering the region at `index`, which was specified as a successor of
         this operation by `getSuccessorRegions`. These operands should
InterfaceMethod<[{
       }],
       "void", "getRegionInvocationBounds",
       (ins "::mlir::ArrayRef<::mlir::Attribute>":$operands,
            "::llvm::SmallVectorImpl<::mlir::InvocationBounds> &"
              :$invocationBounds), [{}],
        [{ invocationBounds.append($_op->getNumRegions(),
                                   ::mlir::InvocationBounds::getUnknown()); }]
     >,
+    InterfaceMethod<[{
+        This method is called to compare types along control-flow edges. By
+        default, the types are checked as equal.
+      }],
+      "bool", "areTypesCompatible",
+      (ins "::mlir::Type":$lhs, "::mlir::Type":$rhs), [{}],
+       [{ return lhs == rhs; }]
+    >,
   ];
 
   let verify = [{
     static_assert(!ConcreteOp::template hasTrait<OpTrait::ZeroRegion>(),
                   "expected operation to have non-zero regions");
-    return success();
+    return detail::verifyTypesAlongControlFlowEdges($_op);
   }];
 
   let extraClassDeclaration = [{
     /// Convenience helper in case none of the operands is known.
     void getSuccessorRegions(Optional<unsigned> index,
                              SmallVectorImpl<RegionSuccessor> &regions) {
        SmallVector<Attribute, 2> nullAttrs(getOperation()->getNumOperands());
        getSuccessorRegions(index, nullAttrs, regions);
     }
-
-    /// Verify types along control flow edges described by this interface.
-    static LogicalResult verifyTypes(Operation *op) {
-      return detail::verifyTypesAlongControlFlowEdges(op);
-    }
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // RegionBranchTerminatorOpInterface
 //===----------------------------------------------------------------------===//
 
 def RegionBranchTerminatorOpInterface :

mlir/lib/Dialect/Async/IR/Async.cpp

if (getOperandTypes() != types)
     return emitOpError("operand types do not match the types returned from "
                        "the parent ExecuteOp");
 
   return success();
 }
 
 MutableOperandRange
 YieldOp::getMutableSuccessorOperands(Optional<unsigned> index) {
-  assert(!index.hasValue());
   return operandsMutable();
 }
 
 //===----------------------------------------------------------------------===//
 /// ExecuteOp
 //===----------------------------------------------------------------------===//
 
 constexpr char kOperandSegmentSizesAttr[] = "operand_segment_sizes";
 
 OperandRange ExecuteOp::getSuccessorEntryOperands(unsigned index) {
   assert(index == 0 && "invalid region index");
   return operands();
 }
 
+bool ExecuteOp::areTypesCompatible(Type lhs, Type rhs) {
+  const auto getValueOrTokenType = [](Type type) {
+    if (auto value = type.dyn_cast<ValueType>())
+      return value.getValueType();
+    return type;
+  };
+  return getValueOrTokenType(lhs) == getValueOrTokenType(rhs);
+}
+
 void ExecuteOp::getSuccessorRegions(Optional<unsigned> index,
                                     ArrayRef<Attribute>,
                                     SmallVectorImpl<RegionSuccessor> &regions) {
   // The `body` region branch back to the parent operation.
   if (index.hasValue()) {
     assert(*index == 0 && "invalid region index");
     regions.push_back(RegionSuccessor(results()));
     return;

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

ParseResult AllocaScopeOp::parse(OpAsmParser &parser, OperationState &result) {
 
   // Parse the optional attribute list.
   if (parser.parseOptionalAttrDict(result.attributes))
     return failure();
 
   return success();
 }
 
-LogicalResult AllocaScopeOp::verify() {
-  return RegionBranchOpInterface::verifyTypes(*this);
-}
-
 void AllocaScopeOp::getSuccessorRegions(
     Optional<unsigned> index, ArrayRef<Attribute> operands,
     SmallVectorImpl<RegionSuccessor> &regions) {
   if (index.hasValue()) {
     regions.push_back(RegionSuccessor(getResults()));
     return;
   }
 

mlir/lib/Dialect/SCF/SCF.cpp

if (std::get<0>(e).getType() != std::get<2>(e).getType())
       return emitOpError() << "types mismatch between " << i
                            << "th iter operand and defined value";
     if (std::get<1>(e).getType() != std::get<2>(e).getType())
       return emitOpError() << "types mismatch between " << i
                            << "th iter region arg and defined value";
 
     i++;
   }
-
-  return RegionBranchOpInterface::verifyTypes(*this);
+  return success();
 }
 
 Optional<Value> ForOp::getSingleInductionVar() { return getInductionVar(); }
 
 Optional<OpFoldResult> ForOp::getSingleLowerBound() {
   return OpFoldResult(getLowerBound());
 }
 
void IfOp::build(OpBuilder &builder, OperationState &result, Value cond,
                  function_ref<void(OpBuilder &, Location)> thenBuilder,
                  function_ref<void(OpBuilder &, Location)> elseBuilder) {
   build(builder, result, TypeRange(), cond, thenBuilder, elseBuilder);
 }
 
 LogicalResult IfOp::verify() {
   if (getNumResults() != 0 && getElseRegion().empty())
     return emitOpError("must have an else block if defining values");
-
-  return RegionBranchOpInterface::verifyTypes(*this);
+  return success();
 }
 
 ParseResult IfOp::parse(OpAsmParser &parser, OperationState &result) {
   // Create the regions for 'then'.
   result.regions.reserve(2);
   Region *thenRegion = result.addRegion();
   Region *elseRegion = result.addRegion();
 
static TerminatorTy verifyAndGetTerminator(scf::WhileOp op, Region &region,
 
   auto diag = op.emitOpError(errorMessage);
   if (terminatorOperation)
     diag.attachNote(terminatorOperation->getLoc()) << "terminator here";
   return nullptr;
 }
 
 LogicalResult scf::WhileOp::verify() {
-  if (failed(RegionBranchOpInterface::verifyTypes(*this)))
-    return failure();
-
   auto beforeTerminator = verifyAndGetTerminator<scf::ConditionOp>(
       *this, getBefore(),
       "expects the 'before' region to terminate with 'scf.condition'");
   if (!beforeTerminator)
     return failure();
 
   auto afterTerminator = verifyAndGetTerminator<scf::YieldOp>(
       *this, getAfter(),

mlir/lib/Dialect/Shape/IR/Shape.cpp

void AssumingOp::build(
   builder.create<AssumingYieldOp>(result.location, yieldValues);
 
   SmallVector<Type, 2> assumingTypes;
   for (Value v : yieldValues)
     assumingTypes.push_back(v.getType());
   result.addTypes(assumingTypes);
 }
 
-LogicalResult AssumingOp::verify() {
-  return RegionBranchOpInterface::verifyTypes(*this);
-}
-
 //===----------------------------------------------------------------------===//
 // AddOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult mlir::shape::AddOp::inferReturnTypes(
     MLIRContext *context, Optional<Location> location, ValueRange operands,
     DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<Type> &inferredReturnTypes) {

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

 
 /// Verify that the given operands match those of the given successor block.
 LogicalResult
 detail::verifyBranchSuccessorOperands(Operation *op, unsigned succNo,
                                       Optional<OperandRange> operands) {
   if (!operands)
     return success();
 
+  auto branchInterface = cast<BranchOpInterface>(op);
+
   // Check the count.
   unsigned operandCount = operands->size();
   Block *destBB = op->getSuccessor(succNo);
   if (operandCount != destBB->getNumArguments())
     return op->emitError() << "branch has " << operandCount
                            << " operands for successor #" << succNo
                            << ", but target block has "
                            << destBB->getNumArguments();
 
   // Check the types.
   auto operandIt = operands->begin();
   for (unsigned i = 0; i != operandCount; ++i, ++operandIt) {
-    if ((*operandIt).getType() != destBB->getArgument(i).getType())
+    if (!branchInterface.areTypesCompatible((*operandIt).getType(),
+                                            destBB->getArgument(i).getType()))
       return op->emitError() << "type mismatch for bb argument #" << i
                              << " of successor #" << succNo;
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // RegionBranchOpInterface
if (sourceTypes->size() != succInputsTypes.size()) {
                                  << " operands, but target successor needs "
                                  << succInputsTypes.size();
     }
 
     for (const auto &typesIdx :
          llvm::enumerate(llvm::zip(*sourceTypes, succInputsTypes))) {
       Type sourceType = std::get<0>(typesIdx.value());
       Type inputType = std::get<1>(typesIdx.value());
-      if (sourceType != inputType) {
+      if (!regionInterface.areTypesCompatible(sourceType, inputType)) {
         InFlightDiagnostic diag = op->emitOpError(" along control flow edge ");
         return printEdgeName(diag)
                << ": source type #" << typesIdx.index() << " " << sourceType
                << " should match input type #" << typesIdx.index() << " "
                << inputType;
       }
     }
   }
LogicalResult detail::verifyTypesAlongControlFlowEdges(Operation *op) {
   // Verify types along control flow edges originating from the parent.
   if (failed(verifyTypesAlongAllEdges(op, llvm::None, inputTypesFromParent)))
     return failure();
 
   // RegionBranchOpInterface should not be implemented by Ops that do not have
   // attached regions.
   assert(op->getNumRegions() != 0);
 
+  auto areTypesCompatible = [&](TypeRange lhs, TypeRange rhs) {
+    if (lhs.size() != rhs.size())
+      return false;
+    for (auto types : llvm::zip(lhs, rhs)) {
+      if (!regionInterface.areTypesCompatible(std::get<0>(types),
+                                              std::get<1>(types))) {
+        return false;
+      }
+    }
+    return true;
+  };
+
   // Verify types along control flow edges originating from each region.
   for (unsigned regionNo : llvm::seq(0U, op->getNumRegions())) {
     Region &region = op->getRegion(regionNo);
 
     // Since there can be multiple `ReturnLike` terminators or others
     // implementing the `RegionBranchTerminatorOpInterface`, all should have the
     // same operand types when passing them to the same region.
 
     Optional<OperandRange> regionReturnOperands;
     for (Block &block : region) {
       Operation *terminator = block.getTerminator();
       auto terminatorOperands =
           getRegionBranchSuccessorOperands(terminator, regionNo);
       if (!terminatorOperands)
         continue;
 
       if (!regionReturnOperands) {
         regionReturnOperands = terminatorOperands;
         continue;
       }
 
       // Found more than one ReturnLike terminator. Make sure the operand types
       // match with the first one.
-      if (regionReturnOperands->getTypes() != terminatorOperands->getTypes())
+      if (!areTypesCompatible(regionReturnOperands->getTypes(),
+                              terminatorOperands->getTypes()))
         return op->emitOpError("Region #")
                << regionNo
                << " operands mismatch between return-like terminators";
     }
 
     auto inputTypesFromRegion =
         [&](Optional<unsigned> regionNo) -> Optional<TypeRange> {
       // If there is no return-like terminator, the op itself should verify

mlir/test/Transforms/control-flow-sink.mlir

 // RUN: mlir-opt -control-flow-sink %s | FileCheck %s
 
 // Test that operations can be sunk.
 
 // CHECK-LABEL: @test_simple_sink
-// CHECK-SAME:  (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)
+// CHECK-SAME:  (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)
 // CHECK-NEXT: %[[V0:.*]] = arith.subi %[[ARG2]], %[[ARG1]]
-// CHECK-NEXT: %[[V1:.*]] = test.region_if %[[ARG0]]: i1 -> i32 then {
+// CHECK-NEXT: %[[V1:.*]] = test.region_if %[[ARG0]]: i32 -> i32 then {
+// CHECK-NEXT: ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:   %[[V2:.*]] = arith.subi %[[ARG1]], %[[ARG2]]
 // CHECK-NEXT:   test.region_if_yield %[[V2]]
 // CHECK-NEXT: } else {
+// CHECK-NEXT: ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:   %[[V2:.*]] = arith.addi %[[ARG1]], %[[ARG1]]
 // CHECK-NEXT:   %[[V3:.*]] = arith.addi %[[V0]], %[[V2]]
 // CHECK-NEXT:   test.region_if_yield %[[V3]]
 // CHECK-NEXT: } join {
+// CHECK-NEXT: ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:   %[[V2:.*]] = arith.addi %[[ARG2]], %[[ARG2]]
 // CHECK-NEXT:   %[[V3:.*]] = arith.addi %[[V2]], %[[V0]]
 // CHECK-NEXT:   test.region_if_yield %[[V3]]
 // CHECK-NEXT: }
 // CHECK-NEXT: return %[[V1]]
-func @test_simple_sink(%arg0: i1, %arg1: i32, %arg2: i32) -> i32 {
+func @test_simple_sink(%arg0: i32, %arg1: i32, %arg2: i32) -> i32 {
   %0 = arith.subi %arg1, %arg2 : i32
   %1 = arith.subi %arg2, %arg1 : i32
   %2 = arith.addi %arg1, %arg1 : i32
   %3 = arith.addi %arg2, %arg2 : i32
-  %4 = test.region_if %arg0: i1 -> i32 then {
+  %4 = test.region_if %arg0: i32 -> i32 then {
+  ^bb0(%arg3: i32):
     test.region_if_yield %0 : i32
   } else {
+  ^bb0(%arg3: i32):
     %5 = arith.addi %1, %2 : i32
     test.region_if_yield %5 : i32
   } join {
+  ^bb0(%arg3: i32):
     %5 = arith.addi %3, %1 : i32
     test.region_if_yield %5 : i32
   }
   return %4 : i32
 }
 
 // Test that a region op can be sunk.
 
 // CHECK-LABEL: @test_region_sink
-// CHECK-SAME:  (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)
-// CHECK-NEXT: %[[V0:.*]] = test.region_if %[[ARG0]]: i1 -> i32 then {
-// CHECK-NEXT:   %[[V1:.*]] = test.region_if %[[ARG0]]: i1 -> i32 then {
+// CHECK-SAME:  (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)
+// CHECK-NEXT: %[[V0:.*]] = test.region_if %[[ARG0]]: i32 -> i32 then {
+// CHECK-NEXT: ^bb0(%{{.*}}: i32):
+// CHECK-NEXT:   %[[V1:.*]] = test.region_if %[[ARG0]]: i32 -> i32 then {
+// CHECK-NEXT:   ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:     test.region_if_yield %[[ARG1]]
 // CHECK-NEXT:   } else {
+// CHECK-NEXT:   ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:     %[[V2:.*]] = arith.subi %[[ARG1]], %[[ARG2]]
 // CHECK-NEXT:     test.region_if_yield %[[V2]]
 // CHECK-NEXT:   } join {
+// CHECK-NEXT:   ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:     test.region_if_yield %[[ARG2]]
 // CHECK-NEXT:   }
 // CHECK-NEXT:   test.region_if_yield %[[V1]]
 // CHECK-NEXT: } else {
+// CHECK-NEXT: ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:   test.region_if_yield %[[ARG1]]
 // CHECK-NEXT: } join {
+// CHECK-NEXT: ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:   test.region_if_yield %[[ARG2]]
 // CHECK-NEXT: }
 // CHECK-NEXT: return %[[V0]]
-func @test_region_sink(%arg0: i1, %arg1: i32, %arg2: i32) -> i32 {
+func @test_region_sink(%arg0: i32, %arg1: i32, %arg2: i32) -> i32 {
   %0 = arith.subi %arg1, %arg2 : i32
-  %1 = test.region_if %arg0: i1 -> i32 then {
+  %1 = test.region_if %arg0: i32 -> i32 then {
+  ^bb0(%arg3: i32):
     test.region_if_yield %arg1 : i32
   } else {
+  ^bb0(%arg3: i32):
     test.region_if_yield %0 : i32
   } join {
+  ^bb0(%arg3: i32):
     test.region_if_yield %arg2 : i32
   }
-  %2 = test.region_if %arg0: i1 -> i32 then {
+  %2 = test.region_if %arg0: i32 -> i32 then {
+  ^bb0(%arg3: i32):
     test.region_if_yield %1 : i32
   } else {
+  ^bb0(%arg3: i32):
     test.region_if_yield %arg1 : i32
   } join {
+  ^bb0(%arg3: i32):
     test.region_if_yield %arg2 : i32
   }
   return %2 : i32
 }
 
 // Test that an entire subgraph can be sunk.
 
 // CHECK-LABEL: @test_subgraph_sink
-// CHECK-SAME:  (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)
-// CHECK-NEXT: %[[V0:.*]] = test.region_if %[[ARG0]]: i1 -> i32 then {
+// CHECK-SAME:  (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)
+// CHECK-NEXT: %[[V0:.*]] = test.region_if %[[ARG0]]: i32 -> i32 then {
+// CHECK-NEXT: ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:   %[[V1:.*]] = arith.subi %[[ARG1]], %[[ARG2]]
 // CHECK-NEXT:   %[[V2:.*]] = arith.addi %[[ARG1]], %[[ARG2]]
 // CHECK-NEXT:   %[[V3:.*]] = arith.subi %[[ARG2]], %[[ARG1]]
 // CHECK-NEXT:   %[[V4:.*]] = arith.muli %[[V3]], %[[V3]]
 // CHECK-NEXT:   %[[V5:.*]] = arith.muli %[[V2]], %[[V1]]
 // CHECK-NEXT:   %[[V6:.*]] = arith.addi %[[V5]], %[[V4]]
 // CHECK-NEXT:   test.region_if_yield %[[V6]]
 // CHECK-NEXT: } else {
+// CHECK-NEXT: ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:   test.region_if_yield %[[ARG1]]
 // CHECK-NEXT: } join {
+// CHECK-NEXT: ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:   test.region_if_yield %[[ARG2]]
 // CHECK-NEXT: }
 // CHECK-NEXT: return %[[V0]]
-func @test_subgraph_sink(%arg0: i1, %arg1: i32, %arg2: i32) -> i32 {
+func @test_subgraph_sink(%arg0: i32, %arg1: i32, %arg2: i32) -> i32 {
   %0 = arith.addi %arg1, %arg2 : i32
   %1 = arith.subi %arg1, %arg2 : i32
   %2 = arith.subi %arg2, %arg1 : i32
   %3 = arith.muli %0, %1 : i32
   %4 = arith.muli %2, %2 : i32
   %5 = arith.addi %3, %4 : i32
-  %6 = test.region_if %arg0: i1 -> i32 then {
+  %6 = test.region_if %arg0: i32 -> i32 then {
+  ^bb0(%arg3: i32):
     test.region_if_yield %5 : i32
   } else {
+  ^bb0(%arg3: i32):
     test.region_if_yield %arg1 : i32
   } join {
+  ^bb0(%arg3: i32):
     test.region_if_yield %arg2 : i32
   }
   return %6 : i32
 }
 
 // Test that ops can be sunk into regions with multiple blocks.
 
 // CHECK-LABEL: @test_multiblock_region_sink
-// CHECK-SAME:  (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)
+// CHECK-SAME:  (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)
 // CHECK-NEXT: %[[V0:.*]] = arith.addi %[[ARG1]], %[[ARG2]]
 // CHECK-NEXT: %[[V1:.*]] = "test.any_cond"() ({
 // CHECK-NEXT:   %[[V3:.*]] = arith.addi %[[V0]], %[[ARG2]]
 // CHECK-NEXT:   %[[V4:.*]] = arith.addi %[[V3]], %[[ARG1]]
 // CHECK-NEXT:   cf.br ^bb1(%[[V4]] : i32)
 // CHECK-NEXT: ^bb1(%[[V5:.*]]: i32):
 // CHECK-NEXT:   %[[V6:.*]] = arith.addi %[[V5]], %[[V4]]
 // CHECK-NEXT:   "test.yield"(%[[V6]])
 // CHECK-NEXT: })
 // CHECK-NEXT: %[[V2:.*]] = arith.addi %[[V0]], %[[V1]]
 // CHECK-NEXT: return %[[V2]]
-func @test_multiblock_region_sink(%arg0: i1, %arg1: i32, %arg2: i32) -> i32 {
+func @test_multiblock_region_sink(%arg0: i32, %arg1: i32, %arg2: i32) -> i32 {
   %0 = arith.addi %arg1, %arg2 : i32
   %1 = arith.addi %0, %arg2 : i32
   %2 = arith.addi %1, %arg1 : i32
   %3 = "test.any_cond"() ({
     cf.br ^bb1(%2 : i32)
   ^bb1(%5: i32):
     %6 = arith.addi %5, %2 : i32
     "test.yield"(%6) : (i32) -> ()
   }) : () -> i32
   %4 = arith.addi %0, %3 : i32
   return %4 : i32
 }
 
 // Test that ops can be sunk recursively into nested regions.
 
 // CHECK-LABEL: @test_nested_region_sink
-// CHECK-SAME:  (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i32) -> i32 {
-// CHECK-NEXT: %[[V0:.*]] = test.region_if %[[ARG0]]: i1 -> i32 then {
-// CHECK-NEXT:   %[[V1:.*]] = test.region_if %[[ARG0]]: i1 -> i32 then {
+// CHECK-SAME:  (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32) -> i32 {
+// CHECK-NEXT: %[[V0:.*]] = test.region_if %[[ARG0]]: i32 -> i32 then {
+// CHECK-NEXT: ^bb0(%{{.*}}: i32):
+// CHECK-NEXT:   %[[V1:.*]] = test.region_if %[[ARG0]]: i32 -> i32 then {
+// CHECK-NEXT:   ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:     %[[V2:.*]] = arith.addi %[[ARG1]], %[[ARG1]]
 // CHECK-NEXT:     test.region_if_yield %[[V2]]
 // CHECK-NEXT:   } else {
+// CHECK-NEXT:   ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:     test.region_if_yield %[[ARG1]]
 // CHECK-NEXT:   } join {
+// CHECK-NEXT:   ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:     test.region_if_yield %[[ARG1]]
 // CHECK-NEXT:   }
 // CHECK-NEXT:   test.region_if_yield %[[V1]]
 // CHECK-NEXT: } else {
+// CHECK-NEXT: ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:   test.region_if_yield %[[ARG1]]
 // CHECK-NEXT: } join {
+// CHECK-NEXT: ^bb0(%{{.*}}: i32):
 // CHECK-NEXT:   test.region_if_yield %[[ARG1]]
 // CHECK-NEXT: }
 // CHECK-NEXT: return %[[V0]]
-func @test_nested_region_sink(%arg0: i1, %arg1: i32) -> i32 {
+func @test_nested_region_sink(%arg0: i32, %arg1: i32) -> i32 {
   %0 = arith.addi %arg1, %arg1 : i32
-  %1 = test.region_if %arg0: i1 -> i32 then {
-    %2 = test.region_if %arg0: i1 -> i32 then {
+  %1 = test.region_if %arg0: i32 -> i32 then {
+  ^bb0(%arg3: i32):
+    %2 = test.region_if %arg0: i32 -> i32 then {
+    ^bb0(%arg4: i32):
       test.region_if_yield %0 : i32
     } else {
+    ^bb0(%arg4: i32):
       test.region_if_yield %arg1 : i32
     } join {
+    ^bb0(%arg4: i32):
       test.region_if_yield %arg1 : i32
     }
     test.region_if_yield %2 : i32
   } else {
+  ^bb0(%arg3: i32):
     test.region_if_yield %arg1 : i32
   } join {
+  ^bb0(%arg3: i32):
     test.region_if_yield %arg1 : i32
   }
   return %1 : i32
 }
 
 // Test that ops are only moved into the entry block, even when their only uses
 // are further along.
 

mlir/unittests/Interfaces/ControlFlowInterfacesTest.cpp

struct SequentialRegionsOp
   static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static StringRef getOperationName() { return "cftest.sequential_regions_op"; }
 
   // Region 0 has Region 1 as a successor.
   void getSuccessorRegions(Optional<unsigned> index,
                            ArrayRef<Attribute> operands,
                            SmallVectorImpl<RegionSuccessor> &regions) {
-    assert(index.hasValue() && "expected index");
-    if (*index == 0) {
+    if (index.hasValue() && *index == 0) {
       Operation *thisOp = this->getOperation();
       regions.push_back(RegionSuccessor(&thisOp->getRegion(1)));
     }
   }
 };
 
 /// A dialect putting all the above together.
 struct CFTestDialect : Dialect {