[mlir][NFC] Make ReturnLike trait imply RegionBranchTerminatorOpInterface

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.

Differential Revision: https://reviews.llvm.org/D157402

Author: zero9178

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

@@ -397,13 +397,13 @@ class AbstractSparseBackwardDataFlowAnalysis : public DataFlowAnalysis {
   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,
+      RegionBranchOpInterface branch);
 
   /// Get the lattice element for a value, and also set up
   /// dependencies so that the analysis on the given ProgramPoint is re-invoked

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

@@ -229,32 +229,6 @@ Region *getEnclosingRepetitiveRegion(Operation *op);
 /// 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
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

@@ -272,6 +272,19 @@ def RegionBranchTerminatorOpInterface :
 //===----------------------------------------------------------------------===//
 
 // 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

@@ -91,15 +91,14 @@ static void collectUnderlyingAddressValues(RegionBranchOpInterface branch,
   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);

mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp

@@ -337,9 +337,8 @@ void AbstractDenseBackwardDataFlowAnalysis::visitBlock(Block *block) {
     // 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

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

@@ -93,11 +93,9 @@ void LivenessAnalysis::visitBranchOperand(OpOperand &operand) {
   // `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
@@ -161,11 +159,10 @@ void LivenessAnalysis::visitBranchOperand(OpOperand &operand) {
   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;

mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp

@@ -226,9 +226,9 @@ void AbstractSparseForwardDataFlowAnalysis::visitRegionSuccessors(
     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) {
@@ -439,10 +439,9 @@ void AbstractSparseBackwardDataFlowAnalysis::visitOperation(Operation *op) {
   // 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;
     }
   }
@@ -506,12 +505,11 @@ void AbstractSparseBackwardDataFlowAnalysis::visitRegionSuccessors(
 }
 
 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,
+        RegionBranchOpInterface branch) {
+  assert(isa<RegionBranchTerminatorOpInterface>(terminator) &&
+         "expected a `RegionBranchTerminatorOpInterface` op");
   assert(terminator->getParentOp() == branch.getOperation() &&
          "expected `branch` to be the parent op of `terminator`");
 
@@ -527,10 +525,8 @@ void AbstractSparseBackwardDataFlowAnalysis::
   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()),

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

@@ -690,7 +690,7 @@ bool AnalysisState::isTensorYielded(Value tensor) const {
       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.

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

@@ -70,14 +70,15 @@ 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();
     }
   }
@@ -447,23 +448,25 @@ class BufferDeallocation : public BufferPlacementTransformationBase {
       // 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) {
-            // Get the actual mutable operands for this terminator op.
-            auto terminatorOperands = *getMutableRegionBranchSuccessorOperands(
-                terminator, 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.
-            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();
-          })))
+      if (failed(walkReturnOperations(
+              &region, [&](RegionBranchTerminatorOpInterface terminator) {
+                // Get the actual mutable operands for this terminator op.
+                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.
+                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();

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

@@ -75,7 +75,8 @@ leavesAllocationScope(Region *parentRegion,
       // 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;
     }
   }

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

@@ -128,14 +128,11 @@ void BufferViewFlowAnalysis::build(Operation *op) {
             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());
-            }
-          }
         }
       }
 

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

@@ -183,7 +183,8 @@ void OneShotAnalysisState::createAliasInfoEntry(Value v) {
 // 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()) {
@@ -1059,7 +1060,7 @@ static LogicalResult assertNoAllocsReturned(Operation *op,
   LogicalResult status = success();
   DominanceInfo domInfo(op);
   op->walk([&](Operation *returnOp) {
-    if (!isRegionReturnLike(returnOp) ||
+    if (!isa<RegionBranchTerminatorOpInterface>(returnOp) ||
         !state.getOptions().isOpAllowed(returnOp))
       return WalkResult::advance();