[mlir] Remove locking for dialect/operation registration.

[rriddle]

Summary

Moving forward dialects should only be registered in a thread safe context. This matches the existing usage we have today, but it allows for removing quite a bit of expensive locking from the context.

This led to ~.5 a second compile time improvement when running one conversion pass on a very large .mlir file(hundreds of thousands of operations).

Depends On D81883

[mehdi_amini]

Nice!

mlir/include/mlir/IR/Dialect.h

private:
   static void
   registerDialectAllocator(TypeID typeID,
                            const DialectAllocatorFunction &function);
   template <typename ConcreteDialect>
   friend void registerDialect();
 };
 /// Registers all dialects and hooks from the global registries with the
 /// specified MLIRContext.
+/// Note: This method is not thread-safe.
 void registerAllDialects(MLIRContext *context);
 
 /// Utility to register a dialect. Client can register their dialect with the
 /// global registry by calling registerDialect<MyDialect>();
+/// Note: This method is not thread-safe.
 template <typename ConcreteDialect> void registerDialect() {
   Dialect::registerDialectAllocator(TypeID::get<ConcreteDialect>(),
                                     [](MLIRContext *ctx) {
                                       // Just allocate the dialect, the context
                                       // takes ownership of it.
                                       new ConcreteDialect(ctx);
                                     });
 }

mlir/lib/IR/MLIRContext.cpp

public:
 
   /// If the current stack trace should be attached when emitting diagnostics.
   bool printStackTraceOnDiagnostic = false;
 
   //===--------------------------------------------------------------------===//
   // Other
   //===--------------------------------------------------------------------===//
 
