[MLIR] Move eraseArguments and eraseResults to FunctionLike

[mikeurbach]

Summary

Previously, they were only defined for FuncOp.

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 getTypeWithoutArgsAndResults.

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

[mikeurbach]

I took a slightly different approach here compared to the previous implementation. Rather than requiring a new hook on the concrete ops, I just added a member function on FunctionLike and defined it in terms of FunctionType. To me, this isn't much different than how getType and setType work today. By doing it this way, existing FunctionLike ops do not need to change, and can immediately access the new functionality. However, this doesn't really fix the assumption we are using FunctionType, and indeed perpetuates it. I updated the comments to try to be more explicit about this and what concrete ops need to do if they want this functionality and want to use their own function type.

This is a little different that what we originally discussed, so I'm curious what you all think.

@ftynse After doing things this way, I ended up removing all the changes related to LLVMFuncOp. It already doesn't define a setType member function because it doesn't use that functionality. Since it isn't using the functionality extracted from FuncOp, I didn't think it was necessary to define setType or getTypeWithoutArgsAndResults. If it's preferable, I can add back the change that defines these for LLVMFuncOp for compatibility. I also took a look at updating the return type of getType, and that can be done smoothly--call sites are using getFunction* helpers on LLVMType and are straightforward to update... I could submit a separate patch for that.

mlir/include/mlir/IR/FunctionSupportImplementation.h

2

@rriddle I'm not sure if this is what you had in mind, but I put the functions removed from FuncOp here. I wasn't sure about the implementations that were already in FunctionSupport.h, but this lets me share these functions between FuncOp and the ops I'm working with out of tree.

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

168

@mravishankar here is what I ended up adding. I was having a trouble figuring out how to pass this DenseSet<unsigned> without copying the indices into a vector here. I expected the following to work, but the DenseSet const iterator types weren't compatible:

ArrayRef<unsigned> unitDimsToErase(unitDims.begin(), unitDims.end())

Is there a better way to pass the indices in this set as an ArrayRef? In general, is ArrayRef the right API to expose here?

[mikeurbach]

Add missing header guard to new header file.

[ftynse]

The approach LGTM.
Deferring to @rriddle on header layering.

mlir/include/mlir/IR/FunctionSupportImplementation.h

2

I am not convinced by having two headers and requiring the user to include an extra header...

mlir/lib/IR/Types.cpp

67–70

This looks like a repetitive pattern. Would it make sense to define a utility function (somewhere in lib/Support) along the lines of

template <typename Range>
BitVector bitVectorSetAtPositions(unsigned max, Range &&range) {
  BitVector result(max);
  for (unsigned i : range)
    result.set(i);
  return result;
}

[mikeurbach]

mlir/lib/IR/Types.cpp

67–70

I think it would make sense. I can a helper like that and use it here.

[mravishankar]

THanks for moving the method to Block. Other knowledgable people can comment on rest of the patch.

[mikeurbach]

Moved BitVector creation into a helper in Support.

[rriddle]

mlir/include/mlir/IR/FunctionSupportImplementation.h

2

This isn't exactly what I had in mind. I would imagine that we could just have a FunctionSupport.cpp file with type erased methods that are called by the templated version. For example, looking at eraseArguments the only thing that relies on a template is getTypeWithoutArgsAndResults, the rest we could write in a type erased fashion(Updating the attributes and the block arguments don't require anything from the template).

mlir/include/mlir/Support/ADTExtras.h

22

nit: This code block doesn't look large enough to justify a special method. We also prefer to keep mlir/Support as free as possible of things that could go directly into llvm/support.

mlir/lib/IR/Block.cpp

180

Maybe have a bitvector variant so we don't need to create multiple bitvectors?

mlir/lib/IR/Types.cpp

59

Can you avoid all of this if the argIndices or resultIndices are empty?
e.g.

