[mlir] Fix verifier of `RegionBranchOpInterface`

[zero9178]

Summary

The verifier incorrectly passed the region number of the predecessor region instead of the successor region to getSuccessorOperands. This went unnoticed since all upstream RegionBranchTerminatorOpInterface implementations did not make use of the index parameter.
Adding an assert to e.g. scf.condition to make sure the index is valid or adding a region terminator that passes different operands to different successors immediately causes the verifier to fail as it suddenly gets incorrect types.

This patch fixes the implementation to correctly pass the successor region index.

Depends on https://reviews.llvm.org/D157506

[zero9178]

Rebase

[Mogball]

mlir/test/IR/test-region-branch-op-verifier.mlir

2

There's no existing test to merge this withj?

[zero9178]

mlir/test/IR/test-region-branch-op-verifier.mlir

2

I failed to find any tests for region branch op interface :( All changes done on the file have always been tested indirectly by making sure a pass relying on the interface works.
There are no dedicated tests for the verifier as far as I could tell.

mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp

for (Region &region : regions) {
       if (regionSuccessor == successorRegions.end())
         continue;
       // Get the operand index in the context of the current successor input
       // bindings.
       size_t operandIndex =
           llvm::find(regionSuccessor->getSuccessorInputs(), argValue)
               .getIndex();
 
+      std::optional<unsigned> successorRegionNumber;
+      if (Region *successorRegion = regionSuccessor->getSuccessor())
+        successorRegionNumber = successorRegion->getRegionNumber();
+
       // Iterate over all immediate terminator operations to introduce
       // new buffer allocations. Thereby, the appropriate terminator operand
       // will be adjusted to point to the newly allocated buffer instead.
       if (failed(walkReturnOperations(
               &region, [&](RegionBranchTerminatorOpInterface terminator) {
                 // Get the actual mutable operands for this terminator op.
                 auto terminatorOperands =
                     terminator.getMutableSuccessorOperands(
-                        region.getRegionNumber());
+                        successorRegionNumber);
                 // Extract the source value from the current terminator.
                 // This conversion needs to exist on a separate line due to a
                 // bug in GCC conversion analysis.
                 OperandRange immutableTerminatorOperands = terminatorOperands;
                 Value sourceValue = immutableTerminatorOperands[operandIndex];
                 // Create a new clone at the current location of the terminator.
                 auto clone = introduceCloneBuffers(sourceValue, terminator);
                 if (failed(clone))

mlir/lib/Dialect/SCF/IR/SCF.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // ConditionOp
 //===----------------------------------------------------------------------===//
 
 MutableOperandRange
 ConditionOp::getMutableSuccessorOperands(std::optional<unsigned> index) {
+  assert((!index || index == getParentOp().getAfter().getRegionNumber()) &&
+         "condition op can only exit the loop or branch to the after"
+         "region");
   // Pass all operands except the condition to the successor region.
   return getArgsMutable();
 }
 
 void ConditionOp::getSuccessorRegions(
     ArrayRef<Attribute> operands, SmallVectorImpl<RegionSuccessor> &regions) {
   FoldAdaptor adaptor(operands);
 

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

detail::verifyBranchSuccessorOperands(Operation *op, unsigned succNo,
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // RegionBranchOpInterface
 //===----------------------------------------------------------------------===//
 
+static InFlightDiagnostic &
+printRegionEdgeName(InFlightDiagnostic &diag, std::optional<unsigned> sourceNo,
+                    std::optional<unsigned> succRegionNo) {
+  diag << "from ";
+  if (sourceNo)
+    diag << "Region #" << sourceNo.value();
+  else
+    diag << "parent operands";
+
+  diag << " to ";
+  if (succRegionNo)
+    diag << "Region #" << succRegionNo.value();
+  else
+    diag << "parent results";
+  return diag;
+}
+
 /// Verify that types match along all region control flow edges originating from
 /// `sourceNo` (region # if source is a region, std::nullopt if source is parent
 /// op). `getInputsTypesForRegion` is a function that returns the types of the
 /// inputs that flow from `sourceIndex' to the given region, or std::nullopt if
 /// the exact type match verification is not necessary (e.g., if the Op verifies
 /// the match itself).
 static LogicalResult verifyTypesAlongAllEdges(
     Operation *op, std::optional<unsigned> sourceNo,
-    function_ref<std::optional<TypeRange>(std::optional<unsigned>)>
+    function_ref<FailureOr<TypeRange>(std::optional<unsigned>)>
         getInputsTypesForRegion) {
   auto regionInterface = cast<RegionBranchOpInterface>(op);
 
   SmallVector<RegionSuccessor, 2> successors;
   regionInterface.getSuccessorRegions(sourceNo, successors);
 
   for (RegionSuccessor &succ : successors) {
     std::optional<unsigned> succRegionNo;
     if (!succ.isParent())
       succRegionNo = succ.getSuccessor()->getRegionNumber();
 
-    auto printEdgeName = [&](InFlightDiagnostic &diag) -> InFlightDiagnostic & {
-      diag << "from ";
-      if (sourceNo)
-        diag << "Region #" << sourceNo.value();
-      else
-        diag << "parent operands";
-
-      diag << " to ";
-      if (succRegionNo)
-        diag << "Region #" << succRegionNo.value();
-      else
-        diag << "parent results";
-      return diag;
-    };
-
-    std::optional<TypeRange> sourceTypes =
-        getInputsTypesForRegion(succRegionNo);
-    if (!sourceTypes.has_value())
-      continue;
+    FailureOr<TypeRange> sourceTypes = getInputsTypesForRegion(succRegionNo);
+    if (failed(sourceTypes))
+      return failure();
 
     TypeRange succInputsTypes = succ.getSuccessorInputs().getTypes();
     if (sourceTypes->size() != succInputsTypes.size()) {
       InFlightDiagnostic diag = op->emitOpError(" region control flow edge ");
-      return printEdgeName(diag) << ": source has " << sourceTypes->size()
-                                 << " operands, but target successor needs "
-                                 << succInputsTypes.size();
+      return printRegionEdgeName(diag, sourceNo, succRegionNo)
+             << ": source has " << sourceTypes->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 (!regionInterface.areTypesCompatible(sourceType, inputType)) {
         InFlightDiagnostic diag = op->emitOpError(" along control flow edge ");
-        return printEdgeName(diag)
+        return printRegionEdgeName(diag, sourceNo, succRegionNo)
                << ": source type #" << typesIdx.index() << " " << sourceType
                << " should match input type #" << typesIdx.index() << " "
                << inputType;
       }
     }
   }
   return success();
 }
auto areTypesCompatible = [&](TypeRange lhs, TypeRange rhs) {
     }
     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.
+    // Since there can be multiple terminators implementing the
+    // `RegionBranchTerminatorOpInterface`, all should have the same operand
+    // types when passing them to the same region.
+
+    SmallVector<RegionBranchTerminatorOpInterface> regionReturnOps;
+    for (Block &block : region)
+      if (auto terminator = dyn_cast<RegionBranchTerminatorOpInterface>(
+              block.getTerminator()))
+        regionReturnOps.push_back(terminator);
 
-    std::optional<OperandRange> regionReturnOperands;
-    for (Block &block : region) {
-      auto terminator =
-          dyn_cast<RegionBranchTerminatorOpInterface>(block.getTerminator());
-      if (!terminator)
-        continue;
+    // If there is no return-like terminator, the op itself should verify
+    // type consistency.
+    if (regionReturnOps.empty())
 
-      OperandRange terminatorOperands =
-          terminator.getSuccessorOperands(regionNo);
+    auto inputTypesForRegion =
+        [&](std::optional<unsigned> succRegionNo) -> FailureOr<TypeRange> {
+      std::optional<OperandRange> regionReturnOperands;
+      for (RegionBranchTerminatorOpInterface regionReturnOp : regionReturnOps) {
+        auto terminatorOperands =
+            regionReturnOp.getSuccessorOperands(succRegionNo);
+
         if (!regionReturnOperands) {
           regionReturnOperands = terminatorOperands;
           continue;
         }
 
-      // Found more than one ReturnLike terminator. Make sure the operand types
-      // match with the first one.
+        // Found more than one ReturnLike terminator. Make sure the operand
+        // types match with the first one.
         if (!areTypesCompatible(regionReturnOperands->getTypes(),
-                              terminatorOperands.getTypes()))
-        return op->emitOpError("Region #")
-               << regionNo
+                                terminatorOperands.getTypes())) {
+          InFlightDiagnostic diag = op->emitOpError(" along control flow edge");
+          return printRegionEdgeName(diag, regionNo, succRegionNo)
                  << " operands mismatch between return-like terminators";
         }
-
-    auto inputTypesFromRegion =
-        [&](std::optional<unsigned> regionNo) -> std::optional<TypeRange> {
-      // If there is no return-like terminator, the op itself should verify
-      // type consistency.
-      if (!regionReturnOperands)
-        return std::nullopt;
+      }
 
       // All successors get the same set of operand types.
       return TypeRange(regionReturnOperands->getTypes());
     };
 
-    if (failed(verifyTypesAlongAllEdges(op, regionNo, inputTypesFromRegion)))
+    if (failed(verifyTypesAlongAllEdges(op, regionNo, inputTypesForRegion)))
       return failure();
   }
 
   return success();
 }
 
 /// Return `true` if region `r` is reachable from region `begin` according to
 /// the RegionBranchOpInterface (by taking a branch).

mlir/test/IR/test-region-branch-op-verifier.mlir

+// RUN: mlir-opt %s
+
+func.func @test_ops_verify(%arg: i32) -> f32 {
+  %0 = "test.constant"() { value = 5.3 : f32 } : () -> f32
+  %1 = test.loop_block %arg : (i32) -> f32 {
+  ^bb0(%arg1 : i32):
+    test.loop_block_term iter %arg exit %0
+  }
+  return %1 : f32
+}

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

 
 void AnyCondOp::getRegionInvocationBounds(
     ArrayRef<Attribute> operands,
     SmallVectorImpl<InvocationBounds> &invocationBounds) {
   invocationBounds.emplace_back(1, 1);
 }
 
 //===----------------------------------------------------------------------===//
+// LoopBlockOp
+//===----------------------------------------------------------------------===//
+
+void LoopBlockOp::getSuccessorRegions(
+    std::optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
+  regions.emplace_back(&getBody(), getBody().getArguments());
+  if (!index)
+    return;
+
+  regions.emplace_back((*this)->getResults());
+}
+
+OperandRange
+LoopBlockOp::getEntrySuccessorOperands(std::optional<unsigned> index) {
+  assert(index == 0);
+  return getInitMutable();
+}
+
+//===----------------------------------------------------------------------===//
+// LoopBlockTerminatorOp
+//===----------------------------------------------------------------------===//
+
+MutableOperandRange LoopBlockTerminatorOp::getMutableSuccessorOperands(
+    std::optional<unsigned> index) {
+  assert(!index || index == 0);
+  if (!index)
+    return getExitArgMutable();
+  return getNextIterArgMutable();
+}
+
+//===----------------------------------------------------------------------===//
 // SingleNoTerminatorCustomAsmOp
 //===----------------------------------------------------------------------===//
 
 ParseResult SingleNoTerminatorCustomAsmOp::parse(OpAsmParser &parser,
                                                  OperationState &state) {
   Region *body = state.addRegion();
   if (parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{}))
     return failure();

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

 def AnyCondOp : TEST_Op<"any_cond",
       [DeclareOpInterfaceMethods<RegionBranchOpInterface,
                                  ["getRegionInvocationBounds"]>,
        RecursiveMemoryEffects]> {
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region AnyRegion:$region);
 }
 
+def LoopBlockOp : TEST_Op<"loop_block",
+    [DeclareOpInterfaceMethods<RegionBranchOpInterface,
+        ["getEntrySuccessorOperands"]>, RecursiveMemoryEffects]> {
+
+  let results = (outs F32:$floatResult);
+  let arguments = (ins I32:$init);
+  let regions = (region SizedRegion<1>:$body);
+
+  let assemblyFormat = [{
+    $init `:` functional-type($init, $floatResult) $body
+    attr-dict-with-keyword
+  }];
+}
+
+def LoopBlockTerminatorOp : TEST_Op<"loop_block_term",
+    [DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface>, Pure,
+     Terminator]> {
+  let arguments = (ins I32:$nextIterArg, F32:$exitArg);
+
+  let assemblyFormat = [{
+    `iter` $nextIterArg `exit` $exitArg attr-dict
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // Test TableGen generated build() methods
 //===----------------------------------------------------------------------===//
 
 def TableGenConstant : TEST_Op<"tblgen_constant"> {
   let results = (outs AnyType);
 }