[mlir][SymbolTable] Use Identifier instead of StringRef when looking …

…up symbol name attributes

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

Differential Revision: https://reviews.llvm.org/D89660

mlir/lib/IR/SymbolTable.cpp

@@ -27,6 +27,11 @@ static Optional<StringRef> getNameIfSymbol(Operation *symbol) {
       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
@@ -49,12 +54,15 @@ collectValidReferencesFor(Operation *symbol, StringRef symbolName,
 
   // 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));
@@ -106,8 +114,10 @@ SymbolTable::SymbolTable(Operation *symbolTableOp)
   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;
 
@@ -269,8 +279,10 @@ Operation *SymbolTable::lookupSymbolIn(Operation *symbolTableOp,
   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;
 }