ArrayRef<Type> inputTypes = getInputs();
SmallVector<Type, 4> newInputTypeBuffer;
if (!argIndices.empty()) {
  // Compute the internals of newInputTypeBuffer.
  inputTypes = newInputTypeBuffer;
}

59–64

Same below.

[rriddle]

Apologies for the delay.

[mikeurbach]

mlir/include/mlir/IR/FunctionSupportImplementation.h

2

I see, this makes sense. I'll update this.

I was thinking about getTypeWithoutArgsAndResults... Maybe we can just inline the call to getType().getWithoutArgsAndResults() inside the type erased methods? It should just work for all the concrete ops using FunctionType.

Any ops using another type still have the same "Note that the concrete class must define a method with the same name" caveat they already have for getType. This adds an additional restriction that the type has to have a getWithoutArgsAndResults method.

As was mentioned on Discourse, it would be nice to eventually enforce this kind of thing through a type interface. Perhaps for now, going this route and documenting it in the comments would be sufficient?

mlir/include/mlir/Support/ADTExtras.h

22

I can remove it from here. In the places where this pattern is repeated, I guess I'll just define this small helper there.

mlir/lib/IR/Block.cpp

180

Sounds good.

mlir/lib/IR/Types.cpp

59

Yep, I can do that.

[mikeurbach]

mlir/include/mlir/IR/FunctionSupportImplementation.h

2

Ok I realize how the getTypeWithoutArgsAndResults part has to be templated but the rest of the implementation can be done with the type erased. I will keep the getTypeWithoutArgsAndResults helper as it is now.

[mikeurbach]

Refactor implementations into FunctionSupport.cpp, and address other comments.

[mikeurbach]

Remove unused file from previous revision.

[mikeurbach]

@rriddle when you get a chance, please take another look. I think I've addressed your comments from the previous revision.

mlir/include/mlir/IR/FunctionSupport.h

272

@rriddle I hope this is more along the lines of what you had in mind regarding "type erased methods that are called by the templated version". I didn't spend any time refactoring pieces out of FunctionSupport.h into FunctionSupport.cpp, but I put the bulk of the new implementation there.

mlir/lib/IR/Block.cpp

180

@rriddle I added an eraseArguments below that takes a bitvector. Is this what you had in mind? I know you also have concerns with adding more dependencies to headers, and I added BitVector.h for this.

[rriddle]

LGTM after resolving final comments.

mlir/include/mlir/IR/Block.h

18

You don't need to include this here, you can forward declare BitVector and move the include to Block.cpp.

mlir/lib/IR/FunctionSupport.cpp

12 ↗

