[mlir][IR] Add an Operation::eraseOperands that supports batch erasure

[rriddle]

Summary

This method allows for removing multiple disjoint operands at once, reducing the need to erase operands individually (which results in shifting the operand list).

Depends On D98071

[jpienaar]

Nice, thanks

mlir/lib/IR/OperationSupport.cpp

307

This is to avoid needing to disambiguate between only prefix specified or not? (meaning, alternative could have been considering indexes from one end and this is up to the Nth operand, instead require all operands to be specified).

Did you consider a list of indices? esp for wher you have large number of operands but small number of operands to remove

[rriddle]

mlir/lib/IR/OperationSupport.cpp

307

This is to avoid needing to disambiguate between only prefix specified or not? (meaning, alternative could have been considering indexes from one end and this is up to the Nth operand, instead require all operands to be specified).

Essentially yes. I don't have a valid use case that is worth the differentiation behavior, so I don't see a need yet in supporting it.

Did you consider a list of indices? esp for wher you have large number of operands but small number of operands to remove

The list approach only really works if you know for a fact that a) you have a very large number of operands, and b) you have a small number of operands to remove. If either of those are false, you don't gain anything by the list approach and often end up using much more memory. For all of the use cases I've encountered so far that needed this, one of the above has been false. We can add another API in the future if there is a need for one.

mlir/include/mlir/IR/Operation.h

public:
   void eraseOperand(unsigned idx) { eraseOperands(idx); }
 
   /// Erase the operands starting at position `idx` and ending at position
   /// 'idx'+'length'.
   void eraseOperands(unsigned idx, unsigned length = 1) {
     getOperandStorage().eraseOperands(idx, length);
   }
 
+  /// Erases the operands that have their corresponding bit set in
+  /// `eraseIndices` and removes them from the operand list.
+  void eraseOperands(const llvm::BitVector &eraseIndices) {
+    getOperandStorage().eraseOperands(eraseIndices);
+  }
+
   // Support operand iteration.
   using operand_range = OperandRange;
   using operand_iterator = operand_range::iterator;
 
   operand_iterator operand_begin() { return getOperands().begin(); }
   operand_iterator operand_end() { return getOperands().end(); }
 
   /// Returns an iterator on the underlying Value's.

mlir/include/mlir/IR/OperationSupport.h

 #include "mlir/IR/Value.h"
 #include "mlir/Support/InterfaceSupport.h"
 #include "llvm/ADT/BitmaskEnum.h"
 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 #include "llvm/Support/TrailingObjects.h"
 #include <memory>
 
