[mlir][core] Slightly improved attribute lookup

[Mogball]

Summary

• String binary search does 1 less string comparison
• Identifier linear scan on large attribute list is switched to string binary search

[Mogball]

Fix commit message

[rriddle]

mlir/include/mlir/IR/OperationSupport.h

305–311

This class is only ever used by NamedAttrList. Is there a reason for this to be separate?

347

438

[Mogball]

mlir/include/mlir/IR/OperationSupport.h

305–311

I was going to add it to DictionaryAttr but since that's tied up in ODS I didn't bother...

[Mogball]

Deleted AttributeDictionaryLike class

[rriddle]

mlir/include/mlir/IR/OperationSupport.h

260–264

271–273

Is the bool return important? What is the performance if you just return end() if the name isn't found?

274–275

Just use ptrdiff_t.

280

Spell out auto here.

295

Drop the 16, it could change. Prefer being somewhat ambiguous on exact numbers of heuristics in comments.

[Mogball]

mlir/include/mlir/IR/OperationSupport.h

271–273

For a linear search, the only gain is that you skip one extra it == end() check.

For a binary search, it's important to know whether the attribute was actually found, because, if not, then it will be the sorted insertion point. If just end() was returned:

• Attribute lookups will need to check it != end() && it->first == name to know if the attribute was actually found (which is no different than using lower_bound)
• Setting an attribute when the attribute doesn't exist will have to re-search the sorted insertion point.

[Mogball]

mlir/include/mlir/IR/OperationSupport.h

271–273

And if the attribute was actually found (with a StringRef lookup) then it->first == name will compare the entire string (think operand_segment_sizes...)

[rriddle]

mlir/include/mlir/IR/OperationSupport.h

271–273

Why would the binary search be different? Wouldn't we just return end in all of the cases where you currently return false, and the current iterator whenever you return true?

[Mogball]

mlir/include/mlir/IR/OperationSupport.h

271–273

Case 1: attribute is found. Returning just an iterator:

c++
auto it = findAttrSorted(begin(), end(), name);
return it == end() ? Attribute() : it->second;

// vs
auto it = findAttrSorted(begin(), end(), name);
return it.second ? it.first->second : Attribute();

All that happens is that the == end() check is skipped, which isn't much, but:

Case 2: setting an attribute

c++
auto it = findAttrSorted(begin(), end(), name);
// a: attribute is found
if (it != end()) {
  it->second = value;
  return;
}
// b: attribute is not found, where should the attribute be inserted?
it = llvm::lower_bound(begin(), end(), name); // re-search the list
attrs.insert(it, {name, value});

// vs
auto it = findAttrSorted(begin(), end(), name);
if (it.second) {
  it.first->second = value;
  return;
}
// it.first will be the sorted insertion point
attrs.insert(it.first, {name, value});

// vs lower_bound
auto it = llvm::lower_bound(begin(), end(), name);
if (it != end() && it->first == name) { // re-compare the whole string (if found)
  it->second = value;
  return; 
}
attrs.insert(it, {name, value});

Essentially, the bool contains the result of it->first == name so that it doesn't have to be re-checked.

[rriddle]

mlir/include/mlir/IR/OperationSupport.h

271–273

Right, I missed where this was being used for insertion.

[Mogball]

Review comments

[rriddle]

mlir/lib/IR/OperationSupport.cpp

110

nit: Add a comment on the use of swap here.

122–128

Why is this the case? Can't we use the insertion point from the findAttr above directly?

[Mogball]

mlir/lib/IR/OperationSupport.cpp

122–128

Since this one is set(Identifier, Attribute), the identifier lookup doesn't return the sorted insertion point (either end() or the found attribute). You can either:

1. insert wherever and unset the sorted flag dictionarySorted.setInt(false)
2. lookup the sorted insertion point and insert maintaining sortedness

I chose (2) because that was the previous behaviour of the function.

[Mogball]

Add swap comment

mlir/include/mlir/IR/BuiltinAttributes.td

