Added new RegionBranchTerminatorOpInterface and adapted uses of hasTrait<ReturnLike>.

[dfki-mako]

Summary

This CL adds a new RegionBranchTerminatorOpInterface to query information about operands that can be
passed to successor regions. Similar to the BranchOpInterface, it allows to freely define the
involved operands. However, in contrast to the BranchOpInterface, it expects an additional region
number to distinguish between various use cases which might require different operands passed to
different regions.

Moreover, we added new utility functions (namely getMutableRegionBranchSuccessorOperands and
getRegionBranchSuccessorOperands) to query (mutable) operand ranges for operations equiped with the
ReturnLike trait and/or implementing the newly added interface. This simplifies reasoning about
terminators in the scope of the nested regions.

We also adjusted the SCF.ConditionOp to benefit from the newly added capabilities.

See also: discussion on Discourse

[rriddle]

Thanks, this seems like a natural evolution.

mlir/include/mlir/Dialect/SCF/SCFOps.td

38–45

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

182–183

Can you fill this in?

190

Can you reference the specific parameters of the method here? So that it is clear what the represent?

mlir/lib/Analysis/BufferViewFlowAnalysis.cpp

105–109

115

You'll also need to update many other usages of the control flow interfaces to use this. Have you considered adding a utility/<something-else> for this? So that users don't need to check multiple interfaces/traits? e.g. if ReturnLike in ODS also implied this interface somehow.

121

If one side of the branch has braces, both sides should.

[mehdi_amini]

mlir/lib/Analysis/BufferViewFlowAnalysis.cpp

121

Alternatively, add continue to the then branch and eliminate the else and reduce indentation afterward.

[dfki-mako]

Addressed reviewer comments and refined RegionBranchTerminatorOpInterface.
Adapted uses of hasTrait<ReturnLike> to use newly added functions (if possible).

[dfki-mako]

@rriddle there is also another use of the ReturnLike trait in the scope of the DataFlowAnalysis. Currently, you check whether a terminator is ReturnLike. If not, you mark the following region successors to have reached a pessimistic fixpoint. This raises the question of whether we want to change this behavior to propagate lattice information in the context of "more return-like" operations.

The following patch would change this behavior to include all operations that are either marked as ReturnLike or implement the RegionBranchTerminatorOpInterface:

@@ -562,26 +562,28 @@ void ForwardDataFlowSolver::visitTerminatorOperation(
     if (regionSuccessors.empty())
       return;
 
+    // Try to get "region-like" successor operands if possible in order to
+    // propagate the operand states to the successors.
+    if (isRegionReturnLike(op))
+      return visitRegionSuccessors(
+          parentOp, regionSuccessors, [&](Optional<unsigned> regionIndex) {
+            // Determine the individual region successor operands for the given
+            // region index (if any).
+            return *getRegionBranchSuccessorOperands(op, regionIndex);
+          });
+

However, I am completely fine with *not* changing this code to make sure that lattice states are propagated to region successors in the scope of "real return-like" operations only.

[rriddle]

In D105018#2873291, @dfki-mako wrote:

@rriddle there is also another use of the ReturnLike trait in the scope of the DataFlowAnalysis. Currently, you check whether a terminator is ReturnLike. If not, you mark the following region successors to have reached a pessimistic fixpoint. This raises the question of whether we want to change this behavior to propagate lattice information in the context of "more return-like" operations.

The following patch would change this behavior to include all operations that are either marked as ReturnLike or implement the RegionBranchTerminatorOpInterface:

@@ -562,26 +562,28 @@ void ForwardDataFlowSolver::visitTerminatorOperation(
     if (regionSuccessors.empty())
       return;
 
+    // Try to get "region-like" successor operands if possible in order to
+    // propagate the operand states to the successors.
+    if (isRegionReturnLike(op))
+      return visitRegionSuccessors(
+          parentOp, regionSuccessors, [&](Optional<unsigned> regionIndex) {
+            // Determine the individual region successor operands for the given
+            // region index (if any).
+            return *getRegionBranchSuccessorOperands(op, regionIndex);
+          });
+

However, I am completely fine with *not* changing this code to make sure that lattice states are propagated to region successors in the scope of "real return-like" operations only.

I would include RegionBranchTerminatorOpInterface handling in that code as well. Using RegionBranchTerminatorOpInterface should allow for providing a more refined result in cases that use it.

[dfki-mako]

Adatped ForwardDataFlowSolver to support the newly added RegionBranchTerminatorOpInterface.

[herhut]

What is the longer term plan to get rid of ReturnLike? This can land in the meantime as a first step in the right direction, though.

[rriddle]

In D105018#2886846, @herhut wrote:

What is the longer term plan to get rid of ReturnLike? This can land in the meantime as a first step in the right direction, though.

I'd prefer if you can wait until I have time to look it over. I'm currently OOO for the next few days.

[zero9178]

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

196–206

Could the types here be fully qualified (eg. ::mlir::OperandRange instead of OperandRange)? Otherwise it is impossible to use this interface in a downstream project without also putting using namespace mlir; into the header.

[dfki-mako]

@herhut I think it might be beneficial to either remove the ReturnLike trait or to change it in a way such that it automatically provides a "default implementation" of the newly added RegionBranchTerminatorOpInterface.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

196–206

Hmm, yes we can do that :) However, there are more occurrences in the same file without explicitly qualified OperandRange/MutableOperandRange types.
@rriddle Should we also explicitly qualify all other uses?

[herhut]

In D105018#2892631, @dfki-mako wrote:

@herhut I think it might be beneficial to either remove the ReturnLike trait or to change it in a way such that it automatically provides a "default implementation" of the newly added RegionBranchTerminatorOpInterface.

I think ReturnLike is common enough to have it as a default implementation. Otherwise we would have to implement this again and again. I see it similar to FunctionLike.

Added benefit, it provides a transition path.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

196–206

The other cases likely are an oversight and would have the same effect. So if you fix them, too, that would be great.

[dfki-mako]

Changed namespace qualifiers to explicitly reference the llvm/mlir namespaces.

[rriddle]

Thanks a lot for this! Looks really great.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

49

Only the interface method params need to have things fully qualified. Everything else (method bodies/extraTraitDecls/extraClassDecls/verify/etc.) only ever get placed in the mlir:: namespace.

208–209

You should be able to use $_op for this instead.

mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp

104

106

mlir/lib/Analysis/BufferViewFlowAnalysis.cpp

106

nit: Cache the result as a variable and reuse it below.

115

mlir/lib/Analysis/DataFlowAnalysis.cpp

567

Can you add braces here? The body spans quite a few lines.

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

198–199

212

230

When is operation ever null? dyn_cast (well isa) asserts on null, so this can't happen.

249

I mentioned it earlier, but I'd love if ReturnLike could imply RegionBranchTerminatorOpInterface somehow. That make this code much simpler to deal with (i.e. users could just interface with RegionBranchTerminatorOpInterface directly). Can you add a TODO here for exploring that?

[dfki-mako]

Addressed reviewer comments from @rriddle.

mlir/include/mlir/Dialect/SCF/SCFOps.td

class SCF_Op<string mnemonic, list<OpTrait> traits = []> :
   //   * LogicalResult verify(${C++ class of Op} op)
   //   * ParseResult parse${C++ class of Op}(OpAsmParser &parser,
   //                                         OperationState &result)
   // functions.
   let printer = [{ return ::print(p, *this); }];
   let verifier = [{ return ::verify(*this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
-def ConditionOp : SCF_Op<"condition",
-                         [HasParent<"WhileOp">, NoSideEffect, Terminator]> {
+def ConditionOp : SCF_Op<"condition", [
+  HasParent<"WhileOp">,
+  DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface>,
+  NoSideEffect,
+  Terminator
+]> {
   let summary = "loop continuation condition";
   let description = [{
     This operation accepts the continuation (i.e., inverse of exit) condition
     of the `scf.while` construct. If its first argument is true, the "after"
     region of `scf.while` is executed, with the remaining arguments forwarded
     to the entry block of the region. Otherwise, the loop terminates.
   }];
 
   let arguments = (ins I1:$condition, Variadic<AnyType>:$args);

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

public:
   ValueRange getSuccessorInputs() const { return inputs; }
 
 private:
   Region *region;
   ValueRange inputs;
 };
 
 //===----------------------------------------------------------------------===//
+// RegionBranchTerminatorOpInterface
+//===----------------------------------------------------------------------===//
+
+/// Returns true if the given operation is either annotated with the
+/// `ReturnLike` trait or implements the `RegionBranchTerminatorOpInterface`.
+bool isRegionReturnLike(Operation *operation);
+
+/// Returns the mutable operands that are passed to the region with the given
+/// `regionIndex`. If the operation does not implement the
+/// `RegionBranchTerminatorOpInterface` and is not marked as `ReturnLike`, the
+/// result will be `llvm::None`. In all other cases, the resulting
+/// `OperandRange` represents all operands that are passed to the specified
+/// successor region. If `regionIndex` is `llvm::None`, all operands that are
+/// passed to the parent operation will be returned.
+Optional<MutableOperandRange>
+getMutableRegionBranchSuccessorOperands(Operation *operation,
+                                        Optional<unsigned> regionIndex);
+
+/// Returns the read only operands that are passed to the region with the given
+/// `regionIndex`. See `getMutableRegionBranchSuccessorOperands` for more
+/// information.
+Optional<OperandRange>
+getRegionBranchSuccessorOperands(Operation *operation,
+                                 Optional<unsigned> regionIndex);
+
+//===----------------------------------------------------------------------===//
 // ControlFlow Traits
 //===----------------------------------------------------------------------===//
 
 namespace OpTrait {
 /// This trait indicates that a terminator operation is "return-like". This
 /// means that it exits its current region and forwards its operands as "exit"
 /// values to the parent region. Operations with this trait are not permitted to
 /// contain successors or produce results.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

let methods = [
     InterfaceMethod<[{
         Returns a 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
         operation.
       }],
       "::mlir::Optional<::mlir::OperandRange>", "getSuccessorOperands",
       (ins "unsigned":$index), [{}], [{
-        ConcreteOp *op = static_cast<ConcreteOp *>(this);
-        auto operands = op->getMutableSuccessorOperands(index);
+        auto operands = $_op.getMutableSuccessorOperands(index);
         return operands ? ::mlir::Optional<::mlir::OperandRange>(*operands) : ::llvm::None;
       }]
     >,
     InterfaceMethod<[{
         Returns the `BlockArgument` corresponding to operand `operandIndex` in
         some successor, or None if `operandIndex` isn't a successor operand
         index.
       }],
       "::mlir::Optional<::mlir::BlockArgument>", "getSuccessorBlockArgument",
let extraClassDeclaration = [{
     /// Verify types along control flow edges described by this interface.
     static LogicalResult verifyTypes(Operation *op) {
       return detail::verifyTypesAlongControlFlowEdges(op);
     }
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// RegionBranchTerminatorOpInterface
+//===----------------------------------------------------------------------===//
+
+def RegionBranchTerminatorOpInterface :
+  OpInterface<"RegionBranchTerminatorOpInterface"> {
+  let description = [{
+    This interface provides information for branching terminator operations
+    in the presence of a parent RegionBranchOpInterface implementation. It
+    specifies which operands are passed to which successor region.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<[{
+        Returns a mutable range of operands that are semantically "returned" by
+        passing them to the region successor given by `index`.  If `index` is
+        None, this function returns the operands that are passed as a result to
+        the parent operation.
+      }],
+      "::mlir::MutableOperandRange", "getMutableSuccessorOperands",
+      (ins "::mlir::Optional<unsigned>":$index)
+    >,
+    InterfaceMethod<[{
+        Returns a range of operands that are semantically "returned" by passing
+        them to the region successor given by `index`.  If `index` is None, this
+        function returns the operands that are passed as a result to the parent
+        operation.
+      }],
+      "::mlir::OperandRange", "getSuccessorOperands",
+      (ins "::mlir::Optional<unsigned>":$index), [{}],
+      /*defaultImplementation=*/[{
+        return $_op.getMutableSuccessorOperands(index);
+      }]
+    >
+  ];
+
+  let verify = [{
+    static_assert(ConcreteOp::template hasTrait<OpTrait::IsTerminator>(),
+                  "expected operation to be a terminator");
+    static_assert(ConcreteOp::template hasTrait<OpTrait::ZeroResult>(),
+                  "expected operation to have zero results");
+    static_assert(ConcreteOp::template hasTrait<OpTrait::ZeroSuccessor>(),
+                  "expected operation to have zero successors");
+    return success();
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // ControlFlow Traits
 //===----------------------------------------------------------------------===//
 
 // Op is "return-like".
 def ReturnLike : NativeOpTrait<"ReturnLike">;
 
 #endif // MLIR_INTERFACES_CONTROLFLOWINTERFACES

mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp

for (RegionSuccessor &successor : successors) {
         break;
       }
       return inputIndex - firstInputIndex;
     }
     return llvm::None;
   };
 
   // Check branches from the parent operation.
+  Optional<unsigned> regionIndex;
   if (region) {
+    // Determine the actual region number from the passed region.
+    regionIndex = region->getRegionNumber();
     if (Optional<unsigned> operandIndex =
             getOperandIndexIfPred(/*predIndex=*/llvm::None)) {
       collectUnderlyingAddressValues(
-          branch.getSuccessorEntryOperands(
-              region->getRegionNumber())[*operandIndex],
+          branch.getSuccessorEntryOperands(*regionIndex)[*operandIndex],
           maxDepth, visited, output);
     }
   }
   // Check branches from each child region.
   Operation *op = branch.getOperation();
   for (int i = 0, e = op->getNumRegions(); i != e; ++i) {
     if (Optional<unsigned> operandIndex = getOperandIndexIfPred(i)) {
       for (Block &block : op->getRegion(i)) {
         Operation *term = block.getTerminator();
-        if (term->hasTrait<OpTrait::ReturnLike>()) {
-          collectUnderlyingAddressValues(term->getOperand(*operandIndex),
+        // Try to determine possible region-branch successor operands for the
+        // current region.
+        auto successorOperands =
+            getRegionBranchSuccessorOperands(term, regionIndex);
+        if (successorOperands) {
+          collectUnderlyingAddressValues((*successorOperands)[*operandIndex],
                                          maxDepth, visited, output);
         } else if (term->getNumSuccessors()) {
           // Otherwise, if this terminator may exit the region we can't make
           // any assumptions about which values get passed.
           output.push_back(inputValue);
           return;
         }
       }
     }
   }
 }
 
 /// Given a result, collect all of the underlying values being addressed.
 static void collectUnderlyingAddressValues(OpResult result, unsigned maxDepth,
                                            DenseSet<Value> &visited,

mlir/lib/Analysis/BufferViewFlowAnalysis.cpp

op->walk([&](RegionBranchOpInterface regionInterface) {
     // Wire flow between regions and from region exits.
     for (Region &region : regionInterface->getRegions()) {
       // Iterate over all successor region entries that are reachable from the
       // current region.
       SmallVector<RegionSuccessor, 2> successorRegions;
       regionInterface.getSuccessorRegions(region.getRegionNumber(),
                                           successorRegions);
       for (RegionSuccessor &successorRegion : successorRegions) {
+        // Determine the current region index (if any).
+        Optional<unsigned> regionIndex;
+        Region *regionSuccessor = successorRegion.getSuccessor();
+        if (regionSuccessor)
+          regionIndex = regionSuccessor->getRegionNumber();
         // Iterate over all immediate terminator operations and wire the
-        // successor inputs with the operands of each terminator.
+        // successor inputs with the successor operands of each terminator.
         for (Block &block : region) {
-          Operation &terminator = *block.getTerminator();
-          if (terminator.hasTrait<OpTrait::ReturnLike>()) {
-            registerDependencies(terminator.getOperands(),
+          auto successorOperands = getRegionBranchSuccessorOperands(
+              block.getTerminator(), regionIndex);
+          if (successorOperands) {
+            registerDependencies(*successorOperands,
                                  successorRegion.getSuccessorInputs());
           }
         }
       }
     }
   });
 }

mlir/lib/Analysis/DataFlowAnalysis.cpp

if (op->getNumSuccessors() == 0) {
     // optimistic lattice state.
     SmallVector<RegionSuccessor, 1> regionSuccessors;
     analysis.getSuccessorsForOperands(regionInterface,
                                       parentRegion->getRegionNumber(),
                                       operandLattices, regionSuccessors);
     if (regionSuccessors.empty())
       return;
 
+    // Try to get "region-like" successor operands if possible in order to
+    // propagate the operand states to the successors.
+    if (isRegionReturnLike(op)) {
+      return visitRegionSuccessors(
+          parentOp, regionSuccessors, [&](Optional<unsigned> regionIndex) {
+            // Determine the individual region successor operands for the given
+            // region index (if any).
+            return *getRegionBranchSuccessorOperands(op, regionIndex);
+          });
+    }
+
     // If this terminator is not "region-like", conservatively mark all of the
     // successor values as having reached the pessimistic fixpoint.
-    if (!op->hasTrait<OpTrait::ReturnLike>()) {
-      for (auto &it : regionSuccessors) {
-        // If the successor is a region, mark the entry block as executable so
-        // that we visit operations defined within. If the successor is the
-        // parent operation, we simply mark the control flow results as having
-        // reached the pessimistic state.
-        if (Region *region = it.getSuccessor())
-          markEntryBlockExecutable(region, /*markPessimisticFixpoint=*/true);
-        else
-          markAllPessimisticFixpointAndVisitUsers(it.getSuccessorInputs());
-      }
-      return;
-    }
-
-    // Otherwise, propagate the operand lattice states to the successors.
-    OperandRange operands = op->getOperands();
-    return visitRegionSuccessors(parentOp, regionSuccessors,
-                                 [&](Optional<unsigned>) { return operands; });
   }
 
   // Try to resolve to a specific set of successors with the current optimistic
   // lattice state.
   Block *block = op->getBlock();
   if (auto branch = dyn_cast<BranchOpInterface>(op)) {
     SmallVector<Block *> successors;
     if (succeeded(analysis.getSuccessorsForOperands(branch, operandLattices,

mlir/lib/Dialect/SCF/SCF.cpp

 };
 
 void ExecuteRegionOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
   results.add<SingleBlockExecuteInliner, MultiBlockExecuteInliner>(context);
 }
 
 //===----------------------------------------------------------------------===//
+// ConditionOp
+//===----------------------------------------------------------------------===//
+
+MutableOperandRange
+ConditionOp::getMutableSuccessorOperands(Optional<unsigned> index) {
+  // Pass all operands except the condition to the successor region.
+  return argsMutable();
+}
+
+//===----------------------------------------------------------------------===//
 // ForOp
 //===----------------------------------------------------------------------===//
 
 void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
                   Value ub, Value step, ValueRange iterArgs,
                   BodyBuilderFn bodyBuilder) {
   result.addOperands({lb, ub, step});
   result.addOperands(iterArgs);
if (failed(RegionBranchOpInterface::verifyTypes(op)))
     return failure();
 
   auto beforeTerminator = verifyAndGetTerminator<scf::ConditionOp>(
       op, op.before(),
       "expects the 'before' region to terminate with 'scf.condition'");
   if (!beforeTerminator)
     return failure();
 
-  TypeRange trailingTerminatorOperands = beforeTerminator.args().getTypes();
-  if (failed(verifyTypeRangesMatch(op, trailingTerminatorOperands,
-                                   op.after().getArgumentTypes(),
-                                   "trailing operands of the 'before' block "
-                                   "terminator and 'after' region arguments")))
-    return failure();
-
-  if (failed(verifyTypeRangesMatch(
-          op, trailingTerminatorOperands, op.getResultTypes(),
-          "trailing operands of the 'before' block terminator and op results")))
-    return failure();
-
   auto afterTerminator = verifyAndGetTerminator<scf::YieldOp>(
       op, op.after(),
       "expects the 'after' region to terminate with 'scf.yield'");
   return success(afterTerminator != nullptr);
 }
 
 namespace {
 /// Replace uses of the condition within the do block with true, since otherwise

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

LogicalResult detail::verifyTypesAlongControlFlowEdges(Operation *op) {
   // RegionBranchOpInterface should not be implemented by Ops that do not have
   // attached regions.
   assert(op->getNumRegions() != 0);
 
   // Verify types along control flow edges originating from each region.
   for (unsigned regionNo : llvm::seq(0U, op->getNumRegions())) {
     Region &region = op->getRegion(regionNo);
 
-    // Since the interface cannot distinguish between different ReturnLike
-    // ops within the region branching to different successors, all ReturnLike
-    // ops in this region should have the same operand types. We will then use
-    // one of them as the representative for type matching.
+    // 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.
 
-    Operation *regionReturn = nullptr;
+    Optional<OperandRange> regionReturnOperands;
     for (Block &block : region) {
       Operation *terminator = block.getTerminator();
-      if (!terminator->hasTrait<OpTrait::ReturnLike>())
+      auto terminatorOperands =
+          getRegionBranchSuccessorOperands(terminator, regionNo);
+      if (!terminatorOperands)
         continue;
 
-      if (!regionReturn) {
-        regionReturn = terminator;
+      if (!regionReturnOperands) {
+        regionReturnOperands = terminatorOperands;
         continue;
       }
 
       // Found more than one ReturnLike terminator. Make sure the operand types
       // match with the first one.
-      if (regionReturn->getOperandTypes() != terminator->getOperandTypes())
+      if (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
       // type consistency.
-      if (!regionReturn)
+      if (!regionReturnOperands)
         return llvm::None;
 
-      // All successors get the same set of operands.
-      return TypeRange(regionReturn->getOperands().getTypes());
+      // All successors get the same set of operand types.
+      return TypeRange(regionReturnOperands->getTypes());
     };
 
     if (failed(verifyTypesAlongAllEdges(op, regionNo, inputTypesFromRegion)))
       return failure();
   }
 
   return success();
 }
+
+//===----------------------------------------------------------------------===//
+// RegionBranchTerminatorOpInterface
+//===----------------------------------------------------------------------===//
+
+/// Returns true if the given operation is either annotated with the
+/// `ReturnLike` trait or implements the `RegionBranchTerminatorOpInterface`.
+bool mlir::isRegionReturnLike(Operation *operation) {
+  return dyn_cast<RegionBranchTerminatorOpInterface>(operation) ||
+         operation->hasTrait<OpTrait::ReturnLike>();
+}
+
+/// Returns the mutable operands that are passed to the region with the given
+/// `regionIndex`. If the operation does not implement the
+/// `RegionBranchTerminatorOpInterface` and is not marked as `ReturnLike`, the
+/// result will be `llvm::None`. In all other cases, the resulting
+/// `OperandRange` represents all operands that are passed to the specified
+/// successor region. If `regionIndex` is `llvm::None`, all operands that are
+/// passed to the parent operation will be returned.
+Optional<MutableOperandRange>
+mlir::getMutableRegionBranchSuccessorOperands(Operation *operation,
+                                              Optional<unsigned> regionIndex) {
+  // Try to query a RegionBranchTerminatorOpInterface to determine
+  // all successor operands that will be passed to the successor
+  // input arguments.
+  if (auto regionTerminatorInterface =
+          dyn_cast<RegionBranchTerminatorOpInterface>(operation))
+    return regionTerminatorInterface.getMutableSuccessorOperands(regionIndex);
+
+  // TODO: The ReturnLike trait should imply a default implementation of the
+  // RegionBranchTerminatorOpInterface. This would make this code significantly
+  // easier. Furthermore, this may even make this function obsolete.
+  if (operation->hasTrait<OpTrait::ReturnLike>())
+    return MutableOperandRange(operation);
+  return llvm::None;
+}
+
+/// Returns the read only operands that are passed to the region with the given
+/// `regionIndex`. See `getMutableRegionBranchSuccessorOperands` for more
+/// information.
+Optional<OperandRange>
+mlir::getRegionBranchSuccessorOperands(Operation *operation,
+                                       Optional<unsigned> regionIndex) {
+  auto range = getMutableRegionBranchSuccessorOperands(operation, regionIndex);
+  return range ? Optional<OperandRange>(*range) : llvm::None;
+}

mlir/lib/Transforms/BufferDeallocation.cpp

 
 using namespace mlir;
 
 /// Walks over all immediate return-like terminators in the given region.
 static void walkReturnOperations(Region *region,
                                  std::function<void(Operation *)> func) {
   for (Block &block : *region) {
     Operation *terminator = block.getTerminator();
-    // Skip non-return-like terminators.
-    if (terminator->hasTrait<OpTrait::ReturnLike>())
+    // Skip non region-return-like terminators.
+    if (isRegionReturnLike(terminator))
       func(terminator);
   }
 }
 
 /// Checks if all operations in a given region that have at least one attached
 /// region implement the RegionBranchOpInterface. This is not required in edge
 /// cases, where we have a single attached region and the parent operation has
 /// no results.
for (Region &region : regions) {
       size_t operandIndex =
           llvm::find(regionSuccessor->getSuccessorInputs(), argValue)
               .getIndex();
 
       // 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.
       walkReturnOperations(&region, [&](Operation *terminator) {
+        // Get the actual mutable operands for this terminator op.
+        auto terminatorOperands = *getMutableRegionBranchSuccessorOperands(
+            terminator, region.getRegionNumber());
         // Extract the source value from the current terminator.
-        Value sourceValue = terminator->getOperand(operandIndex);
+        Value sourceValue = ((OperandRange)terminatorOperands)[operandIndex];
         // Create a new clone at the current location of the terminator.
         Value clone = introduceCloneBuffers(sourceValue, terminator);
         // Wire clone and terminator operand.
-        terminator->setOperand(operandIndex, clone);
+        terminatorOperands.slice(operandIndex, 1).assign(clone);
       });
     }
   }
 
   /// Creates a new memory allocation for the given source value and clones
   /// its content into the newly allocated buffer. The terminator operation is
   /// used to insert the clone operation at the right place.
   Value introduceCloneBuffers(Value sourceValue, Operation *terminator) {

mlir/lib/Transforms/BufferOptimizations.cpp

 static bool
 leavesAllocationScope(Region *parentRegion,
                       const BufferViewFlowAnalysis::ValueSetT &aliases) {
   for (Value alias : aliases) {
     for (auto *use : alias.getUsers()) {
       // If there is at least one alias that leaves the parent region, we know
       // that this alias escapes the whole region and hence the associated
       // allocation leaves allocation scope.
-      if (use->hasTrait<OpTrait::ReturnLike>() &&
-          use->getParentRegion() == parentRegion)
+      if (isRegionReturnLike(use) && use->getParentRegion() == parentRegion)
         return true;
     }
   }
   return false;
 }
 
 /// Checks, if an automated allocation scope for a given alloc value exists.
 static bool hasAllocationScope(Value alloc,

mlir/test/Dialect/SCF/invalid.mlir

scf.while : () -> () {
     scf.yield
   }
 }
 
 // -----
 
 func @while_cross_region_type_mismatch() {
   %true = constant true
-  // expected-error@+1 {{expects the same number of trailing operands of the 'before' block terminator and 'after' region arguments}}
+  // expected-error@+1 {{'scf.while' op  region control flow edge from Region #0 to Region #1: source has 0 operands, but target successor needs 1}}
   scf.while : () -> () {
     scf.condition(%true)
   } do {
   ^bb0(%arg0: i32):
     scf.yield
   }
 }
 
 // -----
 
 func @while_cross_region_type_mismatch() {
   %true = constant true
-  // expected-error@+2 {{expects the same types for trailing operands of the 'before' block terminator and 'after' region arguments}}
-  // expected-note@+1 {{for argument 0, found 'i1' and 'i32}}
+  // expected-error@+1 {{'scf.while' op  along control flow edge from Region #0 to Region #1: source type #0 'i1' should match input type #0 'i32'}}
   scf.while : () -> () {
     scf.condition(%true) %true : i1
   } do {
   ^bb0(%arg0: i32):
     scf.yield
   }
 }
 
 // -----
 
 func @while_result_type_mismatch() {
   %true = constant true
-  // expected-error@+1 {{expects the same number of trailing operands of the 'before' block terminator and op results}}
+  // expected-error@+1 {{'scf.while' op  region control flow edge from Region #0 to parent results: source has 1 operands, but target successor needs 0}}
   scf.while : () -> () {
     scf.condition(%true) %true : i1
   } do {
   ^bb0(%arg0: i1):
     scf.yield
   }
 }
 

mlir/test/Transforms/sccp-structured.mlir

%inner_res = scf.if %cond -> (i32) {
       %si_inc = addi %si, %cst_1 : i32
       scf.yield %si_inc : i32
     }
     scf.yield %inner_res : i32
   }
   return %result : i32
 }
 
-/// Test that we can properly visit region successors when the terminator is not
-/// return-like.
+/// Test that we can properly visit region successors when the terminator
+/// implements the RegionBranchTerminatorOpInterface.
 
-// CHECK-LABEL: func @overdefined_non_returnlike(
-func @overdefined_non_returnlike(%arg1 : i32) {
-  // CHECK: scf.while (%[[ARG:.*]] = %[[INPUT:.*]])
-  // CHECK-NEXT: %[[COND:.*]] = cmpi slt, %[[ARG]], %{{.*}} : i32
-  // CHECK-NEXT: scf.condition(%[[COND]]) %[[ARG]] : i32
+// CHECK-LABEL: func @loop_region_branch_terminator_op(
+func @loop_region_branch_terminator_op(%arg1 : i32) {
+  // CHECK:      %c2_i32 = constant 2 : i32
+  // CHECK-NEXT: return
 
   %c2_i32 = constant 2 : i32
    %0 = scf.while (%arg2 = %c2_i32) : (i32) -> (i32) {
     %1 = cmpi slt, %arg2, %arg1 : i32
     scf.condition(%1) %arg2 : i32
   } do {
   ^bb0(%arg2: i32):
     scf.yield %arg2 : i32
   }
   return
 }