(On Diff #302659)

Namespaces in source files should only be used for classes, everything else should be top-level:

https://llvm.org/docs/CodingStandards.html#use-namespace-qualifiers-to-implement-previously-declared-functions
https://llvm.org/docs/CodingStandards.html#anonymous-namespaces

MLIR makes use of using namespace mlir to remove the need to fully qualify everything.

mlir/lib/IR/FunctionSupportDetail.h

24 ↗

(On Diff #302659)

I'm a bit apprehensive about adding new files for little helpers like this, can you just duplicate locally in each .cpp for now?

[mikeurbach]

Awesome, thanks for taking a look. Those final comments all make sense.. I will clean those up.

[mikeurbach]

Updates after review.

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:
public:
   /// Add one argument to the argument list for each type specified in the list.
   iterator_range<args_iterator> addArguments(TypeRange types);
 
   /// 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.
   void eraseArgument(unsigned index);
+  /// Erases the arguments listed in `argIndices` and removes them from the
+  /// argument list.
+  /// `argIndices` is allowed to have duplicates and can be in any order.
+  void eraseArguments(ArrayRef<unsigned> argIndices);
 
   unsigned getNumArguments() { return arguments.size(); }
   BlockArgument getArgument(unsigned i) { return arguments[i]; }
 
   //===--------------------------------------------------------------------===//
   // Operation list management
   //===--------------------------------------------------------------------===//
 

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
   /// provided (leaving them alone if no entry is present). If the mapper
   /// contains entries for function arguments, these arguments are not included
   /// in the new function. Replaces references to cloned sub-values with the
   /// corresponding value that is copied, and adds those mappings to the mapper.
   FuncOp clone(BlockAndValueMapping &mapper);
   FuncOp clone();

mlir/include/mlir/IR/FunctionSupport.h

 /// - Ops must be symbols, i.e. also have the `Symbol` trait;
 /// - Ops have a single region with multiple blocks that corresponds to the body
 ///   of the function;
 /// - An op with a single empty region corresponds to an external function;
 /// - leading arguments of the first block of the region are treated as function
 ///   arguments;
 /// - they can have argument attributes that are stored in a dictionary
 ///   attribute on the Op itself.
-/// This trait does *NOT* provide type support for the functions, meaning that
-/// concrete Ops must handle the type of the declared or defined function.
-/// `getTypeAttrName()` is a convenience function that returns the name of the
-/// attribute that can be used to store the function type, but the trait makes
-/// no assumption based on it.
 ///
+/// This trait provides limited type support for the declared or defined
+/// functions. The convenience function `getTypeAttrName()` returns the name of
+/// an attribute that can be used to store the function type. In addition, this
+/// trait provides `getType` and `setType` helpers to store a `FunctionType` in
+/// the attribute named by `getTypeAttrName()`.
+///
+/// In general, this trait assumes concrete ops use `FunctionType` under the
+/// hood. If this is not the case, in order to use the function type support,
+/// concrete ops must define the following methods, using the same name, to hide
+/// the ones defined for `FunctionType`: `addBodyBlock`, `getType`,
+/// `getTypeWithoutArgsAndResults` and `setType`.
+///
+/// Besides the requirements above, concrete ops must interact with this trait
+/// using the following functions:
 /// - Concrete ops *must* define a member function `getNumFuncArguments()` that
 ///   returns the number of function arguments based exclusively on type (so
 ///   that it can be called on function declarations).
 /// - Concrete ops *must* define a member function `getNumFuncResults()` that
 ///   returns the number of function results based exclusively on type (so that
 ///   it can be called on function declarations).
 /// - To verify that the type respects op-specific invariants, concrete ops may
 ///   redefine the `verifyType()` hook that will be called after verifying the
public:
   ///
   /// 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 getType() {
     return getTypeAttr().getValue().template cast<FunctionType>();
   }
 
+  /// 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.
+  ///
+  /// 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 getTypeWithoutArgsAndResults(ArrayRef<unsigned> argIndices,
+                                            ArrayRef<unsigned> resultIndices);
+
   bool isTypeAttrValid() {
     auto typeAttr = getTypeAttr();
     if (!typeAttr)
       return false;
     return typeAttr.getValue() != Type{};
   }
 
   /// Change the type of this function in place. This is an extremely dangerous
   /// operation and it is up to the caller to ensure that this is legal for this
   /// function, and to restore invariants:
   ///  - the entry block args must be updated to match the function params.
   ///  - the argument/result attributes may need an update: if the new type
   ///    has less parameters we drop the extra attributes, if there are more
   ///    parameters they won't have any attributes.
   ///
   /// 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.
   void setType(FunctionType newType);
 
   //===--------------------------------------------------------------------===//
-  // Argument Handling
+  // Argument and Result Handling
   //===--------------------------------------------------------------------===//
   using BlockArgListType = Region::BlockArgListType;
 
   unsigned getNumArguments() {
     return static_cast<ConcreteType *>(this)->getNumFuncArguments();
   }
 
   unsigned getNumResults() {
     return static_cast<ConcreteType *>(this)->getNumFuncResults();
   }
 
   /// 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);
