From 25007b4d7e094c569d512770bd2397d8667fd3db Mon Sep 17 00:00:00 2001
From: Rahul Joshi <jurahul@google.com>
Date: Mon, 28 Dec 2020 13:34:11 -0800
Subject: [PATCH] [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
---
 mlir/include/mlir/IR/FunctionSupport.h | 52 +++++++++++++++++++++-----
 1 file changed, 42 insertions(+), 10 deletions(-)

diff --git a/mlir/include/mlir/IR/FunctionSupport.h b/mlir/include/mlir/IR/FunctionSupport.h
index fd50e0dbb51219..481dc5a6986f89 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>