[mlir][SymbolTable] Use Identifier instead of StringRef when looking up symbol name attributes

[rriddle]

Summary

Using an Identifier is much more efficient for attribute lookups because it uses pointer comparison as opposed to string comparison.

Depends On D89659

[jpienaar]

Nice, thanks

mlir/lib/IR/SymbolTable.cpp

 
 /// Returns the string name of the given symbol, or None if this is not a
 /// symbol.
 static Optional<StringRef> getNameIfSymbol(Operation *symbol) {
   auto nameAttr =
       symbol->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName());
   return nameAttr ? nameAttr.getValue() : Optional<StringRef>();
 }
+static Optional<StringRef> getNameIfSymbol(Operation *symbol,
+                                           Identifier symbolAttrNameId) {
+  auto nameAttr = symbol->getAttrOfType<StringAttr>(symbolAttrNameId);
+  return nameAttr ? nameAttr.getValue() : Optional<StringRef>();
+}
 
 /// Computes the nested symbol reference attribute for the symbol 'symbolName'
 /// that are usable within the symbol table operations from 'symbol' as far up
 /// to the given operation 'within', where 'within' is an ancestor of 'symbol'.
 /// Returns success if all references up to 'within' could be computed.
 static LogicalResult
 collectValidReferencesFor(Operation *symbol, StringRef symbolName,
                           Operation *within,
                           SmallVectorImpl<SymbolRefAttr> &results) {
   assert(within->isAncestor(symbol) && "expected 'within' to be an ancestor");
   MLIRContext *ctx = symbol->getContext();
 
   auto leafRef = FlatSymbolRefAttr::get(symbolName, ctx);
   results.push_back(leafRef);
 
   // Early exit for when 'within' is the parent of 'symbol'.
   Operation *symbolTableOp = symbol->getParentOp();
   if (within == symbolTableOp)
     return success();
 
   // Collect references until 'symbolTableOp' reaches 'within'.
   SmallVector<FlatSymbolRefAttr, 1> nestedRefs(1, leafRef);
+  Identifier symbolNameId =
+      Identifier::get(SymbolTable::getSymbolAttrName(), ctx);
   do {
     // Each parent of 'symbol' should define a symbol table.
     if (!symbolTableOp->hasTrait<OpTrait::SymbolTable>())
       return failure();
     // Each parent of 'symbol' should also be a symbol.
-    Optional<StringRef> symbolTableName = getNameIfSymbol(symbolTableOp);
+    Optional<StringRef> symbolTableName =
+        getNameIfSymbol(symbolTableOp, symbolNameId);
     if (!symbolTableName)
       return failure();
     results.push_back(SymbolRefAttr::get(*symbolTableName, nestedRefs, ctx));
 
     symbolTableOp = symbolTableOp->getParentOp();
     if (symbolTableOp == within)
       break;
     nestedRefs.insert(nestedRefs.begin(),
SymbolTable::SymbolTable(Operation *symbolTableOp)
     : symbolTableOp(symbolTableOp) {
   assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>() &&
          "expected operation to have SymbolTable trait");
   assert(symbolTableOp->getNumRegions() == 1 &&
          "expected operation to have a single region");
   assert(llvm::hasSingleElement(symbolTableOp->getRegion(0)) &&
          "expected operation to have a single block");
 
+  Identifier symbolNameId = Identifier::get(SymbolTable::getSymbolAttrName(),
+                                            symbolTableOp->getContext());
   for (auto &op : symbolTableOp->getRegion(0).front()) {
-    Optional<StringRef> name = getNameIfSymbol(&op);
+    Optional<StringRef> name = getNameIfSymbol(&op, symbolNameId);
     if (!name)
       continue;
 
     auto inserted = symbolTable.insert({*name, &op});
     (void)inserted;
     assert(inserted.second &&
            "expected region to contain uniquely named symbol operations");
   }

 /// regions of 'symbolTableOp'. 'symbolTableOp' is required to be an operation
 /// with the 'OpTrait::SymbolTable' trait. Returns nullptr if no valid symbol
 /// was found.
 Operation *SymbolTable::lookupSymbolIn(Operation *symbolTableOp,
                                        StringRef symbol) {
   assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>());
 
   // Look for a symbol with the given name.
+  Identifier symbolNameId = Identifier::get(SymbolTable::getSymbolAttrName(),
+                                            symbolTableOp->getContext());
   for (auto &op : symbolTableOp->getRegion(0).front().without_terminator())
-    if (getNameIfSymbol(&op) == symbol)
+    if (getNameIfSymbol(&op, symbolNameId) == symbol)
       return &op;
   return nullptr;
 }
 Operation *SymbolTable::lookupSymbolIn(Operation *symbolTableOp,
                                        SymbolRefAttr symbol) {
   SmallVector<Operation *, 4> resolvedSymbols;
   if (failed(lookupSymbolIn(symbolTableOp, symbol, resolvedSymbols)))
     return nullptr;