+namespace llvm {
+class BitVector;
+} // end namespace llvm
+
 namespace mlir {
 class Dialect;
 class DictionaryAttr;
 class ElementsAttr;
 class Operation;
 struct OperationState;
 class OpAsmParser;
 class OpAsmParserResult;
public:
   /// with the ones provided in 'operands'. 'operands' may be smaller or larger
   /// than the range pointed to by 'start'+'length'.
   void setOperands(Operation *owner, unsigned start, unsigned length,
                    ValueRange operands);
 
   /// Erase the operands held by the storage within the given range.
   void eraseOperands(unsigned start, unsigned length);
 
+  /// Erase the operands held by the storage that have their corresponding bit
+  /// set in `eraseIndices`.
+  void eraseOperands(const llvm::BitVector &eraseIndices);
+
   /// Get the operation operands held by the storage.
   MutableArrayRef<OpOperand> getOperands() {
     return getStorage().getOperands();
   }
 
   /// Return the number of operands held in the storage.
   unsigned size() { return getStorage().numOperands; }
 

mlir/lib/IR/OperationSupport.cpp

 //===- OperationSupport.cpp -----------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains out-of-line implementations of the support types that
 // Operation and related classes build on top of.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/OperationSupport.h"
-#include "mlir/IR/Block.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
-#include "mlir/IR/Operation.h"
+#include "llvm/ADT/BitVector.h"
+
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // NamedAttrList
 //===----------------------------------------------------------------------===//
 
 NamedAttrList::NamedAttrList(ArrayRef<NamedAttribute> attributes) {
   assign(attributes.begin(), attributes.end());
void detail::OperandStorage::eraseOperands(unsigned start, unsigned length) {
   if (start != storage.numOperands) {
     auto *indexIt = std::next(operands.begin(), start);
     std::rotate(indexIt, std::next(indexIt, length), operands.end());
   }
   for (unsigned i = 0; i != length; ++i)
     operands[storage.numOperands + i].~OpOperand();
 }
 
+void detail::OperandStorage::eraseOperands(
+    const llvm::BitVector &eraseIndices) {
+  TrailingOperandStorage &storage = getStorage();
+  MutableArrayRef<OpOperand> operands = storage.getOperands();
+  assert(eraseIndices.size() == operands.size());
+
+  // Check that at least one operand is erased.
+  int firstErasedIndice = eraseIndices.find_first();
+  if (firstErasedIndice == -1)
+    return;
+
+  // Shift all of the removed operands to the end, and destroy them.
+  storage.numOperands = firstErasedIndice;
+  for (unsigned i = firstErasedIndice + 1, e = operands.size(); i < e; ++i)
+    if (!eraseIndices.test(i))
+      operands[storage.numOperands++] = std::move(operands[i]);
+  for (OpOperand &operand : operands.drop_front(storage.numOperands))
+    operand.~OpOperand();
+}
+
 /// Resize the storage to the given size. Returns the array containing the new
 /// operands.
 MutableArrayRef<OpOperand> detail::OperandStorage::resize(Operation *owner,
                                                           unsigned newSize) {
   TrailingOperandStorage &storage = getStorage();
 
   // If the number of operands is less than or equal to the current amount, we
   // can just update in place.

mlir/unittests/IR/OperationSupportTest.cpp

 //===- OperationSupportTest.cpp - Operation support unit tests ------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "llvm/ADT/BitVector.h"
 #include "gtest/gtest.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 static Operation *createOp(MLIRContext *context,
                            ArrayRef<Value> operands = llvm::None,
                            ArrayRef<Type> resultTypes = llvm::None,
TEST(OperandStorageTest, MutableRange) {
   EXPECT_EQ(mutableOperands.size(), 0u);
   EXPECT_EQ(user->getNumOperands(), 0u);
 
   // Destroy the operations.
   user->destroy();
   useOp->destroy();
 }
 
+TEST(OperandStorageTest, RangeErase) {
+  MLIRContext context;
+  Builder builder(&context);
+
+  Type type = builder.getNoneType();
+  Operation *useOp = createOp(&context, /*operands=*/llvm::None, {type, type});
+  Value operand1 = useOp->getResult(0);
+  Value operand2 = useOp->getResult(1);
+
+  // Create an operation with operands to erase.
+  Operation *user =
+      createOp(&context, {operand2, operand1, operand2, operand1});
+  llvm::BitVector eraseIndices(user->getNumOperands());
+
+  // Check erasing no operands.
+  user->eraseOperands(eraseIndices);
+  EXPECT_EQ(user->getNumOperands(), 4u);
+
+  // Check erasing disjoint operands.
+  eraseIndices.set(0);
+  eraseIndices.set(3);
+  user->eraseOperands(eraseIndices);
+  EXPECT_EQ(user->getNumOperands(), 2u);
+  EXPECT_EQ(user->getOperand(0), operand1);
+  EXPECT_EQ(user->getOperand(1), operand2);
+
+  // Destroy the operations.
+  user->destroy();
+  useOp->destroy();
+}
+
 TEST(OperationOrderTest, OrderIsAlwaysValid) {
   MLIRContext context;
   Builder builder(&context);
 
   Operation *containerOp =
       createOp(&context, /*operands=*/llvm::None, /*resultTypes=*/llvm::None,
                /*numRegions=*/1);
   Region &region = containerOp->getRegion(0);