[mlir] Simplify BranchOpInterface by using MutableOperandRange

[rriddle]

Summary

This range allows for performing many different operations on successor operands, including erasing/adding/setting. This removes the need for the explicit canEraseSuccessorOperand and eraseSuccessorOperand methods.

Depends On D79030

[jpienaar]

Nice, thanks

flang/include/flang/Optimizer/Dialect/FIROps.td

string extraSwitchClassDeclaration = [{
     unsigned getNumDest() { return getOperation()->getNumSuccessors(); }
 
     llvm::Optional<mlir::OperandRange> getCompareOperands(unsigned cond);
     llvm::Optional<llvm::ArrayRef<mlir::Value>> getCompareOperands(
         llvm::ArrayRef<mlir::Value> operands, unsigned cond);
 
     llvm::Optional<llvm::ArrayRef<mlir::Value>> getSuccessorOperands(
         llvm::ArrayRef<mlir::Value> operands, unsigned cond);
+    using BranchOpInterfaceTrait::getSuccessorOperands;
 
     // Helper function to deal with Optional operand forms
     void printSuccessorAtIndex(mlir::OpAsmPrinter &p, unsigned i) {
       auto *succ = getSuccessor(i);
       auto ops = getSuccessorOperands(i);
       if (ops.hasValue())
         p.printSuccessorAndUseList(succ, ops.getValue());
       else

flang/lib/Optimizer/Dialect/FIROps.cpp

 static constexpr llvm::StringRef getCompareOffsetAttr() {
   return "compare_operand_offsets";
 }
 
 static constexpr llvm::StringRef getTargetOffsetAttr() {
   return "target_operand_offsets";
 }
 
-template <typename A>
+template <typename A, typename... AdditionalArgs>
 static A getSubOperands(unsigned pos, A allArgs,
-                        mlir::DenseIntElementsAttr ranges) {
+                        mlir::DenseIntElementsAttr ranges,
+                        AdditionalArgs &&... additionalArgs) {
   unsigned start = 0;
   for (unsigned i = 0; i < pos; ++i)
     start += (*(ranges.begin() + i)).getZExtValue();
-  unsigned end = start + (*(ranges.begin() + pos)).getZExtValue();
-  return {std::next(allArgs.begin(), start), std::next(allArgs.begin(), end)};
+  return allArgs.slice(start, (*(ranges.begin() + pos)).getZExtValue(),
+                       std::forward<AdditionalArgs>(additionalArgs)...);
+}
+
+static mlir::MutableOperandRange
+getMutableSuccessorOperands(unsigned pos, mlir::MutableOperandRange operands,
+                            StringRef offsetAttr) {
+  Operation *owner = operands.getOwner();
+  NamedAttribute targetOffsetAttr =
+      *owner->getMutableAttrDict().getNamed(offsetAttr);
+  return getSubOperands(
+      pos, operands, targetOffsetAttr.second.cast<DenseIntElementsAttr>(),
+      mlir::MutableOperandRange::OperandSegment(pos, targetOffsetAttr));
 }
 
 static unsigned denseElementsSize(mlir::DenseIntElementsAttr attr) {
   return attr.getNumElements();
 }
 
 llvm::Optional<mlir::OperandRange> fir::SelectOp::getCompareOperands(unsigned) {
   return {};
 }
 
 llvm::Optional<llvm::ArrayRef<mlir::Value>>
 fir::SelectOp::getCompareOperands(llvm::ArrayRef<mlir::Value>, unsigned) {
   return {};
 }
 
-llvm::Optional<mlir::OperandRange>
-fir::SelectOp::getSuccessorOperands(unsigned oper) {
-  auto a = getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
-  return {getSubOperands(oper, targetArgs(), a)};
+llvm::Optional<mlir::MutableOperandRange>
+fir::SelectOp::getMutableSuccessorOperands(unsigned oper) {
+  return ::getMutableSuccessorOperands(oper, targetArgsMutable(),
+                                       getTargetOffsetAttr());
 }
 
 llvm::Optional<llvm::ArrayRef<mlir::Value>>
 fir::SelectOp::getSuccessorOperands(llvm::ArrayRef<mlir::Value> operands,
                                     unsigned oper) {
   auto a = getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
   auto segments =
       getAttrOfType<mlir::DenseIntElementsAttr>(getOperandSegmentSizeAttr());
   return {getSubOperands(oper, getSubOperands(2, operands, segments), a)};
 }
 
-bool fir::SelectOp::canEraseSuccessorOperand() { return true; }
-
 unsigned fir::SelectOp::targetOffsetSize() {
   return denseElementsSize(
       getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr()));
 }
 
 //===----------------------------------------------------------------------===//
 // SelectCaseOp
 //===----------------------------------------------------------------------===//
 
 llvm::Optional<mlir::OperandRange>
 fir::SelectCaseOp::getCompareOperands(unsigned cond) {
   auto a = getAttrOfType<mlir::DenseIntElementsAttr>(getCompareOffsetAttr());
   return {getSubOperands(cond, compareArgs(), a)};
 }
 
 llvm::Optional<llvm::ArrayRef<mlir::Value>>
 fir::SelectCaseOp::getCompareOperands(llvm::ArrayRef<mlir::Value> operands,
                                       unsigned cond) {
   auto a = getAttrOfType<mlir::DenseIntElementsAttr>(getCompareOffsetAttr());
   auto segments =
       getAttrOfType<mlir::DenseIntElementsAttr>(getOperandSegmentSizeAttr());
   return {getSubOperands(cond, getSubOperands(1, operands, segments), a)};
 }
 
-llvm::Optional<mlir::OperandRange>
-fir::SelectCaseOp::getSuccessorOperands(unsigned oper) {
-  auto a = getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
-  return {getSubOperands(oper, targetArgs(), a)};
+llvm::Optional<mlir::MutableOperandRange>
+fir::SelectCaseOp::getMutableSuccessorOperands(unsigned oper) {
+  return ::getMutableSuccessorOperands(oper, targetArgsMutable(),
+                                       getTargetOffsetAttr());
 }
 
 llvm::Optional<llvm::ArrayRef<mlir::Value>>
 fir::SelectCaseOp::getSuccessorOperands(llvm::ArrayRef<mlir::Value> operands,
                                         unsigned oper) {
   auto a = getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
   auto segments =
       getAttrOfType<mlir::DenseIntElementsAttr>(getOperandSegmentSizeAttr());
   return {getSubOperands(oper, getSubOperands(2, operands, segments), a)};
 }
 
-bool fir::SelectCaseOp::canEraseSuccessorOperand() { return true; }
-
 // parser for fir.select_case Op
 static mlir::ParseResult parseSelectCase(mlir::OpAsmParser &parser,
                                          mlir::OperationState &result) {
   mlir::OpAsmParser::OperandType selector;
   mlir::Type type;
   if (parseSelector(parser, result, selector, type))
     return mlir::failure();
 
fir::SelectRankOp::getCompareOperands(unsigned) {
   return {};
 }
 
 llvm::Optional<llvm::ArrayRef<mlir::Value>>
 fir::SelectRankOp::getCompareOperands(llvm::ArrayRef<mlir::Value>, unsigned) {
   return {};
 }
 
-llvm::Optional<mlir::OperandRange>
-fir::SelectRankOp::getSuccessorOperands(unsigned oper) {
-  auto a = getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
-  return {getSubOperands(oper, targetArgs(), a)};
+llvm::Optional<mlir::MutableOperandRange>
+fir::SelectRankOp::getMutableSuccessorOperands(unsigned oper) {
+  return ::getMutableSuccessorOperands(oper, targetArgsMutable(),
+                                       getTargetOffsetAttr());
 }
 
 llvm::Optional<llvm::ArrayRef<mlir::Value>>
 fir::SelectRankOp::getSuccessorOperands(llvm::ArrayRef<mlir::Value> operands,
                                         unsigned oper) {
   auto a = getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
   auto segments =
       getAttrOfType<mlir::DenseIntElementsAttr>(getOperandSegmentSizeAttr());
   return {getSubOperands(oper, getSubOperands(2, operands, segments), a)};
 }
 
-bool fir::SelectRankOp::canEraseSuccessorOperand() { return true; }
-
 unsigned fir::SelectRankOp::targetOffsetSize() {
   return denseElementsSize(
       getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr()));
 }
 
 //===----------------------------------------------------------------------===//
 // SelectTypeOp
 //===----------------------------------------------------------------------===//
 
 llvm::Optional<mlir::OperandRange>
 fir::SelectTypeOp::getCompareOperands(unsigned) {
   return {};
 }
 
 llvm::Optional<llvm::ArrayRef<mlir::Value>>
 fir::SelectTypeOp::getCompareOperands(llvm::ArrayRef<mlir::Value>, unsigned) {
   return {};
 }
 
-llvm::Optional<mlir::OperandRange>
-fir::SelectTypeOp::getSuccessorOperands(unsigned oper) {
-  auto a = getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
-  return {getSubOperands(oper, targetArgs(), a)};
+llvm::Optional<mlir::MutableOperandRange>
+fir::SelectTypeOp::getMutableSuccessorOperands(unsigned oper) {
+  return ::getMutableSuccessorOperands(oper, targetArgsMutable(),
+                                       getTargetOffsetAttr());
 }
 
 llvm::Optional<llvm::ArrayRef<mlir::Value>>
 fir::SelectTypeOp::getSuccessorOperands(llvm::ArrayRef<mlir::Value> operands,
                                         unsigned oper) {
   auto a = getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
   auto segments =
       getAttrOfType<mlir::DenseIntElementsAttr>(getOperandSegmentSizeAttr());
   return {getSubOperands(oper, getSubOperands(2, operands, segments), a)};
 }
 
-bool fir::SelectTypeOp::canEraseSuccessorOperand() { return true; }
-
 static ParseResult parseSelectType(OpAsmParser &parser,
                                    OperationState &result) {
   mlir::OpAsmParser::OperandType selector;
   mlir::Type type;
   if (parseSelector(parser, result, selector, type))
     return mlir::failure();
 
   llvm::SmallVector<mlir::Attribute, 8> attrs;

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let extraClassDeclaration = [{
     }
 
     operand_range getTrueOperands() { return trueDestOperands(); }
 
     unsigned getNumTrueOperands()  { return getTrueOperands().size(); }
 
     /// Erase the operand at 'index' from the true operand list.
     void eraseTrueOperand(unsigned index)  {
-      eraseSuccessorOperand(trueIndex, index);
+      trueDestOperandsMutable().erase(index);
     }
 
     // Accessors for operands to the 'false' destination.
     Value getFalseOperand(unsigned idx) {
       assert(idx < getNumFalseOperands());
       return getOperand(getFalseDestOperandIndex() + idx);
     }
     void setFalseOperand(unsigned idx, Value value) {
       assert(idx < getNumFalseOperands());
       setOperand(getFalseDestOperandIndex() + idx, value);
     }
 
     operand_range getFalseOperands() { return falseDestOperands(); }
 
     unsigned getNumFalseOperands() { return getFalseOperands().size(); }
 
     /// Erase the operand at 'index' from the false operand list.
     void eraseFalseOperand(unsigned index) {
-      eraseSuccessorOperand(falseIndex, index);
+      falseDestOperandsMutable().erase(index);
     }
 
   private:
     /// Get the index of the first true destination operand.
     unsigned getTrueDestOperandIndex() { return 1; }
 
     /// Get the index of the first false destination operand.
     unsigned getFalseDestOperandIndex() {

mlir/include/mlir/IR/OperationSupport.h

public:
 
   /// Construct a new mutable range from the given operand, operand start index,
   /// and range length. `operandSegments` is an optional set of operand segments
   /// to be updated when mutating the operand list.
   MutableOperandRange(Operation *owner, unsigned start, unsigned length,
                       ArrayRef<OperandSegment> operandSegments = llvm::None);
   MutableOperandRange(Operation *owner);
 
+  /// Slice this range into a sub range, with the additional operand segment.
+  MutableOperandRange slice(unsigned subStart, unsigned subLen,
+                            Optional<OperandSegment> segment = llvm::None);
+
   /// Append the given values to the range.
   void append(ValueRange values);
 
   /// Assign this range to the given values.
   void assign(ValueRange values);
 
   /// Assign the range to the given value.
   void assign(Value value);
 
   /// Erase the operands within the given sub-range.
   void erase(unsigned subStart, unsigned subLen = 1);
 
   /// Clear this range and erase all of the operands.
   void clear();
 
   /// Returns the current size of the range.
   unsigned size() const { return length; }
 
   /// Allow implicit conversion to an OperandRange.
   operator OperandRange() const;
 
+  /// Returns the owning operation.
+  Operation *getOwner() const { return owner; }
+
 private:
   /// Update the length of this range to the one provided.
   void updateLength(unsigned newLength);
 
   /// The owning operation of this range.
   Operation *owner;
 
   /// The start index of the operand range within the owner operand list, and

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

 namespace mlir {
 class BranchOpInterface;
 
 //===----------------------------------------------------------------------===//
 // BranchOpInterface
 //===----------------------------------------------------------------------===//
 
 namespace detail {
-/// Erase an operand from a branch operation that is used as a successor
-/// operand. `operandIndex` is the operand within `operands` to be erased.
-void eraseBranchSuccessorOperand(OperandRange operands, unsigned operandIndex,
-                                 Operation *op);
-
 /// Return the `BlockArgument` corresponding to operand `operandIndex` in some
 /// successor if `operandIndex` is within the range of `operands`, or None if
 /// `operandIndex` isn't a successor operand index.
 Optional<BlockArgument>
 getBranchSuccessorArgument(Optional<OperandRange> operands,
                            unsigned operandIndex, Block *successor);
 
 /// Verify that the given operands match those of the given successor block.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

 
 def BranchOpInterface : OpInterface<"BranchOpInterface"> {
   let description = [{
     This interface provides information for branching terminator operations,
     i.e. terminator operations with successors.
   }];
   let methods = [
     InterfaceMethod<[{
-        Returns a set of values that correspond to the arguments to the
+        Returns a mutable range of operands that correspond to the arguments of
         successor at the given index. Returns None if the operands to the
         successor are non-materialized values, i.e. they are internal to the
         operation.
       }],
-      "Optional<OperandRange>", "getSuccessorOperands", (ins "unsigned":$index)
+      "Optional<MutableOperandRange>", "getMutableSuccessorOperands",
+      (ins "unsigned":$index)
     >,
     InterfaceMethod<[{
-        Return true if this operation can erase an operand to a successor block.
-      }],
-      "bool", "canEraseSuccessorOperand"
-    >,
-    InterfaceMethod<[{
-        Erase the operand at `operandIndex` from the `index`-th successor. This
-        should only be called if `canEraseSuccessorOperand` returns true.
+        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.
       }],
-      "void", "eraseSuccessorOperand",
-      (ins "unsigned":$index, "unsigned":$operandIndex), [{}],
-      /*defaultImplementation=*/[{
+      "Optional<OperandRange>", "getSuccessorOperands",
+      (ins "unsigned":$index), [{}], [{
         ConcreteOp *op = static_cast<ConcreteOp *>(this);
-        Optional<OperandRange> operands = op->getSuccessorOperands(index);
-        assert(operands && "unable to query operands for successor");
-        detail::eraseBranchSuccessorOperand(*operands, operandIndex, *op);
+        auto operands = op->getMutableSuccessorOperands(index);
+        return operands ? Optional<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.
       }],
       "Optional<BlockArgument>", "getSuccessorBlockArgument",

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

static ParseResult parseAllocaOp(OpAsmParser &parser, OperationState &result) {
   result.addTypes({funcType.getResult(0)});
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // LLVM::BrOp
 //===----------------------------------------------------------------------===//
 
-Optional<OperandRange> BrOp::getSuccessorOperands(unsigned index) {
+Optional<MutableOperandRange>
+BrOp::getMutableSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
-  return getOperands();
+  return destOperandsMutable();
 }
 
-bool BrOp::canEraseSuccessorOperand() { return true; }
-
 //===----------------------------------------------------------------------===//
 // LLVM::CondBrOp
 //===----------------------------------------------------------------------===//
 
-Optional<OperandRange> CondBrOp::getSuccessorOperands(unsigned index) {
+Optional<MutableOperandRange>
+CondBrOp::getMutableSuccessorOperands(unsigned index) {
   assert(index < getNumSuccessors() && "invalid successor index");
-  return index == 0 ? trueDestOperands() : falseDestOperands();
+  return index == 0 ? trueDestOperandsMutable() : falseDestOperandsMutable();
 }
 
-bool CondBrOp::canEraseSuccessorOperand() { return true; }
-
 //===----------------------------------------------------------------------===//
 // Printing/parsing for LLVM::LoadOp.
 //===----------------------------------------------------------------------===//
 
 static void printLoadOp(OpAsmPrinter &p, LoadOp &op) {
   p << op.getOperationName() << ' ' << op.addr();
   p.printOptionalAttrDict(op.getAttrs());
   p << " : " << op.addr().getType();
static ParseResult parseStoreOp(OpAsmParser &parser, OperationState &result) {
 
   return success();
 }
 
 ///===---------------------------------------------------------------------===//
 /// LLVM::InvokeOp
 ///===---------------------------------------------------------------------===//
 
-Optional<OperandRange> InvokeOp::getSuccessorOperands(unsigned index) {
+Optional<MutableOperandRange>
+InvokeOp::getMutableSuccessorOperands(unsigned index) {
   assert(index < getNumSuccessors() && "invalid successor index");
-  return index == 0 ? normalDestOperands() : unwindDestOperands();
+  return index == 0 ? normalDestOperandsMutable() : unwindDestOperandsMutable();
 }
 
-bool InvokeOp::canEraseSuccessorOperand() { return true; }
-
 static LogicalResult verify(InvokeOp op) {
   if (op.getNumResults() > 1)
     return op.emitOpError("must have 0 or 1 result");
 
   Block *unwindDest = op.unwindDest();
   if (unwindDest->empty())
     return op.emitError(
         "must have at least one operation in unwind destination");

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

static LogicalResult verify(spirv::BitcastOp bitcastOp) {
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // spv.BranchOp
 //===----------------------------------------------------------------------===//
 
-Optional<OperandRange> spirv::BranchOp::getSuccessorOperands(unsigned index) {
+Optional<MutableOperandRange>
+spirv::BranchOp::getMutableSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
-  return getOperands();
+  return targetOperandsMutable();
 }
 
-bool spirv::BranchOp::canEraseSuccessorOperand() { return true; }
-
 //===----------------------------------------------------------------------===//
 // spv.BranchConditionalOp
 //===----------------------------------------------------------------------===//
 
-Optional<OperandRange>
-spirv::BranchConditionalOp::getSuccessorOperands(unsigned index) {
+Optional<MutableOperandRange>
+spirv::BranchConditionalOp::getMutableSuccessorOperands(unsigned index) {
   assert(index < 2 && "invalid successor index");
-  return index == kTrueIndex ? getTrueBlockArguments()
-                             : getFalseBlockArguments();
+  return index == kTrueIndex ? trueTargetOperandsMutable()
+                             : falseTargetOperandsMutable();
 }
 
-bool spirv::BranchConditionalOp::canEraseSuccessorOperand() { return true; }
-
 static ParseResult parseBranchConditionalOp(OpAsmParser &parser,
                                             OperationState &state) {
   auto &builder = parser.getBuilder();
   OpAsmParser::OperandType condInfo;
   Block *dest;
 
   // Parse the condition.
   Type boolTy = builder.getI1Type();

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 }
 
 void BranchOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                            MLIRContext *context) {
   results.insert<SimplifyBrToBlockWithSinglePred, SimplifyPassThroughBr>(
       context);
 }
 
-Optional<OperandRange> BranchOp::getSuccessorOperands(unsigned index) {
+Optional<MutableOperandRange>
+BranchOp::getMutableSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
-  return getOperands();
+  return destOperandsMutable();
 }
 
-bool BranchOp::canEraseSuccessorOperand() { return true; }
-
 Block *BranchOp::getSuccessorForOperands(ArrayRef<Attribute>) { return dest(); }
 
 //===----------------------------------------------------------------------===//
 // CallOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(CallOp op) {
   // Check that the callee attribute was specified.

 } // end anonymous namespace
 
 void CondBranchOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<SimplifyConstCondBranchPred, SimplifyPassThroughCondBranch,
                  SimplifyCondBranchIdenticalSuccessors>(context);
 }
 
-Optional<OperandRange> CondBranchOp::getSuccessorOperands(unsigned index) {
+Optional<MutableOperandRange>
+CondBranchOp::getMutableSuccessorOperands(unsigned index) {
   assert(index < getNumSuccessors() && "invalid successor index");
-  return index == trueIndex ? getTrueOperands() : getFalseOperands();
+  return index == trueIndex ? trueDestOperandsMutable()
+                            : falseDestOperandsMutable();
 }
 
-bool CondBranchOp::canEraseSuccessorOperand() { return true; }
-
 Block *CondBranchOp::getSuccessorForOperands(ArrayRef<Attribute> operands) {
   if (BoolAttr condAttr = operands.front().dyn_cast_or_null<BoolAttr>())
     return condAttr.getValue() ? trueDest() : falseDest();
   if (IntegerAttr condAttr = operands.front().dyn_cast_or_null<IntegerAttr>())
     return condAttr.getValue().isOneValue() ? trueDest() : falseDest();
   return nullptr;
 }
 

mlir/lib/IR/OperationSupport.cpp

MutableOperandRange::MutableOperandRange(
     ArrayRef<OperandSegment> operandSegments)
     : owner(owner), start(start), length(length),
       operandSegments(operandSegments.begin(), operandSegments.end()) {
   assert((start + length) <= owner->getNumOperands() && "invalid range");
 }
 MutableOperandRange::MutableOperandRange(Operation *owner)
     : MutableOperandRange(owner, /*start=*/0, owner->getNumOperands()) {}
 
+/// Slice this range into a sub range, with the additional operand segment.
+MutableOperandRange
+MutableOperandRange::slice(unsigned subStart, unsigned subLen,
+                           Optional<OperandSegment> segment) {
+  assert((subStart + subLen) <= length && "invalid sub-range");
+  MutableOperandRange subSlice(owner, start + subStart, subLen,
+                               operandSegments);
+  if (segment)
+    subSlice.operandSegments.push_back(*segment);
+  return subSlice;
+}
+
 /// Append the given values to the range.
 void MutableOperandRange::append(ValueRange values) {
   if (values.empty())
     return;
   owner->insertOperands(start + length, values);
   updateLength(length + values.size());
 }
 

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/ControlFlowInterfaces.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // BranchOpInterface
 //===----------------------------------------------------------------------===//
 
-/// Erase an operand from a branch operation that is used as a successor
-/// operand. 'operandIndex' is the operand within 'operands' to be erased.
-void mlir::detail::eraseBranchSuccessorOperand(OperandRange operands,
-                                               unsigned operandIndex,
-                                               Operation *op) {
-  assert(operandIndex < operands.size() &&
-         "invalid index for successor operands");
-
-  // Erase the operand from the operation.
-  size_t fullOperandIndex = operands.getBeginOperandIndex() + operandIndex;
-  op->eraseOperand(fullOperandIndex);
-
-  // If this operation has an OperandSegmentSizeAttr, keep it up to date.
-  auto operandSegmentAttr =
-      op->getAttrOfType<DenseElementsAttr>("operand_segment_sizes");
-  if (!operandSegmentAttr)
-    return;
-
-  // Find the segment containing the full operand index and decrement it.
-  // TODO: This seems like a general utility that could be added somewhere.
-  SmallVector<int32_t, 4> values(operandSegmentAttr.getValues<int32_t>());
-  unsigned currentSize = 0;
-  for (unsigned i = 0, e = values.size(); i != e; ++i) {
-    currentSize += values[i];
-    if (fullOperandIndex < currentSize) {
-      --values[i];
-      break;
-    }
-  }
-  op->setAttr("operand_segment_sizes",
-              DenseIntElementsAttr::get(operandSegmentAttr.getType(), values));
-}
-
 /// Returns the `BlockArgument` corresponding to operand `operandIndex` in some
 /// successor if 'operandIndex' is within the range of 'operands', or None if
 /// `operandIndex` isn't a successor operand index.
 Optional<BlockArgument> mlir::detail::getBranchSuccessorArgument(
     Optional<OperandRange> operands, unsigned operandIndex, Block *successor) {
   // Check that the operands are valid.
   if (!operands || operands->empty())
     return llvm::None;

mlir/lib/Transforms/Utils/RegionUtils.cpp

 static void propagateLiveness(Region &region, LiveMap &liveMap);
 
 static void propagateTerminatorLiveness(Operation *op, LiveMap &liveMap) {
   // Terminators are always live.
   liveMap.setProvedLive(op);
 
   // Check to see if we can reason about the successor operands and mutate them.
   BranchOpInterface branchInterface = dyn_cast<BranchOpInterface>(op);
-  if (!branchInterface || !branchInterface.canEraseSuccessorOperand()) {
+  if (!branchInterface) {
     for (Block *successor : op->getSuccessors())
       for (BlockArgument arg : successor->getArguments())
         liveMap.setProvedLive(arg);
     return;
   }
 
   // If we can't reason about the operands to a successor, conservatively mark
   // all arguments as live.
   for (unsigned i = 0, e = op->getNumSuccessors(); i != e; ++i) {
-    if (!branchInterface.getSuccessorOperands(i))
+    if (!branchInterface.getMutableSuccessorOperands(i))
       for (BlockArgument arg : op->getSuccessor(i)->getArguments())
         liveMap.setProvedLive(arg);
   }
 }
 
 static void propagateLiveness(Operation *op, LiveMap &liveMap) {
   // All Value's are either a block argument or an op result.
   // We call processValue on those cases.
static void eraseTerminatorSuccessorOperands(Operation *terminator,
 
   for (unsigned succI = 0, succE = terminator->getNumSuccessors();
        succI < succE; succI++) {
     // Iterating successors in reverse is not strictly needed, since we
     // aren't erasing any successors. But it is slightly more efficient
     // since it will promote later operands of the terminator being erased
     // first, reducing the quadratic-ness.
     unsigned succ = succE - succI - 1;
-    Optional<OperandRange> succOperands = branchOp.getSuccessorOperands(succ);
+    Optional<MutableOperandRange> succOperands =
+        branchOp.getMutableSuccessorOperands(succ);
     if (!succOperands)
       continue;
     Block *successor = terminator->getSuccessor(succ);
 
     for (unsigned argI = 0, argE = succOperands->size(); argI < argE; ++argI) {
       // Iterating args in reverse is needed for correctness, to avoid
       // shifting later args when earlier args are erased.
       unsigned arg = argE - argI - 1;
       if (!liveMap.wasProvenLive(successor->getArgument(arg)))
-        branchOp.eraseSuccessorOperand(succ, arg);
+        succOperands->erase(arg);
     }
   }
 }
 
 static LogicalResult deleteDeadness(MutableArrayRef<Region> regions,
                                     LiveMap &liveMap) {
   bool erasedAnything = false;
   for (Region &region : regions) {

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

if (namedAttr.first == "test.invalid_attr")
     return op->emitError() << "invalid to use 'test.invalid_attr'";
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // TestBranchOp
 //===----------------------------------------------------------------------===//
 
-Optional<OperandRange> TestBranchOp::getSuccessorOperands(unsigned index) {
+Optional<MutableOperandRange>
+TestBranchOp::getMutableSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
-  return getOperands();
+  return targetOperandsMutable();
 }
 
-bool TestBranchOp::canEraseSuccessorOperand() { return true; }
-
 //===----------------------------------------------------------------------===//
 // Test IsolatedRegionOp - parse passthrough region arguments.
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseIsolatedRegionOp(OpAsmParser &parser,
                                          OperationState &result) {
   OpAsmParser::OperandType argInfo;
   Type argType = parser.getBuilder().getIndexType();

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

static void emitModelDecl(OpInterface &interface, raw_ostream &os) {
   }
   os << "  };\n";
 }
 
 static void emitTraitDecl(OpInterface &interface, raw_ostream &os,
                           StringRef interfaceName,
                           StringRef interfaceTraitsName) {
   os << "  template <typename ConcreteOp>\n  "
-     << llvm::formatv("struct Trait : public OpInterface<{0},"
+     << llvm::formatv("struct {0}Trait : public OpInterface<{0},"
                       " detail::{1}>::Trait<ConcreteOp> {{\n",
                       interfaceName, interfaceTraitsName);
 
   // Insert the default implementation for any methods.
   for (auto &method : interface.getMethods()) {
     // Flag interface methods named verifyTrait.
     if (method.getName() == "verifyTrait")
       PrintFatalError(
           formatv("'verifyTrait' method cannot be specified as interface "
                   "method for '{0}'; set 'verify' on OpInterfaceTrait instead",
                   interfaceName));
     auto defaultImpl = method.getDefaultImplementation();
     if (!defaultImpl)
       continue;
 
     os << "  " << (method.isStatic() ? "static " : "") << method.getReturnType()
        << " ";
     emitMethodNameAndArgs(method, os, /*addOperationArg=*/false);
     os << " {\n" << defaultImpl.getValue() << "  }\n";
   }
 
   tblgen::FmtContext traitCtx;
   traitCtx.withOp("op");
   if (auto verify = interface.getVerify()) {
-    os << "  static LogicalResult verifyTrait(Operation* op) {\n"
+    os << "    static LogicalResult verifyTrait(Operation* op) {\n"
        << std::string(tblgen::tgfmt(*verify, &traitCtx)) << "\n  }\n";
   }
   if (auto extraTraitDecls = interface.getExtraTraitClassDeclaration())
     os << extraTraitDecls << "\n";
 
   os << "  };\n";
+
+  // Emit a utility using directive for the trait class.
+  os << "    template <typename ConcreteOp>\n    "
+     << llvm::formatv("using Trait = {0}Trait<ConcreteOp>;\n", interfaceName);
 }
 
 static void emitInterfaceDecl(OpInterface &interface, raw_ostream &os) {
   StringRef interfaceName = interface.getName();
   auto interfaceTraitsName = (interfaceName + "InterfaceTraits").str();
 
   // Emit the traits struct containing the concept and model declarations.
   os << "namespace detail {\n"