[llvm-branch-commits] [mlir] 25007b4 - [MLIR][NFC] Change FunctionLike::setAllArgAttrs/setAllResultAttrs to do a one-shot attribute update.

[Rahul Joshi via llvm-branch-commits]

Author: Rahul Joshi
Date: 2020-12-28T14:15:47-08:00
New Revision: 25007b4d7e094c569d512770bd2397d8667fd3db

URL: https://github.com/llvm/llvm-project/commit/25007b4d7e094c569d512770bd2397d8667fd3db
DIFF: https://github.com/llvm/llvm-project/commit/25007b4d7e094c569d512770bd2397d8667fd3db.diff

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

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

Differential Revision: https://reviews.llvm.org/D93870

Added: 
    

Modified: 
    mlir/include/mlir/IR/FunctionSupport.h

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/FunctionSupport.h b/mlir/include/mlir/IR/FunctionSupport.h
index fd50e0dbb512..481dc5a6986f 100644
--- a/mlir/include/mlir/IR/FunctionSupport.h
+++ b/mlir/include/mlir/IR/FunctionSupport.h
@@ -333,11 +333,7 @@ class FunctionLike : public OpTrait::TraitBase<ConcreteType, FunctionLike> {
   /// 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.
@@ -400,11 +396,7 @@ class FunctionLike : public OpTrait::TraitBase<ConcreteType, FunctionLike> {
   /// 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.
@@ -591,6 +583,26 @@ void FunctionLike<ConcreteType>::setArgAttrs(unsigned index,
                                          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>
@@ -648,6 +660,26 @@ void FunctionLike<ConcreteType>::setResultAttrs(unsigned index,
                                   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>