[mlir][DictionaryAttr] Add a new getWithSorted and use it when possible

[rriddle]

Summary

The elements of a DictionaryAttr are sorted by name. In many situations, e.g NamedAttributeList, we can guarantee that the elements are sorted on construction and remove the need to perform extra checks. In places with lots of calls to attribute methods, this leads to a good performance improvement.

[jpienaar]

Thanks, this will fit the bill well :-)

mlir/include/mlir/IR/Attributes.h

281

Set confuses me here, is it meant to convey unordered? (an alternative is just to say array or vector as by default there is no assumption that it would be sorted specially IMHO)

[mehdi_amini]

mlir/include/mlir/IR/Attributes.h

281

I would document here that is the client can guarantee that the array is sorted they should use the other method for performance

mlir/lib/IR/Attributes.cpp

101

If Attr is a non zero terminated string that is a prefix of name, won’t this read out of bound?

[rriddle]

mlir/lib/IR/Attributes.cpp

101

Attr guarantees null termination because it uses Identifier.

[rriddle]

Resolve comments

[mehdi_amini]

mlir/lib/IR/Attributes.cpp

101

Oh I was on my phone and mis-read This is correct even when attr.first.data() is not a zero terminated string

Actually the comment is adding more confusion than it helps for me: the comparison is obviously correct as long as attr is null terminated.

mlir/include/mlir/IR/Attributes.h

 /// be unique within the collection.
 class DictionaryAttr
     : public Attribute::AttrBase<DictionaryAttr, Attribute,
                                  detail::DictionaryAttributeStorage> {
 public:
   using Base::Base;
   using ValueType = ArrayRef<NamedAttribute>;
 
+  /// Construct a dictionary attribute with the provided list of named
+  /// attributes. This method assumes that the provided list is unordered. If
+  /// the caller can guarantee that the attributes are ordered by name,
+  /// getWithSorted should be used instead.
   static DictionaryAttr get(ArrayRef<NamedAttribute> value,
                             MLIRContext *context);
 
+  /// Construct a dictionary with an array of values that is known to already be
+  /// sorted by name and uniqued.
+  static DictionaryAttr getWithSorted(ArrayRef<NamedAttribute> value,
+                                      MLIRContext *context);
+
   ArrayRef<NamedAttribute> getValue() const;
 
   /// Return the specified attribute if present, null otherwise.
   Attribute get(StringRef name) const;
   Attribute get(Identifier name) const;
 
   /// Support range iteration.
   using iterator = llvm::ArrayRef<NamedAttribute>::iterator;

 inline ::llvm::hash_code hash_value(Attribute arg) {
   return ::llvm::hash_value(arg.impl);
 }
 
 //===----------------------------------------------------------------------===//
 // NamedAttributeList
 //===----------------------------------------------------------------------===//
 
-/// A NamedAttributeList is used to manage a list of named attributes. This
-/// provides simple interfaces for adding/removing/finding attributes from
+/// A NamedAttributeList is a mutable wrapper around a DictionaryAttr. It
+/// provides additional interfaces for adding, removing, replacing attributes
 /// within a DictionaryAttr.
 ///
 /// We assume there will be relatively few attributes on a given operation
 /// (maybe a dozen or so, but not hundreds or thousands) so we use linear
 /// searches for everything.
 class NamedAttributeList {
 public:
   NamedAttributeList(DictionaryAttr attrs = nullptr)

mlir/lib/IR/Attributes.cpp

 
 /// Perform a three-way comparison between the names of the specified
 /// NamedAttributes.
 static int compareNamedAttributes(const NamedAttribute *lhs,
                                   const NamedAttribute *rhs) {
   return strcmp(lhs->first.data(), rhs->first.data());
 }
 
+/// Returns if the name of the given attribute precedes that of 'name'.
+static bool compareNamedAttributeWithName(const NamedAttribute &attr,
+                                          StringRef name) {
+  // This is correct even when attr.first.data()[name.size()] is not a zero
+  // string terminator, because we only care about a less than comparison.
+  // This can't use memcmp, because it doesn't guarantee that it will stop
+  // reading both buffers if one is shorter than the other, even if there is
+  // a difference.
+  return strncmp(attr.first.data(), name.data(), name.size()) < 0;
+}
+
 DictionaryAttr DictionaryAttr::get(ArrayRef<NamedAttribute> value,
                                    MLIRContext *context) {
   assert(llvm::all_of(value,
                       [](const NamedAttribute &attr) { return attr.second; }) &&
          "value cannot have null entries");
 
   // We need to sort the element list to canonicalize it, but we also don't want
   // to do a ton of work in the super common case where the element list is
assert(std::adjacent_find(value.begin(), value.end(),
                                 return l.first == r.first;
                               }) == value.end() &&
            "DictionaryAttr element names must be unique");
   }
 
   return Base::get(context, StandardAttributes::Dictionary, value);
 }
 
