[MLIR] Move eraseArguments helper from FuncOp to FunctionLike

[mikeurbach]

Summary

To support this, FunctionLike needs a way to get an updated type
from the concrete operation. This adds a new hook for that purpose,
called getTypeWithoutArguments.

For now, FunctionLike continues to assume the type is
FunctionType, and concrete operations that use another type can hide
the getType, setType, and getTypeWithoutArguments methods.

[mikeurbach]

I'm not sure if there is a way to mark this as a draft, but please consider it as such. I'm trying to illustrate the ideas discussed here: https://llvm.discourse.group/t/moving-erasearguments-from-funcop-to-functionlike/2016.

[ftynse]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

761

This should really be LLVMFunctionType. The LLVMType is used in many places due to legcacy.

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

122

Let's not pollute LLVMType with more methods. The existing methods are mirroring LLVM's API, but new methods are much better off on specific types.

mlir/include/mlir/IR/Function.h

126–127

Nit: this look like irrelevant formatting changes. Does git-clang-format force you to do this because you changed the file?

[mravishankar]

Can this logic be generalized for all Regions? In the patch https://reviews.llvm.org/D90118 I need to drop arguments for region of an linalg::IndexedGenericOp. Having this method there would be more ergonomic.

[mikeurbach]

In D90076#2354042, @mravishankar wrote:

Can this logic be generalized for all Regions? In the patch https://reviews.llvm.org/D90118 I need to drop arguments for region of an linalg::IndexedGenericOp. Having this method there would be more ergonomic.

I am thinking about how to generalize this to something useful for your case. In addition to erasing the requisite block arguments, this helper also handles attributes that are specific to FunctionLike ops. The only part that really generalizes to your situation is looping through a set of indices and calling eraseArgument on the Block.

I guess we could add that part to an eraseArguments(ArrayRef<unsigned> indices) method on the Block class. That could be used by eraseArguments on the FunctionLike trait as well as your case. Is that what you mean?

[rriddle]

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

336

Comments here should be ///

mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td

266

Comments should be /// here.

mlir/include/mlir/IR/FunctionSupport.h

523

Can you please move as much of this out of line into the .cpp file as possible? Templates end up generating a lot of code, and also require including otherwise unnecessary headers (like BitVector here). (I know this comment also applies to the existing code, but code size is something we should really start paying more attention to).

mlir/include/mlir/IR/Types.h

108

Can you please revert the unrelated formatting changes? Using git-clang-format should help here.

255

Please add a much more detailed comment here, and also use /// for comments.

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

1465

nit: Spell out the type instead of using auto here.

[mikeurbach]

Thanks for the feedback. A couple questions to inline comments, otherwise I'm addressing the rest of the comments.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

761

Sounds good. Should I also change getType to use LLVMFunctionType at this point or leave it alone?

mlir/include/mlir/IR/FunctionSupport.h

523

I didn't see a .cpp file at lib/IR/FunctionSupport.cpp. Should I create one and put these definitions there?

[ftynse]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

761

It would be a nice cleanup if it doesn't require a lot of changes in its call sites..

[mikeurbach]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

761

I will try it out, and if it isn't a big mess, include it in the patch.

[mravishankar]

I guess we could add that part to an eraseArguments(ArrayRef<unsigned> indices) method on the Block class. That could be used by eraseArguments on the FunctionLike trait as well as your case. Is that what you mean?

Yes. That would be useful!

[mikeurbach]

mlir/include/mlir/IR/FunctionSupport.h

523

I started looking into this, and now it is making more sense. I'm not a C++ expert, but it looks like these definitions will need to be available to the various places FunctionLike ops use them. I think a good way to achieve that would be to move them into a header here in include/mlir/IR so they can be included in the various .cpp files that need them. Like how there is DialectImplementation.h, OpImplementation.h, etc., there could be a FunctionSupportImplementation.h or something of the sort.

I think this will work, but is that what you had in mind? Is there a better way to do it?

[mikeurbach]

Well, I tried to share my updates and ended up creating a new revision, that is here now: https://reviews.llvm.org/D90363

[mikeurbach]

This revision was superseded by https://reviews.llvm.org/D90363.

mlir/include/mlir/Dialect/GPU/GPUOps.td

