[Mlir-commits] [mlir] 8dbdbd3 - [mlir] Support thread-safe context access without parallel execution (#236551)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Nov 1 03:32:17 PST 2026
Author: Mehdi Amini
Date: 2026-11-01T11:32:12Z
New Revision: 8dbdbd3b146d576f270931a56dd93ec2c2cdd5b3
URL: https://github.com/llvm/llvm-project/commit/8dbdbd3b146d576f270931a56dd93ec2c2cdd5b3
DIFF: https://github.com/llvm/llvm-project/commit/8dbdbd3b146d576f270931a56dd93ec2c2cdd5b3.diff
LOG: [mlir] Support thread-safe context access without parallel execution (#236551)
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. Enabling multithreading also enables locking.
Follow-up to #235294.
Assisted-by: Codex
Added:
mlir/unittests/IR/MLIRContextTest.cpp
Modified:
mlir/docs/ReleaseNotes.md
mlir/include/mlir/IR/MLIRContext.h
mlir/include/mlir/Support/StorageUniquer.h
mlir/lib/IR/MLIRContext.cpp
mlir/lib/Support/StorageUniquer.cpp
mlir/unittests/IR/CMakeLists.txt
Removed:
################################################################################
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..703d9c00fc28de 100644
--- a/mlir/include/mlir/IR/MLIRContext.h
+++ b/mlir/include/mlir/IR/MLIRContext.h
@@ -61,6 +61,13 @@ class RemarkEngine;
/// processRequest(request, cxt);
/// }
///
+/// Clients that manage their own threads can enable thread-safe context access
+/// without creating a thread pool or enabling parallel execution within MLIR.
+/// Required dialects must be loaded before accessing the context concurrently:
+///
+/// MLIRContext ctx(registry, MLIRContext::Threading::DISABLED);
+/// ctx.enableThreadSafeLocking();
+///
class MLIRContext {
public:
enum class Threading { DISABLED, ENABLED };
@@ -190,6 +197,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 +206,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 +216,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 f5a878edd53612..d8abff14be05f2 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.
@@ -291,14 +294,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)
@@ -636,7 +640,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.
@@ -657,6 +661,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");
@@ -927,8 +941,8 @@ OperationName::OperationName(StringRef name, MLIRContext *context) {
}
// Check for an existing name in read-only mode.
- bool isMultithreadingEnabled = context->isMultithreadingEnabled();
- if (isMultithreadingEnabled) {
+ bool shouldLock = context->isThreadSafeLockingEnabled();
+ if (shouldLock) {
llvm::sys::SmartScopedReader<true> contextLock(ctxImpl.operationInfoMutex);
auto it = ctxImpl.unregisteredOperationsByName.find(name);
if (it != ctxImpl.unregisteredOperationsByName.end()) {
@@ -938,7 +952,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.unregisteredOperationsByName.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