[mlir] Refactor the representation of function-like argument/result attributes.

[rriddle]

Summary

The current design uses a unique entry for each argument/result attribute, with the name of the entry being something like "arg0". This provides for a somewhat sparse design, but ends up being much more expensive (from a runtime perspective) in-practice. The design requires building a string every time we lookup the dictionary for a specific arg/result, and also requires N attribute lookups when collecting all of the arg/result attribute dictionaries.

This revision restructures the design to instead have an ArrayAttr that contains all of the attribute dictionaries for arguments and another for results. This design reduces the number of attribute name lookups to 1, and allows for O(1) lookup for individual element dictionaries. The major downside is that we can end up with larger memory usage, as the ArrayAttr contains an entry for each element even if that element has no attributes. If the memory usage becomes too problematic, we can experiment with a more sparse structure that still provides a lot of the wins in this revision.

This dropped the compilation time of a somewhat large TensorFlow model from ~650 seconds to ~400 seconds.

[mehdi_amini]

mlir/include/mlir/IR/FunctionSupport.h

47

Looks like a candidate to be in a "detail" namespace, this really looks like a low-level helper for the other methods.

49

Add that the ArrayRef size must matching the number of arguments to the function

[rriddle]

update

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

def GPU_GPUFuncOp : GPU_Op<"func", [HasParent<"GPUModuleOp">,
   let extraClassDeclaration = [{
     /// Returns `true` if the GPU function defined by this Op is a kernel, i.e.
     /// it is intended to be launched from host.
     bool isKernel() {
       return (*this)->getAttrOfType<UnitAttr>(
           GPUDialect::getKernelFuncAttrName()) != nullptr;
     }
 
-    /// 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.
-    // TODO: consider removing this function thanks to rewrite patterns.
-    void setType(FunctionType newType);
-
     /// Returns the number of buffers located in the workgroup memory.
     unsigned getNumWorkgroupAttributions() {
       return (*this)->getAttrOfType<IntegerAttr>(
           getNumWorkgroupAttributionsAttrName()).getInt();
     }
 
     /// Returns a list of block arguments that correspond to buffers located in
     /// the workgroup memory

mlir/include/mlir/IR/BuiltinAttributes.td

let description = [{
     ```mlir
     {}
     {attr_name = "string attribute"}
     {int_attr = 10, "string attr name" = "string attribute"}
     ```
   }];
   let parameters = (ins ArrayRefParameter<"NamedAttribute", "">:$value);
   let builders = [
-    AttrBuilder<(ins "ArrayRef<NamedAttribute>":$value)>
+    AttrBuilder<(ins CArg<"ArrayRef<NamedAttribute>", "llvm::None">:$value)>
   ];
   let extraClassDeclaration = [{
     using ValueType = ArrayRef<NamedAttribute>;
 
     /// Construct a dictionary with an array of values that is known to already
     /// be sorted by name and uniqued.
     static DictionaryAttr getWithSorted(MLIRContext *context,
                                         ArrayRef<NamedAttribute> value);

mlir/include/mlir/IR/FunctionImplementation.h

 #ifndef MLIR_IR_FUNCTIONIMPLEMENTATION_H_
 #define MLIR_IR_FUNCTIONIMPLEMENTATION_H_
 
 #include "mlir/IR/FunctionSupport.h"
 #include "mlir/IR/OpImplementation.h"
 
 namespace mlir {
 
-namespace impl {
+namespace function_like_impl {
 
 /// A named class for passing around the variadic flag.
 class VariadicFlag {
 public:
   explicit VariadicFlag(bool variadic) : variadic(variadic) {}
   bool isVariadic() const { return variadic; }
 
 private:
   /// Underlying storage.
   bool variadic;
 };
 
 /// Adds argument and result attributes, provided as `argAttrs` and
 /// `resultAttrs` arguments, to the list of operation attributes in `result`.
 /// Internally, argument and result attributes are stored as dict attributes
 /// with special names given by getResultAttrName, getArgumentAttrName.
 void addArgAndResultAttrs(Builder &builder, OperationState &result,
+                          ArrayRef<DictionaryAttr> argAttrs,
+                          ArrayRef<DictionaryAttr> resultAttrs);
+void addArgAndResultAttrs(Builder &builder, OperationState &result,
                           ArrayRef<NamedAttrList> argAttrs,
                           ArrayRef<NamedAttrList> resultAttrs);
 
 /// Callback type for `parseFunctionLikeOp`, the callback should produce the
 /// type that will be associated with a function-like operation from lists of
 /// function arguments and results, VariadicFlag indicates whether the function
 /// should have variadic arguments; in case of error, it may populate the last
 /// argument with a message.

 /// attributes with names listed in "elided" as well as those used by the
 /// function-like operation internally are not printed. Nothing is printed
 /// if all attributes are elided. Assumes `op` has the `FunctionLike` trait and
 /// passed the verification.
 void printFunctionAttributes(OpAsmPrinter &p, Operation *op, unsigned numInputs,
                              unsigned numResults,
                              ArrayRef<StringRef> elided = {});
 
-} // namespace impl
+} // namespace function_like_impl
 
 } // namespace mlir
 
 #endif // MLIR_IR_FUNCTIONIMPLEMENTATION_H_

mlir/include/mlir/IR/FunctionSupport.h

 #define MLIR_IR_FUNCTIONSUPPORT_H
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
 #include "llvm/ADT/SmallString.h"
 
 namespace mlir {
 
-namespace impl {
+namespace function_like_impl {
 
 /// Return the name of the attribute used for function types.
 inline StringRef getTypeAttrName() { return "type"; }
 
-/// Return the name of the attribute used for function arguments.
-inline StringRef getArgAttrName(unsigned arg, SmallVectorImpl<char> &out) {
-  out.clear();
-  return ("arg" + Twine(arg)).toStringRef(out);
-}
-
-/// Returns true if the given name is a valid argument attribute name.
-inline bool isArgAttrName(StringRef name) {
-  APInt unused;
-  return name.startswith("arg") &&
-         !name.drop_front(3).getAsInteger(/*Radix=*/10, unused);
-}
-
-/// Return the name of the attribute used for function results.
-inline StringRef getResultAttrName(unsigned arg, SmallVectorImpl<char> &out) {
-  out.clear();
-  return ("result" + Twine(arg)).toStringRef(out);
-}
+/// Return the name of the attribute used for function argument attributes.
+inline StringRef getArgDictAttrName() { return "arg_attrs"; }
+
+/// Return the name of the attribute used for function argument attributes.
+inline StringRef getResultDictAttrName() { return "res_attrs"; }
 
 /// Returns the dictionary attribute corresponding to the argument at 'index'.
 /// If there are no argument attributes at 'index', a null attribute is
 /// returned.
-inline DictionaryAttr getArgAttrDict(Operation *op, unsigned index) {
-  SmallString<8> nameOut;
-  return op->getAttrOfType<DictionaryAttr>(getArgAttrName(index, nameOut));
-}
+DictionaryAttr getArgAttrDict(Operation *op, unsigned index);
 
 /// Returns the dictionary attribute corresponding to the result at 'index'.
 /// If there are no result attributes at 'index', a null attribute is
 /// returned.
-inline DictionaryAttr getResultAttrDict(Operation *op, unsigned index) {
-  SmallString<8> nameOut;
-  return op->getAttrOfType<DictionaryAttr>(getResultAttrName(index, nameOut));
-}
+DictionaryAttr getResultAttrDict(Operation *op, unsigned index);
+
+namespace detail {
+/// Update the given index into an argument or result attribute dictionary.
+void setArgResAttrDict(Operation *op, StringRef attrName,
+                       unsigned numTotalIndices, unsigned index,
+                       DictionaryAttr attrs);
+} // namespace detail
+
+/// Set all of the argument or result attribute dictionaries for a function. The
+/// size of `attrs` is expected to match the number of arguments/results of the
+/// given `op`.
+void setAllArgAttrDicts(Operation *op, ArrayRef<DictionaryAttr> attrs);
+void setAllArgAttrDicts(Operation *op, ArrayRef<Attribute> attrs);
+void setAllResultAttrDicts(Operation *op, ArrayRef<DictionaryAttr> attrs);
+void setAllResultAttrDicts(Operation *op, ArrayRef<Attribute> attrs);
 
 /// Return all of the attributes for the argument at 'index'.
 inline ArrayRef<NamedAttribute> getArgAttrs(Operation *op, unsigned index) {
   auto argDict = getArgAttrDict(op, index);
   return argDict ? argDict.getValue() : llvm::None;
 }
 
 /// Return all of the attributes for the result at 'index'.

 
 /// Get and set a FunctionLike operation's type signature.
 FunctionType getFunctionType(Operation *op);
 void setFunctionType(Operation *op, FunctionType newType);
 
 /// Get a FunctionLike operation's body.
 Region &getFunctionBody(Operation *op);
 
-} // namespace impl
+} // namespace function_like_impl
 
 namespace OpTrait {
 
 /// This trait provides APIs for Ops that behave like functions.  In particular:
 /// - 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;
public:
   //===--------------------------------------------------------------------===//
   // Body Handling
   //===--------------------------------------------------------------------===//
 
   /// Returns true if this function is external, i.e. it has no body.
   bool isExternal() { return empty(); }
 
   Region &getBody() {
-    return ::mlir::impl::getFunctionBody(this->getOperation());
+    return function_like_impl::getFunctionBody(this->getOperation());
   }
 
   /// Delete all blocks from this function.
   void eraseBody() {
     getBody().dropAllReferences();
     getBody().getBlocks().clear();
   }
 
public:
   /// region exists.
   LogicalResult verifyBody();
 
   //===--------------------------------------------------------------------===//
   // Type Attribute Handling
   //===--------------------------------------------------------------------===//
 
   /// Return the name of the attribute used for function types.
-  static StringRef getTypeAttrName() { return ::mlir::impl::getTypeAttrName(); }
+  static StringRef getTypeAttrName() {
+    return function_like_impl::getTypeAttrName();
+  }
 
   TypeAttr getTypeAttr() {
     return this->getOperation()->template getAttrOfType<TypeAttr>(
         getTypeAttrName());
   }
 
   /// Return the type of this function.
   ///
   /// 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 ::mlir::impl::getFunctionType(this->getOperation());
+    return function_like_impl::getFunctionType(this->getOperation());
   }
 
   /// 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
public:
   /// 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) {
     unsigned originalNumArgs = getNumArguments();
     Type newType = getTypeWithoutArgsAndResults(argIndices, {});
-    ::mlir::impl::eraseFunctionArguments(this->getOperation(), argIndices,
+    function_like_impl::eraseFunctionArguments(this->getOperation(), argIndices,
                                                originalNumArgs, newType);
   }
 
   /// 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) {
     unsigned originalNumResults = getNumResults();
     Type newType = getTypeWithoutArgsAndResults({}, resultIndices);
-    ::mlir::impl::eraseFunctionResults(this->getOperation(), resultIndices,
-                                       originalNumResults, newType);
+    function_like_impl::eraseFunctionResults(
+        this->getOperation(), resultIndices, originalNumResults, newType);
   }
 
   //===--------------------------------------------------------------------===//
   // 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 function_like_impl::getArgAttrs(this->getOperation(), index);
   }
 
-  /// Return all argument attributes of this function. If an argument does not
-  /// have any attributes, the corresponding entry in `result` is nullptr.
+  /// Return an ArrayAttr containing all argument attribute dictionaries of this
+  /// function, or nullptr if no arguments have attributes.
+  ArrayAttr getAllArgAttrs() {
+    return this->getOperation()->template getAttrOfType<ArrayAttr>(
+        function_like_impl::getArgDictAttrName());
+  }
+  /// Return all argument attributes of this function.
   void getAllArgAttrs(SmallVectorImpl<DictionaryAttr> &result) {
-    for (unsigned i = 0, e = getNumArguments(); i != e; ++i)
-      result.emplace_back(getArgAttrDict(i));
+    if (ArrayAttr argAttrs = getAllArgAttrs()) {
+      auto argAttrRange = argAttrs.template getAsRange<DictionaryAttr>();
+      result.append(argAttrRange.begin(), argAttrRange.end());
+    } else {
+      result.resize(getNumArguments());
+    }
   }
 
   /// Return the specified attribute, if present, for the argument at 'index',
   /// null otherwise.
   Attribute getArgAttr(unsigned index, Identifier name) {
     auto argDict = getArgAttrDict(index);
     return argDict ? argDict.get(name) : nullptr;
   }
public:
   }
 
   /// Set the attributes held by the argument at 'index'.
   void setArgAttrs(unsigned index, ArrayRef<NamedAttribute> attributes);
 
   /// Set the attributes held by the argument at 'index'. `attributes` may be
   /// null, in which case any existing argument attributes are removed.
   void setArgAttrs(unsigned index, DictionaryAttr attributes);
-  void setAllArgAttrs(ArrayRef<DictionaryAttr> attributes);
+  void setAllArgAttrs(ArrayRef<DictionaryAttr> attributes) {
+    assert(attributes.size() == getNumArguments());
+    function_like_impl::setAllArgAttrDicts(this->getOperation(), attributes);
+  }
+  void setAllArgAttrs(ArrayRef<Attribute> attributes) {
+    assert(attributes.size() == getNumArguments());
+    function_like_impl::setAllArgAttrDicts(this->getOperation(), attributes);
+  }
+  void setAllArgAttrs(ArrayAttr attributes) {
+    assert(attributes.size() == getNumArguments());
+    this->getOperation()->setAttr(function_like_impl::getArgDictAttrName(),
+                                  attributes);
+  }
 
   /// If the an attribute exists with the specified name, change it to the new
   /// value. Otherwise, add a new attribute with the specified name/value.
   void setArgAttr(unsigned index, Identifier name, Attribute value);
   void setArgAttr(unsigned index, StringRef name, Attribute value) {
     setArgAttr(index, Identifier::get(name, this->getOperation()->getContext()),
                value);
   }
public:
   /// respective function results. These result attributes are stored as
   /// DictionaryAttrs in the main operation attribute dictionary. The name of
   /// these entries is `result` followed by the index of the result. These
   /// result attribute dictionaries are optional, and will generally only
   /// exist if they are non-empty.
 
   /// Return all of the attributes for the result at 'index'.
   ArrayRef<NamedAttribute> getResultAttrs(unsigned index) {
-    return ::mlir::impl::getResultAttrs(this->getOperation(), index);
+    return function_like_impl::getResultAttrs(this->getOperation(), index);
   }
 
-  /// Return all result attributes of this function. If a result does not have
-  /// any attributes, the corresponding entry in `result` is nullptr.
+  /// Return an ArrayAttr containing all result attribute dictionaries of this
+  /// function, or nullptr if no result have attributes.
+  ArrayAttr getAllResultAttrs() {
+    return this->getOperation()->template getAttrOfType<ArrayAttr>(
+        function_like_impl::getResultDictAttrName());
+  }
+  /// Return all result attributes of this function.
   void getAllResultAttrs(SmallVectorImpl<DictionaryAttr> &result) {
-    for (unsigned i = 0, e = getNumResults(); i != e; ++i)
-      result.emplace_back(getResultAttrDict(i));
+    if (ArrayAttr argAttrs = getAllResultAttrs()) {
+      auto argAttrRange = argAttrs.template getAsRange<DictionaryAttr>();
+      result.append(argAttrRange.begin(), argAttrRange.end());
+    } else {
+      result.resize(getNumResults());
+    }
   }
 
   /// Return the specified attribute, if present, for the result at 'index',
   /// null otherwise.
   Attribute getResultAttr(unsigned index, Identifier name) {
     auto argDict = getResultAttrDict(index);
     return argDict ? argDict.get(name) : nullptr;
   }
   Attribute getResultAttr(unsigned index, StringRef name) {
     auto argDict = getResultAttrDict(index);
     return argDict ? argDict.get(name) : nullptr;
   }
 
   template <typename AttrClass>
   AttrClass getResultAttrOfType(unsigned index, Identifier name) {
     return getResultAttr(index, name).template dyn_cast_or_null<AttrClass>();
   }
   template <typename AttrClass>
   AttrClass getResultAttrOfType(unsigned index, StringRef name) {
     return getResultAttr(index, name).template dyn_cast_or_null<AttrClass>();
   }
 
   /// Set the attributes held by the result at 'index'.
   void setResultAttrs(unsigned index, ArrayRef<NamedAttribute> attributes);
+
   /// Set the attributes held by the result at 'index'. `attributes` may be
   /// null, in which case any existing argument attributes are removed.
   void setResultAttrs(unsigned index, DictionaryAttr attributes);
-  void setAllResultAttrs(ArrayRef<DictionaryAttr> attributes);
+  void setAllResultAttrs(ArrayRef<DictionaryAttr> attributes) {
+    assert(attributes.size() == getNumResults());
+    function_like_impl::setAllResultAttrDicts(this->getOperation(), attributes);
+  }
+  void setAllResultAttrs(ArrayRef<Attribute> attributes) {
+    assert(attributes.size() == getNumResults());
+    function_like_impl::setAllResultAttrDicts(this->getOperation(), attributes);
+  }
+  void setAllResultAttrs(ArrayAttr attributes) {
+    assert(attributes.size() == getNumResults());
+    this->getOperation()->setAttr(function_like_impl::getResultDictAttrName(),
+                                  attributes);
+  }
 
   /// If the an attribute exists with the specified name, change it to the new
   /// value. Otherwise, add a new attribute with the specified name/value.
   void setResultAttr(unsigned index, Identifier name, Attribute value);
   void setResultAttr(unsigned index, StringRef name, Attribute value) {
     setResultAttr(index,
                   Identifier::get(name, this->getOperation()->getContext()),
                   value);
   }
 
   /// Remove the attribute 'name' from the result at 'index'. Return the
   /// attribute that was erased, or nullptr if there was no attribute with such
   /// name.
   Attribute removeResultAttr(unsigned index, Identifier name);
 
 protected:
-  /// Returns the attribute entry name for the set of argument attributes at
-  /// 'index'.
-  static StringRef getArgAttrName(unsigned index, SmallVectorImpl<char> &out) {
-    return ::mlir::impl::getArgAttrName(index, out);
-  }
-
   /// Returns the dictionary attribute corresponding to the argument at 'index'.
   /// If there are no argument attributes at 'index', a null attribute is
   /// returned.
   DictionaryAttr getArgAttrDict(unsigned index) {
     assert(index < getNumArguments() && "invalid argument number");
-    return ::mlir::impl::getArgAttrDict(this->getOperation(), index);
-  }
-
-  /// Returns the attribute entry name for the set of result attributes at
-  /// 'index'.
-  static StringRef getResultAttrName(unsigned index,
-                                     SmallVectorImpl<char> &out) {
-    return ::mlir::impl::getResultAttrName(index, out);
+    return function_like_impl::getArgAttrDict(this->getOperation(), index);
   }
 
   /// Returns the dictionary attribute corresponding to the result at 'index'.
   /// If there are no result attributes at 'index', a null attribute is
   /// returned.
   DictionaryAttr getResultAttrDict(unsigned index) {
     assert(index < getNumResults() && "invalid result number");
-    return ::mlir::impl::getResultAttrDict(this->getOperation(), index);
+    return function_like_impl::getResultAttrDict(this->getOperation(), index);
   }
 
   /// Hook for concrete classes to verify that the type attribute respects
   /// op-specific invariants.  Default implementation always succeeds.
   LogicalResult verifyType() { return success(); }
 };
 
 /// Default verifier checks that if the entry block exists, it has the same
if (funcOp.front().getNumArguments() != numArguments)
     return funcOp.emitOpError("entry block must have ")
            << numArguments << " arguments to match function signature";
 
   return success();
 }
 
 template <typename ConcreteType>
 LogicalResult FunctionLike<ConcreteType>::verifyTrait(Operation *op) {
-  MLIRContext *ctx = op->getContext();
   auto funcOp = cast<ConcreteType>(op);
-
   if (!funcOp.isTypeAttrValid())
     return funcOp.emitOpError("requires a type attribute '")
            << getTypeAttrName() << '\'';
 
   if (failed(funcOp.verifyType()))
     return failure();
 
-  for (unsigned i = 0, e = funcOp.getNumArguments(); i != e; ++i) {
+  if (ArrayAttr allArgAttrs = funcOp.getAllArgAttrs()) {
+    unsigned numArgs = funcOp.getNumArguments();
+    if (allArgAttrs.size() != numArgs) {
+      return funcOp.emitOpError()
+             << "expects argument attribute array `"
+             << function_like_impl::getArgDictAttrName()
+             << "` to have the same number of elements as the number of "
+                "function arguments, got "
+             << allArgAttrs.size() << ", but expected " << numArgs;
+    }
+    for (unsigned i = 0; i != numArgs; ++i) {
+      DictionaryAttr argAttrs = allArgAttrs[i].dyn_cast<DictionaryAttr>();
+      if (!argAttrs) {
+        return funcOp.emitOpError() << "expects argument attribute dictionary "
+                                       "to be a DictionaryAttr, but got `"
+                                    << allArgAttrs[i] << "`";
+      }
+
       // Verify that all of the argument attributes are dialect attributes, i.e.
       // that they contain a dialect prefix in their name.  Call the dialect, if
       // registered, to verify the attributes themselves.
-    for (auto attr : funcOp.getArgAttrs(i)) {
+      for (auto attr : argAttrs) {
         if (!attr.first.strref().contains('.'))
-        return funcOp.emitOpError("arguments may only have dialect attributes");
-      auto dialectNamePair = attr.first.strref().split('.');
-      if (auto *dialect = ctx->getLoadedDialect(dialectNamePair.first)) {
+          return funcOp.emitOpError(
+              "arguments may only have dialect attributes");
+        if (Dialect *dialect = attr.first.getDialect()) {
           if (failed(dialect->verifyRegionArgAttribute(op, /*regionIndex=*/0,
                                                        /*argIndex=*/i, attr)))
             return failure();
         }
       }
     }
+  }
+  if (ArrayAttr allResultAttrs = funcOp.getAllResultAttrs()) {
+    unsigned numResults = funcOp.getNumResults();
+    if (allResultAttrs.size() != numResults) {
+      return funcOp.emitOpError()
+             << "expects result attribute array `"
+             << function_like_impl::getResultDictAttrName()
+             << "` to have the same number of elements as the number of "
+                "function results, got "
+             << allResultAttrs.size() << ", but expected " << numResults;
+    }
+    for (unsigned i = 0; i != numResults; ++i) {
+      DictionaryAttr resultAttrs = allResultAttrs[i].dyn_cast<DictionaryAttr>();
+      if (!resultAttrs) {
+        return funcOp.emitOpError() << "expects result attribute dictionary "
+                                       "to be a DictionaryAttr, but got `"
+                                    << allResultAttrs[i] << "`";
+      }
 
-  for (unsigned i = 0, e = funcOp.getNumResults(); i != e; ++i) {
       // Verify that all of the result attributes are dialect attributes, i.e.
       // that they contain a dialect prefix in their name.  Call the dialect, if
       // registered, to verify the attributes themselves.
-    for (auto attr : funcOp.getResultAttrs(i)) {
+      for (auto attr : resultAttrs) {
         if (!attr.first.strref().contains('.'))
           return funcOp.emitOpError("results may only have dialect attributes");
-      auto dialectNamePair = attr.first.strref().split('.');
-      if (auto *dialect = ctx->getLoadedDialect(dialectNamePair.first)) {
+        if (Dialect *dialect = attr.first.getDialect()) {
           if (failed(dialect->verifyRegionResultAttribute(op, /*regionIndex=*/0,
                                                           /*resultIndex=*/i,
                                                           attr)))
             return failure();
         }
       }
     }
+  }
 
   // Check that the op has exactly one region for the body.
   if (op->getNumRegions() != 1)
     return funcOp.emitOpError("expects one region");
 
   return funcOp.verifyBody();
 }
 

 }
 
 //===----------------------------------------------------------------------===//
 // Function Type Attribute.
 //===----------------------------------------------------------------------===//
 
 template <typename ConcreteType>
 void FunctionLike<ConcreteType>::setType(FunctionType newType) {
-  ::mlir::impl::setFunctionType(this->getOperation(), newType);
+  function_like_impl::setFunctionType(this->getOperation(), newType);
 }
 
 //===----------------------------------------------------------------------===//
 // 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");
-  SmallString<8> nameOut;
-  getArgAttrName(index, nameOut);
-
   Operation *op = this->getOperation();
-  if (attributes.empty())
-    return (void)op->removeAttr(nameOut);
-  op->setAttr(nameOut, DictionaryAttr::get(op->getContext(), attributes));
+  return function_like_impl::detail::setArgResAttrDict(
+      op, function_like_impl::getArgDictAttrName(), getNumArguments(), index,
+      DictionaryAttr::get(op->getContext(), attributes));
 }
 
 template <typename ConcreteType>
 void FunctionLike<ConcreteType>::setArgAttrs(unsigned index,
                                              DictionaryAttr attributes) {
-  assert(index < getNumArguments() && "invalid argument number");
-  SmallString<8> nameOut;
-  if (!attributes || attributes.empty())
-    this->getOperation()->removeAttr(getArgAttrName(index, nameOut));
-  else
-    return this->getOperation()->setAttr(getArgAttrName(index, nameOut),
-                                         attributes);
-}
-
-template <typename ConcreteType>
-void FunctionLike<ConcreteType>::setAllArgAttrs(
-    ArrayRef<DictionaryAttr> attributes) {
-  assert(attributes.size() == getNumArguments());
-  NamedAttrList attrs = this->getOperation()->getAttrs();
-
-  // Instead of calling setArgAttrs() multiple times, which rebuild the
-  // attribute dictionary every time, build a new list of attributes for the
-  // operation so that we rebuild the attribute dictionary in one shot.
-  SmallString<8> argAttrName;
-  for (unsigned i = 0, e = attributes.size(); i != e; ++i) {
-    StringRef attrName = getArgAttrName(i, argAttrName);
-    if (!attributes[i] || attributes[i].empty())
-      attrs.erase(attrName);
-    else
-      attrs.set(attrName, attributes[i]);
-  }
-  this->getOperation()->setAttrs(attrs);
+  Operation *op = this->getOperation();
+  return function_like_impl::detail::setArgResAttrDict(
+      op, function_like_impl::getArgDictAttrName(), getNumArguments(), index,
+      attributes ? attributes : DictionaryAttr::get(op->getContext()));
 }
 
 /// If the an attribute exists with the specified name, change it to the new
 /// value. Otherwise, add a new attribute with the specified name/value.
 template <typename ConcreteType>
 void FunctionLike<ConcreteType>::setArgAttr(unsigned index, Identifier name,
                                             Attribute value) {
   NamedAttrList attributes(getArgAttrDict(index));

 // Function Result Attribute.
 //===----------------------------------------------------------------------===//
 
 /// Set the attributes held by the result at 'index'.
 template <typename ConcreteType>
 void FunctionLike<ConcreteType>::setResultAttrs(
     unsigned index, ArrayRef<NamedAttribute> attributes) {
   assert(index < getNumResults() && "invalid result number");
-  SmallString<8> nameOut;
-  getResultAttrName(index, nameOut);
-
-  if (attributes.empty())
-    return (void)this->getOperation()->removeAttr(nameOut);
   Operation *op = this->getOperation();
-  op->setAttr(nameOut, DictionaryAttr::get(op->getContext(), attributes));
+  return function_like_impl::detail::setArgResAttrDict(
+      op, function_like_impl::getResultDictAttrName(), getNumResults(), index,
+      DictionaryAttr::get(op->getContext(), attributes));
 }
 
 template <typename ConcreteType>
 void FunctionLike<ConcreteType>::setResultAttrs(unsigned index,
                                                 DictionaryAttr attributes) {
   assert(index < getNumResults() && "invalid result number");
-  SmallString<8> nameOut;
-  if (!attributes || attributes.empty())
-    this->getOperation()->removeAttr(getResultAttrName(index, nameOut));
-  else
-    this->getOperation()->setAttr(getResultAttrName(index, nameOut),
-                                  attributes);
-}
-
-template <typename ConcreteType>
-void FunctionLike<ConcreteType>::setAllResultAttrs(
-    ArrayRef<DictionaryAttr> attributes) {
-  assert(attributes.size() == getNumResults());
-  NamedAttrList attrs = this->getOperation()->getAttrs();
-
-  // Instead of calling setResultAttrs() multiple times, which rebuild the
-  // attribute dictionary every time, build a new list of attributes for the
-  // operation so that we rebuild the attribute dictionary in one shot.
-  SmallString<8> resultAttrName;
-  for (unsigned i = 0, e = attributes.size(); i != e; ++i) {
-    StringRef attrName = getResultAttrName(i, resultAttrName);
-    if (!attributes[i] || attributes[i].empty())
-      attrs.erase(attrName);
-    else
-      attrs.set(attrName, attributes[i]);
-  }
-  this->getOperation()->setAttrs(attrs);
+  Operation *op = this->getOperation();
+  return function_like_impl::detail::setArgResAttrDict(
+      op, function_like_impl::getResultDictAttrName(), getNumResults(), index,
+      attributes ? attributes : DictionaryAttr::get(op->getContext()));
 }
 
 /// If the an attribute exists with the specified name, change it to the new
 /// value. Otherwise, add a new attribute with the specified name/value.
 template <typename ConcreteType>
 void FunctionLike<ConcreteType>::setResultAttr(unsigned index, Identifier name,
                                                Attribute value) {
   NamedAttrList attributes(getResultAttrDict(index));

mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp

GPUFuncOpLowering::matchAndRewrite(gpu::GPUFuncOp gpuFuncOp,
   getTypeConverter()->convertFunctionSignature(
       gpuFuncOp.getType(), /*isVariadic=*/false, signatureConversion);
 
   // Create the new function operation. Only copy those attributes that are
   // not specific to function modeling.
   SmallVector<NamedAttribute, 4> attributes;
   for (const auto &attr : gpuFuncOp->getAttrs()) {
     if (attr.first == SymbolTable::getSymbolAttrName() ||
-        attr.first == impl::getTypeAttrName() ||
+        attr.first == function_like_impl::getTypeAttrName() ||
         attr.first == gpu::GPUFuncOp::getNumWorkgroupAttributionsAttrName())
       continue;
     attributes.push_back(attr);
   }
   // Add a dialect specific kernel attribute in addition to GPU kernel
   // attribute. The former is necessary for further translation while the
   // latter is expected by gpu.launch_func.
   if (gpuFuncOp.isKernel())

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp

for (auto argType : enumerate(funcOp.getType().getInputs())) {
       signatureConverter.addInputs(argType.index(), convertedType);
     }
   }
   auto newFuncOp = rewriter.create<spirv::FuncOp>(
       funcOp.getLoc(), funcOp.getName(),
       rewriter.getFunctionType(signatureConverter.getConvertedTypes(),
                                llvm::None));
   for (const auto &namedAttr : funcOp->getAttrs()) {
-    if (namedAttr.first == impl::getTypeAttrName() ||
+    if (namedAttr.first == function_like_impl::getTypeAttrName() ||
         namedAttr.first == SymbolTable::getSymbolAttrName())
       continue;
     newFuncOp->setAttr(namedAttr.first, namedAttr.second);
   }
 
   rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(),
                               newFuncOp.end());
   if (failed(rewriter.convertRegionTypes(&newFuncOp.getBody(), typeConverter,

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 /// Only retain those attributes that are not constructed by
 /// `LLVMFuncOp::build`. If `filterArgAttrs` is set, also filter out argument
 /// attributes.
 static void filterFuncAttributes(ArrayRef<NamedAttribute> attrs,
                                  bool filterArgAttrs,
                                  SmallVectorImpl<NamedAttribute> &result) {
   for (const auto &attr : attrs) {
     if (attr.first == SymbolTable::getSymbolAttrName() ||
-        attr.first == impl::getTypeAttrName() || attr.first == "std.varargs" ||
-        (filterArgAttrs && impl::isArgAttrName(attr.first.strref())))
+        attr.first == function_like_impl::getTypeAttrName() ||
+        attr.first == "std.varargs" ||
+        (filterArgAttrs &&
+         attr.first == function_like_impl::getArgDictAttrName()))
       continue;
     result.push_back(attr);
   }
 }
 
 /// Creates an auxiliary function with pointer-to-memref-descriptor-struct
 /// arguments instead of unpacked arguments. This function can be called from C
 /// by passing a pointer to a C struct corresponding to a memref descriptor.
convertFuncOpToLLVMFuncOp(FuncOp funcOp,
     if (!llvmType)
       return nullptr;
 
     // Propagate argument attributes to all converted arguments obtained after
     // converting a given original argument.
     SmallVector<NamedAttribute, 4> attributes;
     filterFuncAttributes(funcOp->getAttrs(), /*filterArgAttrs=*/true,
                          attributes);
+    if (ArrayAttr argAttrDicts = funcOp.getAllArgAttrs()) {
+      SmallVector<Attribute, 4> newArgAttrs(
+          llvmType.cast<LLVM::LLVMFunctionType>().getNumParams());
       for (unsigned i = 0, e = funcOp.getNumArguments(); i < e; ++i) {
-      auto attr = impl::getArgAttrDict(funcOp, i);
-      if (!attr)
-        continue;
-
         auto mapping = result.getInputMapping(i);
-      assert(mapping.hasValue() && "unexpected deletion of function argument");
-
-      SmallString<8> name;
-      for (size_t j = 0; j < mapping->size; ++j) {
-        impl::getArgAttrName(mapping->inputNo + j, name);
-        attributes.push_back(rewriter.getNamedAttr(name, attr));
-      }
+        assert(mapping.hasValue() &&
+               "unexpected deletion of function argument");
+        for (size_t j = 0; j < mapping->size; ++j)
+          newArgAttrs[mapping->inputNo + j] = argAttrDicts[i];
+      }
+      attributes.push_back(
+          rewriter.getNamedAttr(function_like_impl::getArgDictAttrName(),
+                                rewriter.getArrayAttr(newArgAttrs)));
     }
 
     // Create an LLVM function, use external linkage by default until MLIR
     // functions have linkage.
     auto newFuncOp = rewriter.create<LLVM::LLVMFuncOp>(
         funcOp.getLoc(), funcOp.getName(), llvmType, LLVM::Linkage::External,
         attributes);
     rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(),

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

 static ParseResult
 parseLaunchFuncOperands(OpAsmParser &parser,
                         SmallVectorImpl<OpAsmParser::OperandType> &argNames,
                         SmallVectorImpl<Type> &argTypes) {
   if (parser.parseOptionalKeyword("args"))
     return success();
   SmallVector<NamedAttrList, 4> argAttrs;
   bool isVariadic = false;
-  return impl::parseFunctionArgumentList(parser, /*allowAttributes=*/false,
-                                         /*allowVariadic=*/false, argNames,
-                                         argTypes, argAttrs, isVariadic);
+  return function_like_impl::parseFunctionArgumentList(
+      parser, /*allowAttributes=*/false,
+      /*allowVariadic=*/false, argNames, argTypes, argAttrs, isVariadic);
 }
 
 static void printLaunchFuncOperands(OpAsmPrinter &printer, Operation *,
                                     OperandRange operands, TypeRange types) {
   if (operands.empty())
     return;
   printer << "args(";
   llvm::interleaveComma(llvm::zip(operands, types), printer,
static ParseResult parseGPUFuncOp(OpAsmParser &parser, OperationState &result) {
 
   // Parse the function name.
   StringAttr nameAttr;
   if (parser.parseSymbolName(nameAttr, ::mlir::SymbolTable::getSymbolAttrName(),
                              result.attributes))
     return failure();
 
   auto signatureLocation = parser.getCurrentLocation();
-  if (failed(impl::parseFunctionSignature(
+  if (failed(function_like_impl::parseFunctionSignature(
           parser, /*allowVariadic=*/false, entryArgs, argTypes, argAttrs,
           isVariadic, resultTypes, resultAttrs)))
     return failure();
 
   if (entryArgs.empty() && !argTypes.empty())
     return parser.emitError(signatureLocation)
            << "gpu.func requires named arguments";
 
static ParseResult parseGPUFuncOp(OpAsmParser &parser, OperationState &result) {
   // Parse the kernel attribute if present.
   if (succeeded(parser.parseOptionalKeyword(GPUFuncOp::getKernelKeyword())))
     result.addAttribute(GPUDialect::getKernelFuncAttrName(),
                         builder.getUnitAttr());
 
   // Parse attributes.
   if (failed(parser.parseOptionalAttrDictWithKeyword(result.attributes)))
     return failure();
-  mlir::impl::addArgAndResultAttrs(builder, result, argAttrs, resultAttrs);
+  function_like_impl::addArgAndResultAttrs(builder, result, argAttrs,
+                                           resultAttrs);
 
   // Parse the region. If no argument names were provided, take all names
   // (including those of attributions) from the entry block.
   auto *body = result.addRegion();
   return parser.parseRegion(*body, entryArgs, argTypes);
 }
 
 static void printAttributions(OpAsmPrinter &p, StringRef keyword,
                               ArrayRef<BlockArgument> values) {
   if (values.empty())
     return;
 
   p << ' ' << keyword << '(';
   llvm::interleaveComma(
       values, p, [&p](BlockArgument v) { p << v << " : " << v.getType(); });
   p << ')';
 }
 
 /// Prints a GPU Func op.
 static void printGPUFuncOp(OpAsmPrinter &p, GPUFuncOp op) {
   p << GPUFuncOp::getOperationName() << ' ';
   p.printSymbolName(op.getName());
 
   FunctionType type = op.getType();
-  impl::printFunctionSignature(p, op.getOperation(), type.getInputs(),
-                               /*isVariadic=*/false, type.getResults());
+  function_like_impl::printFunctionSignature(
+      p, op.getOperation(), type.getInputs(),
 
   printAttributions(p, op.getWorkgroupKeyword(), op.getWorkgroupAttributions());
   printAttributions(p, op.getPrivateKeyword(), op.getPrivateAttributions());
   if (op.isKernel())
     p << ' ' << op.getKernelKeyword();
 
-  impl::printFunctionAttributes(p, op.getOperation(), type.getNumInputs(),
-                                type.getNumResults(),
-                                {op.getNumWorkgroupAttributionsAttrName(),
-                                 GPUDialect::getKernelFuncAttrName()});
+  function_like_impl::printFunctionAttributes(
+      p, op.getOperation(), type.getNumInputs(), type.getNumResults(),
   p.printRegion(op.getBody(), /*printEntryBlockArgs=*/false);
 }
 
-void GPUFuncOp::setType(FunctionType newType) {
-  auto oldType = getType();
-  assert(newType.getNumResults() == oldType.getNumResults() &&
-         "unimplemented: changes to the number of results");
-
-  SmallVector<char, 16> nameBuf;
-  for (int i = newType.getNumInputs(), e = oldType.getNumInputs(); i < e; i++)
-    (*this)->removeAttr(getArgAttrName(i, nameBuf));
-
-  (*this)->setAttr(getTypeAttrName(), TypeAttr::get(newType));
-}
-
 /// Hook for FunctionLike verifier.
 LogicalResult GPUFuncOp::verifyType() {
   Type type = getTypeAttr().getValue();
   if (!type.isa<FunctionType>())
     return emitOpError("requires '" + getTypeAttrName() +
                        "' attribute of function type");
 
   if (isKernel() && getType().getNumResults() != 0)

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

result.addAttribute(SymbolTable::getSymbolAttrName(),
                       builder.getStringAttr(name));
   result.addAttribute("type", TypeAttr::get(type));
   result.addAttribute(getLinkageAttrName(),
                       builder.getI64IntegerAttr(static_cast<int64_t>(linkage)));
   result.attributes.append(attrs.begin(), attrs.end());
   if (argAttrs.empty())
     return;
 
-  unsigned numInputs = type.cast<LLVMFunctionType>().getNumParams();
-  assert(numInputs == argAttrs.size() &&
+  assert(type.cast<LLVMFunctionType>().getNumParams() == argAttrs.size() &&
          "expected as many argument attribute lists as arguments");
-  SmallString<8> argAttrName;
-  for (unsigned i = 0; i < numInputs; ++i)
-    if (DictionaryAttr argDict = argAttrs[i])
-      result.addAttribute(getArgAttrName(i, argAttrName), argDict);
+  function_like_impl::addArgAndResultAttrs(builder, result, argAttrs,
+                                           /*resultAttrs=*/llvm::None);
 }
 
 // Builds an LLVM function type from the given lists of input and output types.
 // Returns a null type if any of the types provided are non-LLVM types, or if
 // there is more than one output type.
-static Type buildLLVMFunctionType(OpAsmParser &parser, llvm::SMLoc loc,
-                                  ArrayRef<Type> inputs, ArrayRef<Type> outputs,
-                                  impl::VariadicFlag variadicFlag) {
+static Type
+buildLLVMFunctionType(OpAsmParser &parser, llvm::SMLoc loc,
+                      function_like_impl::VariadicFlag variadicFlag) {
   Builder &b = parser.getBuilder();
   if (outputs.size() > 1) {
     parser.emitError(loc, "failed to construct function type: expected zero or "
                           "one function result");
     return {};
   }
 
   // Convert inputs to LLVM types, exit early on error.
static ParseResult parseLLVMFuncOp(OpAsmParser &parser,
   SmallVector<NamedAttrList, 1> resultAttrs;
   SmallVector<Type, 8> argTypes;
   SmallVector<Type, 4> resultTypes;
   bool isVariadic;
 
   auto signatureLocation = parser.getCurrentLocation();
   if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(),
                              result.attributes) ||
-      impl::parseFunctionSignature(parser, /*allowVariadic=*/true, entryArgs,
-                                   argTypes, argAttrs, isVariadic, resultTypes,
-                                   resultAttrs))
+      function_like_impl::parseFunctionSignature(
+          parser, /*allowVariadic=*/true, entryArgs, argTypes, argAttrs,
+          isVariadic, resultTypes, resultAttrs))
     return failure();
 
   auto type =
       buildLLVMFunctionType(parser, signatureLocation, argTypes, resultTypes,
-                            impl::VariadicFlag(isVariadic));
+                            function_like_impl::VariadicFlag(isVariadic));
   if (!type)
     return failure();
-  result.addAttribute(impl::getTypeAttrName(), TypeAttr::get(type));
+  result.addAttribute(function_like_impl::getTypeAttrName(),
+                      TypeAttr::get(type));
 
   if (failed(parser.parseOptionalAttrDictWithKeyword(result.attributes)))
     return failure();
-  impl::addArgAndResultAttrs(parser.getBuilder(), result, argAttrs,
-                             resultAttrs);
+  function_like_impl::addArgAndResultAttrs(parser.getBuilder(), result,
+                                           argAttrs, resultAttrs);
 
   auto *body = result.addRegion();
   OptionalParseResult parseResult = parser.parseOptionalRegion(
       *body, entryArgs, entryArgs.empty() ? ArrayRef<Type>() : argTypes);
   return failure(parseResult.hasValue() && failed(*parseResult));
 }
 
 // Print the LLVMFuncOp. Collects argument and result types and passes them to
static void printLLVMFuncOp(OpAsmPrinter &p, LLVMFuncOp op) {
   argTypes.reserve(fnType.getNumParams());
   for (unsigned i = 0, e = fnType.getNumParams(); i < e; ++i)
     argTypes.push_back(fnType.getParamType(i));
 
   Type returnType = fnType.getReturnType();
   if (!returnType.isa<LLVMVoidType>())
     resTypes.push_back(returnType);
 
-  impl::printFunctionSignature(p, op, argTypes, op.isVarArg(), resTypes);
-  impl::printFunctionAttributes(p, op, argTypes.size(), resTypes.size(),
-                                {getLinkageAttrName()});
+  function_like_impl::printFunctionSignature(p, op, argTypes, op.isVarArg(),
+                                             resTypes);
+  function_like_impl::printFunctionAttributes(
+      p, op, argTypes.size(), resTypes.size(), {getLinkageAttrName()});
 
   // Print the body if this is not an external function.
   Region &body = op.body();
   if (!body.empty())
     p.printRegion(body, /*printEntryBlockArgs=*/false,
                   /*printBlockTerminators=*/true);
 }
 

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

FunctionNonEntryBlockConversion(StringRef functionLikeOpName,
                                   DenseSet<BlockArgument> blockArgsToDetensor)
       : ConversionPattern(converter, functionLikeOpName, /*benefit=*/1, ctx),
         blockArgsToDetensor(blockArgsToDetensor) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.startRootUpdate(op);
-    Region &region = mlir::impl::getFunctionBody(op);
+    Region &region = function_like_impl::getFunctionBody(op);
     SmallVector<TypeConverter::SignatureConversion, 2> conversions;
 
     for (Block &block : llvm::drop_begin(region, 1)) {
       conversions.emplace_back(block.getNumArguments());
       TypeConverter::SignatureConversion &back = conversions.back();
 
       for (BlockArgument blockArgument : block.getArguments()) {
         int idx = blockArgument.getArgNumber();

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

static ParseResult parseFuncOp(OpAsmParser &parser, OperationState &state) {
   // Parse the name as a symbol.
   StringAttr nameAttr;
   if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(),
                              state.attributes))
     return failure();
 
   // Parse the function signature.
   bool isVariadic = false;
-  if (impl::parseFunctionSignature(parser, /*allowVariadic=*/false, entryArgs,
-                                   argTypes, argAttrs, isVariadic, resultTypes,
-                                   resultAttrs))
+  if (function_like_impl::parseFunctionSignature(
+          parser, /*allowVariadic=*/false, entryArgs, argTypes, argAttrs,
+          isVariadic, resultTypes, resultAttrs))
     return failure();
 
   auto fnType = builder.getFunctionType(argTypes, resultTypes);
-  state.addAttribute(impl::getTypeAttrName(), TypeAttr::get(fnType));
+  state.addAttribute(function_like_impl::getTypeAttrName(),
+                     TypeAttr::get(fnType));
 
   // Parse the optional function control keyword.
   spirv::FunctionControl fnControl;
   if (parseEnumStrAttr(fnControl, parser, state))
     return failure();
 
   // If additional attributes are present, parse them.
   if (parser.parseOptionalAttrDictWithKeyword(state.attributes))
     return failure();
 
   // Add the attributes to the function arguments.
   assert(argAttrs.size() == argTypes.size());
   assert(resultAttrs.size() == resultTypes.size());
-  impl::addArgAndResultAttrs(builder, state, argAttrs, resultAttrs);
+  function_like_impl::addArgAndResultAttrs(builder, state, argAttrs,
+                                           resultAttrs);
 
   // Parse the optional function body.
   auto *body = state.addRegion();
   OptionalParseResult result = parser.parseOptionalRegion(
       *body, entryArgs, entryArgs.empty() ? ArrayRef<Type>() : argTypes);
   return failure(result.hasValue() && failed(*result));
 }
 
 static void print(spirv::FuncOp fnOp, OpAsmPrinter &printer) {
   // Print function name, signature, and control.
   printer << spirv::FuncOp::getOperationName() << " ";
   printer.printSymbolName(fnOp.sym_name());
   auto fnType = fnOp.getType();
-  impl::printFunctionSignature(printer, fnOp, fnType.getInputs(),
-                               /*isVariadic=*/false, fnType.getResults());
+  function_like_impl::printFunctionSignature(printer, fnOp, fnType.getInputs(),
+                                             /*isVariadic=*/false,
+                                             fnType.getResults());
   printer << " \"" << spirv::stringifyFunctionControl(fnOp.function_control())
           << "\"";
-  impl::printFunctionAttributes(
+  function_like_impl::printFunctionAttributes(
       printer, fnOp, fnType.getNumInputs(), fnType.getNumResults(),
       {spirv::attributeName<spirv::FunctionControl>()});
 
   // Print the body if this is not an external function.
   Region &body = fnOp.body();
   if (!body.empty())
     printer.printRegion(body, /*printEntryBlockArgs=*/false,
                         /*printBlockTerminators=*/true);

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

FuncOpConversion::matchAndRewrite(FuncOp funcOp, ArrayRef<Value> operands,
   auto newFuncOp = rewriter.create<spirv::FuncOp>(
       funcOp.getLoc(), funcOp.getName(),
       rewriter.getFunctionType(signatureConverter.getConvertedTypes(),
                                resultType ? TypeRange(resultType)
                                           : TypeRange()));
 
   // Copy over all attributes other than the function name and type.
   for (const auto &namedAttr : funcOp->getAttrs()) {
-    if (namedAttr.first != impl::getTypeAttrName() &&
+    if (namedAttr.first != function_like_impl::getTypeAttrName() &&
         namedAttr.first != SymbolTable::getSymbolAttrName())
       newFuncOp->setAttr(namedAttr.first, namedAttr.second);
   }
 
   rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(),
                               newFuncOp.end());
   if (failed(rewriter.convertRegionTypes(
           &newFuncOp.getBody(), *getTypeConverter(), &signatureConverter)))

mlir/lib/IR/BuiltinDialect.cpp

state.addAttribute(SymbolTable::getSymbolAttrName(),
                      builder.getStringAttr(name));
   state.addAttribute(getTypeAttrName(), TypeAttr::get(type));
   state.attributes.append(attrs.begin(), attrs.end());
   state.addRegion();
 
   if (argAttrs.empty())
     return;
   assert(type.getNumInputs() == argAttrs.size());
-  SmallString<8> argAttrName;
-  for (unsigned i = 0, e = type.getNumInputs(); i != e; ++i)
-    if (DictionaryAttr argDict = argAttrs[i])
-      state.addAttribute(getArgAttrName(i, argAttrName), argDict);
+  function_like_impl::addArgAndResultAttrs(builder, state, argAttrs,
+                                           /*resultAttrs=*/llvm::None);
 }
 
 static ParseResult parseFuncOp(OpAsmParser &parser, OperationState &result) {
   auto buildFuncType = [](Builder &builder, ArrayRef<Type> argTypes,
-                          ArrayRef<Type> results, impl::VariadicFlag,
-                          std::string &) {
+                          ArrayRef<Type> results,
+                          function_like_impl::VariadicFlag, std::string &) {
     return builder.getFunctionType(argTypes, results);
   };
 
-  return impl::parseFunctionLikeOp(parser, result, /*allowVariadic=*/false,
-                                   buildFuncType);
+  return function_like_impl::parseFunctionLikeOp(
+      parser, result, /*allowVariadic=*/false, buildFuncType);
 }
 
 static void print(FuncOp op, OpAsmPrinter &p) {
   FunctionType fnType = op.getType();
-  impl::printFunctionLikeOp(p, op, fnType.getInputs(), /*isVariadic=*/false,
-                            fnType.getResults());
+  function_like_impl::printFunctionLikeOp(
+      p, op, fnType.getInputs(), /*isVariadic=*/false, fnType.getResults());
 }
 
 static LogicalResult verify(FuncOp op) {
   // If this function is external there is nothing to do.
   if (op.isExternal())
     return success();
 
   // Verify that the argument list of the function and the arg list of the entry

 }
 
 /// 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). Replaces references
 /// to cloned sub-values with the corresponding value that is copied, and adds
 /// those mappings to the mapper.
 FuncOp FuncOp::clone(BlockAndValueMapping &mapper) {
-  FunctionType newType = getType();
+  // Create the new function.
+  FuncOp newFunc = cast<FuncOp>(getOperation()->cloneWithoutRegions());
 
   // If the function has a body, then the user might be deleting arguments to
   // the function by specifying them in the mapper. If so, we don't add the
   // argument to the input type vector.
-  bool isExternalFn = isExternal();
-  if (!isExternalFn) {
-    SmallVector<Type, 4> inputTypes;
-    inputTypes.reserve(newType.getNumInputs());
-    for (unsigned i = 0, e = getNumArguments(); i != e; ++i)
-      if (!mapper.contains(getArgument(i)))
-        inputTypes.push_back(newType.getInput(i));
-    newType = FunctionType::get(getContext(), inputTypes, newType.getResults());
-  }
+  if (!isExternal()) {
+    FunctionType oldType = getType();
 
-  // Create the new function.
-  FuncOp newFunc = cast<FuncOp>(getOperation()->cloneWithoutRegions());
-  newFunc.setType(newType);
+    unsigned oldNumArgs = oldType.getNumInputs();
+    SmallVector<Type, 4> newInputs;
+    newInputs.reserve(oldNumArgs);
+    for (unsigned i = 0; i != oldNumArgs; ++i)
+      if (!mapper.contains(getArgument(i)))
+        newInputs.push_back(oldType.getInput(i));
 
-  /// Set the argument attributes for arguments that aren't being replaced.
-  for (unsigned i = 0, e = getNumArguments(), destI = 0; i != e; ++i)
-    if (isExternalFn || !mapper.contains(getArgument(i)))
-      newFunc.setArgAttrs(destI++, getArgAttrs(i));
+    /// If any of the arguments were dropped, update the type and drop any
+    /// necessary argument attributes.
+    if (newInputs.size() != oldNumArgs) {
+      newFunc.setType(FunctionType::get(oldType.getContext(), newInputs,
+                                        oldType.getResults()));
+
+      if (ArrayAttr argAttrs = getAllArgAttrs()) {
+        SmallVector<Attribute> newArgAttrs;
+        newArgAttrs.reserve(newInputs.size());
+        for (unsigned i = 0; i != oldNumArgs; ++i)
+          if (!mapper.contains(getArgument(i)))
+            newArgAttrs.push_back(argAttrs[i]);
+        newFunc.setAllArgAttrs(newArgAttrs);
+      }
+    }
+  }
 
   /// Clone the current function into the new one and return it.
   cloneInto(newFunc, mapper);
   return newFunc;
 }
 FuncOp FuncOp::clone() {
   BlockAndValueMapping mapper;
   return clone(mapper);

mlir/lib/IR/FunctionImplementation.cpp

 //===- FunctionImplementation.cpp - Utilities 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/FunctionImplementation.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/FunctionSupport.h"
 #include "mlir/IR/SymbolTable.h"
 
 using namespace mlir;
 
-ParseResult mlir::impl::parseFunctionArgumentList(
+ParseResult mlir::function_like_impl::parseFunctionArgumentList(
     OpAsmParser &parser, bool allowAttributes, bool allowVariadic,
     SmallVectorImpl<OpAsmParser::OperandType> &argNames,
     SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
     bool &isVariadic) {
   if (parser.parseLParen())
     return failure();
 
   // The argument list either has to consistently have ssa-id's followed by
parseFunctionResultList(OpAsmParser &parser, SmallVectorImpl<Type> &resultTypes,
   } while (succeeded(parser.parseOptionalComma()));
   return parser.parseRParen();
 }
 
 /// Parses a function signature using `parser`. The `allowVariadic` argument
 /// indicates whether functions with variadic arguments are supported. The
 /// trailing arguments are populated by this function with names, types and
 /// attributes of the arguments and those of the results.
-ParseResult mlir::impl::parseFunctionSignature(
+ParseResult mlir::function_like_impl::parseFunctionSignature(
     OpAsmParser &parser, bool allowVariadic,
     SmallVectorImpl<OpAsmParser::OperandType> &argNames,
     SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
     bool &isVariadic, SmallVectorImpl<Type> &resultTypes,
     SmallVectorImpl<NamedAttrList> &resultAttrs) {
   bool allowArgAttrs = true;
   if (parseFunctionArgumentList(parser, allowArgAttrs, allowVariadic, argNames,
                                 argTypes, argAttrs, isVariadic))
     return failure();
   if (succeeded(parser.parseOptionalArrow()))
     return parseFunctionResultList(parser, resultTypes, resultAttrs);
   return success();
 }
 
-void mlir::impl::addArgAndResultAttrs(Builder &builder, OperationState &result,
-                                      ArrayRef<NamedAttrList> argAttrs,
-                                      ArrayRef<NamedAttrList> resultAttrs) {
-  // Add the attributes to the function arguments.
-  SmallString<8> attrNameBuf;
-  for (unsigned i = 0, e = argAttrs.size(); i != e; ++i)
-    if (!argAttrs[i].empty())
-      result.addAttribute(getArgAttrName(i, attrNameBuf),
-                          builder.getDictionaryAttr(argAttrs[i]));
+/// Implementation of `addArgAndResultAttrs` that is attribute list type
+/// agnostic.
+template <typename AttrListT, typename AttrArrayBuildFnT>
+static void addArgAndResultAttrsImpl(Builder &builder, OperationState &result,
+                                     ArrayRef<AttrListT> argAttrs,
+                                     ArrayRef<AttrListT> resultAttrs,
+                                     AttrArrayBuildFnT &&buildAttrArrayFn) {
+  auto nonEmptyAttrsFn = [](const AttrListT &attrs) { return !attrs.empty(); };
 
+  // Add the attributes to the function arguments.
+  if (!argAttrs.empty() && llvm::any_of(argAttrs, nonEmptyAttrsFn)) {
+    ArrayAttr attrDicts = builder.getArrayAttr(buildAttrArrayFn(argAttrs));
+    result.addAttribute(function_like_impl::getArgDictAttrName(), attrDicts);
+  }
   // Add the attributes to the function results.
-  for (unsigned i = 0, e = resultAttrs.size(); i != e; ++i)
-    if (!resultAttrs[i].empty())
-      result.addAttribute(getResultAttrName(i, attrNameBuf),
-                          builder.getDictionaryAttr(resultAttrs[i]));
+  if (!resultAttrs.empty() && llvm::any_of(resultAttrs, nonEmptyAttrsFn)) {
+    ArrayAttr attrDicts = builder.getArrayAttr(buildAttrArrayFn(resultAttrs));
+    result.addAttribute(function_like_impl::getResultDictAttrName(), attrDicts);
+  }
+}
+
+void mlir::function_like_impl::addArgAndResultAttrs(
+    Builder &builder, OperationState &result, ArrayRef<DictionaryAttr> argAttrs,
+    ArrayRef<DictionaryAttr> resultAttrs) {
+  auto buildFn = [](ArrayRef<DictionaryAttr> attrs) {
+    return ArrayRef<Attribute>(attrs.data(), attrs.size());
+  };
+  addArgAndResultAttrsImpl(builder, result, argAttrs, resultAttrs, buildFn);
+}
+void mlir::function_like_impl::addArgAndResultAttrs(
+    Builder &builder, OperationState &result, ArrayRef<NamedAttrList> argAttrs,
+    ArrayRef<NamedAttrList> resultAttrs) {
+  MLIRContext *context = builder.getContext();
+  auto buildFn = [=](ArrayRef<NamedAttrList> attrs) {
+    return llvm::to_vector<8>(
+        llvm::map_range(attrs, [=](const NamedAttrList &attrList) -> Attribute {
+          return attrList.getDictionary(context);
+        }));
+  };
+  addArgAndResultAttrsImpl(builder, result, argAttrs, resultAttrs, buildFn);
 }
 
 /// Parser implementation for function-like operations.  Uses `funcTypeBuilder`
 /// to construct the custom function type given lists of input and output types.
-ParseResult
-mlir::impl::parseFunctionLikeOp(OpAsmParser &parser, OperationState &result,
-                                bool allowVariadic,
-                                mlir::impl::FuncTypeBuilder funcTypeBuilder) {
+ParseResult mlir::function_like_impl::parseFunctionLikeOp(
+    OpAsmParser &parser, OperationState &result, bool allowVariadic,
+    FuncTypeBuilder funcTypeBuilder) {
   SmallVector<OpAsmParser::OperandType, 4> entryArgs;
   SmallVector<NamedAttrList, 4> argAttrs;
   SmallVector<NamedAttrList, 4> resultAttrs;
   SmallVector<Type, 4> argTypes;
   SmallVector<Type, 4> resultTypes;
   auto &builder = parser.getBuilder();
 
   // Parse visibility.
   impl::parseOptionalVisibilityKeyword(parser, result.attributes);
 
   // Parse the name as a symbol.
   StringAttr nameAttr;
   if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(),
                              result.attributes))
     return failure();
 
   // Parse the function signature.
   llvm::SMLoc signatureLocation = parser.getCurrentLocation();
   bool isVariadic = false;
   if (parseFunctionSignature(parser, allowVariadic, entryArgs, argTypes,
                              argAttrs, isVariadic, resultTypes, resultAttrs))
     return failure();
 
   std::string errorMessage;
-  if (auto type = funcTypeBuilder(builder, argTypes, resultTypes,
-                                  impl::VariadicFlag(isVariadic), errorMessage))
-    result.addAttribute(getTypeAttrName(), TypeAttr::get(type));
-  else
+  Type type = funcTypeBuilder(builder, argTypes, resultTypes,
+                              VariadicFlag(isVariadic), errorMessage);
+  if (!type) {
     return parser.emitError(signatureLocation)
            << "failed to construct function type"
            << (errorMessage.empty() ? "" : ": ") << errorMessage;
+  }
+  result.addAttribute(getTypeAttrName(), TypeAttr::get(type));
 
   // If function attributes are present, parse them.
   NamedAttrList parsedAttributes;
   llvm::SMLoc attributeDictLocation = parser.getCurrentLocation();
   if (parser.parseOptionalAttrDictWithKeyword(parsedAttributes))
     return failure();
 
   // Disallow attributes that are inferred from elsewhere in the attribute
if (failed(*parseResult))
       return failure();
     // Function body was parsed, make sure its not empty.
     if (body->empty())
       return parser.emitError(loc, "expected non-empty function body");
   }
   return success();
 }
 
-// Print a function result list.
+/// Print a function result list. The provided `attrs` must either be null, or
+/// contain a set of DictionaryAttrs of the same arity as `types`.
 static void printFunctionResultList(OpAsmPrinter &p, ArrayRef<Type> types,
-                                    ArrayRef<ArrayRef<NamedAttribute>> attrs) {
+                                    ArrayAttr attrs) {
   assert(!types.empty() && "Should not be called for empty result list.");
+  assert((!attrs || attrs.size() == types.size()) &&
+         "Invalid number of attributes.");
+
   auto &os = p.getStream();
-  bool needsParens =
-      types.size() > 1 || types[0].isa<FunctionType>() || !attrs[0].empty();
+  bool needsParens = types.size() > 1 || types[0].isa<FunctionType>() ||
+                     (attrs && !attrs[0].cast<DictionaryAttr>().empty());
   if (needsParens)
     os << '(';
-  llvm::interleaveComma(
-      llvm::zip(types, attrs), os,
-      [&](const std::tuple<Type, ArrayRef<NamedAttribute>> &t) {
-        p.printType(std::get<0>(t));
-        p.printOptionalAttrDict(std::get<1>(t));
-      });
+  llvm::interleaveComma(llvm::seq<size_t>(0, types.size()), os, [&](size_t i) {
+    p.printType(types[i]);
+    if (attrs)
+      p.printOptionalAttrDict(attrs[i].cast<DictionaryAttr>().getValue());
   if (needsParens)
     os << ')';
 }
 
 /// Print the signature of the function-like operation `op`.  Assumes `op` has
 /// the FunctionLike trait and passed the verification.
-void mlir::impl::printFunctionSignature(OpAsmPrinter &p, Operation *op,
-                                        ArrayRef<Type> argTypes,
-                                        bool isVariadic,
-                                        ArrayRef<Type> resultTypes) {
+void mlir::function_like_impl::printFunctionSignature(
+    OpAsmPrinter &p, Operation *op, ArrayRef<Type> argTypes, bool isVariadic,
   Region &body = op->getRegion(0);
   bool isExternal = body.empty();
 
   p << '(';
+  ArrayAttr argAttrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName());
   for (unsigned i = 0, e = argTypes.size(); i < e; ++i) {
     if (i > 0)
       p << ", ";
 
     if (!isExternal) {
       p.printOperand(body.getArgument(i));
       p << ": ";
     }
 
     p.printType(argTypes[i]);
-    p.printOptionalAttrDict(::mlir::impl::getArgAttrs(op, i));
+    if (argAttrs)
+      p.printOptionalAttrDict(argAttrs[i].cast<DictionaryAttr>().getValue());
   }
 
   if (isVariadic) {
     if (!argTypes.empty())
       p << ", ";
     p << "...";
   }
 
   p << ')';
 
   if (!resultTypes.empty()) {
     p.getStream() << " -> ";
-    SmallVector<ArrayRef<NamedAttribute>, 4> resultAttrs;
-    for (int i = 0, e = resultTypes.size(); i < e; ++i)
-      resultAttrs.push_back(::mlir::impl::getResultAttrs(op, i));
+    auto resultAttrs = op->getAttrOfType<ArrayAttr>(getResultDictAttrName());
     printFunctionResultList(p, resultTypes, resultAttrs);
   }
 }
 
 /// Prints the list of function prefixed with the "attributes" keyword. The
 /// attributes with names listed in "elided" as well as those used by the
 /// function-like operation internally are not printed. Nothing is printed
 /// if all attributes are elided. Assumes `op` has the `FunctionLike` trait and
 /// passed the verification.
-void mlir::impl::printFunctionAttributes(OpAsmPrinter &p, Operation *op,
-                                         unsigned numInputs,
-                                         unsigned numResults,
-                                         ArrayRef<StringRef> elided) {
+void mlir::function_like_impl::printFunctionAttributes(
+    OpAsmPrinter &p, Operation *op, unsigned numInputs, unsigned numResults,
   // Print out function attributes, if present.
   SmallVector<StringRef, 2> ignoredAttrs = {
-      ::mlir::SymbolTable::getSymbolAttrName(), getTypeAttrName()};
+      ::mlir::SymbolTable::getSymbolAttrName(), getTypeAttrName(),
+      getArgDictAttrName(), getResultDictAttrName()};
   ignoredAttrs.append(elided.begin(), elided.end());
 
-  SmallString<8> attrNameBuf;
-
-  // Ignore any argument attributes.
-  std::vector<SmallString<8>> argAttrStorage;
-  for (unsigned i = 0; i != numInputs; ++i)
-    if (op->getAttr(getArgAttrName(i, attrNameBuf)))
-      argAttrStorage.emplace_back(attrNameBuf);
-  ignoredAttrs.append(argAttrStorage.begin(), argAttrStorage.end());
-
-  // Ignore any result attributes.
-  std::vector<SmallString<8>> resultAttrStorage;
-  for (unsigned i = 0; i != numResults; ++i)
-    if (op->getAttr(getResultAttrName(i, attrNameBuf)))
-      resultAttrStorage.emplace_back(attrNameBuf);
-  ignoredAttrs.append(resultAttrStorage.begin(), resultAttrStorage.end());
-
   p.printOptionalAttrDictWithKeyword(op->getAttrs(), ignoredAttrs);
 }
 
 /// Printer implementation for function-like operations.  Accepts lists of
 /// argument and result types to use while printing.
-void mlir::impl::printFunctionLikeOp(OpAsmPrinter &p, Operation *op,
-                                     ArrayRef<Type> argTypes, bool isVariadic,
+void mlir::function_like_impl::printFunctionLikeOp(OpAsmPrinter &p,
+                                                   Operation *op,
+                                                   ArrayRef<Type> argTypes,
+                                                   bool isVariadic,
                                                    ArrayRef<Type> resultTypes) {
   // Print the operation and the function name.
   auto funcName =
       op->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName())
           .getValue();
   p << op->getName() << ' ';
 
   StringRef visibilityAttrName = SymbolTable::getVisibilityAttrName();
   if (auto visibility = op->getAttrOfType<StringAttr>(visibilityAttrName))

mlir/lib/IR/FunctionSupport.cpp

for (unsigned i = 0; i < totalIndices; ++i)
     if (!skipIndices.test(i))
       callback(i);
 }
 
 //===----------------------------------------------------------------------===//
 // Function Arguments and Results.
 //===----------------------------------------------------------------------===//
 
-void mlir::impl::eraseFunctionArguments(Operation *op,
-                                        ArrayRef<unsigned> argIndices,
-                                        unsigned originalNumArgs,
-                                        Type newType) {
+static bool isEmptyAttrDict(Attribute attr) {
+  return attr.cast<DictionaryAttr>().empty();
+}
+
+DictionaryAttr mlir::function_like_impl::getArgAttrDict(Operation *op,
+                                                        unsigned index) {
+  ArrayAttr attrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName());
+  DictionaryAttr argAttrs =
+      attrs ? attrs[index].cast<DictionaryAttr>() : DictionaryAttr();
+  return (argAttrs && !argAttrs.empty()) ? argAttrs : DictionaryAttr();
+}
+
+DictionaryAttr mlir::function_like_impl::getResultAttrDict(Operation *op,
+                                                           unsigned index) {
+  ArrayAttr attrs = op->getAttrOfType<ArrayAttr>(getResultDictAttrName());
+  DictionaryAttr resAttrs =
+      attrs ? attrs[index].cast<DictionaryAttr>() : DictionaryAttr();
+  return (resAttrs && !resAttrs.empty()) ? resAttrs : DictionaryAttr();
+}
+
+void mlir::function_like_impl::detail::setArgResAttrDict(
+    Operation *op, StringRef attrName, unsigned numTotalIndices, unsigned index,
+    DictionaryAttr attrs) {
+  ArrayAttr allAttrs = op->getAttrOfType<ArrayAttr>(attrName);
+  if (!allAttrs) {
+    if (attrs.empty())
+      return;
+
+    // If this attribute is not empty, we need to create a new attribute array.
+    SmallVector<Attribute, 8> newAttrs(numTotalIndices,
+                                       DictionaryAttr::get(op->getContext()));
+    newAttrs[index] = attrs;
+    op->setAttr(attrName, ArrayAttr::get(op->getContext(), newAttrs));
+    return;
+  }
+  // Check to see if the attribute is different from what we already have.
+  if (allAttrs[index] == attrs)
+    return;
+
+  // If it is, check to see if the attribute array would now contain only empty
+  // dictionaries.
+  ArrayRef<Attribute> rawAttrArray = allAttrs.getValue();
+  if (attrs.empty() &&
+      llvm::all_of(rawAttrArray.take_front(index), isEmptyAttrDict) &&
+      llvm::all_of(rawAttrArray.drop_front(index + 1), isEmptyAttrDict)) {
+    op->removeAttr(attrName);
+    return;
+  }
+
+  // Otherwise, create a new attribute array with the updated dictionary.
+  SmallVector<Attribute, 8> newAttrs(rawAttrArray.begin(), rawAttrArray.end());
+  newAttrs[index] = attrs;
+  op->setAttr(attrName, ArrayAttr::get(op->getContext(), newAttrs));
+}
+
+/// Set all of the argument or result attribute dictionaries for a function.
+static void setAllArgResAttrDicts(Operation *op, StringRef attrName,
+                                  ArrayRef<Attribute> attrs) {
+  if (llvm::all_of(attrs, isEmptyAttrDict))
+    op->removeAttr(attrName);
+  else
+    op->setAttr(attrName, ArrayAttr::get(op->getContext(), attrs));
+}
+
+void mlir::function_like_impl::setAllArgAttrDicts(
+    Operation *op, ArrayRef<DictionaryAttr> attrs) {
+  setAllArgAttrDicts(op, ArrayRef<Attribute>(attrs.data(), attrs.size()));
+}
+void mlir::function_like_impl::setAllArgAttrDicts(Operation *op,
+                                                  ArrayRef<Attribute> attrs) {
+  auto wrappedAttrs = llvm::map_range(attrs, [op](Attribute attr) -> Attribute {
+    return !attr ? DictionaryAttr::get(op->getContext()) : attr;
+  });
+  setAllArgResAttrDicts(op, getArgDictAttrName(),
+                        llvm::to_vector<8>(wrappedAttrs));
+}
+
+void mlir::function_like_impl::setAllResultAttrDicts(
+    Operation *op, ArrayRef<DictionaryAttr> attrs) {
+  setAllResultAttrDicts(op, ArrayRef<Attribute>(attrs.data(), attrs.size()));
+}
+void mlir::function_like_impl::setAllResultAttrDicts(
+    Operation *op, ArrayRef<Attribute> attrs) {
+  auto wrappedAttrs = llvm::map_range(attrs, [op](Attribute attr) -> Attribute {
+    return !attr ? DictionaryAttr::get(op->getContext()) : attr;
+  });
+  setAllArgResAttrDicts(op, getResultDictAttrName(),
+                        llvm::to_vector<8>(wrappedAttrs));
+}
+
+void mlir::function_like_impl::eraseFunctionArguments(
+    Operation *op, ArrayRef<unsigned> argIndices, unsigned originalNumArgs,
   // There are 3 things that need to be updated:
   // - Function type.
   // - Arg attrs.
   // - Block arguments of entry block.
   Block &entry = op->getRegion(0).front();
-  SmallString<8> nameBuf;
 
-  // Collect arg attrs to set.
+  // 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) {
-    newArgAttrs.emplace_back(getArgAttrDict(op, i));
+      newArgAttrs.emplace_back(argAttrs[i].cast<DictionaryAttr>());
     });
-
-  // Remove any arg attrs that are no longer needed.
-  for (unsigned i = newArgAttrs.size(), e = originalNumArgs; i < e; ++i)
-    op->removeAttr(getArgAttrName(i, nameBuf));
-
-  // Set the function type.
-  op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
-
-  // Set the new arg attrs, or remove them if empty.
-  for (unsigned i = 0, e = newArgAttrs.size(); i != e; ++i) {
-    auto nameAttr = getArgAttrName(i, nameBuf);
-    if (newArgAttrs[i] && !newArgAttrs[i].empty())
-      op->setAttr(nameAttr, newArgAttrs[i]);
-    else
-      op->removeAttr(nameAttr);
+    setAllArgAttrDicts(op, newArgAttrs);
   }
 
-  // Update the entry block's arguments.
+  // Update the function type and any entry block arguments.
+  op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
   entry.eraseArguments(argIndices);
 }
 
-void mlir::impl::eraseFunctionResults(Operation *op,
-                                      ArrayRef<unsigned> resultIndices,
-                                      unsigned originalNumResults,
-                                      Type newType) {
+void mlir::function_like_impl::eraseFunctionResults(
+    Operation *op, ArrayRef<unsigned> resultIndices,
+    unsigned originalNumResults, Type newType) {
   // There are 2 things that need to be updated:
   // - Function type.
   // - Result attrs.
-  SmallString<8> nameBuf;
 
-  // Collect result attrs to set.
+  // 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) {
-    newResultAttrs.emplace_back(getResultAttrDict(op, i));
+      newResultAttrs.emplace_back(resAttrs[i].cast<DictionaryAttr>());
     });
+    setAllResultAttrDicts(op, newResultAttrs);
+  }
 
-  // Remove any result attrs that are no longer needed.
-  for (unsigned i = newResultAttrs.size(), e = originalNumResults; i < e; ++i)
-    op->removeAttr(getResultAttrName(i, nameBuf));
-
-  // Set the function type.
+  // Update the function type.
   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
-
-  // Set the new result attrs, or remove them if empty.
-  for (unsigned i = 0, e = newResultAttrs.size(); i != e; ++i) {
-    auto nameAttr = getResultAttrName(i, nameBuf);
-    if (newResultAttrs[i] && !newResultAttrs[i].empty())
-      op->setAttr(nameAttr, newResultAttrs[i]);
-    else
-      op->removeAttr(nameAttr);
-  }
 }
 
 //===----------------------------------------------------------------------===//
 // Function type signature.
 //===----------------------------------------------------------------------===//
 
-FunctionType mlir::impl::getFunctionType(Operation *op) {
+FunctionType mlir::function_like_impl::getFunctionType(Operation *op) {
   assert(op->hasTrait<OpTrait::FunctionLike>());
-  return op->getAttrOfType<TypeAttr>(mlir::impl::getTypeAttrName())
+  return op->getAttrOfType<TypeAttr>(getTypeAttrName())
       .getValue()
       .cast<FunctionType>();
 }
 
-void mlir::impl::setFunctionType(Operation *op, FunctionType newType) {
+void mlir::function_like_impl::setFunctionType(Operation *op,
+                                               FunctionType newType) {
   assert(op->hasTrait<OpTrait::FunctionLike>());
-  SmallVector<char, 16> nameBuf;
   FunctionType oldType = getFunctionType(op);
-
-  for (int i = newType.getNumInputs(), e = oldType.getNumInputs(); i < e; i++)
-    op->removeAttr(getArgAttrName(i, nameBuf));
-  for (int i = newType.getNumResults(), e = oldType.getNumResults(); i < e; i++)
-    op->removeAttr(getResultAttrName(i, nameBuf));
   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
+
+  // Functor used to update the argument and result attributes of the function.
+  auto updateAttrFn = [&](StringRef attrName, unsigned oldCount,
+                          unsigned newCount, auto setAttrFn) {
+    if (oldCount == newCount)
+      return;
+    // The new type has no arguments/results, just drop the attribute.
+    if (newCount == 0) {
+      op->removeAttr(attrName);
+      return;
+    }
+    ArrayAttr attrs = op->getAttrOfType<ArrayAttr>(attrName);
+    if (!attrs)
+      return;
+
+    // The new type has less arguments/results, take the first N attributes.
+    if (newCount < oldCount)
+      return setAttrFn(op, attrs.getValue().take_front(newCount));
+
+    // Otherwise, the new type has more arguments/results. Initialize the new
+    // arguments/results with empty attributes.
+    SmallVector<Attribute> newAttrs(attrs.begin(), attrs.end());
+    newAttrs.resize(newCount);
+    setAttrFn(op, newAttrs);
+  };
+
+  // Update the argument and result attributes.
+  updateAttrFn(function_like_impl::getArgDictAttrName(), oldType.getNumInputs(),
+               newType.getNumInputs(), [&](Operation *op, auto &&attrs) {
+                 setAllArgAttrDicts(op, attrs);
+               });
+  updateAttrFn(
+      function_like_impl::getResultDictAttrName(), oldType.getNumResults(),
+      newType.getNumResults(),
+      [&](Operation *op, auto &&attrs) { setAllResultAttrDicts(op, attrs); });
 }
 
 //===----------------------------------------------------------------------===//
 // Function body.
 //===----------------------------------------------------------------------===//
 
-Region &mlir::impl::getFunctionBody(Operation *op) {
+Region &mlir::function_like_impl::getFunctionBody(Operation *op) {
   assert(op->hasTrait<OpTrait::FunctionLike>());
   return op->getRegion(0);
 }

mlir/lib/Transforms/Utils/DialectConversion.cpp

struct FunctionLikeSignatureConversion : public ConversionPattern {
   FunctionLikeSignatureConversion(StringRef functionLikeOpName,
                                   MLIRContext *ctx, TypeConverter &converter)
       : ConversionPattern(converter, functionLikeOpName, /*benefit=*/1, ctx) {}
 
   /// Hook to implement combined matching and rewriting for FunctionLike ops.
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    FunctionType type = mlir::impl::getFunctionType(op);
+    FunctionType type = function_like_impl::getFunctionType(op);
 
     // Convert the original function types.
     TypeConverter::SignatureConversion result(type.getNumInputs());
     SmallVector<Type, 1> newResults;
     if (failed(typeConverter->convertSignatureArgs(type.getInputs(), result)) ||
         failed(typeConverter->convertTypes(type.getResults(), newResults)) ||
-        failed(rewriter.convertRegionTypes(&mlir::impl::getFunctionBody(op),
-                                           *typeConverter, &result)))
+        failed(rewriter.convertRegionTypes(
+            &function_like_impl::getFunctionBody(op), *typeConverter, &result)))
       return failure();
 
     // Update the function signature in-place.
     auto newType = FunctionType::get(rewriter.getContext(),
                                      result.getConvertedTypes(), newResults);
 
     rewriter.updateRootInPlace(
-        op, [&] { mlir::impl::setFunctionType(op, newType); });
+        op, [&] { function_like_impl::setFunctionType(op, newType); });
 
     return success();
   }
 };
 } // end anonymous namespace
 
 void mlir::populateFunctionLikeTypeConversionPattern(
     StringRef functionLikeOpName, RewritePatternSet &patterns,

mlir/test/Dialect/LLVMIR/func.mlir

module {
   // GENERIC-SAME: type = !llvm.func<i64 (i64)>
   // GENERIC-SAME: () -> ()
   }) {sym_name = "baz", type = !llvm.func<i64 (i64)>} : () -> ()
 
   // CHECK: llvm.func @qux(!llvm.ptr<i64> {llvm.noalias = true}, i64)
   // CHECK: attributes {xxx = {yyy = 42 : i64}}
   "llvm.func"() ({
   }) {sym_name = "qux", type = !llvm.func<void (ptr<i64>, i64)>,
-      arg0 = {llvm.noalias = true}, xxx = {yyy = 42}} : () -> ()
+      arg_attrs = [{llvm.noalias = true}, {}], xxx = {yyy = 42}} : () -> ()
 
   // CHECK: llvm.func @roundtrip1()
   llvm.func @roundtrip1()
 
   // CHECK: llvm.func @roundtrip2(i64, f32) -> f64
   llvm.func @roundtrip2(i64, f32) -> f64
 
   // CHECK: llvm.func @roundtrip3(i32, i1)

mlir/test/IR/invalid-func-op.mlir

 // expected-error@+1 {{'sym_name' is an inferred attribute and should not be specified in the explicit attribute dictionary}}
 func private @invalid_symbol_name_attr() attributes { sym_name = "x" }
 
 // -----
 
 // expected-error@+1 {{'type' is an inferred attribute and should not be specified in the explicit attribute dictionary}}
 func private @invalid_symbol_type_attr() attributes { type = "x" }
 
+// -----
+
+// expected-error@+1 {{argument attribute array `arg_attrs` to have the same number of elements as the number of function arguments}}
+func private @invalid_arg_attrs() attributes { arg_attrs = [{}] }
+
+// -----
+
+// expected-error@+1 {{expects argument attribute dictionary to be a DictionaryAttr, but got `10 : i64`}}
+func private @invalid_arg_attrs(i32) attributes { arg_attrs = [10] }
+
+// -----
+
+// expected-error@+1 {{result attribute array `res_attrs` to have the same number of elements as the number of function results}}
+func private @invalid_res_attrs() attributes { res_attrs = [{}] }
+
+// -----
+
+// expected-error@+1 {{expects result attribute dictionary to be a DictionaryAttr, but got `10 : i64`}}
+func private @invalid_res_attrs() -> i32 attributes { res_attrs = [10] }

mlir/test/IR/test-func-set-type.mlir

 // RUN: mlir-opt %s -test-func-set-type -split-input-file | FileCheck %s
 
 // It's currently not possible to have an attribute with a function type due to
 // parser ambiguity. So instead we reference a function declaration to take the
 // type from.
 
 // -----
 
 // Test case: The setType call needs to erase some arg attrs.
 
 // CHECK: func private @erase_arg(f32 {test.A})
-// CHECK-NOT: attributes{{.*arg[0-9]}}
 func private @t(f32)
 func private @erase_arg(%arg0: f32 {test.A}, %arg1: f32 {test.B})
 attributes {test.set_type_from = @t}
 
 // -----
 
 // Test case: The setType call needs to erase some result attrs.
 
 // CHECK: func private @erase_result() -> (f32 {test.A})
-// CHECK-NOT: attributes{{.*result[0-9]}}
 func private @t() -> (f32)
 func private @erase_result() -> (f32 {test.A}, f32 {test.B})
 attributes {test.set_type_from = @t}