[Mlir-commits] [mlir] 092634f - [mlir] Store registered op info only in the by name map (#236949)

[llvmlistbot at llvm.org]

Author: Arseniy Obolenskiy
Date: 2026-10-30T10:19:57Z
New Revision: 092634f93290d162acf38c0f64269aea45404394

URL: https://github.com/llvm/llvm-project/commit/092634f93290d162acf38c0f64269aea45404394
DIFF: https://github.com/llvm/llvm-project/commit/092634f93290d162acf38c0f64269aea45404394.diff

LOG: [mlir] Store registered op info only in the by name map (#236949)

Loading a dialect inserted each op into two string maps

Insert it into one and take its entries from a bump allocator instead of
one malloc each

This speeds up context setup, which is noticable for short runs on small
inputs

Added: 
    

Modified: 
    mlir/lib/IR/MLIRContext.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 350680692a6e0..f5a878edd5361 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -182,12 +182,15 @@ class MLIRContextImpl {
   /// An allocator used for AbstractAttribute and AbstractType objects.
   llvm::BumpPtrAllocator abstractDialectSymbolAllocator;
 
-  /// This is a mapping from operation name to the operation info describing it.
-  llvm::StringMap<std::unique_ptr<OperationName::Impl>> operations;
+  /// This is a mapping from unregistered operation name to the operation info
+  /// describing it.
+  llvm::StringMap<std::unique_ptr<OperationName::Impl>>
+      unregisteredOperationsByName;
 
   /// A vector of operation info specifically for registered operations.
   llvm::DenseMap<TypeID, RegisteredOperationName> registeredOperations;
-  llvm::StringMap<RegisteredOperationName> registeredOperationsByName;
+  llvm::StringMap<std::unique_ptr<OperationName::Impl>, llvm::BumpPtrAllocator>
+      registeredOperationsByName;
 
   /// This is a sorted container of registered operations for a deterministic
   /// and efficient `getRegisteredOperations` implementation.
@@ -271,7 +274,8 @@ class MLIRContextImpl {
 
   /// Bundled state dynamically allocated when in a transient scope.
   struct TransientScopeState {
-    /// Set of operation names in `operations` at snapshot time.
+    /// Set of operation names in `unregisteredOperationsByName` at snapshot
+    /// time.
     llvm::DenseSet<StringRef> baseOperations;
 
     /// Number of entries in `dialectReferencingStrAttrs` per dialect at
@@ -705,10 +709,10 @@ void MLIRContext::beginTransientScope() {
   ctxImpl.affineUniquer.beginTransientScope();
   ctxImpl.distinctAttributeAllocator.beginTransientScope();
 
-  // Record base operations in operations map.
+  // Record base unregistered operations.
   {
     llvm::sys::SmartScopedReader<true> contextLock(ctxImpl.operationInfoMutex);
-    for (const auto &entry : ctxImpl.operations)
+    for (const auto &entry : ctxImpl.unregisteredOperationsByName)
       ctxImpl.transientState->baseOperations.insert(entry.first());
   }
 
@@ -735,13 +739,11 @@ void MLIRContext::endTransientScope() {
   {
     llvm::sys::SmartScopedWriter<true> contextLock(ctxImpl.operationInfoMutex);
     SmallVector<StringRef> opsToErase;
-    for (const auto &entry : ctxImpl.operations) {
-      if (!entry.second->isRegistered() &&
-          !ctxImpl.transientState->baseOperations.contains(entry.first()))
+    for (const auto &entry : ctxImpl.unregisteredOperationsByName)
+      if (!ctxImpl.transientState->baseOperations.contains(entry.first()))
         opsToErase.push_back(entry.first());
-    }
     for (StringRef op : opsToErase)
-      ctxImpl.operations.erase(op);
+      ctxImpl.unregisteredOperationsByName.erase(op);
   }
 
   // Restore dialect referencing string attributes before destroying the
@@ -915,21 +917,21 @@ OperationName::Impl::Impl(StringRef name, Dialect *dialect, TypeID typeID,
 OperationName::OperationName(StringRef name, MLIRContext *context) {
   MLIRContextImpl &ctxImpl = context->getImpl();
 
+  // Check the registered info map first. In the overwhelmingly common case,
+  // the entry will be in here and it also removes the need to acquire any
+  // locks.
+  auto registeredIt = ctxImpl.registeredOperationsByName.find(name);
+  if (LLVM_LIKELY(registeredIt != ctxImpl.registeredOperationsByName.end())) {
+    impl = registeredIt->second.get();
+    return;
+  }
+
   // Check for an existing name in read-only mode.
   bool isMultithreadingEnabled = context->isMultithreadingEnabled();
   if (isMultithreadingEnabled) {
-    // Check the registered info map first. In the overwhelmingly common case,
-    // the entry will be in here and it also removes the need to acquire any
-    // locks.
-    auto registeredIt = ctxImpl.registeredOperationsByName.find(name);
-    if (LLVM_LIKELY(registeredIt != ctxImpl.registeredOperationsByName.end())) {
-      impl = registeredIt->second.impl;
-      return;
-    }
-
     llvm::sys::SmartScopedReader<true> contextLock(ctxImpl.operationInfoMutex);
-    auto it = ctxImpl.operations.find(name);
-    if (it != ctxImpl.operations.end()) {
+    auto it = ctxImpl.unregisteredOperationsByName.find(name);
+    if (it != ctxImpl.unregisteredOperationsByName.end()) {
       impl = it->second.get();
       return;
     }
@@ -938,7 +940,7 @@ OperationName::OperationName(StringRef name, MLIRContext *context) {
   // Acquire a writer-lock so that we can safely create the new instance.
   ScopedWriterLock lock(ctxImpl.operationInfoMutex, isMultithreadingEnabled);
 
-  auto it = ctxImpl.operations.try_emplace(name);
+  auto it = ctxImpl.unregisteredOperationsByName.try_emplace(name);
   if (it.second) {
     auto nameAttr = StringAttr::get(context, name);
     it.first->second = std::make_unique<UnregisteredOpModel>(
@@ -1065,7 +1067,7 @@ RegisteredOperationName::lookup(StringRef name, MLIRContext *ctx) {
   auto &impl = ctx->getImpl();
   auto it = impl.registeredOperationsByName.find(name);
   if (it != impl.registeredOperationsByName.end())
-    return it->getValue();
+    return RegisteredOperationName(it->second.get());
   return std::nullopt;
 }
 
@@ -1091,15 +1093,13 @@ void RegisteredOperationName::insert(
     impl->attributeNames = cachedAttrNames;
   }
   StringRef name = impl->getName().strref();
-  // Insert the operation info if it doesn't exist yet.
-  ctxImpl.operations[name] = std::move(ownedImpl);
 
   // Update the registered info for this operation.
   auto emplaced = ctxImpl.registeredOperations.try_emplace(
       impl->getTypeID(), RegisteredOperationName(impl));
   assert(emplaced.second && "operation name registration must be successful");
   auto emplacedByName = ctxImpl.registeredOperationsByName.try_emplace(
-      name, RegisteredOperationName(impl));
+      name, std::move(ownedImpl));
   (void)emplacedByName;
   assert(emplacedByName.second &&
          "operation name registration must be successful");