[mlir][AsmPrinter] Don't use string comparison when filtering list attributes

[rriddle]

Summary

In .mlir modules with large amounts of attributes, e.g. a function with a larger number of argument attributes, the string comparison filtering greatly affects compile time. This revision switches to using a SmallDenseSet in these situations, resulting in over a 10x speed up in some situations.

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

451

NamedAttribute are keyed on identifier, which are allowing pointer hashing/comparison instead of string, can you use these in the DenseSet?

[rriddle]

mlir/lib/IR/AsmPrinter.cpp

451

elidedAttrs is an array of StringRef, so we would have to unique identifiers for each of them. The additional overhead of uniquing elidedAttrs is noticeable, and DenseSet<StringRef> ends up being faster(~200-250 ms).

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

451

Surprising to me, it is because we have a lock in the context when retrieving the Identifiers from the elidedAttr?

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

451

Actually there would be two hashes using Identifier: one to convert the StringRef to an Identifier, and then another one to hash the Identifier pointer to insert in the DenseSet...

[rriddle]

mlir/lib/IR/AsmPrinter.cpp

451

I think it is a combo of all of the above, multiple hashes+locking+etc.

mlir/lib/IR/AsmPrinter.cpp

private:
   void printAttributeWithoutType(Attribute attr) override {
     printAttribute(attr);
   }
 
   /// Print the given set of attributes with names not included within
   /// 'elidedAttrs'.
   void printOptionalAttrDict(ArrayRef<NamedAttribute> attrs,
                              ArrayRef<StringRef> elidedAttrs = {}) override {
-    // Filter out any attributes that shouldn't be included.
-    SmallVector<NamedAttribute, 8> filteredAttrs(
-        llvm::make_filter_range(attrs, [&](NamedAttribute attr) {
-          return !llvm::is_contained(elidedAttrs, attr.first.strref());
-        }));
-    for (const NamedAttribute &attr : filteredAttrs)
+    if (attrs.empty())
+      return;
+    if (elidedAttrs.empty()) {
+      for (const NamedAttribute &attr : attrs)
+        printAttribute(attr.second);
+      return;
+    }
+    llvm::SmallDenseSet<StringRef> elidedAttrsSet(elidedAttrs.begin(),
+                                                  elidedAttrs.end());
+    for (const NamedAttribute &attr : attrs)
+      if (!elidedAttrsSet.contains(attr.first.strref()))
         printAttribute(attr.second);
   }
   void printOptionalAttrDictWithKeyword(
       ArrayRef<NamedAttribute> attrs,
       ArrayRef<StringRef> elidedAttrs = {}) override {
     printOptionalAttrDict(attrs, elidedAttrs);
   }
 
   /// Return 'nulls' as the output stream, this will ignore any data fed to it.

 
 void ModulePrinter::printOptionalAttrDict(ArrayRef<NamedAttribute> attrs,
                                           ArrayRef<StringRef> elidedAttrs,
                                           bool withKeyword) {
   // If there are no attributes, then there is nothing to be done.
   if (attrs.empty())
     return;
 
-  // Filter out any attributes that shouldn't be included.
-  SmallVector<NamedAttribute, 8> filteredAttrs(
-      llvm::make_filter_range(attrs, [&](NamedAttribute attr) {
-        return !llvm::is_contained(elidedAttrs, attr.first.strref());
-      }));
-
-  // If there are no attributes left to print after filtering, then we're done.
-  if (filteredAttrs.empty())
-    return;
-
+  // Functor used to print a filtered attribute list.
+  auto printFilteredAttributesFn = [&](auto filteredAttrs) {
     // Print the 'attributes' keyword if necessary.
     if (withKeyword)
       os << " attributes";
 
     // Otherwise, print them all out in braces.
     os << " {";
     interleaveComma(filteredAttrs,
                     [&](NamedAttribute attr) { printNamedAttribute(attr); });
     os << '}';
+  };
+
+  // If no attributes are elided, we can directly print with no filtering.
+  if (elidedAttrs.empty())
+    return printFilteredAttributesFn(attrs);
+
+  // Otherwise, filter out any attributes that shouldn't be included.
+  llvm::SmallDenseSet<StringRef> elidedAttrsSet(elidedAttrs.begin(),
+                                                elidedAttrs.end());
+  auto filteredAttrs = llvm::make_filter_range(attrs, [&](NamedAttribute attr) {
+    return !elidedAttrsSet.contains(attr.first.strref());
+  });
+  if (!filteredAttrs.empty())
+    printFilteredAttributesFn(filteredAttrs);
 }
 
 void ModulePrinter::printNamedAttribute(NamedAttribute attr) {
   if (isBareIdentifier(attr.first)) {
     os << attr.first;
   } else {
     os << '"';
     printEscapedString(attr.first.strref(), os);