let extraClassDeclaration = [{
     }
 
     // FunctionLike trait needs access to the functions below.
     friend class OpTrait::FunctionLike<GPUFuncOp>;
 
     /// Hooks for the input/output type enumeration in FunctionLike .
     unsigned getNumFuncArguments() { return getType().getNumInputs(); }
     unsigned getNumFuncResults() { return getType().getNumResults(); }
+    Type getTypeWithoutArguments(ArrayRef<unsigned> argIndices) {
+      return getType().getWithoutArguments(argIndices);
+    }
 
     /// Returns the keywords used in the custom syntax for this Op.
     static StringRef getWorkgroupKeyword() { return "workgroup"; }
     static StringRef getPrivateKeyword() { return "private"; }
     static StringRef getKernelKeyword() { return "kernel"; }
 
     /// Hook for FunctionLike verifier.
     LogicalResult verifyType();

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

let extraClassDeclaration = [{
     // Add an entry block to an empty function, and set up the block arguments
     // to match the signature of the function.
     Block *addEntryBlock();
 
     LLVMType getType() {
       return getAttrOfType<TypeAttr>(getTypeAttrName())
           .getValue().cast<LLVMType>();
     }
+
+    void setType(LLVMType newType);
+
     bool isVarArg() {
       return getType().isFunctionVarArg();
     }
 
     // Hook for OpTrait::FunctionLike, returns the number of function arguments`.
     // Depends on the type attribute being correct as checked by verifyType.
     unsigned getNumFuncArguments();
 
     // Hook for OpTrait::FunctionLike, returns the number of function results.
     // Depends on the type attribute being correct as checked by verifyType.
     unsigned getNumFuncResults();
 
+    // Hook for OpTrait::FunctionLike, returns the signature type without the
+    // specified arguments. Depends on the type attribute being correct as
+    // checked by verifyType.
+    Type getTypeWithoutArguments(ArrayRef<unsigned> argIndices);
+
     // Hook for OpTrait::FunctionLike, called after verifying that the 'type'
     // attribute is present.  This can check for preconditions of the
     // getNumArguments hook not failing.
     LogicalResult verifyType();
   }];
 
   let verifier = [{ return ::verify(*this); }];
   let printer = [{ printLLVMFuncOp(p, *this); }];

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

public:
   unsigned getVectorNumElements();
   llvm::ElementCount getVectorElementCount();
   bool isVectorTy();
 
   /// Function type utilities.
   LLVMType getFunctionParamType(unsigned argIdx);
   unsigned getFunctionNumParams();
   LLVMType getFunctionResultType();
+  LLVMType getFunctionTypeWithoutArguments(ArrayRef<unsigned> argIndices);
   bool isFunctionTy();
   bool isFunctionVarArg();
 
   /// Pointer type utilities.
   LLVMType getPointerTo(unsigned addrSpace = 0);
   LLVMType getPointerElementTy();
   bool isPointerTy();
 
public:
 
   /// Returns `i`-th argument of the function. Asserts on out-of-bounds.
   LLVMType getParamType(unsigned i);
 
   /// Returns a list of argument types of the function.
   ArrayRef<LLVMType> getParams();
   ArrayRef<LLVMType> params() { return getParams(); }
 
+  // Helper for returning modified type without arguments.
+  LLVMType getWithoutArguments(ArrayRef<unsigned> argIndices);
+
   /// Verifies that the type about to be constructed is well-formed.
   static LogicalResult
   verifyConstructionInvariants(Location loc, LLVMType result,
                                ArrayRef<LLVMType> arguments, bool);
 };
 
 //===----------------------------------------------------------------------===//
 // LLVMIntegerType.

mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td

private:
     friend class OpTrait::FunctionLike<FuncOp>;
 
     /// Returns the number of arguments. Hook for OpTrait::FunctionLike.
     unsigned getNumFuncArguments() { return getType().getNumInputs(); }
 
     /// Returns the number of results. Hook for OpTrait::FunctionLike.
     unsigned getNumFuncResults() { return getType().getNumResults(); }
 