-  /// A general purpose mutex to lock access to parts of the context that do not
-  /// have a more specific mutex, e.g. registry operations.
-  llvm::sys::SmartRWMutex<true> contextMutex;
-
   /// This is a list of dialects that are created referring to this context.
   /// The MLIRContext owns the objects.
   std::vector<std::unique_ptr<Dialect>> dialects;
 
   /// This is a mapping from operation name to AbstractOperation for registered
   /// operations.
   llvm::StringMap<AbstractOperation> registeredOperations;
 

 DiagnosticEngine &MLIRContext::getDiagEngine() { return getImpl().diagEngine; }
 
 //===----------------------------------------------------------------------===//
 // Dialect and Operation Registration
 //===----------------------------------------------------------------------===//
 
 /// Return information about all registered IR dialects.
 std::vector<Dialect *> MLIRContext::getRegisteredDialects() {
-  // Lock access to the context registry.
-  ScopedReaderLock registryLock(impl->contextMutex, impl->threadingIsEnabled);
   std::vector<Dialect *> result;
   result.reserve(impl->dialects.size());
   for (auto &dialect : impl->dialects)
     result.push_back(dialect.get());
   return result;
 }
 
 /// Get a registered IR dialect with the given namespace. If none is found,
 /// then return nullptr.
 Dialect *MLIRContext::getRegisteredDialect(StringRef name) {
-  // Lock access to the context registry.
-  ScopedReaderLock registryLock(impl->contextMutex, impl->threadingIsEnabled);
-
   // Dialects are sorted by name, so we can use binary search for lookup.
   auto it = llvm::lower_bound(
       impl->dialects, name,
       [](const auto &lhs, StringRef rhs) { return lhs->getNamespace() < rhs; });
   return (it != impl->dialects.end() && (*it)->getNamespace() == name)
              ? (*it).get()
              : nullptr;
 }
 
 /// Register this dialect object with the specified context.  The context
 /// takes ownership of the heap allocated dialect.
 void Dialect::registerDialect(MLIRContext *context) {
   auto &impl = context->getImpl();
   std::unique_ptr<Dialect> dialect(this);
 
-  // Lock access to the context registry.
-  ScopedWriterLock registryLock(impl.contextMutex, impl.threadingIsEnabled);
-
   // Get the correct insertion position sorted by namespace.
   auto insertPt = llvm::lower_bound(
       impl.dialects, dialect, [](const auto &lhs, const auto &rhs) {
         return lhs->getNamespace() < rhs->getNamespace();
       });
 
   // Abort if dialect with namespace has already been registered.
   if (insertPt != impl.dialects.end() &&

 void MLIRContext::printStackTraceOnDiagnostic(bool enable) {
   impl->printStackTraceOnDiagnostic = enable;
 }
 
 /// Return information about all registered operations.  This isn't very
 /// efficient, typically you should ask the operations about their properties
 /// directly.
 std::vector<AbstractOperation *> MLIRContext::getRegisteredOperations() {
-  std::vector<std::pair<StringRef, AbstractOperation *>> opsToSort;
-
-  { // Lock access to the context registry.
-    ScopedReaderLock registryLock(impl->contextMutex, impl->threadingIsEnabled);
-
-    // We just have the operations in a non-deterministic hash table order. Dump
-    // into a temporary array, then sort it by operation name to get a stable
-    // ordering.
-    llvm::StringMap<AbstractOperation> &registeredOps =
-        impl->registeredOperations;
 
-    opsToSort.reserve(registeredOps.size());
-    for (auto &elt : registeredOps)
-      opsToSort.push_back({elt.first(), &elt.second});
-  }
-
-  llvm::array_pod_sort(opsToSort.begin(), opsToSort.end());
+  std::vector<AbstractOperation *> result;
+  result.reserve(registeredOps.size());
+    result.push_back(&elt.second);
+  llvm::array_pod_sort(
+      result.begin(), result.end(),
+      [](AbstractOperation *const *lhs, AbstractOperation *const *rhs) {
+        return (*lhs)->name.compare((*rhs)->name);
+      });
 
-  std::vector<AbstractOperation *> result;
-  result.reserve(opsToSort.size());
-  for (auto &elt : opsToSort)
-    result.push_back(elt.second);
   return result;
 }
 
 bool MLIRContext::isOperationRegistered(StringRef name) {
-  // Lock access to the context registry.
-  ScopedReaderLock registryLock(impl->contextMutex, impl->threadingIsEnabled);
-
   return impl->registeredOperations.count(name);
 }
 
 void Dialect::addOperation(AbstractOperation opInfo) {
   assert((getNamespace().empty() || opInfo.dialect.name == getNamespace()) &&
          "op name doesn't start with dialect namespace");
   assert(&opInfo.dialect == this && "Dialect object mismatch");
   auto &impl = context->getImpl();
-
-  // Lock access to the context registry.
   StringRef opName = opInfo.name;
-  ScopedWriterLock registryLock(impl.contextMutex, impl.threadingIsEnabled);
   if (!impl.registeredOperations.insert({opName, std::move(opInfo)}).second) {
-    llvm::errs() << "error: operation named '" << opName
+    llvm::errs() << "error: operation named '" << opInfo.name
                  << "' is already registered.\n";
     abort();
   }
 }
 
 void Dialect::addType(TypeID typeID, AbstractType &&typeInfo) {
   auto &impl = context->getImpl();
-
-  // Lock access to the context registry.
-  ScopedWriterLock registryLock(impl.contextMutex, impl.threadingIsEnabled);
   auto *newInfo =
       new (impl.abstractDialectSymbolAllocator.Allocate<AbstractType>())
           AbstractType(std::move(typeInfo));
   if (!impl.registeredTypes.insert({typeID, newInfo}).second)
     llvm::report_fatal_error("Dialect Type already registered.");
 }
 
 void Dialect::addAttribute(TypeID typeID, AbstractAttribute &&attrInfo) {
   auto &impl = context->getImpl();
-
-  // Lock access to the context registry.
-  ScopedWriterLock registryLock(impl.contextMutex, impl.threadingIsEnabled);
   auto *newInfo =
       new (impl.abstractDialectSymbolAllocator.Allocate<AbstractAttribute>())
           AbstractAttribute(std::move(attrInfo));
   if (!impl.registeredAttributes.insert({typeID, newInfo}).second)
     llvm::report_fatal_error("Dialect Attribute already registered.");
 }
 
 /// Get the dialect that registered the attribute with the provided typeid.
 const AbstractAttribute &AbstractAttribute::lookup(TypeID typeID,
                                                    MLIRContext *context) {
   auto &impl = context->getImpl();
   auto it = impl.registeredAttributes.find(typeID);
   if (it == impl.registeredAttributes.end())
     llvm::report_fatal_error("Trying to create an Attribute that was not "
                              "registered in this MLIRContext.");
   return *it->second;
 }
 
 /// Look up the specified operation in the operation set and return a pointer
 /// to it if present.  Otherwise, return a null pointer.
 const AbstractOperation *AbstractOperation::lookup(StringRef opName,
                                                    MLIRContext *context) {
   auto &impl = context->getImpl();
-
-  // Lock access to the context registry.
-  ScopedReaderLock registryLock(impl.contextMutex, impl.threadingIsEnabled);
   auto it = impl.registeredOperations.find(opName);
   if (it != impl.registeredOperations.end())
     return &it->second;
   return nullptr;
 }
 
 /// Get the dialect that registered the type with the provided typeid.
 const AbstractType &AbstractType::lookup(TypeID typeID, MLIRContext *context) {