[mlir:Function][NFC] Use BitVector instead of ArrayRef for indices when removing arguments/results

[rriddle]

Summary

We already convert to BitVector internally, and other APIs (namely Operation::eraseOperands)
already use BitVector as well. Switching over provides a common format between
API and also reduces the amount of format conversions necessary.

Fixes #53325

[mehdi_amini]

NFC?

[jpienaar]

Thanks!

mlir/include/mlir/IR/BuiltinTypes.h

16

Should this be in Support/LLVM.h too? Or not as general as the rest there?

[rriddle]

In D118083#3268098, @mehdi_amini wrote:

NFC?

Yep, thanks.

mlir/include/mlir/IR/BuiltinTypes.h

16

We could add it there given how much it's used. I'll do that and a few other common ones in followups.

mlir/include/mlir/IR/BuiltinTypes.h

 //===- BuiltinTypes.h - MLIR Builtin Type Classes ---------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_BUILTINTYPES_H
 #define MLIR_IR_BUILTINTYPES_H
 
 #include "BuiltinAttributeInterfaces.h"
 #include "SubElementInterfaces.h"
 
 namespace llvm {
+class BitVector;
 struct fltSemantics;
 } // namespace llvm
 
 //===----------------------------------------------------------------------===//
 // Tablegen Interface Declarations
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/BuiltinTypeInterfaces.h.inc"

mlir/include/mlir/IR/BuiltinTypes.td

