[Mlir-commits] [mlir] [mlir] Support thread-safe context access without parallel execution (PR #236551)
Mehdi Amini
llvmlistbot at llvm.org
Thu Oct 29 03:46:43 PDT 2026
https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/236551
>From 8774c20cf780bfba906f72e1b52e73dd05331bc7 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 28 Oct 2026 12:11:15 -0700
Subject: [PATCH] [mlir] Support thread-safe context access without parallel
execution
Add enableThreadSafeLocking() and isThreadSafeLockingEnabled() so
clients using external threads can unique context-owned entities
without enabling MLIR parallel execution or creating a thread pool.
Use the locking mode for operation-name uniquing as well as the
type, attribute, and affine uniquers. Rename the StorageUniquer setter
to setThreadSafeLocking(bool enable) to reflect its purpose.
Enabling multithreading also enables locking. Retain the existing
behavior of disabling both together, and document how to enable
locking independently, including with --mlir-disable-threading.
Add coverage for concurrent uniquing, sequential MLIR dispatch,
threading transitions, and externally supplied thread pools.
Follow-up to #235294.
Assisted-by: Codex
---
mlir/docs/ReleaseNotes.md | 23 ++++
mlir/include/mlir/IR/MLIRContext.h | 26 +++++
mlir/include/mlir/Support/StorageUniquer.h | 4 +-
mlir/lib/IR/MLIRContext.cpp | 34 ++++--
mlir/lib/Support/StorageUniquer.cpp | 6 +-
mlir/unittests/IR/CMakeLists.txt | 1 +
mlir/unittests/IR/MLIRContextTest.cpp | 126 +++++++++++++++++++++
7 files changed, 205 insertions(+), 15 deletions(-)
create mode 100644 mlir/unittests/IR/MLIRContextTest.cpp
diff --git a/mlir/docs/ReleaseNotes.md b/mlir/docs/ReleaseNotes.md
index f57642308f4528..856afb3125ee23 100644
--- a/mlir/docs/ReleaseNotes.md
+++ b/mlir/docs/ReleaseNotes.md
@@ -10,6 +10,29 @@ specifically, it is a snapshot of the MLIR development at the time of the releas
## LLVM 24
+### Context Locking Without Parallel Execution
+
+Clients that access an `MLIRContext` from their own threads can enable
+thread-safe locking without creating a context-owned thread pool or enabling
+parallel execution within MLIR:
+
+```c++
+MLIRContext context(MLIRContext::Threading::DISABLED);
+context.enableThreadSafeLocking();
+```
+
+This enables locking for uniquing types, attributes, affine expressions, and
+operation names. Dialects must be loaded before accessing the context
+concurrently. `isMultithreadingEnabled()` remains false, so MLIR's parallel
+helpers execute sequentially. Enabling multithreading also enables locking;
+disabling multithreading disables locking, which can then be re-enabled
+independently. The `--mlir-disable-threading` debugging flag does not override
+explicit calls to `enableThreadSafeLocking()`.
+
+`StorageUniquer::disableMultithreading(disable)` has been renamed to
+`StorageUniquer::setThreadSafeLocking(!disable)` to reflect that it controls
+locking. The new setter takes an explicit boolean argument.
+
### Strict Property Assembly Formats
Property-backed operations now use strict declarative assembly formats by
diff --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h
index d63c863de6544c..72d3c72f6a1525 100644
--- a/mlir/include/mlir/IR/MLIRContext.h
+++ b/mlir/include/mlir/IR/MLIRContext.h
@@ -190,6 +190,8 @@ class MLIRContext {
};
/// Return true if multi-threading is enabled by the context.
+ /// This controls parallel execution within MLIR. Use
+ /// isThreadSafeLockingEnabled() to query locking for context access.
#if LLVM_ENABLE_THREADS
bool isMultithreadingEnabled();
#else
@@ -197,6 +199,9 @@ class MLIRContext {
#endif
/// Set the flag specifying if multi-threading is disabled by the context.
+ /// Enabling multi-threading also enables thread-safe locking. Disabling it
+ /// also disables thread-safe locking, which can be re-enabled independently
+ /// with enableThreadSafeLocking().
/// The command line debugging flag `--mlir-disable-threading` is overriding
/// this call and making it a no-op!
void disableMultithreading(bool disable = true);
@@ -204,6 +209,27 @@ class MLIRContext {
disableMultithreading(!enable);
}
+ /// Return true if thread-safe locking is enabled for context access.
+ /// This does not imply that parallel execution within MLIR is enabled.
+#if LLVM_ENABLE_THREADS
+ bool isThreadSafeLockingEnabled();
+#else
+ bool isThreadSafeLockingEnabled() { return false; }
+#endif
+
+ /// Enable thread-safe locking for context access, including
+ /// uniquing types, attributes, affine expressions, and operation names.
+ /// This does not create a thread pool or change isMultithreadingEnabled().
+ /// For example, clients using their own threads can construct a context with
+ /// Threading::DISABLED and then call enableThreadSafeLocking().
+ ///
+ /// Disabling multi-threading also disables locking. This method must be
+ /// called when no other threads are accessing the context. Dialects
+ /// must still be loaded before accessing the context concurrently.
+ /// Unlike enableMultithreading(), this is not overridden by the debugging
+ /// flag --mlir-disable-threading.
+ void enableThreadSafeLocking();
+
/// Set a new thread pool to be used in this context. This method requires
/// that multithreading is disabled for this context prior to the call. This
/// allows to share a thread pool across multiple contexts, as well as
diff --git a/mlir/include/mlir/Support/StorageUniquer.h b/mlir/include/mlir/Support/StorageUniquer.h
index 39018ebba106f9..1a73446d70401c 100644
--- a/mlir/include/mlir/Support/StorageUniquer.h
+++ b/mlir/include/mlir/Support/StorageUniquer.h
@@ -139,8 +139,8 @@ class StorageUniquer {
StorageUniquer();
~StorageUniquer();
- /// Set the flag specifying if multi-threading is disabled within the uniquer.
- void disableMultithreading(bool disable = true);
+ /// Set whether thread-safe locking is enabled within the uniquer.
+ void setThreadSafeLocking(bool enable);
/// Register a new parametric storage class, this is necessary to create
/// instances of this class type. `id` is the type identifier that will be
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 350680692a6e0f..8cd73471a7f706 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -151,6 +151,9 @@ class MLIRContextImpl {
/// Enable support for multi-threading within MLIR.
bool threadingIsEnabled = true;
+ /// Enable locking for context access independently of parallel execution.
+ bool threadSafeLockingIsEnabled = true;
+
/// Track if we are currently executing in a threaded execution environment
/// (like the pass-manager): this is only a debugging feature to help reducing
/// the chances of data races one some context APIs.
@@ -287,14 +290,15 @@ class MLIRContextImpl {
ownedThreadPool = std::make_unique<llvm::DefaultThreadPool>();
threadPool = ownedThreadPool.get();
} else {
- disableUniquerMultithreading(true);
+ setThreadSafeLocking(false);
}
}
- /// Update the threading mode for each of the uniquers.
- void disableUniquerMultithreading(bool disable) {
- affineUniquer.disableMultithreading(disable);
- attributeUniquer.disableMultithreading(disable);
- typeUniquer.disableMultithreading(disable);
+ /// Update the locking mode for context access and each of the uniquers.
+ void setThreadSafeLocking(bool enable) {
+ threadSafeLockingIsEnabled = enable;
+ affineUniquer.setThreadSafeLocking(enable);
+ attributeUniquer.setThreadSafeLocking(enable);
+ typeUniquer.setThreadSafeLocking(enable);
}
~MLIRContextImpl() {
for (auto typeMapping : registeredTypes)
@@ -632,7 +636,7 @@ void MLIRContext::disableMultithreading(bool disable) {
impl->threadingIsEnabled = !disable;
- impl->disableUniquerMultithreading(disable);
+ impl->setThreadSafeLocking(!disable);
// Destroy thread pool (stop all threads) if it is no longer needed, or create
// a new one if multithreading was re-enabled.
@@ -653,6 +657,16 @@ void MLIRContext::disableMultithreading(bool disable) {
}
}
+#if LLVM_ENABLE_THREADS
+bool MLIRContext::isThreadSafeLockingEnabled() {
+ return impl->threadSafeLockingIsEnabled;
+}
+#endif
+
+void MLIRContext::enableThreadSafeLocking() {
+ impl->setThreadSafeLocking(true);
+}
+
void MLIRContext::setThreadPool(llvm::ThreadPoolInterface &pool) {
assert(!isMultithreadingEnabled() &&
"expected multi-threading to be disabled when setting a ThreadPool");
@@ -916,8 +930,8 @@ OperationName::OperationName(StringRef name, MLIRContext *context) {
MLIRContextImpl &ctxImpl = context->getImpl();
// Check for an existing name in read-only mode.
- bool isMultithreadingEnabled = context->isMultithreadingEnabled();
- if (isMultithreadingEnabled) {
+ bool shouldLock = context->isThreadSafeLockingEnabled();
+ if (shouldLock) {
// 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.
@@ -936,7 +950,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);
+ ScopedWriterLock lock(ctxImpl.operationInfoMutex, shouldLock);
auto it = ctxImpl.operations.try_emplace(name);
if (it.second) {
diff --git a/mlir/lib/Support/StorageUniquer.cpp b/mlir/lib/Support/StorageUniquer.cpp
index 1de660c22328f7..636fc56e03bdf1 100644
--- a/mlir/lib/Support/StorageUniquer.cpp
+++ b/mlir/lib/Support/StorageUniquer.cpp
@@ -528,9 +528,9 @@ struct StorageUniquerImpl {
StorageUniquer::StorageUniquer() : impl(new StorageUniquerImpl()) {}
StorageUniquer::~StorageUniquer() = default;
-/// Set the flag specifying if multi-threading is disabled within the uniquer.
-void StorageUniquer::disableMultithreading(bool disable) {
- impl->threadingIsEnabled = !disable;
+/// Set whether thread-safe locking is enabled within the uniquer.
+void StorageUniquer::setThreadSafeLocking(bool enable) {
+ impl->threadingIsEnabled = enable;
}
void StorageUniquer::beginTransientScope() { impl->beginTransientScope(); }
diff --git a/mlir/unittests/IR/CMakeLists.txt b/mlir/unittests/IR/CMakeLists.txt
index ce9b230222f621..d39ccb5a565b6b 100644
--- a/mlir/unittests/IR/CMakeLists.txt
+++ b/mlir/unittests/IR/CMakeLists.txt
@@ -13,6 +13,7 @@ add_mlir_unittest(MLIRIRTests
InterfaceAttachmentTest.cpp
LocationTest.cpp
MLIRContextResetTest.cpp
+ MLIRContextTest.cpp
MemrefLayoutTest.cpp
OpImplementationTest.cpp
OperationSupportTest.cpp
diff --git a/mlir/unittests/IR/MLIRContextTest.cpp b/mlir/unittests/IR/MLIRContextTest.cpp
new file mode 100644
index 00000000000000..ed57c5bb26310b
--- /dev/null
+++ b/mlir/unittests/IR/MLIRContextTest.cpp
@@ -0,0 +1,126 @@
+//===- MLIRContextTest.cpp - MLIR context tests ---------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/Threading.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/ThreadPool.h"
+#include "gtest/gtest.h"
+#include <array>
+#include <string>
+#include <thread>
+
+using namespace mlir;
+
+namespace {
+TEST(MLIRContextTest, LockingWithoutThreadPool) {
+ MLIRContext context(MLIRContext::Threading::DISABLED);
+ EXPECT_FALSE(context.isMultithreadingEnabled());
+ EXPECT_FALSE(context.isThreadSafeLockingEnabled());
+
+ context.enableThreadSafeLocking();
+ EXPECT_EQ(context.isThreadSafeLockingEnabled(), LLVM_ENABLE_THREADS != 0);
+ EXPECT_FALSE(context.isMultithreadingEnabled());
+ EXPECT_EQ(context.getNumThreads(), 1u);
+
+ // MLIR's parallel helpers must still execute on the calling thread, without
+ // trying to access a thread pool.
+ auto callingThread = std::this_thread::get_id();
+ parallelFor(&context, 0, 16, [&](size_t) {
+ EXPECT_EQ(std::this_thread::get_id(), callingThread);
+ });
+
+ context.disableMultithreading();
+ EXPECT_FALSE(context.isThreadSafeLockingEnabled());
+ EXPECT_FALSE(context.isMultithreadingEnabled());
+}
+
+TEST(MLIRContextTest, MultithreadingEnablesLocking) {
+ MLIRContext context;
+ EXPECT_EQ(context.isMultithreadingEnabled(), LLVM_ENABLE_THREADS != 0);
+ EXPECT_EQ(context.isThreadSafeLockingEnabled(), LLVM_ENABLE_THREADS != 0);
+
+ context.disableMultithreading();
+ EXPECT_FALSE(context.isMultithreadingEnabled());
+ EXPECT_FALSE(context.isThreadSafeLockingEnabled());
+
+ context.enableMultithreading();
+ EXPECT_EQ(context.isMultithreadingEnabled(), LLVM_ENABLE_THREADS != 0);
+ EXPECT_EQ(context.isThreadSafeLockingEnabled(), LLVM_ENABLE_THREADS != 0);
+
+ context.disableMultithreading();
+ context.enableThreadSafeLocking();
+ context.enableMultithreading();
+ EXPECT_EQ(context.isMultithreadingEnabled(), LLVM_ENABLE_THREADS != 0);
+ EXPECT_EQ(context.isThreadSafeLockingEnabled(), LLVM_ENABLE_THREADS != 0);
+
+ context.disableMultithreading();
+ EXPECT_FALSE(context.isThreadSafeLockingEnabled());
+}
+
+TEST(MLIRContextTest, ExternalThreadPoolEnablesLocking) {
+ llvm::SingleThreadExecutor pool;
+ MLIRContext context(MLIRContext::Threading::DISABLED);
+ context.setThreadPool(pool);
+ EXPECT_EQ(context.isMultithreadingEnabled(), LLVM_ENABLE_THREADS != 0);
+ EXPECT_EQ(context.isThreadSafeLockingEnabled(), LLVM_ENABLE_THREADS != 0);
+}
+
+#if LLVM_ENABLE_THREADS
+TEST(MLIRContextTest, ConcurrentAccessWithoutThreadPool) {
+ MLIRContext context(MLIRContext::Threading::DISABLED);
+ context.enableThreadSafeLocking();
+
+ struct UniquedEntities {
+ IntegerType type;
+ StringAttr attribute;
+ AffineExpr expression;
+ OperationName operationName;
+ };
+ constexpr unsigned numThreads = 4;
+ constexpr unsigned numEntities = 256;
+ std::array<SmallVector<UniquedEntities>, numThreads> results;
+
+ // Use an external pool directly: the context itself has no pool and MLIR's
+ // parallel helpers would execute sequentially.
+ llvm::DefaultThreadPool pool(llvm::hardware_concurrency(numThreads));
+ for (unsigned thread = 0; thread < numThreads; ++thread) {
+ pool.async([&, thread] {
+ for (unsigned i = 0; i < numEntities; ++i) {
+ // Use the same keys on every thread, including keys not yet present in
+ // the context, to exercise concurrent insertion as well as lookup.
+ results[thread].push_back(
+ {IntegerType::get(&context, i + 129),
+ StringAttr::get(&context, "external.attr" + std::to_string(i)),
+ getAffineDimExpr(i, &context) +
+ getAffineConstantExpr(i + 1, &context),
+ OperationName("external.op" + std::to_string(i), &context)});
+ }
+ });
+ }
+ pool.wait();
+
+ for (unsigned i = 0; i < numEntities; ++i) {
+ const auto &expected = results[0][i];
+ for (unsigned thread = 1; thread < numThreads; ++thread) {
+ const auto &actual = results[thread][i];
+ EXPECT_EQ(actual.type, expected.type);
+ EXPECT_EQ(actual.attribute, expected.attribute);
+ EXPECT_EQ(actual.expression, expected.expression);
+ EXPECT_EQ(actual.operationName, expected.operationName);
+ }
+ }
+ EXPECT_FALSE(context.isMultithreadingEnabled());
+ EXPECT_TRUE(context.isThreadSafeLockingEnabled());
+}
+#endif
+} // namespace
More information about the Mlir-commits
mailing list