+/// Construct a dictionary with an array of values that is known to already be
+/// sorted by name and uniqued.
+DictionaryAttr DictionaryAttr::getWithSorted(ArrayRef<NamedAttribute> value,
+                                             MLIRContext *context) {
+  // Ensure that the attribute elements are unique and sorted.
+  assert(llvm::is_sorted(value,
+                         [](NamedAttribute l, NamedAttribute r) {
+                           return l.first.strref() < r.first.strref();
+                         }) &&
+         "expected attribute values to be sorted");
+  assert(std::adjacent_find(value.begin(), value.end(),
+                            [](NamedAttribute l, NamedAttribute r) {
+                              return l.first == r.first;
+                            }) == value.end() &&
+         "DictionaryAttr element names must be unique");
+  return Base::get(context, StandardAttributes::Dictionary, value);
+}
+
 ArrayRef<NamedAttribute> DictionaryAttr::getValue() const {
   return getImpl()->getElements();
 }
 
 /// Return the specified attribute if present, null otherwise.
 Attribute DictionaryAttr::get(StringRef name) const {
   ArrayRef<NamedAttribute> values = getValue();
-  auto compare = [](NamedAttribute attr, StringRef name) -> bool {
-    // This is correct even when attr.first.data()[name.size()] is not a zero
-    // string terminator, because we only care about a less than comparison.
-    // This can't use memcmp, because it doesn't guarantee that it will stop
-    // reading both buffers if one is shorter than the other, even if there is
-    // a difference.
-    return strncmp(attr.first.data(), name.data(), name.size()) < 0;
-  };
-  auto it = llvm::lower_bound(values, name, compare);
+  auto it = llvm::lower_bound(values, name, compareNamedAttributeWithName);
   return it != values.end() && it->first == name ? it->second : Attribute();
 }
 Attribute DictionaryAttr::get(Identifier name) const {
   for (auto elt : getValue())
     if (elt.first == name)
       return elt.second;
   return nullptr;
 }
Attribute NamedAttributeList::get(Identifier name) const {
   return attrs ? attrs.get(name) : nullptr;
 }
 
 /// 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 NamedAttributeList::set(Identifier name, Attribute value) {
   assert(value && "attributes may never be null");
 
-  // If we already have this attribute, replace it.
-  auto origAttrs = getAttrs();
-  SmallVector<NamedAttribute, 8> newAttrs(origAttrs.begin(), origAttrs.end());
-  for (auto &elt : newAttrs)
-    if (elt.first == name) {
-      elt.second = value;
-      attrs = DictionaryAttr::get(newAttrs, value.getContext());
-      return;
-    }
+  // Look for an existing value for the given name, and set it in-place.
+  ArrayRef<NamedAttribute> values = getAttrs();
+  auto it = llvm::find_if(
+      values, [name](NamedAttribute attr) { return attr.first == name; });
+  if (it != values.end()) {
+    // Bail out early if the value is the same as what we already have.
+    if (it->second == value)
+      return;
+
+    SmallVector<NamedAttribute, 8> newAttrs(values.begin(), values.end());
+    newAttrs[it - values.begin()].second = value;
+    attrs = DictionaryAttr::getWithSorted(newAttrs, value.getContext());
 
-  // Otherwise, add it.
+  // Otherwise, insert the new attribute into its sorted position.
+  it = llvm::lower_bound(values, name, compareNamedAttributeWithName);
+  SmallVector<NamedAttribute, 8> newAttrs;
+  newAttrs.reserve(values.size() + 1);
+  newAttrs.append(values.begin(), it);
   newAttrs.push_back({name, value});
-  attrs = DictionaryAttr::get(newAttrs, value.getContext());
+  newAttrs.append(it, values.end());
+  attrs = DictionaryAttr::getWithSorted(newAttrs, value.getContext());
 }
 
 /// Remove the attribute with the specified name if it exists.  The return
 /// value indicates whether the attribute was present or not.
 auto NamedAttributeList::remove(Identifier name) -> RemoveResult {
   auto origAttrs = getAttrs();
   for (unsigned i = 0, e = origAttrs.size(); i != e; ++i) {
     if (origAttrs[i].first == name) {
       // Handle the simple case of removing the only attribute in the list.
       if (e == 1) {
         attrs = nullptr;
         return RemoveResult::Removed;
       }
 
       SmallVector<NamedAttribute, 8> newAttrs;
       newAttrs.reserve(origAttrs.size() - 1);
       newAttrs.append(origAttrs.begin(), origAttrs.begin() + i);
       newAttrs.append(origAttrs.begin() + i + 1, origAttrs.end());
-      attrs = DictionaryAttr::get(newAttrs, newAttrs[0].second.getContext());
+      attrs = DictionaryAttr::getWithSorted(newAttrs,
+                                            newAttrs[0].second.getContext());
       return RemoveResult::Removed;
     }
   }
   return RemoveResult::NotFound;
 }