[mlir] Refactor a few users to no longer rely on the successor operand API of Operation.

[rriddle]

Summary

The existing API for successor operands on operations is in the process of being removed. This revision simplifies a later one that completely removes the existing API.

Depends On D75315

[rriddle]

Rebase

[rriddle]

Rebase

[rriddle]

Rebase

mlir/include/mlir/IR/Block.h

 
 #ifndef MLIR_IR_BLOCK_H
 #define MLIR_IR_BLOCK_H
 
 #include "mlir/IR/BlockSupport.h"
 #include "mlir/IR/Visitors.h"
 
 namespace mlir {
+class TypeRange;
+template <typename ValueRangeT> class ValueTypeRange;
+
 /// `Block` represents an ordered list of `Operation`s.
 class Block : public IRObjectWithUseList<BlockOperand>,
               public llvm::ilist_node_with_parent<Block, Region> {
 public:
   explicit Block() {}
   ~Block();
 
   void clear() {
public:
   // Block argument management
   //===--------------------------------------------------------------------===//
 
   // This is the list of arguments to the block.
   using BlockArgListType = MutableArrayRef<BlockArgument>;
 
   BlockArgListType getArguments() { return arguments; }
 
+  /// Return a range containing the types of the arguments for this block.
+  ValueTypeRange<BlockArgListType> getArgumentTypes();
+
   using args_iterator = BlockArgListType::iterator;
   using reverse_args_iterator = BlockArgListType::reverse_iterator;
   args_iterator args_begin() { return getArguments().begin(); }
   args_iterator args_end() { return getArguments().end(); }
   reverse_args_iterator args_rbegin() { return getArguments().rbegin(); }
   reverse_args_iterator args_rend() { return getArguments().rend(); }
 
   bool args_empty() { return arguments.empty(); }
 
   /// Add one value to the argument list.
   BlockArgument addArgument(Type type);
 
   /// Insert one value to the position in the argument list indicated by the
   /// given iterator. The existing arguments are shifted. The block is expected
   /// not to have predecessors.
   BlockArgument insertArgument(args_iterator it, Type type);
 
   /// Add one argument to the argument list for each type specified in the list.
-  iterator_range<args_iterator> addArguments(ArrayRef<Type> types);
+  iterator_range<args_iterator> addArguments(TypeRange types);
 
-  // Add one value to the argument list at the specified position.
+  /// Add one value to the argument list at the specified position.
   BlockArgument insertArgument(unsigned index, Type type);
 
-  /// Erase the argument at 'index' and remove it from the argument list. If
-  /// 'updatePredTerms' is set to true, this argument is also removed from the
-  /// terminators of each predecessor to this block.
-  void eraseArgument(unsigned index, bool updatePredTerms = true);
+  /// Erase the argument at 'index' and remove it from the argument list.
+  void eraseArgument(unsigned index);
 
   unsigned getNumArguments() { return arguments.size(); }
   BlockArgument getArgument(unsigned i) { return arguments[i]; }
 
   //===--------------------------------------------------------------------===//
   // Operation list management
   //===--------------------------------------------------------------------===//
 

mlir/include/mlir/IR/OperationSupport.h

: public detail::indexed_accessor_range_base<
           llvm::PointerUnion<const Value *, const Type *, OpOperand *>, Type,
           Type, Type> {
 public:
   using RangeBaseT::RangeBaseT;
   TypeRange(ArrayRef<Type> types = llvm::None);
   explicit TypeRange(OperandRange values);
   explicit TypeRange(ResultRange values);
   explicit TypeRange(ValueRange values);
+  explicit TypeRange(ArrayRef<Value> values);
+  explicit TypeRange(ArrayRef<BlockArgument> values)
+      : TypeRange(ArrayRef<Value>(values.data(), values.size())) {}
   template <typename ValueRangeT>
   TypeRange(ValueTypeRange<ValueRangeT> values)
       : TypeRange(ValueRangeT(values.begin().getCurrent(),
                               values.end().getCurrent())) {}
+  template <typename Arg,
+            typename = typename std::enable_if_t<
+                std::is_constructible<ArrayRef<Type>, Arg>::value>>
+  TypeRange(Arg &&arg) : TypeRange(ArrayRef<Type>(std::forward<Arg>(arg))) {}
+  TypeRange(std::initializer_list<Type> types)
+      : TypeRange(ArrayRef<Type>(types)) {}
 
 private:
   /// The owner of the range is either:
   /// * A pointer to the first element of an array of values.
   /// * A pointer to the first element of an array of types.
   /// * A pointer to the first element of an array of operands.
   using OwnerT = llvm::PointerUnion<const Value *, const Type *, OpOperand *>;
 

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
 
     return matchSuccess();
   }
 };
 
 } // namespace
 
 static void ensureDistinctSuccessors(Block &bb) {
-  auto *terminator = bb.getTerminator();
+  Operation *terminator = bb.getTerminator();
+  if (terminator->getNumSuccessors() < 2)
+    return;
 
   // Find repeated successors with arguments.
   llvm::SmallDenseMap<Block *, SmallVector<int, 4>> successorPositions;
   for (int i = 0, e = terminator->getNumSuccessors(); i < e; ++i) {
     Block *successor = terminator->getSuccessor(i);
     // Blocks with no arguments are safe even if they appear multiple times
     // because they don't need PHI nodes.
     if (successor->getNumArguments() == 0)
       continue;
     successorPositions[successor].push_back(i);
   }
 
   // If a successor appears for the second or more time in the terminator,
   // create a new dummy block that unconditionally branches to the original
   // destination, and retarget the terminator to branch to this new block.
   // There is no need to pass arguments to the dummy block because it will be
   // dominated by the original block and can therefore use any values defined in
   // the original block.
+  OpBuilder builder(terminator->getContext());
   for (const auto &successor : successorPositions) {
-    const auto &positions = successor.second;
     // Start from the second occurrence of a block in the successor list.
-    for (auto position = std::next(positions.begin()), end = positions.end();
-         position != end; ++position) {
-      auto *dummyBlock = new Block();
-      bb.getParent()->push_back(dummyBlock);
-      auto builder = OpBuilder(dummyBlock);
-      SmallVector<Value, 8> operands(
-          terminator->getSuccessorOperands(*position));
-      builder.create<BranchOp>(terminator->getLoc(), successor.first, operands);
-      terminator->setSuccessor(dummyBlock, *position);
-      for (int i = 0, e = terminator->getNumSuccessorOperands(*position); i < e;
-           ++i)
-        terminator->eraseSuccessorOperand(*position, i);
+    for (int position : llvm::drop_begin(successor.second, 1)) {
+      Block *dummyBlock = builder.createBlock(bb.getParent());
+      terminator->setSuccessor(dummyBlock, position);
+      dummyBlock->addArguments(successor.first->getArgumentTypes());
+      builder.create<BranchOp>(terminator->getLoc(), successor.first,
+                               dummyBlock->getArguments());
     }
   }
 }
 
 void mlir::LLVM::ensureDistinctSuccessors(ModuleOp m) {
   for (auto f : m.getOps<FuncOp>()) {
     for (auto &bb : f.getBlocks()) {
       ::ensureDistinctSuccessors(bb);

mlir/lib/Dialect/AffineOps/AffineOps.cpp

static LogicalResult verify(AffineIfOp op) {
   // Verify that there are enough operands for the condition.
   IntegerSet condition = conditionAttr.getValue();
   if (op.getNumOperands() != condition.getNumInputs())
     return op.emitOpError(
         "operand count and condition integer set dimension and "
         "symbol count must match");
 
   // Verify that the operands are valid dimension/symbols.
-  if (failed(verifyDimAndSymbolIdentifiers(
-          op, op.getOperation()->getNonSuccessorOperands(),
+  if (failed(verifyDimAndSymbolIdentifiers(op, op.getOperands(),
                                            condition.getNumDims())))
     return failure();
 
   // Verify that the entry of each child region does not have arguments.
   for (auto &region : op.getOperation()->getRegions()) {
     for (auto &b : region)
       if (b.getNumArguments() != 0)
         return op.emitOpError(
             "requires that child entry blocks have no arguments");

mlir/lib/IR/Block.cpp

void Block::recomputeOpOrder() {
   for (auto &op : *this)
     op.orderIndex = (orderIndex += Operation::kOrderStride);
 }
 
 //===----------------------------------------------------------------------===//
 // Argument list management.
 //===----------------------------------------------------------------------===//
 
+/// Return a range containing the types of the arguments for this block.
+auto Block::getArgumentTypes() -> ValueTypeRange<BlockArgListType> {
+  return ValueTypeRange<BlockArgListType>(getArguments());
+}
+
 BlockArgument Block::addArgument(Type type) {
   BlockArgument arg = BlockArgument::create(type, this);
   arguments.push_back(arg);
   return arg;
 }
 
 /// Add one argument to the argument list for each type specified in the list.
-auto Block::addArguments(ArrayRef<Type> types)
-    -> iterator_range<args_iterator> {
-  arguments.reserve(arguments.size() + types.size());
-  auto initialSize = arguments.size();
-  for (auto type : types) {
+auto Block::addArguments(TypeRange types) -> iterator_range<args_iterator> {
+  size_t initialSize = arguments.size();
+  arguments.reserve(initialSize + types.size());
+  for (auto type : types)
     addArgument(type);
-  }
   return {arguments.data() + initialSize, arguments.data() + arguments.size()};
 }
 
 BlockArgument Block::insertArgument(unsigned index, Type type) {
   auto arg = BlockArgument::create(type, this);
   assert(index <= arguments.size());
   arguments.insert(arguments.begin() + index, arg);
   return arg;
 }
 
-void Block::eraseArgument(unsigned index, bool updatePredTerms) {
+void Block::eraseArgument(unsigned index) {
   assert(index < arguments.size());
-
-  // If requested, update predecessors. We do this first since this block might
-  // be a predecessor of itself and use this block argument as a successor
-  // operand.
-  if (updatePredTerms) {
-    // Erase this argument from each of the predecessor's terminator.
-    for (auto predIt = pred_begin(), predE = pred_end(); predIt != predE;
-         ++predIt) {
-      auto *predTerminator = (*predIt)->getTerminator();
-      predTerminator->eraseSuccessorOperand(predIt.getSuccessorIndex(), index);
-    }
-  }
-
-  // Delete the argument.
   arguments[index].destroy();
   arguments.erase(arguments.begin() + index);
 }
 
 /// Insert one value to the given position of the argument list. The existing
 /// arguments are shifted. The block is expected not to have predecessors.
 BlockArgument Block::insertArgument(args_iterator it, Type type) {
   assert(llvm::empty(getPredecessors()) &&

mlir/lib/IR/OperationSupport.cpp

 
 TypeRange::TypeRange(ArrayRef<Type> types)
     : TypeRange(types.data(), types.size()) {}
 TypeRange::TypeRange(OperandRange values)
     : TypeRange(values.begin().getBase(), values.size()) {}
 TypeRange::TypeRange(ResultRange values)
     : TypeRange(values.getBase()->getResultTypes().slice(values.getStartIndex(),
                                                          values.size())) {}
+TypeRange::TypeRange(ArrayRef<Value> values)
+    : TypeRange(values.data(), values.size()) {}
 TypeRange::TypeRange(ValueRange values) : TypeRange(OwnerT(), values.size()) {
   detail::ValueRangeOwner owner = values.begin().getBase();
   if (auto *op = reinterpret_cast<Operation *>(owner.ptr.dyn_cast<void *>()))
     this->base = &op->getResultTypes()[owner.startIndex];
   else if (auto *operand = owner.ptr.dyn_cast<OpOperand *>())
     this->base = operand;
   else
     this->base = owner.ptr.get<const Value *>();

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

static Value getPHISourceValue(Block *current, Block *pred,
   assert(condBranchOp &&
          "only branch operations can be terminators of a block that "
          "has successors");
   assert((condBranchOp.getSuccessor(0) != condBranchOp.getSuccessor(1)) &&
          "successors with arguments in LLVM conditional branches must be "
          "different blocks");
 
   return condBranchOp.getSuccessor(0) == current
-             ? terminator.getSuccessorOperand(0, index)
-             : terminator.getSuccessorOperand(1, index);
+             ? condBranchOp.trueDestOperands()[index]
+             : condBranchOp.falseDestOperands()[index];
 }
 
 void ModuleTranslation::connectPHINodes(LLVMFuncOp func) {
   // Skip the first block, it cannot be branched to and its arguments correspond
   // to the arguments of the LLVM function.
   for (auto it = std::next(func.begin()), eit = func.end(); it != eit; ++it) {
     Block *bb = &*it;
     llvm::BasicBlock *llvmBB = blockMapping.lookup(bb);

mlir/lib/Transforms/Utils/RegionUtils.cpp

for (Region &region : regions) {
     // Delete block arguments.
     // The entry block has an unknown contract with their enclosing block, so
     // skip it.
     for (Block &block : llvm::drop_begin(region.getBlocks(), 1)) {
       // Iterate in reverse to avoid shifting later arguments when deleting
       // earlier arguments.
       for (unsigned i = 0, e = block.getNumArguments(); i < e; i++)
         if (!liveMap.wasProvenLive(block.getArgument(e - i - 1))) {
-          block.eraseArgument(e - i - 1, /*updatePredTerms=*/false);
+          block.eraseArgument(e - i - 1);
           erasedAnything = true;
         }
     }
   }
   return success(erasedAnything);
 }
 
 // This function performs a simple dead code elimination algorithm over the

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir

 
 // CHECK-NEXT: ^bb2:
 ^bb2:
 // CHECK-NEXT:  llvm.br ^bb1
   br ^bb1
 }
 // CHECK-LABEL: func @cond_br_same_target(%arg0: !llvm.i1, %arg1: !llvm.i32, %arg2: !llvm.i32)
 func @cond_br_same_target(%arg0: i1, %arg1: i32, %arg2 : i32) -> (i32) {
-// CHECK-NEXT:  llvm.cond_br %arg0, ^[[origBlock:bb[0-9]+]](%arg1 : !llvm.i32), ^[[dummyBlock:bb[0-9]+]]
+// CHECK-NEXT:  llvm.cond_br %arg0, ^[[origBlock:bb[0-9]+]](%arg1 : !llvm.i32), ^[[dummyBlock:bb[0-9]+]](%arg2 : !llvm.i32)
   cond_br %arg0, ^bb1(%arg1 : i32), ^bb1(%arg2 : i32)
 
-// CHECK:      ^[[origBlock]](%0: !llvm.i32):
-// CHECK-NEXT:  llvm.return %0 : !llvm.i32
+// CHECK:      ^[[origBlock]](%[[BLOCKARG1:.*]]: !llvm.i32):
+// CHECK-NEXT:  llvm.return %[[BLOCKARG1]] : !llvm.i32
 ^bb1(%0 : i32):
   return %0 : i32
 
-// CHECK:      ^[[dummyBlock]]:
-// CHECK-NEXT:  llvm.br ^[[origBlock]](%arg2 : !llvm.i32)
+// CHECK:      ^[[dummyBlock]](%[[BLOCKARG2:.*]]: !llvm.i32):
+// CHECK-NEXT:  llvm.br ^[[origBlock]](%[[BLOCKARG2]] : !llvm.i32)
 }
 
 // CHECK-LABEL: func @fcmp(%arg0: !llvm.float, %arg1: !llvm.float) {
 func @fcmp(f32, f32) -> () {
 ^bb0(%arg0: f32, %arg1: f32):
   // CHECK:      llvm.fcmp "oeq" %arg0, %arg1 : !llvm.float
   // CHECK-NEXT: llvm.fcmp "ogt" %arg0, %arg1 : !llvm.float
   // CHECK-NEXT: llvm.fcmp "oge" %arg0, %arg1 : !llvm.float