+  /// 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);
+  /// Erases the results listed in `resultIndices`.
+  /// `resultIndices` is allowed to have duplicates and can be in any order.
+  void eraseResults(ArrayRef<unsigned> resultIndices);
+
   //===--------------------------------------------------------------------===//
   // 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
   /// argument attribute dictionaries are optional, and will generally only
   /// exist if they are non-empty.
 
   /// Return all of the attributes for the argument at 'index'.
   ArrayRef<NamedAttribute> getArgAttrs(unsigned index) {
     return ::mlir::impl::getArgAttrs(this->getOperation(), index);
   }
 
   /// Return all argument attributes of this function.

mlir/include/mlir/IR/FunctionSupportImplementation.h

+//===- FunctionSupportImplementation.h - Implementations for function-like --=//
+//
+// 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_FUNCTIONSUPPORTIMPLEMENTATION_H
+#define MLIR_IR_FUNCTIONSUPPORTIMPLEMENTATION_H
+
+#include "mlir/IR/FunctionSupport.h"
+#include "mlir/Support/ADTExtras.h"
+
+namespace mlir {
+
+namespace OpTrait {
+
+//===----------------------------------------------------------------------===//
+// Function Type Attribute.
+//===----------------------------------------------------------------------===//
+
+template <typename ConcreteType>
+FunctionType FunctionLike<ConcreteType>::getTypeWithoutArgsAndResults(
+    ArrayRef<unsigned> argIndices, ArrayRef<unsigned> resultIndices) {
+  return getType().getWithoutArgsAndResults(argIndices, resultIndices);
+}
+
+//===----------------------------------------------------------------------===//
+// Function Arguments and Results.
+//===----------------------------------------------------------------------===//
+
+template <typename ConcreteType>
+void FunctionLike<ConcreteType>::eraseArgument(unsigned argIndex) {
+  eraseArguments({argIndex});
+}
+
+template <typename ConcreteType>
+void FunctionLike<ConcreteType>::eraseArguments(ArrayRef<unsigned> argIndices) {
+  size_t originalNumArgs = getNumArguments();
+  llvm::BitVector eraseIndices =
+      bitVectorSetAtPositions(originalNumArgs, argIndices);
+
+  // 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 (eraseIndices.test(i))
+      continue;
+    newArgAttrs.emplace_back(getArgAttrDict(i));
+  }
+
+  // Update the function type.
+  setType(getTypeWithoutArgsAndResults(argIndices, {}));
+
+  // Update the arg attrs.
+  setAllArgAttrs(newArgAttrs);
+
+  // Update the entry block's arguments.
+  Block &entry = front();
+  entry.eraseArguments(argIndices);
+}
+
+template <typename ConcreteType>
+void FunctionLike<ConcreteType>::eraseResult(unsigned resultIndex) {
+  eraseResults({resultIndex});
+}
+
+template <typename ConcreteType>
+void FunctionLike<ConcreteType>::eraseResults(
+    ArrayRef<unsigned> resultIndices) {
+  size_t originalNumResults = getNumResults();
+  llvm::BitVector eraseIndices =
+      bitVectorSetAtPositions(originalNumResults, resultIndices);
+
+  // There are 2 things that need to be updated:
+  // - Function type.
+  // - Result attrs.
+
+  // Collect result attrs to set.
+  SmallVector<MutableDictionaryAttr, 4> newResultAttrs;
+  for (size_t i = 0; i < originalNumResults; i++) {
+    if (eraseIndices.test(i))
+      continue;
+    newResultAttrs.emplace_back(getResultAttrDict(i));
+  }
+
+  // Update the function type.
+  setType(getTypeWithoutArgsAndResults({}, resultIndices));
+
+  // Update the result attrs.
+  setAllResultAttrs(newResultAttrs);
+}
+
+} // end namespace OpTrait
+
+} // end namespace mlir
+
+#endif // MLIR_IR_FUNCTIONSUPPORTIMPLEMENTATION_H

