[mlir][NFC] Make `ReturnLike` trait imply `RegionBranchTerminatorOpInterface`

[zero9178]

Summary

This implication was already done de-facto and there were plenty of users and wrapper functions specifically used to handle the "return-like or RegionBranchTerminatorOpInterface" case. These simply existed due to up until recently missing features in ODS.

With the new capabilities of traits, we can make ReturnLike imply RegionBranchTerminatorOpInterface and auto generate proper definitions for its methods.
Various occurrences and wrapper methods used for isa<RegionBranchTerminatorOpInterface>() || hasTrait<ReturnLike>() have all been removed.

[jpienaar]

I like all the red :-)

[mehdi_amini]

Nice!

[olegshyshkov]

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

281–286

This commit is causing errors in our codebase and I'm not sure how to best approach it. We have ReturnLike ops with custom implementations of getMutableSuccessorOperands. https://github.com/tensorflow/tensorflow/blob/675d73c284396a100555300fe2eab5d02f4bd847/tensorflow/core/ir/ops.cc#L1299.

Is it intentional that now there is not way to define custom getMutableSuccessorOperands and add ReturnLike trait at the same time? Or am I missing something and there is a workaround?

[zero9178]

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

281–286

Hi!

Is it intentional that now there is not way to define custom getMutableSuccessorOperands and add ReturnLike trait at the same > time?

Admittedly no, it was somewhat of a misconception of mine that ReturnLike would only be used for ops that simpyl forward all its operands like the ReturnLike ops upstream.

Or am I missing something and there is a workaround?

I would first try to reevaluate whether you need ReturnLike on the op. At least the upstream ops that have the trait essentially use it as a default implementation for RegionBranchTerminatorOpInterface. Upstream ops like scf.condition, which seems similar to the op you linked, are not marked ReturnLike e.g..

If you do end up really requiring it, a workaround would be to instead use NativeOpTrait<"ReturnLike"> instead of just ReturnLike in the trait list of the ops.
Do tell if there is something fundamentally wrong with this patch that needs upstream rather than downstream changes.

[maerhart]

