Avoid string comparisons on the fast path of MLIR Identifier lookup (NFC)

[rriddle]

mlir/lib/IR/MLIRContext.cpp

733

This is still the fast path when in single threaded mode, unless we don't really care about that one?

[mehdi_amini]

Avoid the string comparison on the single-thread fast-path

mlir/lib/IR/MLIRContext.cpp

 #endif
     dialect = ctor();
     assert(dialect && "dialect ctor failed");
 
     // Refresh all the identifiers dialect field, this catches cases where a
     // dialect may be loaded after identifier prefixed with this dialect name
     // were already created.
     for (auto &identifierEntry : impl.identifiers)
-      if (identifierEntry.first().startswith(dialectNamespace))
+      if (!identifierEntry.second &&
+          identifierEntry.first().startswith(dialectNamespace))
         identifierEntry.second = dialect.get();
 
     return dialect.get();
   }
 
   // Abort if dialect with namespace has already been registered.
   if (dialect->getTypeID() != dialectID)
     llvm::report_fatal_error("a dialect with namespace '" + dialectNamespace +

 Identifier Identifier::get(StringRef str, MLIRContext *context) {
   // Check invariants after seeing if we already have something in the
   // identifier table - if we already had it in the table, then it already
   // passed invariant checks.
   assert(!str.empty() && "Cannot create an empty identifier");
   assert(str.find('\0') == StringRef::npos &&
          "Cannot create an identifier with a nul character");
 
+  auto getDialectOrContext = [&]() {
     PointerUnion<Dialect *, MLIRContext *> dialectOrContext = context;
     auto dialectNamePair = str.split('.');
     if (!dialectNamePair.first.empty())
       if (Dialect *dialect = context->getLoadedDialect(dialectNamePair.first))
         dialectOrContext = dialect;
+    return dialectOrContext;
+  };
 
   auto &impl = context->getImpl();
-  if (!context->isMultithreadingEnabled())
-    return Identifier(&*impl.identifiers.insert({str, dialectOrContext}).first);
+  if (!context->isMultithreadingEnabled()) {
+    auto insertedIt = impl.identifiers.insert({str, nullptr});
+    if (insertedIt.second)
+      insertedIt.first->second = getDialectOrContext();
+    return Identifier(&*insertedIt.first);
+  }
 
   // Check for an existing instance in the local cache.
   auto *&localEntry = (*impl.localIdentifierCache)[str];
   if (localEntry)
     return Identifier(localEntry);
 
   // Check for an existing identifier in read-only mode.
   {
     llvm::sys::SmartScopedReader<true> contextLock(impl.identifierMutex);
     auto it = impl.identifiers.find(str);
     if (it != impl.identifiers.end()) {
       localEntry = &*it;
       return Identifier(localEntry);
     }
   }
 
   // Acquire a writer-lock so that we can safely create the new instance.
   llvm::sys::SmartScopedWriter<true> contextLock(impl.identifierMutex);
-  auto it = impl.identifiers.insert({str, dialectOrContext}).first;
+  auto it = impl.identifiers.insert({str, getDialectOrContext()}).first;
   localEntry = &*it;
   return Identifier(localEntry);
 }
 
 Dialect *Identifier::getDialect() {
   return entry->second.dyn_cast<Dialect *>();
 }