mlir/include/mlir/IR/Types.h

 class FunctionType
     : public Type::TypeBase<FunctionType, Type, detail::FunctionTypeStorage> {
 public:
   using Base::Base;
 
   static FunctionType get(TypeRange inputs, TypeRange results,
                           MLIRContext *context);
 
-  // Input types.
+  /// Input types.
   unsigned getNumInputs() const;
   Type getInput(unsigned i) const { return getInputs()[i]; }
   ArrayRef<Type> getInputs() const;
 
-  // Result types.
+  /// Result types.
   unsigned getNumResults() const;
   Type getResult(unsigned i) const { return getResults()[i]; }
   ArrayRef<Type> getResults() const;
+
+  /// Returns a new function type without the specified arguments and results.
+  FunctionType getWithoutArgsAndResults(ArrayRef<unsigned> argIndices,
+                                        ArrayRef<unsigned> resultIndices);
 };
 
 //===----------------------------------------------------------------------===//
 // 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

mlir/include/mlir/Support/ADTExtras.h

+//===- ADTExtras.h - ADT functions relevant to MLIR -------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains abstract data type functions relevant to MLIR to augment
+// what is already provided by `llvm/ADT/`.
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_SUPPORT_ADTEXTRAS_H_
+#define MLIR_SUPPORT_ADTEXTRAS_H_
+
+#include "llvm/ADT/BitVector.h"
+
+namespace mlir {
+
+template <typename RangeT>
+inline llvm::BitVector bitVectorSetAtPositions(unsigned width,
+                                               RangeT &&positions) {
+  llvm::BitVector result(width);
+  for (unsigned i : positions)
+    result.set(i);
+
+  return result;
+}
+
+} // end namespace mlir
+
+#endif // MLIR_SUPPORT_ADTEXTRAS_H_

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
-#include <set>
-
 #define DEBUG_TYPE "linalg-drop-unit-dims"
 
 using namespace mlir;
 using namespace mlir::edsc;
 using namespace mlir::edsc::intrinsics;
 using namespace mlir::linalg;
 
 /// Implements a pass that canonicalizes the uses of unit-extent dimensions for