Is it intentional that operations like the upstream func.return now implement RegionBranchTerminatorOpInterface but func.func does not (and cannot, at least in its current form) implement the RegionBranchOpInterface? As far as I understand, the parent of the op implementing the RegionBranchTerminatorOpInterface has to implement RegionBranchOpInterface which is also assumed in the default implementation of the getSuccessorRegions method.
Should the RegionBranchOpInterface be modified such that it can be implemented by func.func (or is it already possible and I'm misunderstanding something?) or should ReturnLike be dropped from func.return and NativeOpTrait<"ReturnLike"> or something else (e.g., some new trait specific to FunctionOpInterface) be used instead to determine somewhat generically in a pass that a terminator has function-return-like properties?

[zero9178]

In D157402#4595233, @maerhart wrote:

Is it intentional that operations like the upstream func.return now implement RegionBranchTerminatorOpInterface but func.func does not (and cannot, at least in its current form) implement the RegionBranchOpInterface? As far as I understand, the parent of the op implementing the RegionBranchTerminatorOpInterface has to implement RegionBranchOpInterface which is also assumed in the default implementation of the getSuccessorRegions method.
Should the RegionBranchOpInterface be modified such that it can be implemented by func.func (or is it already possible and I'm misunderstanding something?) or should ReturnLike be dropped from func.return and NativeOpTrait<"ReturnLike"> or something else (e.g., some new trait specific to FunctionOpInterface) be used instead to determine somewhat generically in a pass that a terminator has function-return-like properties?

With the concerns you and @olegshyshkov raised, it sounds to me like there should be some kind of split of what is currently "ReturnLike", one being RegionBranchReturnLike and the other possibly CallableReturnLike (to mirror CallableOpInterface). The latter can even live in the TableGen file defining CallableOpInterface while the former remains in ControlFlowInterfaces.
Would that make sense to you?

The documentation of the ReturnLike trait is not very descriptive IMO and does not properly describe what it entails. Just doing a grep through upstream shows it essentially being used as return for CallableOpInterfaces (the uses of it as region branch terminators having been removed in this patch).
I suppose this might be worth an RFC.

[zero9178]

In D157402#4595233, @maerhart wrote:

Is it intentional that operations like the upstream func.return now implement RegionBranchTerminatorOpInterface but func.func does not (and cannot, at least in its current form) implement the RegionBranchOpInterface? As far as I understand, the parent of the op implementing the RegionBranchTerminatorOpInterface has to implement RegionBranchOpInterface which is also assumed in the default implementation of the getSuccessorRegions method.
Should the RegionBranchOpInterface be modified such that it can be implemented by func.func (or is it already possible and I'm misunderstanding something?) or should ReturnLike be dropped from func.return and NativeOpTrait<"ReturnLike"> or something else (e.g., some new trait specific to FunctionOpInterface) be used instead to determine somewhat generically in a pass that a terminator has function-return-like properties?

To directly answer your questions. It was intentional, but I am not a big fan of it at all. The assumption was that one would never call the methods of RegionBranchTerminatorOpInterface if it wasn't nested inside a RegionBranchOpInterface. I believe that RegionBranchOpInterface is only supposed to be used for ops that can only ever directly execute their regions, as opposed to something like func.func or a dialect that has a "lambda" like nested-function op, where the region execution is deffered to a call op.
These ops should be implementing the CallableOpInterface instead.

mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h

private:
 
   /// Visit a block.
   void visitBlock(Block *block);
 
   /// Visit an op with regions (like e.g. `scf.while`)
   void visitRegionSuccessors(RegionBranchOpInterface branch,
                              ArrayRef<AbstractSparseLattice *> operands);
 
-  /// Visit a terminator (an op implementing `RegionBranchTerminatorOpInterface`
-  /// or a return-like op) to compute the lattice values of its operands, given
-  /// its parent op `branch`. The lattice value of an operand is determined
-  /// based on the corresponding arguments in `terminator`'s region
-  /// successor(s).
-  void visitRegionSuccessorsFromTerminator(Operation *terminator,
-                                           RegionBranchOpInterface branch);
+  /// Visit a `RegionBranchTerminatorOpInterface` to compute the lattice values
+  /// of its operands, given its parent op `branch`. The lattice value of an
+  /// operand is determined based on the corresponding arguments in
+  /// `terminator`'s region successor(s).
+  void visitRegionSuccessorsFromTerminator(
+      RegionBranchTerminatorOpInterface terminator,
 
   /// Get the lattice element for a value, and also set up
   /// dependencies so that the analysis on the given ProgramPoint is re-invoked
   /// if the value changes.
   const AbstractSparseLattice *getLatticeElementFor(ProgramPoint point,
                                                     Value value);
 
   /// Get the lattice elements for a range of values, and also set up

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

 Region *getEnclosingRepetitiveRegion(Operation *op);
 
 /// Return the first enclosing region of the given Value that may be executed
 /// repetitively as per RegionBranchOpInterface or `nullptr` if no such region
 /// exists.
 Region *getEnclosingRepetitiveRegion(Value value);
 
 //===----------------------------------------------------------------------===//
-// 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 `std::nullopt`. In all other cases, the resulting
-/// `OperandRange` represents all operands that are passed to the specified
-/// successor region. If `regionIndex` is `std::nullopt`, all operands that are
-/// passed to the parent operation will be returned.
-std::optional<MutableOperandRange>
-getMutableRegionBranchSuccessorOperands(Operation *operation,
-                                        std::optional<unsigned> regionIndex);
-
-/// Returns the read only operands that are passed to the region with the given
-/// `regionIndex`. See `getMutableRegionBranchSuccessorOperands` for more
-/// information.
-std::optional<OperandRange>
-getRegionBranchSuccessorOperands(Operation *operation,
-                                 std::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

def RegionBranchTerminatorOpInterface :
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // ControlFlow Traits
 //===----------------------------------------------------------------------===//
 
 // Op is "return-like".
-def ReturnLike : NativeOpTrait<"ReturnLike">;
+def ReturnLike : TraitList<[
+    DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface>,
+    NativeOpTrait<
+        /*name=*/"ReturnLike",
+        /*traits=*/[],
+        /*extraOpDeclaration=*/"",
+        /*extraOpDefinition=*/[{
+          ::mlir::MutableOperandRange $cppClass::getMutableSuccessorOperands(
+            ::std::optional<unsigned> index) {
+            return ::mlir::MutableOperandRange(*this);
+          }
+        }]
+    >
+]>;
 
 #endif // MLIR_INTERFACES_CONTROLFLOWINTERFACES

mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp

collectUnderlyingAddressValues(
         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 (std::optional<unsigned> operandIndex = getOperandIndexIfPred(i)) {
       for (Block &block : op->getRegion(i)) {
-        Operation *term = block.getTerminator();
         // 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()) {
+        if (auto term = dyn_cast<RegionBranchTerminatorOpInterface>(
+                block.getTerminator())) {
+          collectUnderlyingAddressValues(
+              term.getSuccessorOperands(regionIndex)[*operandIndex], maxDepth,
+              visited, output);
+        } else if (block.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;
         }
       }
     }
   }

mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp

void AbstractDenseBackwardDataFlowAnalysis::visitBlock(Block *block) {
   auto isExitBlock = [](Block *b) {
     // Treat empty and terminator-less blocks as exit blocks.
     if (b->empty() || !b->back().mightHaveTrait<OpTrait::IsTerminator>())
       return true;
 
     // There may be a weird case where a terminator may be transferring control
     // either to the parent or to another block, so exit blocks and successors
     // are not mutually exclusive.
-    Operation *terminator = b->getTerminator();
-    return terminator && (terminator->hasTrait<OpTrait::ReturnLike>() ||
-                          isa<RegionBranchTerminatorOpInterface>(terminator));
+    return isa_and_nonnull<RegionBranchTerminatorOpInterface>(
+        b->getTerminator());
   };
   if (isExitBlock(block)) {
     // If this block is exiting from a callable, the successors of exiting from
     // a callable are the successors of all call sites. And the call sites
     // themselves are predecessors of the callable.
     auto callable = dyn_cast<CallableOpInterface>(block->getParentOp());
     if (callable && callable.getCallableRegion() == block->getParent()) {
       const auto *callsites = getOrCreateFor<PredecessorState>(block, callable);

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

 }
 
 void LivenessAnalysis::visitBranchOperand(OpOperand &operand) {
   // We know (at the moment) and assume (for the future) that `operand` is a
   // non-forwarded branch operand of a `RegionBranchOpInterface`,
   // `BranchOpInterface`, `RegionBranchTerminatorOpInterface` or return-like op.
   Operation *op = operand.getOwner();
   assert((isa<RegionBranchOpInterface>(op) || isa<BranchOpInterface>(op) ||
-          isa<RegionBranchTerminatorOpInterface>(op) ||
-          op->hasTrait<OpTrait::ReturnLike>()) &&
+          isa<RegionBranchTerminatorOpInterface>(op)) &&
          "expected the op to be `RegionBranchOpInterface`, "
-         "`BranchOpInterface`, `RegionBranchTerminatorOpInterface`, or "
-         "return-like");
+         "`BranchOpInterface` or `RegionBranchTerminatorOpInterface`");
 
   // The lattices of the non-forwarded branch operands don't get updated like
   // the forwarded branch operands or the non-branch operands. Thus they need
   // to be handled separately. This is where we handle them.
 
   // This marks values of type (1.b) liveness as "live". A non-forwarded
   // branch operand will be live if a block where its op could take the control
   // has an op with memory effects.
void LivenessAnalysis::visitBranchOperand(OpOperand &operand) {
   SmallVector<Liveness *, 4> operandLiveness;
   operandLiveness.push_back(getLatticeElement(operand.get()));
   SmallVector<const Liveness *, 4> resultsLiveness;
   for (const Value result : op->getResults())
     resultsLiveness.push_back(getLatticeElement(result));
   visitOperation(op, operandLiveness, resultsLiveness);
 
   // We also visit the parent op with the parent's results and this operand if
-  // `op` is a `RegionBranchTerminatorOpInterface` or return-like because its
-  // non-forwarded operand depends on not only its memory effects/results but
-  // also on those of its parent's.
-  if (!isa<RegionBranchTerminatorOpInterface>(op) &&
-      !op->hasTrait<OpTrait::ReturnLike>())
+  // `op` is a `RegionBranchTerminatorOpInterface` because its non-forwarded
+  // operand depends on not only its memory effects/results but also on those of
+  // its parent's.
+  if (!isa<RegionBranchTerminatorOpInterface>(op))
     return;
   Operation *parentOp = op->getParentOp();
   SmallVector<const Liveness *, 4> parentResultsLiveness;
   for (const Value parentResult : parentOp->getResults())
     parentResultsLiveness.push_back(getLatticeElement(parentResult));
   visitOperation(parentOp, operandLiveness, parentResultsLiveness);
 }
 

mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp

void AbstractSparseForwardDataFlowAnalysis::visitRegionSuccessors(
   for (Operation *op : predecessors->getKnownPredecessors()) {
     // Get the incoming successor operands.
     std::optional<OperandRange> operands;
 
     // Check if the predecessor is the parent op.
     if (op == branch) {
       operands = branch.getSuccessorEntryOperands(successorIndex);
       // Otherwise, try to deduce the operands from a region return-like op.
-    } else {
-      if (isRegionReturnLike(op))
-        operands = getRegionBranchSuccessorOperands(op, successorIndex);
+    } else if (auto regionTerminator =
+                   dyn_cast<RegionBranchTerminatorOpInterface>(op)) {
+      operands = regionTerminator.getSuccessorOperands(successorIndex);
     }
 
     if (!operands) {
       // We can't reason about the data-flow.
       return setAllToEntryStates(lattices);
     }
 
     ValueRange inputs = predecessors->getSuccessorInputs(op);
void AbstractSparseBackwardDataFlowAnalysis::visitOperation(Operation *op) {
   // return-like terminator, the region's successors' arguments flow back into
   // the "successor operands" of this terminator.
   //
   // A successor operand with respect to an op implementing
   // `RegionBranchOpInterface` is an operand that is forwarded to a region
   // successor's input. There are two types of successor operands: the operands
   // of this op itself and the operands of the terminators of the regions of
   // this op.
-  if (isa<RegionBranchTerminatorOpInterface>(op) ||
-      op->hasTrait<OpTrait::ReturnLike>()) {
+  if (auto terminator = dyn_cast<RegionBranchTerminatorOpInterface>(op)) {
     if (auto branch = dyn_cast<RegionBranchOpInterface>(op->getParentOp())) {
-      visitRegionSuccessorsFromTerminator(op, branch);
+      visitRegionSuccessorsFromTerminator(terminator, branch);
       return;
     }
   }
 
   if (op->hasTrait<OpTrait::ReturnLike>()) {
     // Going backwards, the operands of the return are derived from the
     // results of all CallOps calling this CallableOp.
     if (auto callable = dyn_cast<CallableOpInterface>(op->getParentOp())) {
void AbstractSparseBackwardDataFlowAnalysis::visitRegionSuccessors(
   // All operands not forwarded to regions are typically parameters of the
   // branch operation itself (for example the boolean for if/else).
   for (int index : unaccounted.set_bits()) {
     visitBranchOperand(op->getOpOperand(index));
   }
 }
 
 void AbstractSparseBackwardDataFlowAnalysis::
-    visitRegionSuccessorsFromTerminator(Operation *terminator,
-                                        RegionBranchOpInterface branch) {
-  assert(isa<RegionBranchTerminatorOpInterface>(terminator) ||
-         terminator->hasTrait<OpTrait::ReturnLike>() &&
-             "expected a `RegionBranchTerminatorOpInterface` op or a "
-             "return-like op");
+    visitRegionSuccessorsFromTerminator(
+        RegionBranchTerminatorOpInterface terminator,
+  assert(isa<RegionBranchTerminatorOpInterface>(terminator) &&
+         "expected a `RegionBranchTerminatorOpInterface` op");
   assert(terminator->getParentOp() == branch.getOperation() &&
          "expected `branch` to be the parent op of `terminator`");
 
   SmallVector<Attribute> operandAttributes(terminator->getNumOperands(),
                                            nullptr);
   SmallVector<RegionSuccessor> successors;
   branch.getSuccessorRegions(terminator->getParentRegion()->getRegionNumber(),
                              operandAttributes, successors);
   // All operands not forwarded to any successor. This set can be
   // non-contiguous in the presence of multiple successors.
   BitVector unaccounted(terminator->getNumOperands(), true);
 
   for (const RegionSuccessor &successor : successors) {
     ValueRange inputs = successor.getSuccessorInputs();
     Region *region = successor.getSuccessor();
-    OperandRange operands =
-        region ? *getRegionBranchSuccessorOperands(terminator,
-                                                   region->getRegionNumber())
-               : *getRegionBranchSuccessorOperands(terminator, {});
+    OperandRange operands = terminator.getSuccessorOperands(
+        region ? region->getRegionNumber() : std::optional<unsigned>{});
     MutableArrayRef<OpOperand> opOperands = operandsToOpOperands(operands);
     for (auto [opOperand, input] : llvm::zip(opOperands, inputs)) {
       meet(getLatticeElement(opOperand.get()),
            *getLatticeElementFor(terminator, input));
       unaccounted.reset(const_cast<OpOperand &>(opOperand).getOperandNumber());
     }
   }
   // Visit operands of the branch op not forwarded to the next region.

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

if (!options.dynCastBufferizableOp(op))
       continue;
 
     // We cannot analyze through ToMemrefOps, so we have to conservatively
     // assume that the value is yielded.
     if (isa<ToMemrefOp>(op))
       return true;
 
     // Check if the op is returning/yielding.
-    if (isRegionReturnLike(op))
+    if (isa<RegionBranchTerminatorOpInterface>(op))
       return true;
 
     // Add all aliasing OpResults to the worklist.
     // Note: In the absence of detailed analysis information (e.g., there may be
     // no function call analysis information), this `getAliasingOpResult` is
     // conservative and may report additional OpResults as potentially aliasing.
     for (AliasingOpResult alias : getAliasingOpResults(*operand))
       for (OpOperand &use : alias.opResult.getUses())

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

 #include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
 } // namespace bufferization
 } // namespace mlir
 
 using namespace mlir;
 using namespace mlir::bufferization;
 
 /// Walks over all immediate return-like terminators in the given region.
-static LogicalResult
-walkReturnOperations(Region *region,
-                     llvm::function_ref<LogicalResult(Operation *)> func) {
+static LogicalResult walkReturnOperations(
+    Region *region,
+    llvm::function_ref<LogicalResult(RegionBranchTerminatorOpInterface)> func) {
   for (Block &block : *region) {
     Operation *terminator = block.getTerminator();
     // Skip non region-return-like terminators.
-    if (isRegionReturnLike(terminator)) {
-      if (failed(func(terminator)))
+    if (auto regionTerminator =
+            dyn_cast<RegionBranchTerminatorOpInterface>(terminator)) {
+      if (failed(func(regionTerminator)))
         return failure();
     }
   }
   return success();
 }
 
 /// Checks if all operations that have at least one attached region implement
 /// the RegionBranchOpInterface. This is not required in edge cases, where we
for (Region &region : regions) {
       // bindings.
       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.
-      if (failed(walkReturnOperations(&region, [&](Operation *terminator) {
+      if (failed(walkReturnOperations(
+              &region, [&](RegionBranchTerminatorOpInterface terminator) {
                 // Get the actual mutable operands for this terminator op.
-            auto terminatorOperands = *getMutableRegionBranchSuccessorOperands(
-                terminator, region.getRegionNumber());
+                auto terminatorOperands =
+                    terminator.getMutableSuccessorOperands(
+                        region.getRegionNumber());
                 // 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.
+                // 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))
                   return failure();
                 // Wire clone and terminator operand.
                 terminatorOperands.slice(operandIndex, 1).assign(*clone);
                 return success();
               })))
         return failure();
     }
     return success();
   }
 
   /// 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.

mlir/lib/Dialect/Bufferization/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 (isRegionReturnLike(use) && use->getParentRegion() == parentRegion)
+      if (isa<RegionBranchTerminatorOpInterface>(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/lib/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.cpp

if (auto regionInterface = dyn_cast<RegionBranchOpInterface>(op)) {
         for (RegionSuccessor &successorRegion : successorRegions) {
           // Determine the current region index (if any).
           std::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 successor operands of each terminator.
-          for (Block &block : region) {
-            auto successorOperands = getRegionBranchSuccessorOperands(
-                block.getTerminator(), regionIndex);
-            if (successorOperands) {
-              registerDependencies(*successorOperands,
+          for (Block &block : region)
+            if (auto terminator = dyn_cast<RegionBranchTerminatorOpInterface>(
+                    block.getTerminator()))
+              registerDependencies(terminator.getSuccessorOperands(regionIndex),
                                    successorRegion.getSuccessorInputs());
-            }
-          }
-        }
-      }
 
       return WalkResult::advance();
     }
 
     // Unknown op: Assume that all operands alias with all results.
     for (Value operand : op->getOperands()) {
       if (!isa<BaseMemRefType>(operand.getType()))
         continue;

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

 }
 
 // Gather yielded tensors in `yieldedTensors` by querying all aliases. This is
 // to ensure that such information is available during bufferization time.
 // Alias information can no longer be queried once we have started modifying
 // the IR.
 void OneShotAnalysisState::gatherYieldedTensors(Operation *op) {
   op->walk([&](Operation *returnOp) {
-    if (!isRegionReturnLike(returnOp) || !getOptions().isOpAllowed(returnOp))
+    if (!isa<RegionBranchTerminatorOpInterface>(returnOp) ||
+        !getOptions().isOpAllowed(returnOp))
       return WalkResult::advance();
 
     for (OpOperand &returnValOperand : returnOp->getOpOperands()) {
       Value returnVal = returnValOperand.get();
       // Skip non-tensor values.
       if (!isa<TensorType>(returnVal.getType()))
         continue;
 

 /// and bufferizes to a new allocation.
 // TODO: Remove buffer deallocation from One-Shot Bufferize and fix the buffer
 // deallocation pass.
 static LogicalResult assertNoAllocsReturned(Operation *op,
                                             const OneShotAnalysisState &state) {
   LogicalResult status = success();
   DominanceInfo domInfo(op);
   op->walk([&](Operation *returnOp) {
-    if (!isRegionReturnLike(returnOp) ||
+    if (!isa<RegionBranchTerminatorOpInterface>(returnOp) ||
         !state.getOptions().isOpAllowed(returnOp))
       return WalkResult::advance();
 
     for (OpOperand &returnValOperand : returnOp->getOpOperands()) {
       Value returnVal = returnValOperand.get();
       // Skip non-tensor values.
       if (!isa<TensorType>(returnVal.getType()))
         continue;

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

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.
 
     std::optional<OperandRange> regionReturnOperands;
     for (Block &block : region) {
-      Operation *terminator = block.getTerminator();
-      auto terminatorOperands =
-          getRegionBranchSuccessorOperands(terminator, regionNo);
-      if (!terminatorOperands)
+      auto terminator =
+          dyn_cast<RegionBranchTerminatorOpInterface>(block.getTerminator());
+      if (!terminator)
         continue;
 
+      OperandRange terminatorOperands =
+          terminator.getSuccessorOperands(regionNo);
       if (!regionReturnOperands) {
         regionReturnOperands = terminatorOperands;
         continue;
       }
 
       // Found more than one ReturnLike terminator. Make sure the operand types
       // match with the first one.
       if (!areTypesCompatible(regionReturnOperands->getTypes(),
-                              terminatorOperands->getTypes()))
+                              terminatorOperands.getTypes()))
         return op->emitOpError("Region #")
                << regionNo
                << " 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

 void RegionBranchOpInterface::getSuccessorRegions(
     std::optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
   unsigned numInputs = 0;
   if (index) {
     // If the predecessor is a region, get the number of operands from an
     // exiting terminator in the region.
     for (Block &block : getOperation()->getRegion(*index)) {
       Operation *terminator = block.getTerminator();
-      if (getRegionBranchSuccessorOperands(terminator, *index)) {
+      if (isa<RegionBranchTerminatorOpInterface>(terminator)) {
         numInputs = terminator->getNumOperands();
         break;
       }
     }
   } else {
     // Otherwise, use the number of parent operation operands.
     numInputs = getOperation()->getNumOperands();
   }
while (region) {
     Operation *op = region->getParentOp();
     if (auto branchOp = dyn_cast<RegionBranchOpInterface>(op))
       if (branchOp.isRepetitiveRegion(region->getRegionNumber()))
         return region;
     region = op->getParentRegion();
   }
   return nullptr;
 }
-
-//===----------------------------------------------------------------------===//
-// 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 `std::nullopt`. In all other cases, the resulting
-/// `OperandRange` represents all operands that are passed to the specified
-/// successor region. If `regionIndex` is `std::nullopt`, all operands that are
-/// passed to the parent operation will be returned.
-std::optional<MutableOperandRange>
-mlir::getMutableRegionBranchSuccessorOperands(
-    Operation *operation, std::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 std::nullopt;
-}
-
-/// Returns the read only operands that are passed to the region with the given
-/// `regionIndex`. See `getMutableRegionBranchSuccessorOperands` for more
-/// information.
-std::optional<OperandRange>
-mlir::getRegionBranchSuccessorOperands(Operation *operation,
-                                       std::optional<unsigned> regionIndex) {
-  auto range = getMutableRegionBranchSuccessorOperands(operation, regionIndex);
-  if (range)
-    return range->operator OperandRange();
-  return std::nullopt;
-}