[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

471

[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

94

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

99–104

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

[Mogball]

mlir/lib/IR/OperationSupport.cpp

99–104

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/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 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) {
+  auto length = std::distance(first, last);
+  using difference_type = decltype(length);
+
+  while (length > 0) {
+    difference_type half = length / 2;
+    IteratorT mid = first + half;
+    auto 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 when the size of
+/// the attribute list exceeds 16.
+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);
+}
+
+/// CRTP class to implement a dictionary-like interface on a collection of
+/// named attributes. Derived classes are required simply to implement `begin`,
+/// `end`, and `isSorted`, which returns true if the underlying representation
+/// has been sorted by the attributes' string names.
+template <typename DerivedT>
+class AttributeDictionaryLike {
+protected:
+  /// Common attribute lookup functions.
+  template <typename DerivedPtrT, typename NameT>
+  static auto findAttr(DerivedPtrT derived, NameT name) {
+    return derived->isSorted()
+               ? findAttrSorted(derived->begin(), derived->end(), name)
+               : findAttrUnsorted(derived->begin(), derived->end(), name);
+  }
+  template <typename NameT>
+  auto findAttr(NameT name) {
+    return findAttr(static_cast<DerivedT *>(this), name);
+  }
+  template <typename NameT>
+  auto findAttr(NameT name) const {
+    return findAttr(static_cast<const DerivedT *>(this), name);
+  }
+
+public:
+  using BaseT = AttributeDictionaryLike<DerivedT>;
+
+  /// Lookup the value of an attribute. Returns null if it was not found.
+  template <typename NameT>
+  Attribute get(NameT name) const {
+    auto it = findAttr(name);
+    return it.second ? it.first->second : Attribute();
+  }
+
+  /// Lookup an attribute and return it as a named attribute if one was found.
+  template <typename NameT>
+  Optional<NamedAttribute> getNamed(NameT name) const {
+    auto it = findAttr(name);
+    return it.second ? *it.first : Optional<NamedAttribute>();
+  }
+
+  /// Returns whether an attribute exists.
+  template <typename NameT>
+  bool has(NameT name) const {
+    return findAttr(name).second;
+  }
+};
+
+} // 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 {
+class NamedAttrList : public impl::AttributeDictionaryLike<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:
 
   /// Return a dictionary attribute for the underlying dictionary. This will
   /// return an empty dictionary attribute if empty rather than null.
   DictionaryAttr getDictionary(MLIRContext *context) const;
 
   /// Return all of the attributes on this operation.
   ArrayRef<NamedAttribute> getAttrs() const;
 
-  /// Return the specified attribute if present, null otherwise.
-  Attribute get(Identifier name) const;
-  Attribute get(StringRef name) const;
-
-  /// Return the specified named attribute if present, None otherwise.
-  Optional<NamedAttribute> getNamed(StringRef name) const;
-  Optional<NamedAttribute> getNamed(Identifier name) const;
-
   /// 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.
   /// Returns the previous attribute value of `name`, or null if no
   /// attribute previously existed with `name`.
   Attribute set(Identifier name, Attribute value);
   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);
 
   // 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;
+
+  /// Give AttributeDictionaryLike access to `isSorted.
+  friend class impl::AttributeDictionaryLike<NamedAttrList>;
 };
 
 //===----------------------------------------------------------------------===//
 // OperationName
 //===----------------------------------------------------------------------===//
 
 class OperationName {
 public:

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>();
 }
 
 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;
-}
-
-/// 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;
-}
-
-/// 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>();
-}
-Optional<NamedAttribute> NamedAttrList::getNamed(Identifier name) const {
-  auto *it = findAttr(attrs, name, isSorted());
-  return it != attrs.end() ? *it : 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()) {
+  auto it = findAttr(name);
+  if (it.second) {
     // Only update if the value is different from the existing.
-    Attribute oldValue = it->second;
-    if (oldValue != value) {
+    if (it.first->second != value) {
+      std::swap(it.first->second, value);
       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(name.strref());
+  attrs.insert(it.first, {name, 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(name);
+  return it.second ? eraseImpl(it.first) : Attribute();
 }
 
 Attribute NamedAttrList::erase(StringRef name) {
-  return eraseImpl(findAttr(attrs, name, isSorted()));
+  auto it = findAttr(name);
+  return it.second ? eraseImpl(it.first) : Attribute();
 }
 
 NamedAttrList &
 NamedAttrList::operator=(const SmallVectorImpl<NamedAttribute> &rhs) {
   assign(rhs.begin(), rhs.end());
   return *this;
 }