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

[Arseniy Obolenskiy]

https://github.com/aobolensk updated https://github.com/llvm/llvm-project/pull/236949

>From 1a502a7fd0820b591f33a991c7bad8a43de8af7b Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Thu, 29 Oct 2026 22:08:09 +0100
Subject: [PATCH 1/2] [mlir] Store registered op info only in the by name map

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
---
 mlir/lib/IR/MLIRContext.cpp | 27 +++++++++++++--------------
 1 file changed, 13 insertions(+), 14 deletions(-)

diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 350680692a6e0f..cd474d18d08858 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -187,7 +187,8 @@ class MLIRContextImpl {
 
   /// 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.
@@ -915,18 +916,18 @@ 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()) {
@@ -1065,7 +1066,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 +1092,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");

>From f8c10567a6a54ec5c5285d39d41f1a6c832af699 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Fri, 30 Oct 2026 11:03:12 +0100
Subject: [PATCH 2/2] Address comments

---
 mlir/lib/IR/MLIRContext.cpp | 27 ++++++++++++++-------------
 1 file changed, 14 insertions(+), 13 deletions(-)

diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index cd474d18d08858..f5a878edd53612 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -182,8 +182,10 @@ 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;
@@ -272,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
@@ -706,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());
   }
 
@@ -736,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
@@ -929,8 +930,8 @@ OperationName::OperationName(StringRef name, MLIRContext *context) {
   bool isMultithreadingEnabled = context->isMultithreadingEnabled();
   if (isMultithreadingEnabled) {
     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;
     }
@@ -939,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>(