Add functionality for erasing terminator successor operands and basic block arguments.

PiperOrigin-RevId: 222303233

diff --git a/mlir/include/mlir/IR/BasicBlock.h b/mlir/include/mlir/IR/BasicBlock.h
index 7220be524df27f3adbfe5fcc52cc80b49f9a73ec..72161f5b58f4db1573b2c27d54dc3c6dbf7dddf7 100644 (file)
--- a/mlir/include/mlir/IR/BasicBlock.h
+++ b/mlir/include/mlir/IR/BasicBlock.h
@@ -63,12 +63,15 @@ public:
 
   bool args_empty() const { return arguments.empty(); }
 
-  /// Add one value to the operand list.
+  /// Add one value to the argument list.
   BBArgument *addArgument(Type type);
 
   /// Add one argument to the argument list for each type specified in the list.
   llvm::iterator_range<args_iterator> addArguments(ArrayRef<Type> types);
 
+  /// Erase the argument at 'index' and remove it from the argument list.
+  void eraseArgument(unsigned index);
+
   unsigned getNumArguments() const { return arguments.size(); }
   BBArgument *getArgument(unsigned i) { return arguments[i]; }
   const BBArgument *getArgument(unsigned i) const { return arguments[i]; }
@@ -258,6 +261,11 @@ public:
     return *this;
   }
 
+  /// Get the successor number in the predecessor terminator.
+  unsigned getSuccessorIndex() const {
+    return bbUseIterator->getOperandNumber();
+  }
+
 private:
   using BBUseIterator = SSAValueUseIterator<BasicBlockOperand, Instruction>;
   BBUseIterator bbUseIterator;
@@ -316,6 +324,9 @@ public:
   BlockType *operator*() const {
     return this->object->getSuccessor(this->index);
   }
+
+  /// Get the successor number in the terminator.
+  unsigned getSuccessorIndex() const { return this->index; }
 };
 
 inline auto BasicBlock::succ_begin() const -> const_succ_iterator {

diff --git a/mlir/include/mlir/IR/BuiltinOps.h b/mlir/include/mlir/IR/BuiltinOps.h
index dec77ae7fe9730f9f2035ec53d37579ab98b40cd..4916f0e5f134f39f811498a31cc2f3a4e7cae1e3 100644 (file)
--- a/mlir/include/mlir/IR/BuiltinOps.h
+++ b/mlir/include/mlir/IR/BuiltinOps.h
@@ -119,8 +119,8 @@ public:
   /// Add a list of values to the operand list.
   void addOperands(ArrayRef<SSAValue *> values);
 
-  /// Erase a specific argument from the arg list.
-  // TODO: void eraseArgument(int Index);
+  /// Erase the operand at 'index' from the operand list.
+  void eraseOperand(unsigned index);
 
 private:
   friend class Operation;
@@ -207,6 +207,9 @@ public:
   /// Add a list of values to the operand list.
   void addTrueOperands(ArrayRef<SSAValue *> values);
 
+  /// Erase the operand at 'index' from the true operand list.
+  void eraseTrueOperand(unsigned index);
+
   // Accessors for operands to the 'false' destination.
   SSAValue *getFalseOperand(unsigned idx) {
     assert(idx < getNumFalseOperands());
@@ -246,6 +249,9 @@ public:
   /// Add a list of values to the false operand list.
   void addFalseOperands(ArrayRef<SSAValue *> values);
 
+  /// Erase the operand at 'index' from the false operand list.
+  void eraseFalseOperand(unsigned index);
+
 private:
   /// Get the index of the first true destination operand.
   unsigned getTrueDestOperandIndex() { return 1; }

diff --git a/mlir/include/mlir/IR/Instructions.h b/mlir/include/mlir/IR/Instructions.h
index 1397e6fd5ae08ad48a95a37e55cd212cf79294e7..5e436acf11c5c0f1fc6c303339f7aa1898a1dd43 100644 (file)
--- a/mlir/include/mlir/IR/Instructions.h
+++ b/mlir/include/mlir/IR/Instructions.h
@@ -227,8 +227,14 @@ public:
   /// index.
   void addSuccessorOperands(unsigned index, ArrayRef<CFGValue *> values);
 
-  /// Erase a specific argument from the arg list.
-  // TODO: void eraseSuccessorOperand(unsigned index, unsigned argIndex);
+  /// Erase a specific operand from the operand list of the successor at
+  /// 'index'.
+  void eraseSuccessorOperand(unsigned succIndex, unsigned opIndex) {
+    assert(succIndex < getNumSuccessors());
+    assert(opIndex < getNumSuccessorOperands(succIndex));
+    eraseOperand(getSuccessorOperandIndex(succIndex) + opIndex);
+    --getTrailingObjects<unsigned>()[succIndex];
+  }
 
   /// Get the index of the first operand of the successor at the provided
   /// index.
@@ -328,6 +334,9 @@ private:
   // does not have a virtual destructor.
   ~Instruction();
 
+  /// Erase the operand at 'index'.
+  void eraseOperand(unsigned index);
+
   friend struct llvm::ilist_traits<Instruction>;
   friend class BasicBlock;
 

diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index 06108c60462eec72d1329885bbcbe540bb7c91f5..622ed553d8aa8042a488048a050235664493d6d0 100644 (file)
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -113,6 +113,7 @@ public:
     for (auto *value : values)
       addSuccessorOperand(index, value);
   }
+  void eraseSuccessorOperand(unsigned succIndex, unsigned opIndex);
   llvm::iterator_range<const_operand_iterator>
   getSuccessorOperands(unsigned index) const;
   llvm::iterator_range<operand_iterator> getSuccessorOperands(unsigned index);

diff --git a/mlir/lib/IR/BasicBlock.cpp b/mlir/lib/IR/BasicBlock.cpp
index cf677b34747f9022e950cfa3c7879467461780ad..8ba457af11d7165a017208c37b4bed11b2a18eaa 100644 (file)
--- a/mlir/lib/IR/BasicBlock.cpp
+++ b/mlir/lib/IR/BasicBlock.cpp
@@ -50,6 +50,21 @@ auto BasicBlock::addArguments(ArrayRef<Type> types)
   return {arguments.data() + initialSize, arguments.data() + arguments.size()};
 }
 