+    /// Returns the signature type without the specified arguments. This is a
+    // hook for OpTrait::FunctionLike.
+    Type getTypeWithoutArguments(ArrayRef<unsigned> argIndices) {
+      return getType().getWithoutArguments(argIndices);
+    }
+
     /// Hook for OpTrait::FunctionLike, called after verifying that the 'type'
     /// attribute is present and checks if it holds a function type. Ensures
     /// getType, getNumFuncArguments, and getNumFuncResults can be called safely
     LogicalResult verifyType();
 
     /// Hook for OpTrait::FunctionLike, called after verifying the function
     /// type and the presence of the (potentially empty) function body.
     /// Ensures SPIR-V specific semantics.

mlir/include/mlir/IR/Function.h

static void build(OpBuilder &builder, OperationState &result, StringRef name,
                     FunctionType type, ArrayRef<NamedAttribute> attrs = {},
                     ArrayRef<MutableDictionaryAttr> argAttrs = {});
 
   /// Operation hooks.
   static ParseResult parse(OpAsmParser &parser, OperationState &result);
   void print(OpAsmPrinter &p);
   LogicalResult verify();
 
-  /// Erase a single argument at `argIndex`.
-  void eraseArgument(unsigned argIndex) { eraseArguments({argIndex}); }
-  /// Erases the arguments listed in `argIndices`.
-  /// `argIndices` is allowed to have duplicates and can be in any order.
-  void eraseArguments(ArrayRef<unsigned> argIndices);
-
   /// Erase a single result at `resultIndex`.
   void eraseResult(unsigned resultIndex) { eraseResults({resultIndex}); }
   /// Erases the results listed in `resultIndices`.
   /// `resultIndices` is allowed to have duplicates and can be in any order.
   void eraseResults(ArrayRef<unsigned> resultIndices);
 
   /// Create a deep copy of this function and all of its blocks, remapping
   /// any operands that use values outside of the function using the map that is
private:
   friend class OpTrait::FunctionLike<FuncOp>;
 
   /// Returns the number of arguments. This is a hook for OpTrait::FunctionLike.
   unsigned getNumFuncArguments() { return getType().getInputs().size(); }
 
   /// Returns the number of results. This is a hook for OpTrait::FunctionLike.
   unsigned getNumFuncResults() { return getType().getResults().size(); }
 