LogicalResult replaceBlockArgForUnitDimLoops<IndexedGenericOp>(
     PatternRewriter &rewriter) {
   OpBuilder::InsertionGuard guard(rewriter);
   Block *entryBlock = &op.getOperation()->getRegion(0).front();
   rewriter.setInsertionPointToStart(entryBlock);
   Value zero = rewriter.create<ConstantIndexOp>(op.getLoc(), 0);
   for (unsigned unitDimLoop : unitDims) {
     entryBlock->getArgument(unitDimLoop).replaceAllUsesWith(zero);
   }
-  std::set<unsigned> orderedUnitDims(unitDims.begin(), unitDims.end());
-  for (unsigned i : llvm::reverse(orderedUnitDims))
-    entryBlock->eraseArgument(i);
+  SmallVector<unsigned, 8> unitDimsToErase(unitDims.begin(), unitDims.end());
+  entryBlock->eraseArguments(unitDimsToErase);
   return success();
 }
 
 namespace {
 /// Pattern to fold unit-trip count loops in GenericOps.
 // TODO: Generalize this to indexed-generic as well by modifying the region args
 // as well.
 template <typename GenericOpTy>

mlir/lib/IR/Block.cpp

 //===- Block.cpp - MLIR Block Class ---------------------------------------===//
 //
 // 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/Block.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/Support/ADTExtras.h"
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // BlockArgument
 //===----------------------------------------------------------------------===//
 
 /// Returns the number of this argument.
 unsigned BlockArgument::getArgNumber() const {

 }
 
 void Block::eraseArgument(unsigned index) {
   assert(index < arguments.size());
   arguments[index].destroy();
   arguments.erase(arguments.begin() + index);
 }
 
+void Block::eraseArguments(ArrayRef<unsigned> argIndices) {
+  size_t originalNumArgs = getNumArguments();
+  llvm::BitVector eraseIndices =
+      bitVectorSetAtPositions(originalNumArgs, argIndices);
+
+  // We do this in reverse so that we erase later indices before earlier
+  // indices, to avoid shifting the later indices.
+  for (size_t i = 0; i < originalNumArgs; ++i)
+    if (eraseIndices.test(originalNumArgs - i - 1))
+      eraseArgument(originalNumArgs - i - 1);
+}
+
 /// 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()) &&
          "cannot insert arguments to blocks with predecessors");
 
   // Use the args_iterator (on the BlockArgListType) to compute the insertion
   // iterator in the underlying argument storage.

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); };
-
-  // There are 2 things that need to be updated:
-  // - Function type.
-  // - Result attrs.
-
-  // Update the function type and result attrs.
-  SmallVector<Type, 4> newResultTypes;
-  SmallVector<MutableDictionaryAttr, 4> newResultAttrs;
-  for (int i = 0; i < originalNumResults; i++) {
-    if (shouldEraseResult(i))
-      continue;
-    newResultTypes.emplace_back(oldType.getResult(i));
-    newResultAttrs.emplace_back(getResultAttrDict(i));
-  }
-  setType(FunctionType::get(oldType.getInputs(), newResultTypes, getContext()));
-  setAllResultAttrs(newResultAttrs);
-}
-
 /// Clone the internal blocks from this function into dest and all attributes
 /// from this function to dest.
 void FuncOp::cloneInto(FuncOp dest, BlockAndValueMapping &mapper) {
   // Add the attributes of this function to dest.
   llvm::MapVector<Identifier, Attribute> newAttrs;
   for (auto &attr : dest.getAttrs())
     newAttrs.insert(attr);
   for (auto &attr : getAttrs())

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 "mlir/Support/ADTExtras.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();
 }
 
+// Returns a new function type without the specified arguments and results.
+FunctionType
+FunctionType::getWithoutArgsAndResults(ArrayRef<unsigned> argIndices,
+                                       ArrayRef<unsigned> resultIndices) {
+  unsigned originalNumArgs = getNumInputs();
+  llvm::BitVector skipArgIndices =
+      bitVectorSetAtPositions(originalNumArgs, argIndices);
+
+  SmallVector<Type, 4> newInputTypes;
+  for (unsigned i = 0; i < originalNumArgs; ++i) {
+    if (skipArgIndices.test(i))
+      continue;
+
+    newInputTypes.emplace_back(getInput(i));
+  }
+
+  unsigned originalNumResults = getNumResults();
+  llvm::BitVector skipResultIndices =
+      bitVectorSetAtPositions(originalNumResults, resultIndices);
+
+  SmallVector<Type, 4> newOutputTypes;
+  for (unsigned i = 0; i < originalNumResults; ++i) {
+    if (skipResultIndices.test(i))
+      continue;
+
+    newOutputTypes.emplace_back(getResult(i));
+  }
+
+  return get(newInputTypes, newOutputTypes, getContext());
+}
+
 //===----------------------------------------------------------------------===//
 // OpaqueType
 //===----------------------------------------------------------------------===//
 
 OpaqueType OpaqueType::get(Identifier dialect, StringRef typeData,
                            MLIRContext *context) {
   return Base::get(context, dialect, typeData);
 }

mlir/test/lib/IR/TestFunc.cpp

 //===- TestFunctionLike.cpp - Pass to test helpers on FunctionLike --------===//
 //
 // 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/Function.h"
+#include "mlir/IR/FunctionSupportImplementation.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 namespace {
 /// This is a test pass for verifying FuncOp's eraseArgument method.
 struct TestFuncEraseArg
     : public PassWrapper<TestFuncEraseArg, OperationPass<ModuleOp>> {