+void BasicBlock::eraseArgument(unsigned index) {
+  assert(index < arguments.size());
+
+  // Delete the argument.
+  delete arguments[index];
+  arguments.erase(arguments.begin() + index);
+
+  // 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);
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // Terminator management
 //===----------------------------------------------------------------------===//

diff --git a/mlir/lib/IR/BuiltinOps.cpp b/mlir/lib/IR/BuiltinOps.cpp
index fe1c4024c5e7e20f5efaa8bf542403fe15586fd4..8e512dc3adb011577f4f04fc6932460a427703a7 100644 (file)
--- a/mlir/lib/IR/BuiltinOps.cpp
+++ b/mlir/lib/IR/BuiltinOps.cpp
@@ -211,6 +211,10 @@ void BranchOp::addOperands(ArrayRef<SSAValue *> values) {
   return getOperation()->addSuccessorOperands(0, values);
 }
 
+void BranchOp::eraseOperand(unsigned index) {
+  getOperation()->eraseSuccessorOperand(0, index);
+}
+
 //===----------------------------------------------------------------------===//
 // CondBranchOp
 //===----------------------------------------------------------------------===//
@@ -290,6 +294,10 @@ void CondBranchOp::addTrueOperands(ArrayRef<SSAValue *> values) {
   return getOperation()->addSuccessorOperands(trueIndex, values);
 }
 
+void CondBranchOp::eraseTrueOperand(unsigned index) {
+  getOperation()->eraseSuccessorOperand(trueIndex, index);
+}
+
 unsigned CondBranchOp::getNumFalseOperands() const {
   return getOperation()->getNumSuccessorOperands(falseIndex);
 }
@@ -302,6 +310,10 @@ void CondBranchOp::addFalseOperands(ArrayRef<SSAValue *> values) {
   return getOperation()->addSuccessorOperands(falseIndex, values);
 }
 
+void CondBranchOp::eraseFalseOperand(unsigned index) {
+  getOperation()->eraseSuccessorOperand(falseIndex, index);
+}
+
 //===----------------------------------------------------------------------===//
 // Constant*Op
 //===----------------------------------------------------------------------===//

diff --git a/mlir/lib/IR/Instructions.cpp b/mlir/lib/IR/Instructions.cpp
index d1e3aab3d05068e84ebfbeae79ff6e6fb04ee51a..cd9bd594247b40b8002e83ad1ef749f82787486d 100644 (file)
--- a/mlir/lib/IR/Instructions.cpp
+++ b/mlir/lib/IR/Instructions.cpp
@@ -188,6 +188,16 @@ Instruction::~Instruction() {
       successor.~BasicBlockOperand();
 }
 
+void Instruction::eraseOperand(unsigned index) {
+  assert(index < operands.size());
+
+  // Shift all operands down by 1.
+  for (unsigned i = index, e = operands.size() - 1; i != e; ++i)
+    operands[i].set(operands[i + 1].get());
+  // Drop the last operand.
+  operands.pop_back();
+}
+
 /// Destroy this instruction.
 void Instruction::destroy() {
   this->~Instruction();

diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index bdf38745b7094f105e0833b1a0db2eeb768d98f7..728d2362c99d17b010a24205f1958f49d9dce2c0 100644 (file)
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -168,6 +168,12 @@ void Operation::addSuccessorOperand(unsigned index, SSAValue *value) {
   return llvm::cast<Instruction>(this)->addSuccessorOperand(
       index, llvm::cast<CFGValue>(value));
 }
+void Operation::eraseSuccessorOperand(unsigned succIndex, unsigned opIndex) {
+  assert(isTerminator() && "Only terminators have successors.");
+  assert(llvm::isa<Instruction>(this) && "Only instructions have successors.");
+  return llvm::cast<Instruction>(this)->eraseSuccessorOperand(succIndex,
+                                                              opIndex);
+}
 auto Operation::getSuccessorOperands(unsigned index) const
     -> llvm::iterator_range<const_operand_iterator> {
   assert(isTerminator() && "Only terminators have successors.");