+  /// Returns the signature type without the specified arguments. This is a hook
+  /// for OpTrait::FunctionLike.
+  FunctionType getTypeWithoutArguments(ArrayRef<unsigned> argIndices) {
+    return getType().getWithoutArguments(argIndices);
+  }
+
   /// Hook for OpTrait::FunctionLike, called after verifying that the 'type'
   /// attribute is present and checks if it holds a function type.  Ensures
   /// getType, getNumFuncArguments, and getNumFuncResults can be called safely.
   LogicalResult verifyType() {
     auto type = getTypeAttr().getValue();
     if (!type.isa<FunctionType>())
       return emitOpError("requires '" + getTypeAttrName() +
                          "' attribute of function type");
     return success();
   }
 };
 } // end namespace mlir
 
 namespace llvm {
 
 // Functions hash just like pointers.
-template <> struct DenseMapInfo<mlir::FuncOp> {
+template <>
+struct DenseMapInfo<mlir::FuncOp> {
   static mlir::FuncOp getEmptyKey() {
     auto pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
     return mlir::FuncOp::getFromOpaquePointer(pointer);
   }
   static mlir::FuncOp getTombstoneKey() {
     auto pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
     return mlir::FuncOp::getFromOpaquePointer(pointer);
   }
   static unsigned getHashValue(mlir::FuncOp val) {
     return hash_value(val.getAsOpaquePointer());
   }
   static bool isEqual(mlir::FuncOp LHS, mlir::FuncOp RHS) { return LHS == RHS; }
 };
 
 /// Allow stealing the low bits of FuncOp.
-template <> struct PointerLikeTypeTraits<mlir::FuncOp> {
+template <>
+struct PointerLikeTypeTraits<mlir::FuncOp> {
 public:
   static inline void *getAsVoidPointer(mlir::FuncOp I) {
     return const_cast<void *>(I.getAsOpaquePointer());
   }
   static inline mlir::FuncOp getFromVoidPointer(void *P) {
     return mlir::FuncOp::getFromOpaquePointer(P);
   }
   static constexpr int NumLowBitsAvailable = 3;
 };
 
 } // namespace llvm
 
 #endif // MLIR_IR_FUNCTION_H

mlir/include/mlir/IR/FunctionSupport.h

 // constructs to use.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_FUNCTIONSUPPORT_H
 #define MLIR_IR_FUNCTIONSUPPORT_H
 
 #include "mlir/IR/OpDefinition.h"
+#include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/SmallString.h"
 
 namespace mlir {
 
 namespace impl {
 
 /// Return the name of the attribute used for function types.
 inline StringRef getTypeAttrName() { return "type"; }
public:
   unsigned getNumArguments() {
     return static_cast<ConcreteType *>(this)->getNumFuncArguments();
   }
 
   unsigned getNumResults() {
     return static_cast<ConcreteType *>(this)->getNumFuncResults();
   }
 
+  /// Get an updated signature type without the specified arguments.
+  ///
+  /// Note that the concrete class must define a method with the same name to
+  /// hide this one if the concrete class does not use FunctionType for the
+  /// function type under the hood.
+  FunctionType getTypeWithoutArguments(ArrayRef<unsigned> argIndices) {
+    return static_cast<ConcreteType *>(this)->getTypeWithoutArguments(
+        argIndices);
+  }
+
   /// Gets argument.
   BlockArgument getArgument(unsigned idx) { return getBody().getArgument(idx); }
 
   /// Support argument iteration.
   using args_iterator = Region::args_iterator;
   args_iterator args_begin() { return getBody().args_begin(); }
   args_iterator args_end() { return getBody().args_end(); }
   Block::BlockArgListType getArguments() { return getBody().getArguments(); }
 
   ValueTypeRange<BlockArgListType> getArgumentTypes() {
     return getBody().getArgumentTypes();
   }
 
+  /// Erase a single argument at `argIndex`.
+  void eraseArgument(unsigned argIndex) { eraseArguments({argIndex}); }
+  /// Erases the arguments listed in `argIndices`.
+  /// `argIndices` is allowed to have duplicates and can be in any order.
+  void eraseArguments(ArrayRef<unsigned> argIndices);
+
   //===--------------------------------------------------------------------===//
   // Argument Attributes
   //===--------------------------------------------------------------------===//
 
   /// FunctionLike operations allow for attaching attributes to each of the
   /// respective function arguments. These argument attributes are stored as
   /// DictionaryAttrs in the main operation attribute dictionary. The name of
   /// these entries is `arg` followed by the index of the argument. These
void FunctionLike<ConcreteType>::setType(FunctionType newType) {
   for (int i = newType.getNumInputs(), e = oldType.getNumInputs(); i < e; i++)
     concreteOp->removeAttr(getArgAttrName(i, nameBuf));
   for (int i = newType.getNumResults(), e = oldType.getNumResults(); i < e; i++)
     concreteOp->removeAttr(getResultAttrName(i, nameBuf));
   concreteOp->setAttr(getTypeAttrName(), TypeAttr::get(newType));
 }
 
 //===----------------------------------------------------------------------===//
+// Function Arguments.
+//===----------------------------------------------------------------------===//
+
+template <typename ConcreteType>
+void FunctionLike<ConcreteType>::eraseArguments(ArrayRef<unsigned> argIndices) {
+  size_t originalNumArgs = getNumArguments();
+  llvm::BitVector eraseIndices(originalNumArgs);
+  for (auto index : argIndices)
+    eraseIndices.set(index);
+  auto shouldEraseArg = [&](size_t i) { return eraseIndices.test(i); };
+
+  // There are 3 things that need to be updated:
+  // - Function type.
+  // - Arg attrs.
+  // - Block arguments of entry block.
+
+  // Collect arg attrs to set.
+  SmallVector<MutableDictionaryAttr, 4> newArgAttrs;
+  for (size_t i = 0; i < originalNumArgs; ++i) {
+    if (shouldEraseArg(i))
+      continue;
+    newArgAttrs.emplace_back(getArgAttrDict(i));
+  }
+
+  // Update the function type.
+  setType(getTypeWithoutArguments(argIndices));
+
+  // Update the arg attrs.
+  setAllArgAttrs(newArgAttrs);
+
+  // Update the entry block's arguments.
+  // We do this in reverse so that we erase later indices before earlier
+  // indices, to avoid shifting the later indices.
+  Block &entry = front();
+  for (size_t i = 0; i < originalNumArgs; ++i)
+    if (shouldEraseArg(originalNumArgs - i - 1))
+      entry.eraseArgument(originalNumArgs - i - 1);
+}
+
+//===----------------------------------------------------------------------===//
 // Function Argument Attribute.
 //===----------------------------------------------------------------------===//
 
 /// Set the attributes held by the argument at 'index'.
 template <typename ConcreteType>
 void FunctionLike<ConcreteType>::setArgAttrs(
     unsigned index, ArrayRef<NamedAttribute> attributes) {
   assert(index < getNumArguments() && "invalid argument number");

mlir/include/mlir/IR/Types.h

public:
   Type &operator=(const Type &other) = default;
 
   bool operator==(Type other) const { return impl == other.impl; }
   bool operator!=(Type other) const { return !(*this == other); }
   explicit operator bool() const { return impl; }
 
   bool operator!() const { return impl == nullptr; }
 
-  template <typename U> bool isa() const;
+  template <typename U>
+  bool isa() const;
   template <typename First, typename Second, typename... Rest>
   bool isa() const;
-  template <typename U> U dyn_cast() const;
-  template <typename U> U dyn_cast_or_null() const;
-  template <typename U> U cast() const;
+  template <typename U>
+  U dyn_cast() const;
+  template <typename U>
+  U dyn_cast_or_null() const;
+  template <typename U>
+  U cast() const;
 
   // Support type casting Type to itself.
   static bool classof(Type) { return true; }
 
   /// Return a unique identifier for the concrete type. This is used to support
   /// dynamic type casting.
   TypeID getTypeID() { return impl->getAbstractType().getTypeID(); }
 
public:
   unsigned getNumInputs() const;
   Type getInput(unsigned i) const { return getInputs()[i]; }
   ArrayRef<Type> getInputs() const;
 
   // Result types.
   unsigned getNumResults() const;
   Type getResult(unsigned i) const { return getResults()[i]; }
   ArrayRef<Type> getResults() const;
+
+  // Helper for returning modified type without arguments.
+  FunctionType getWithoutArguments(ArrayRef<unsigned> argIndices);
 };
 
 //===----------------------------------------------------------------------===//
 // OpaqueType
 //===----------------------------------------------------------------------===//
 
 /// Opaque types represent types of non-registered dialects. These are types
 /// represented in their raw string form, and can only usefully be tested for
static LogicalResult verifyConstructionInvariants(Location loc,
                                                     StringRef typeData);
 };
 
 // Make Type hashable.
 inline ::llvm::hash_code hash_value(Type arg) {
   return ::llvm::hash_value(arg.impl);
 }
 
-template <typename U> bool Type::isa() const {
+template <typename U>
+bool Type::isa() const {
   assert(impl && "isa<> used on a null type.");
   return U::classof(*this);
 }
 
 template <typename First, typename Second, typename... Rest>
 bool Type::isa() const {
   return isa<First>() || isa<Second, Rest...>();
 }
 
-template <typename U> U Type::dyn_cast() const {
+template <typename U>
+U Type::dyn_cast() const {
   return isa<U>() ? U(impl) : U(nullptr);
 }
-template <typename U> U Type::dyn_cast_or_null() const {
+template <typename U>
+U Type::dyn_cast_or_null() const {
   return (impl && isa<U>()) ? U(impl) : U(nullptr);
 }
-template <typename U> U Type::cast() const {
+template <typename U>
+U Type::cast() const {
   assert(isa<U>());
   return U(impl);
 }
 
 } // end namespace mlir
 
 namespace llvm {
 
 // Type hash just like pointers.
-template <> struct DenseMapInfo<mlir::Type> {
+template <>
+struct DenseMapInfo<mlir::Type> {
   static mlir::Type getEmptyKey() {
     auto pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
     return mlir::Type(static_cast<mlir::Type::ImplType *>(pointer));
   }
   static mlir::Type getTombstoneKey() {
     auto pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
     return mlir::Type(static_cast<mlir::Type::ImplType *>(pointer));
   }
   static unsigned getHashValue(mlir::Type val) { return mlir::hash_value(val); }
   static bool isEqual(mlir::Type LHS, mlir::Type RHS) { return LHS == RHS; }
 };
 
 /// We align TypeStorage by 8, so allow LLVM to steal the low bits.
-template <> struct PointerLikeTypeTraits<mlir::Type> {
+template <>
+struct PointerLikeTypeTraits<mlir::Type> {
 public:
   static inline void *getAsVoidPointer(mlir::Type I) {
     return const_cast<void *>(I.getAsOpaquePointer());
   }
   static inline mlir::Type getFromVoidPointer(void *P) {
     return mlir::Type::getFromOpaquePointer(P);
   }
   static constexpr int NumLowBitsAvailable = 3;
 };
 
 } // namespace llvm
 
 #endif // MLIR_IR_TYPES_H

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

unsigned LLVMFuncOp::getNumFuncResults() {
   // If we modeled a void return as one result, then it would be possible to
   // attach an MLIR result attribute to it, and it isn't clear what semantics we
   // would assign to that.
   if (getType().getFunctionResultType().isVoidTy())
     return 0;
   return 1;
 }
 
+// Hook for OpTrait::FunctionLike, returns the signature type without the
+// specified arguments. Depends on the type attribute being correct as checked
+// by verifyType.
+Type LLVMFuncOp::getTypeWithoutArguments(ArrayRef<unsigned> argIndices) {
+  return getType().getFunctionTypeWithoutArguments(argIndices);
+}
+
+void LLVMFuncOp::setType(LLVMType newType) {
+  SmallVector<char, 16> nameBuf;
+  auto oldType = getType();
+
+  for (size_t i = newType.getFunctionNumParams(),
+              e = oldType.getFunctionNumParams();
+       i < e; ++i)
+    removeAttr(getArgAttrName(i, nameBuf));
+
+  setAttr(getTypeAttrName(), TypeAttr::get(newType));
+}
+
 // Verifies LLVM- and implementation-specific properties of the LLVM func Op:
 // - functions don't have 'common' linkage
 // - external functions have 'external' or 'extern_weak' linkage;
 // - vararg is (currently) only supported for external functions;
 // - entry block arguments are of LLVM types and match the function signature.
 static LogicalResult verify(LLVMFuncOp op) {
   if (op.linkage() == LLVM::Linkage::Common)
     return op.emitOpError()

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

 
 #include "TypeDetail.h"
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/TypeSupport.h"
 
+#include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/TypeSize.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
 
 //===----------------------------------------------------------------------===//
 // LLVMType.

 unsigned LLVMType::getFunctionNumParams() {
   return cast<LLVMFunctionType>().getNumParams();
 }
 
 LLVMType LLVMType::getFunctionResultType() {
   return cast<LLVMFunctionType>().getReturnType();
 }
 
+LLVMType
+LLVMType::getFunctionTypeWithoutArguments(ArrayRef<unsigned> argIndices) {
+  return cast<LLVMFunctionType>().getWithoutArguments(argIndices);
+}
+
 bool LLVMType::isFunctionTy() { return isa<LLVMFunctionType>(); }
 
 bool LLVMType::isFunctionVarArg() {
   return cast<LLVMFunctionType>().isVarArg();
 }
 
 //----------------------------------------------------------------------------//
 // Pointer type utilities.

 }
 
 bool LLVMFunctionType::isVarArg() { return getImpl()->isVariadic(); }
 
 ArrayRef<LLVMType> LLVMFunctionType::getParams() {
   return getImpl()->getArgumentTypes();
 }
 
+LLVMType LLVMFunctionType::getWithoutArguments(ArrayRef<unsigned> argIndices) {
+  unsigned originalNumArgs = getNumParams();
+  llvm::BitVector skipIndices(originalNumArgs);
+  for (auto index : argIndices)
+    skipIndices.set(index);
+
+  SmallVector<LLVMType, 4> newInputTypes;
+  for (unsigned i = 0; i < originalNumArgs; ++i) {
+    if (skipIndices.test(i))
+      continue;
+
+    newInputTypes.emplace_back(getParamType(i));
+  }
+
+  return get(getReturnType(), newInputTypes, isVarArg());
+}
+
 LogicalResult LLVMFunctionType::verifyConstructionInvariants(
     Location loc, LLVMType result, ArrayRef<LLVMType> arguments, bool) {
   if (!isValidResultType(result))
     return emitError(loc, "invalid function result type: ") << result;
 
   for (LLVMType arg : arguments)
     if (!isValidArgumentType(arg))
       return emitError(loc, "invalid function argument type: ") << arg;

mlir/lib/IR/Function.cpp

if (fnInputTypes[i] != entryBlock.getArgument(i).getType())
       return emitOpError("type of entry block argument #")
              << i << '(' << entryBlock.getArgument(i).getType()
              << ") must match the type of the corresponding argument in "
              << "function signature(" << fnInputTypes[i] << ')';
 
   return success();
 }
 
-void FuncOp::eraseArguments(ArrayRef<unsigned> argIndices) {
-  auto oldType = getType();
-  int originalNumArgs = oldType.getNumInputs();
-  llvm::BitVector eraseIndices(originalNumArgs);
-  for (auto index : argIndices)
-    eraseIndices.set(index);
-  auto shouldEraseArg = [&](int i) { return eraseIndices.test(i); };
-
-  // There are 3 things that need to be updated:
-  // - Function type.
-  // - Arg attrs.
-  // - Block arguments of entry block.
-
-  // Update the function type and arg attrs.
-  SmallVector<Type, 4> newInputTypes;
-  SmallVector<MutableDictionaryAttr, 4> newArgAttrs;
-  for (int i = 0; i < originalNumArgs; i++) {
-    if (shouldEraseArg(i))
-      continue;
-    newInputTypes.emplace_back(oldType.getInput(i));
-    newArgAttrs.emplace_back(getArgAttrDict(i));
-  }
-  setType(FunctionType::get(newInputTypes, oldType.getResults(), getContext()));
-  setAllArgAttrs(newArgAttrs);
-
-  // Update the entry block's arguments.
-  // We do this in reverse so that we erase later indices before earlier
-  // indices, to avoid shifting the later indices.
-  Block &entry = front();
-  for (int i = 0; i < originalNumArgs; i++)
-    if (shouldEraseArg(originalNumArgs - i - 1))
-      entry.eraseArgument(originalNumArgs - i - 1);
-}
-
 void FuncOp::eraseResults(ArrayRef<unsigned> resultIndices) {
   auto oldType = getType();
   int originalNumResults = oldType.getNumResults();
   llvm::BitVector eraseIndices(originalNumResults);
   for (auto index : resultIndices)
     eraseIndices.set(index);
   auto shouldEraseResult = [&](int i) { return eraseIndices.test(i); };
 

mlir/lib/IR/Types.cpp

 //===- Types.cpp - MLIR Type Classes --------------------------------------===//
 //
 // 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/Types.h"
 #include "TypeDetail.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
+#include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/Twine.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 //===----------------------------------------------------------------------===//
 // Type
 //===----------------------------------------------------------------------===//

 }
 
 unsigned FunctionType::getNumResults() const { return getImpl()->numResults; }
 
 ArrayRef<Type> FunctionType::getResults() const {
   return getImpl()->getResults();
 }
 
+// Helper for returning modified type without arguments.
+FunctionType FunctionType::getWithoutArguments(ArrayRef<unsigned> argIndices) {
+  unsigned originalNumArgs = getNumInputs();
+  llvm::BitVector skipIndices(originalNumArgs);
+  for (auto index : argIndices)
+    skipIndices.set(index);
+
+  SmallVector<Type, 4> newInputTypes;
+  for (unsigned i = 0; i < originalNumArgs; ++i) {
+    if (skipIndices.test(i))
+      continue;
+
+    newInputTypes.emplace_back(getInput(i));
+  }
+
+  return get(newInputTypes, getResults(), getContext());
+}
+
 //===----------------------------------------------------------------------===//
 // OpaqueType
 //===----------------------------------------------------------------------===//
 
 OpaqueType OpaqueType::get(Identifier dialect, StringRef typeData,
                            MLIRContext *context) {
   return Base::get(context, dialect, typeData);
 }