let extraClassDeclaration = [{
     /// Return the specified attribute if present, null otherwise.
     Attribute get(StringRef name) const;
     Attribute get(Identifier name) const;
 
     /// Return the specified named attribute if present, None otherwise.
     Optional<NamedAttribute> getNamed(StringRef name) const;
     Optional<NamedAttribute> getNamed(Identifier name) const;
 
+    /// Return whether the specified attribute is present.
+    bool contains(StringRef name) const;
+    bool contains(Identifier name) const;
+
     /// Support range iteration.
     using iterator = llvm::ArrayRef<NamedAttribute>::iterator;
     iterator begin() const;
     iterator end() const;
     bool empty() const { return size() == 0; }
     size_t size() const;
 
     /// Sorts the NamedAttributes in the array ordered by name as expected by

mlir/include/mlir/IR/Operation.h

template <typename AttrClass> AttrClass getAttrOfType(Identifier name) {
     return getAttr(name).dyn_cast_or_null<AttrClass>();
   }
   template <typename AttrClass> AttrClass getAttrOfType(StringRef name) {
     return getAttr(name).dyn_cast_or_null<AttrClass>();
   }
 
   /// Return true if the operation has an attribute with the provided name,
   /// false otherwise.
-  bool hasAttr(Identifier name) { return static_cast<bool>(getAttr(name)); }
-  bool hasAttr(StringRef name) { return static_cast<bool>(getAttr(name)); }
+  bool hasAttr(Identifier name) { return attrs.contains(name); }
+  bool hasAttr(StringRef name) { return attrs.contains(name); }
   template <typename AttrClass, typename NameT>
   bool hasAttrOfType(NameT &&name) {
     return static_cast<bool>(
         getAttrOfType<AttrClass>(std::forward<NameT>(name)));
   }
 
   /// 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.

mlir/include/mlir/IR/OperationSupport.h

private:
 
   /// A list of attribute names registered to this operation in identifier form.
   /// This allows for operation classes to use identifiers for attribute
   /// lookup/creation/etc., as opposed to strings.
   ArrayRef<Identifier> attributeNames;
 };
 
 //===----------------------------------------------------------------------===//
+// Attribute Dictionary-Like Interface
+//===----------------------------------------------------------------------===//
+
+/// Attribute collections provide a dictionary-like interface. Define common
+/// lookup functions.
+namespace impl {
+
+/// Unsorted string search or identifier lookups are linear scans.
+template <typename IteratorT, typename NameT>
+std::pair<IteratorT, bool> findAttrUnsorted(IteratorT first, IteratorT last,
+                                            NameT name) {
+  for (auto it = first; it != last; ++it)
+    if (it->first == name)
+      return {it, true};
+  return {last, false};
+}
+
+/// Using llvm::lower_bound requires an extra string comparison to check whether
+/// the returned iterator points to the found element or whether it indicates
+/// the lower bound. Skip this redundant comparison by checking if `compare ==
+/// 0` during the binary search.
+template <typename IteratorT>
+std::pair<IteratorT, bool> findAttrSorted(IteratorT first, IteratorT last,
+                                          StringRef name) {
+  ptrdiff_t length = std::distance(first, last);
+
+  while (length > 0) {
+    ptrdiff_t half = length / 2;
+    IteratorT mid = first + half;
+    int compare = mid->first.strref().compare(name);
+    if (compare < 0) {
+      first = mid + 1;
+      length = length - half - 1;
+    } else if (compare > 0) {
+      length = half;
+    } else {
+      return {mid, true};
+    }
+  }
+  return {first, false};
+}
+
+/// Identifier lookups on large attribute lists will switch to string binary
+/// search. String binary searches become significantly faster than linear scans
+/// with the identifier when the attribute list becomes very large.
+template <typename IteratorT>
+std::pair<IteratorT, bool> findAttrSorted(IteratorT first, IteratorT last,
+                                          Identifier name) {
+  constexpr unsigned kSmallAttributeList = 16;
+  if (std::distance(first, last) > kSmallAttributeList)
+    return findAttrSorted(first, last, name.strref());
+  return findAttrUnsorted(first, last, name);
+}
+
+} // end namespace impl
+
+//===----------------------------------------------------------------------===//
 // NamedAttrList
 //===----------------------------------------------------------------------===//
 
 /// NamedAttrList is array of NamedAttributes that tracks whether it is sorted
 /// and does some basic work to remain sorted.
 class NamedAttrList {
 public:
+  using iterator = SmallVectorImpl<NamedAttribute>::iterator;
   using const_iterator = SmallVectorImpl<NamedAttribute>::const_iterator;
-  using const_reference = const NamedAttribute &;
   using reference = NamedAttribute &;
+  using const_reference = const NamedAttribute &;
   using size_type = size_t;
 
   NamedAttrList() : dictionarySorted({}, true) {}
   NamedAttrList(ArrayRef<NamedAttribute> attributes);
   NamedAttrList(DictionaryAttr attributes);
   NamedAttrList(const_iterator in_start, const_iterator in_end);
 
   bool operator!=(const NamedAttrList &other) const {
public:
   /// Append the given named attribute.
   void append(NamedAttribute attr) { push_back(attr); }
 
   /// Add an array of named attributes.
   template <typename RangeT>
   void append(RangeT &&newAttributes) {
     append(std::begin(newAttributes), std::end(newAttributes));
   }
 
   /// Add a range of named attributes.
   template <typename IteratorT>
   void append(IteratorT in_start, IteratorT in_end) {
     // TODO: expand to handle case where values appended are in order & after
     // end of current list.
     dictionarySorted.setPointerAndInt(nullptr, false);
     attrs.append(in_start, in_end);
   }
public:
   Attribute set(StringRef name, Attribute value);
 
   /// Erase the attribute with the given name from the list. Return the
   /// attribute that was erased, or nullptr if there was no attribute with such
   /// name.
   Attribute erase(Identifier name);
   Attribute erase(StringRef name);
 
+  iterator begin() { return attrs.begin(); }
+  iterator end() { return attrs.end(); }
   const_iterator begin() const { return attrs.begin(); }
   const_iterator end() const { return attrs.end(); }
 
   NamedAttrList &operator=(const SmallVectorImpl<NamedAttribute> &rhs);
   operator ArrayRef<NamedAttribute>() const;
 
 private:
   /// Return whether the attributes are sorted.
   bool isSorted() const { return dictionarySorted.getInt(); }
 
   /// Erase the attribute at the given iterator position.
   Attribute eraseImpl(SmallVectorImpl<NamedAttribute>::iterator it);
 
+  /// Lookup an attribute in the list.
+  template <typename AttrListT, typename NameT>
+  static auto findAttr(AttrListT &attrs, NameT name) {
+    return attrs.isSorted()
+               ? impl::findAttrSorted(attrs.begin(), attrs.end(), name)
+               : impl::findAttrUnsorted(attrs.begin(), attrs.end(), name);
+  }
+
   // These are marked mutable as they may be modified (e.g., sorted)
   mutable SmallVector<NamedAttribute, 4> attrs;
   // Pair with cached DictionaryAttr and status of whether attrs is sorted.
   // Note: just because sorted does not mean a DictionaryAttr has been created
   // but the case where there is a DictionaryAttr but attrs isn't sorted should
   // not occur.
   mutable llvm::PointerIntPair<Attribute, 1, bool> dictionarySorted;
 };
 
 //===----------------------------------------------------------------------===//
 // OperationName
 //===----------------------------------------------------------------------===//
 
 class OperationName {
 public:
   using RepresentationUnion =
       PointerUnion<Identifier, const AbstractOperation *>;

mlir/lib/IR/BuiltinAttributes.cpp

assert(llvm::is_sorted(value,
          "expected attribute values to be sorted");
   assert(!findDuplicateElement(value) &&
          "DictionaryAttr element names must be unique");
   return Base::get(context, value);
 }
 
 /// Return the specified attribute if present, null otherwise.
 Attribute DictionaryAttr::get(StringRef name) const {
-  Optional<NamedAttribute> attr = getNamed(name);
-  return attr ? attr->second : nullptr;
+  auto it = impl::findAttrSorted(begin(), end(), name);
+  return it.second ? it.first->second : Attribute();
 }
 Attribute DictionaryAttr::get(Identifier name) const {
-  Optional<NamedAttribute> attr = getNamed(name);
-  return attr ? attr->second : nullptr;
+  auto it = impl::findAttrSorted(begin(), end(), name);
+  return it.second ? it.first->second : Attribute();
 }
 
 /// Return the specified named attribute if present, None otherwise.
 Optional<NamedAttribute> DictionaryAttr::getNamed(StringRef name) const {
-  ArrayRef<NamedAttribute> values = getValue();
-  const auto *it = llvm::lower_bound(values, name);
-  return it != values.end() && it->first == name ? *it
-                                                 : Optional<NamedAttribute>();
+  auto it = impl::findAttrSorted(begin(), end(), name);
+  return it.second ? *it.first : Optional<NamedAttribute>();
 }
 Optional<NamedAttribute> DictionaryAttr::getNamed(Identifier name) const {
-  for (auto elt : getValue())
-    if (elt.first == name)
-      return elt;
-  return llvm::None;
+  auto it = impl::findAttrSorted(begin(), end(), name);
+  return it.second ? *it.first : Optional<NamedAttribute>();
+}
+
+/// Return whether the specified attribute is present.
+bool DictionaryAttr::contains(StringRef name) const {
+  return impl::findAttrSorted(begin(), end(), name).second;
+}
+bool DictionaryAttr::contains(Identifier name) const {
+  return impl::findAttrSorted(begin(), end(), name).second;
 }
 
 DictionaryAttr::iterator DictionaryAttr::begin() const {
   return getValue().begin();
 }
 DictionaryAttr::iterator DictionaryAttr::end() const {
   return getValue().end();
 }

mlir/lib/IR/OperationSupport.cpp

void NamedAttrList::push_back(NamedAttribute newAttribute) {
   if (isSorted())
     dictionarySorted.setInt(
         attrs.empty() ||
         strcmp(attrs.back().first.data(), newAttribute.first.data()) < 0);
   dictionarySorted.setPointer(nullptr);
   attrs.push_back(newAttribute);
 }
 
-/// Helper function to find attribute in possible sorted vector of
-/// NamedAttributes.
-template <typename T>
-static auto *findAttr(SmallVectorImpl<NamedAttribute> &attrs, T name,
-                      bool sorted) {
-  if (!sorted) {
-    return llvm::find_if(
-        attrs, [name](NamedAttribute attr) { return attr.first == name; });
-  }
-
-  auto *it = llvm::lower_bound(attrs, name);
-  if (it == attrs.end() || it->first != name)
-    return attrs.end();
-  return it;
-}
-
 /// Return the specified attribute if present, null otherwise.
 Attribute NamedAttrList::get(StringRef name) const {
-  auto *it = findAttr(attrs, name, isSorted());
-  return it != attrs.end() ? it->second : nullptr;
+  auto it = findAttr(*this, name);
+  return it.second ? it.first->second : Attribute();
 }
-
-/// Return the specified attribute if present, null otherwise.
 Attribute NamedAttrList::get(Identifier name) const {
-  auto *it = findAttr(attrs, name, isSorted());
-  return it != attrs.end() ? it->second : nullptr;
+  auto it = findAttr(*this, name);
+  return it.second ? it.first->second : Attribute();
 }
 
 /// Return the specified named attribute if present, None otherwise.
 Optional<NamedAttribute> NamedAttrList::getNamed(StringRef name) const {
-  auto *it = findAttr(attrs, name, isSorted());
-  return it != attrs.end() ? *it : Optional<NamedAttribute>();
+  auto it = findAttr(*this, name);
+  return it.second ? *it.first : Optional<NamedAttribute>();
 }
 Optional<NamedAttribute> NamedAttrList::getNamed(Identifier name) const {
-  auto *it = findAttr(attrs, name, isSorted());
-  return it != attrs.end() ? *it : Optional<NamedAttribute>();
+  auto it = findAttr(*this, name);
+  return it.second ? *it.first : Optional<NamedAttribute>();
 }
 
 /// 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.
 Attribute NamedAttrList::set(Identifier name, Attribute value) {
   assert(value && "attributes may never be null");
 
-  // Look for an existing value for the given name, and set it in-place.
-  auto *it = findAttr(attrs, name, isSorted());
-  if (it != attrs.end()) {
-    // Only update if the value is different from the existing.
-    Attribute oldValue = it->second;
-    if (oldValue != value) {
+  // Look for an existing attribute with the given name, and set its value
+  // in-place. Return the previous value of the attribute, if there was one.
+  auto it = findAttr(*this, name);
+  if (it.second) {
+    // Update the existing attribute by swapping out the old value for the new
+    // value. Return the old value.
+    if (it.first->second != value) {
+      std::swap(it.first->second, value);
+      // If the attributes have changed, the dictionary is invalidated.
       dictionarySorted.setPointer(nullptr);
-      it->second = value;
     }
-    return oldValue;
+    return value;
   }
-
-  // Otherwise, insert the new attribute into its sorted position.
-  it = llvm::lower_bound(attrs, name);
+  // Perform a string lookup to insert the new attribute into its sorted
+  // position.
+  if (isSorted())
+    it = findAttr(*this, name.strref());
+  attrs.insert(it.first, {name, value});
+  // Invalidate the dictionary. Return null as there was no previous value.
   dictionarySorted.setPointer(nullptr);
-  attrs.insert(it, {name, value});
   return Attribute();
 }
+
 Attribute NamedAttrList::set(StringRef name, Attribute value) {
-  assert(value && "setting null attribute not supported");
+  assert(value && "attributes may never be null");
   return set(mlir::Identifier::get(name, value.getContext()), value);
 }
 
 Attribute
 NamedAttrList::eraseImpl(SmallVectorImpl<NamedAttribute>::iterator it) {
-  if (it == attrs.end())
-    return nullptr;
-
   // Erasing does not affect the sorted property.
   Attribute attr = it->second;
   attrs.erase(it);
   dictionarySorted.setPointer(nullptr);
   return attr;
 }
 
 Attribute NamedAttrList::erase(Identifier name) {
-  return eraseImpl(findAttr(attrs, name, isSorted()));
+  auto it = findAttr(*this, name);
+  return it.second ? eraseImpl(it.first) : Attribute();
 }
 
 Attribute NamedAttrList::erase(StringRef name) {
-  return eraseImpl(findAttr(attrs, name, isSorted()));
+  auto it = findAttr(*this, name);
+  return it.second ? eraseImpl(it.first) : Attribute();
 }
 
 NamedAttrList &
 NamedAttrList::operator=(const SmallVectorImpl<NamedAttribute> &rhs) {
   assign(rhs.begin(), rhs.end());
   return *this;
 }