[MLIR][NFC] Change FunctionLike::setAllArgAttrs/setAllResultAttrs to do a one-shot attribute update.

[jurahul]

Summary

• Change FunctionLike::setAllArgAttrs() and setAllResultAttrs() to rebuild the new list of function attributes locally and call setAttr() just once instead of calling setArgAttr()/setResultAttrs() for each argument which incrementally build the attribute dictionary and can end up creating a lot of unused DictionaryAttr's (which are uniqued and nor garbage collected).

[jpienaar]

Thanks (this reminds me of wanting some benchmarks/tracking ...)

mlir/include/mlir/IR/FunctionSupport.h

598

An if-then-else seems shorter and simpler here (and below).

[mehdi_amini]

LGTM.

NFC?

[jurahul]

Use if-else

mlir/include/mlir/IR/FunctionSupport.h

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) {
-    assert(attributes.size() == getNumArguments());
-    for (unsigned i = 0, e = attributes.size(); i != e; ++i)
-      setArgAttrs(i, attributes[i]);
-  }
+  void setAllArgAttrs(ArrayRef<DictionaryAttr> 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);
   }
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) {
-    assert(attributes.size() == getNumResults());
-    for (unsigned i = 0, e = attributes.size(); i != e; ++i)
-      setResultAttrs(i, attributes[i]);
-  }
+  void setAllResultAttrs(ArrayRef<DictionaryAttr> 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);
void FunctionLike<ConcreteType>::setArgAttrs(unsigned index,
   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);
+}
+
 /// 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));
   Attribute oldValue = attributes.set(name, value);
 
void FunctionLike<ConcreteType>::setResultAttrs(unsigned index,
   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);
+}
+
 /// 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));
   Attribute oldAttr = attributes.set(name, value);