let extraClassDeclaration = [{
     /// Returns a new function type with the specified arguments and results
     /// inserted.
     FunctionType getWithArgsAndResults(ArrayRef<unsigned> argIndices,
                                        TypeRange argTypes,
                                        ArrayRef<unsigned> resultIndices,
                                        TypeRange resultTypes);
 
     /// Returns a new function type without the specified arguments and results.
-    FunctionType getWithoutArgsAndResults(ArrayRef<unsigned> argIndices,
-                                          ArrayRef<unsigned> resultIndices);
+    FunctionType getWithoutArgsAndResults(const llvm::BitVector &argIndices,
+                                          const llvm::BitVector &resultIndices);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // IndexType
 //===----------------------------------------------------------------------===//
 
 def Builtin_Index : Builtin_Type<"Index"> {

mlir/include/mlir/IR/FunctionInterfaces.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_FUNCTIONINTERFACES_H
 #define MLIR_IR_FUNCTIONINTERFACES_H
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
+#include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/SmallString.h"
 
 namespace mlir {
 
 namespace function_interface_impl {
 
 /// Return the name of the attribute used for function types.
 inline StringRef getTypeAttrName() { return "type"; }

 
 /// Insert the specified results and update the function type attribute.
 void insertFunctionResults(Operation *op, ArrayRef<unsigned> resultIndices,
                            TypeRange resultTypes,
                            ArrayRef<DictionaryAttr> resultAttrs,
                            unsigned originalNumResults, Type newType);
 
 /// Erase the specified arguments and update the function type attribute.
-void eraseFunctionArguments(Operation *op, ArrayRef<unsigned> argIndices,
-                            unsigned originalNumArgs, Type newType);
+void eraseFunctionArguments(Operation *op, const llvm::BitVector &argIndices,
+                            Type newType);
 
 /// Erase the specified results and update the function type attribute.
-void eraseFunctionResults(Operation *op, ArrayRef<unsigned> resultIndices,
-                          unsigned originalNumResults, Type newType);
+void eraseFunctionResults(Operation *op, const llvm::BitVector &resultIndices,
+                          Type newType);
 
 /// Set a FunctionOpInterface operation's type signature.
 void setFunctionType(Operation *op, Type newType);
 
 /// Insert a set of `newTypes` into `oldTypes` at the given `indices`. If any
 /// types are inserted, `storage` is used to hold the new type list. The new
 /// type list is returned. `indices` must be sorted by increasing index.
 TypeRange insertTypesInto(TypeRange oldTypes, ArrayRef<unsigned> indices,
                           TypeRange newTypes, SmallVectorImpl<Type> &storage);
 
 /// Filters out any elements referenced by `indices`. If any types are removed,
 /// `storage` is used to hold the new type list. Returns the new type list.
-TypeRange filterTypesOut(TypeRange types, ArrayRef<unsigned> indices,
+TypeRange filterTypesOut(TypeRange types, const llvm::BitVector &indices,
                          SmallVectorImpl<Type> &storage);
 
 //===----------------------------------------------------------------------===//
 // Function Argument Attribute.
 //===----------------------------------------------------------------------===//
 
 /// Set the attributes held by the argument at 'index'.
 template <typename ConcreteType>

mlir/include/mlir/IR/FunctionInterfaces.td

void insertResults(ArrayRef<unsigned> resultIndices, TypeRange resultTypes,
       Type newType = $_op.getTypeWithArgsAndResults(
         /*argIndices=*/{}, /*argTypes=*/{}, resultIndices, resultTypes);
       function_interface_impl::insertFunctionResults(
           this->getOperation(), resultIndices, resultTypes, resultAttrs,
           originalNumResults, newType);
     }
 
     /// Erase a single argument at `argIndex`.
-    void eraseArgument(unsigned argIndex) { eraseArguments({argIndex}); }
+    void eraseArgument(unsigned argIndex) {
+      llvm::BitVector argsToErase($_op.getNumArguments());
+      argsToErase.set(argIndex);
+      eraseArguments(argsToErase);
+    }
 
     /// Erases the arguments listed in `argIndices`.
-    /// `argIndices` is allowed to have duplicates and can be in any order.
-    void eraseArguments(ArrayRef<unsigned> argIndices) {
-      unsigned originalNumArgs = $_op.getNumArguments();
-      Type newType = $_op.getTypeWithoutArgsAndResults(argIndices, {});
-      function_interface_impl::eraseFunctionArguments(this->getOperation(), argIndices,
-                                                originalNumArgs, newType);
+    void eraseArguments(const llvm::BitVector &argIndices) {
+      Type newType = $_op.getTypeWithoutArgs(argIndices);
+      function_interface_impl::eraseFunctionArguments(
+        this->getOperation(), argIndices, newType);
     }
 
     /// Erase a single result at `resultIndex`.
-    void eraseResult(unsigned resultIndex) { eraseResults({resultIndex}); }
+    void eraseResult(unsigned resultIndex) {
+      llvm::BitVector resultsToErase($_op.getNumResults());
+      resultsToErase.set(resultIndex);
+      eraseResults(resultsToErase);
+    }
 
     /// Erases the results listed in `resultIndices`.
-    /// `resultIndices` is allowed to have duplicates and can be in any order.
-    void eraseResults(ArrayRef<unsigned> resultIndices) {
-      unsigned originalNumResults = $_op.getNumResults();
-      Type newType = $_op.getTypeWithoutArgsAndResults({}, resultIndices);
+    void eraseResults(const llvm::BitVector &resultIndices) {
+      Type newType = $_op.getTypeWithoutResults(resultIndices);
       function_interface_impl::eraseFunctionResults(
-          this->getOperation(), resultIndices, originalNumResults, newType);
+          this->getOperation(), resultIndices, newType);
     }
 
     /// Return the type of this function with the specified arguments and
     /// results inserted. This is used to update the function's signature in
     /// the `insertArguments` and `insertResults` methods. The arrays must be
     /// sorted by increasing index.
     Type getTypeWithArgsAndResults(
       ArrayRef<unsigned> argIndices, TypeRange argTypes,
       ArrayRef<unsigned> resultIndices, TypeRange resultTypes) {
       SmallVector<Type> argStorage, resultStorage;
       TypeRange newArgTypes = function_interface_impl::insertTypesInto(
           $_op.getArgumentTypes(), argIndices, argTypes, argStorage);
       TypeRange newResultTypes = function_interface_impl::insertTypesInto(
           $_op.getResultTypes(), resultIndices, resultTypes, resultStorage);
       return $_op.cloneTypeWith(newArgTypes, newResultTypes);
     }
 
     /// Return the type of this function without the specified arguments and
     /// results. This is used to update the function's signature in the
-    /// `eraseArguments` and `eraseResults` methods. The arrays of indices are
-    /// allowed to have duplicates and can be in any order.
+    /// `eraseArguments` and `eraseResults` methods.
     Type getTypeWithoutArgsAndResults(
-      ArrayRef<unsigned> argIndices, ArrayRef<unsigned> resultIndices) {
+      const llvm::BitVector &argIndices, const llvm::BitVector &resultIndices) {
       SmallVector<Type> argStorage, resultStorage;
       TypeRange newArgTypes = function_interface_impl::filterTypesOut(
           $_op.getArgumentTypes(), argIndices, argStorage);
       TypeRange newResultTypes = function_interface_impl::filterTypesOut(
           $_op.getResultTypes(), resultIndices, resultStorage);
       return $_op.cloneTypeWith(newArgTypes, newResultTypes);
     }
+    Type getTypeWithoutArgs(const llvm::BitVector &argIndices) {
+      SmallVector<Type> argStorage;
+      TypeRange newArgTypes = function_interface_impl::filterTypesOut(
+          $_op.getArgumentTypes(), argIndices, argStorage);
+      return $_op.cloneTypeWith(newArgTypes, $_op.getResultTypes());
+    }
+    Type getTypeWithoutResults(const llvm::BitVector &resultIndices) {
+      SmallVector<Type> resultStorage;
+      TypeRange newResultTypes = function_interface_impl::filterTypesOut(
+          $_op.getResultTypes(), resultIndices, resultStorage);
+      return $_op.cloneTypeWith($_op.getArgumentTypes(), newResultTypes);
+    }
 
     //===------------------------------------------------------------------===//
     // Argument Attributes
     //===------------------------------------------------------------------===//
 
     /// Return all of the attributes for the argument at 'index'.
     ArrayRef<NamedAttribute> getArgAttrs(unsigned index) {
       return function_interface_impl::getArgAttrs(this->getOperation(), index);

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

 //
 // Any args appended to the entry block are added to `appendedEntryArgs`.
 static void updateFuncOp(FuncOp func,
                          SmallVectorImpl<BlockArgument> &appendedEntryArgs) {
   auto functionType = func.getType();
 
   // Collect information about the results will become appended arguments.
   SmallVector<Type, 6> erasedResultTypes;
-  SmallVector<unsigned, 6> erasedResultIndices;
+  llvm::BitVector erasedResultIndices(functionType.getNumResults());
   for (const auto &resultType : llvm::enumerate(functionType.getResults())) {
     if (resultType.value().isa<BaseMemRefType>()) {
-      erasedResultIndices.push_back(resultType.index());
+      erasedResultIndices.set(resultType.index());
       erasedResultTypes.push_back(resultType.value());
     }
   }
 
   // Add the new arguments to the function type.
   auto newArgTypes = llvm::to_vector<6>(
       llvm::concat<const Type>(functionType.getInputs(), erasedResultTypes));
   auto newFunctionType = FunctionType::get(func.getContext(), newArgTypes,
                                            functionType.getResults());
   func.setType(newFunctionType);
 
   // Transfer the result attributes to arg attributes.
-  for (int i = 0, e = erasedResultTypes.size(); i < e; i++)
+  auto erasedIndicesIt = erasedResultIndices.set_bits_begin();
+  for (int i = 0, e = erasedResultTypes.size(); i < e; ++i, ++erasedIndicesIt) {
     func.setArgAttrs(functionType.getNumInputs() + i,
-                     func.getResultAttrs(erasedResultIndices[i]));
+                     func.getResultAttrs(*erasedIndicesIt));
+  }
 
   // Erase the results.
   func.eraseResults(erasedResultIndices);
 
   // Add the new arguments to the entry block if the function is not external.
   if (func.isExternal())
     return;
   Location loc = func.getLoc();

mlir/lib/IR/BuiltinTypes.cpp

TypeRange newArgTypes = function_interface_impl::insertTypesInto(
       getInputs(), argIndices, argTypes, argStorage);
   TypeRange newResultTypes = function_interface_impl::insertTypesInto(
       getResults(), resultIndices, resultTypes, resultStorage);
   return clone(newArgTypes, newResultTypes);
 }
 
 /// Returns a new function type without the specified arguments and results.
 FunctionType
-FunctionType::getWithoutArgsAndResults(ArrayRef<unsigned> argIndices,
-                                       ArrayRef<unsigned> resultIndices) {
+FunctionType::getWithoutArgsAndResults(const llvm::BitVector &argIndices,
+                                       const llvm::BitVector &resultIndices) {
   SmallVector<Type> argStorage, resultStorage;
   TypeRange newArgTypes = function_interface_impl::filterTypesOut(
       getInputs(), argIndices, argStorage);
   TypeRange newResultTypes = function_interface_impl::filterTypesOut(
       getResults(), resultIndices, resultStorage);
   return clone(newArgTypes, newResultTypes);
 }
 

mlir/lib/IR/FunctionInterfaces.cpp

 //===- FunctionSupport.cpp - Utility types for function-like ops ----------===//
 //
 // 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/FunctionInterfaces.h"
-#include "mlir/Support/LLVM.h"
-#include "llvm/ADT/BitVector.h"
 
 using namespace mlir;
 
-/// Helper to call a callback once on each index in the range
-/// [0, `totalIndices`), *except* for the indices given in `indices`.
-/// `indices` is allowed to have duplicates and can be in any order.
-inline static void iterateIndicesExcept(unsigned totalIndices,
-                                        ArrayRef<unsigned> indices,
-                                        function_ref<void(unsigned)> callback) {
-  llvm::BitVector skipIndices(totalIndices);
-  for (unsigned i : indices)
-    skipIndices.set(i);
-
-  for (unsigned i = 0; i < totalIndices; ++i)
-    if (!skipIndices.test(i))
-      callback(i);
-}
-
 //===----------------------------------------------------------------------===//
 // Tablegen Interface Definitions
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/FunctionOpInterfaces.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // Function Arguments and Results.
if (oldResultAttrs || !resultAttrs.empty()) {
     setAllResultAttrDicts(op, newResultAttrs);
   }
 
   // Update the function type.
   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
 }
 
 void mlir::function_interface_impl::eraseFunctionArguments(
-    Operation *op, ArrayRef<unsigned> argIndices, unsigned originalNumArgs,
-    Type newType) {
+    Operation *op, const llvm::BitVector &argIndices, Type newType) {
   // There are 3 things that need to be updated:
   // - Function type.
   // - Arg attrs.
   // - Block arguments of entry block.
   Block &entry = op->getRegion(0).front();
 
   // Update the argument attributes of the function.
   if (auto argAttrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName())) {
     SmallVector<DictionaryAttr, 4> newArgAttrs;
     newArgAttrs.reserve(argAttrs.size());
-    iterateIndicesExcept(originalNumArgs, argIndices, [&](unsigned i) {
+    for (unsigned i = 0, e = argIndices.size(); i < e; ++i)
+      if (!argIndices[i])
         newArgAttrs.emplace_back(argAttrs[i].cast<DictionaryAttr>());
-    });
     setAllArgAttrDicts(op, newArgAttrs);
   }
 
   // Update the function type and any entry block arguments.
   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
   entry.eraseArguments(argIndices);
 }
 
 void mlir::function_interface_impl::eraseFunctionResults(
-    Operation *op, ArrayRef<unsigned> resultIndices,
-    unsigned originalNumResults, Type newType) {
+    Operation *op, const llvm::BitVector &resultIndices, Type newType) {
   // There are 2 things that need to be updated:
   // - Function type.
   // - Result attrs.
 
   // Update the result attributes of the function.
   if (auto resAttrs = op->getAttrOfType<ArrayAttr>(getResultDictAttrName())) {
     SmallVector<DictionaryAttr, 4> newResultAttrs;
     newResultAttrs.reserve(resAttrs.size());
-    iterateIndicesExcept(originalNumResults, resultIndices, [&](unsigned i) {
+    for (unsigned i = 0, e = resultIndices.size(); i < e; ++i)
+      if (!resultIndices[i])
         newResultAttrs.emplace_back(resAttrs[i].cast<DictionaryAttr>());
-    });
     setAllResultAttrDicts(op, newResultAttrs);
   }
 
   // Update the function type.
   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
 }
 
 TypeRange mlir::function_interface_impl::insertTypesInto(
for (auto it : llvm::zip(indices, newTypes)) {
     fromIt = toIt;
   }
   storage.append(fromIt, oldTypes.end());
   return storage;
 }
 
 TypeRange
 mlir::function_interface_impl::filterTypesOut(TypeRange types,
-                                              ArrayRef<unsigned> indices,
+                                              const llvm::BitVector &indices,
                                               SmallVectorImpl<Type> &storage) {
-  if (indices.empty())
+  if (indices.none())
     return types;
-  iterateIndicesExcept(types.size(), indices,
-                       [&](unsigned i) { storage.emplace_back(types[i]); });
+
+  for (unsigned i = 0, e = types.size(); i < e; ++i)
+    if (!indices[i])
+      storage.emplace_back(types[i]);
   return storage;
 }
 
 //===----------------------------------------------------------------------===//
 // Function type signature.
 //===----------------------------------------------------------------------===//
 
 void mlir::function_interface_impl::setFunctionType(Operation *op,

mlir/test/lib/IR/TestFunc.cpp

 struct TestFuncEraseArg
     : public PassWrapper<TestFuncEraseArg, OperationPass<ModuleOp>> {
   StringRef getArgument() const final { return "test-func-erase-arg"; }
   StringRef getDescription() const final { return "Test erasing func args."; }
   void runOnOperation() override {
     auto module = getOperation();
 
     for (FuncOp func : module.getOps<FuncOp>()) {
-      SmallVector<unsigned, 4> indicesToErase;
-      for (auto argIndex : llvm::seq<int>(0, func.getNumArguments())) {
-        if (func.getArgAttr(argIndex, "test.erase_this_arg")) {
-          // Push back twice to test that duplicate arg indices are handled
-          // correctly.
-          indicesToErase.push_back(argIndex);
-          indicesToErase.push_back(argIndex);
-        }
-      }
-      // Reverse the order to test that unsorted index lists are handled
-      // correctly.
-      std::reverse(indicesToErase.begin(), indicesToErase.end());
+      llvm::BitVector indicesToErase(func.getNumArguments());
+      for (auto argIndex : llvm::seq<int>(0, func.getNumArguments()))
+        if (func.getArgAttr(argIndex, "test.erase_this_arg"))
+          indicesToErase.set(argIndex);
       func.eraseArguments(indicesToErase);
     }
   }
 };
 
 /// This is a test pass for verifying FuncOp's eraseResult method.
 struct TestFuncEraseResult
     : public PassWrapper<TestFuncEraseResult, OperationPass<ModuleOp>> {
   StringRef getArgument() const final { return "test-func-erase-result"; }
   StringRef getDescription() const final {
     return "Test erasing func results.";
   }
   void runOnOperation() override {
     auto module = getOperation();
 
     for (FuncOp func : module.getOps<FuncOp>()) {
-      SmallVector<unsigned, 4> indicesToErase;
-      for (auto resultIndex : llvm::seq<int>(0, func.getNumResults())) {
-        if (func.getResultAttr(resultIndex, "test.erase_this_result")) {
-          // Push back twice to test that duplicate indices are handled
-          // correctly.
-          indicesToErase.push_back(resultIndex);
-          indicesToErase.push_back(resultIndex);
-        }
-      }
-      // Reverse the order to test that unsorted index lists are handled
-      // correctly.
-      std::reverse(indicesToErase.begin(), indicesToErase.end());
+      llvm::BitVector indicesToErase(func.getNumResults());
+      for (auto resultIndex : llvm::seq<int>(0, func.getNumResults()))
+        if (func.getResultAttr(resultIndex, "test.erase_this_result"))
+          indicesToErase.set(resultIndex);
       func.eraseResults(indicesToErase);
     }
   }
 };
 
 /// This is a test pass for verifying FuncOp's setType method.
 struct TestFuncSetType
     : public PassWrapper<TestFuncSetType, OperationPass<ModuleOp>> {