[Mlir-commits] [mlir] [MLIR] Fix use-after-frees when accessing DistinctAttr storage (PR #148666)
Artemiy Bulavin
llvmlistbot at llvm.org
Tue Jul 15 08:51:39 PDT 2025
https://github.com/abulavin updated https://github.com/llvm/llvm-project/pull/148666
>From c1c4bd25a1e01999f0ee0f109a38c48640cce506 Mon Sep 17 00:00:00 2001
From: Artemiy Bulavin <artemiyb at graphcore.ai>
Date: Thu, 30 Jan 2025 14:32:03 +0000
Subject: [PATCH 1/8] Use shared storage allocator for DistinctAttr when crash
recovery enabled
---
mlir/include/mlir/IR/MLIRContext.h | 7 ++
mlir/include/mlir/Pass/PassManager.h | 5 +
mlir/lib/IR/AttributeDetail.h | 109 ++++++++++++++++--
mlir/lib/IR/CMakeLists.txt | 2 +
mlir/lib/IR/MLIRContext.cpp | 7 ++
mlir/lib/Pass/PassCrashRecovery.cpp | 26 ++++-
mlir/unittests/IR/CMakeLists.txt | 1 +
.../IR/DistinctAttributeAllocatorTest.cpp | 54 +++++++++
8 files changed, 198 insertions(+), 13 deletions(-)
create mode 100644 mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp
diff --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h
index ef8dab87f131a..a4ac3ab08b53c 100644
--- a/mlir/include/mlir/IR/MLIRContext.h
+++ b/mlir/include/mlir/IR/MLIRContext.h
@@ -153,6 +153,13 @@ class MLIRContext {
disableMultithreading(!enable);
}
+ /// Set the flag specifying if thread-local storage should be used by
+ /// attribute storage allocators in this context.
+ void disableThreadLocalStorage(bool disable = true);
+ void enableThreadLocalStorage(bool enable = true) {
+ disableThreadLocalStorage(!enable);
+ }
+
/// 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/Pass/PassManager.h b/mlir/include/mlir/Pass/PassManager.h
index 6e59b0f32ac6f..e6260415a83e1 100644
--- a/mlir/include/mlir/Pass/PassManager.h
+++ b/mlir/include/mlir/Pass/PassManager.h
@@ -17,6 +17,7 @@
#include "llvm/Support/raw_ostream.h"
#include <functional>
+#include <mutex>
#include <optional>
namespace mlir {
@@ -497,6 +498,10 @@ class PassManager : public OpPassManager {
/// A flag that indicates if the IR should be verified in between passes.
bool verifyPasses : 1;
+
+ /// A mutex used to serialixe thread access when running the pass manager with
+ /// crash reproduction enabled.
+ std::mutex crashRecoveryLock;
};
/// Register a set of useful command-line options that can be used to configure
diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index 26d40ac3a38f6..820fb021d96fe 100644
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/mlir/lib/IR/AttributeDetail.h
@@ -22,8 +22,10 @@
#include "mlir/Support/StorageUniquer.h"
#include "mlir/Support/ThreadLocalCache.h"
#include "llvm/ADT/APFloat.h"
-#include "llvm/ADT/PointerIntPair.h"
-#include "llvm/Support/TrailingObjects.h"
+#include "llvm/Support/Allocator.h"
+#include <atomic>
+#include <cstdint>
+#include <mutex>
namespace mlir {
namespace detail {
@@ -396,27 +398,112 @@ class DistinctAttributeUniquer {
Attribute referencedAttr);
};
-/// An allocator for distinct attribute storage instances. It uses thread local
-/// bump pointer allocators stored in a thread local cache to ensure the storage
-/// is freed after the destruction of the distinct attribute allocator.
-class DistinctAttributeAllocator {
+/// An allocator for distinct attribute storage instances. It uses thread-local
+/// bump pointer allocators stored in a thread-local cache to ensure the storage
+/// is freed after the destruction of the distinct attribute allocator. The way
+/// in which storage is allocated may be changed from a thread-local allocator
+/// to a shared, locked allocator. This can be used to prevent use-after-free
+/// errors if storage is allocated on a thread that has a lifetime less than the
+/// lifetime of the storage.
+class DistinctAttributeAllocator final {
public:
- DistinctAttributeAllocator() = default;
+ /// The allocation strategy for DistinctAttribute storage.
+ enum class AllocationMode : uint8_t {
+ /// Use a thread-local allocator. Lock-free, but storage lifetime is tied to
+ /// a thread, which may have shorter lifetime than the storage allocated
+ /// here.
+ ThreadLocal,
+ /// Use a single, shared allocator protected by a mutex. Slower due to
+ /// locking, but storage persists for the lifetime of the MLIR context.
+ SharedLocked,
+ };
+
+ DistinctAttributeAllocator() {
+ setAllocationMode(AllocationMode::ThreadLocal);
+ };
DistinctAttributeAllocator(DistinctAttributeAllocator &&) = delete;
DistinctAttributeAllocator(const DistinctAttributeAllocator &) = delete;
DistinctAttributeAllocator &
operator=(const DistinctAttributeAllocator &) = delete;
- /// Allocates a distinct attribute storage using a thread local bump pointer
- /// allocator to enable synchronization free parallel allocations.
+ /// Allocates a distinct attribute storage. In most cases this will be
+ /// using a thread-local allocator for synchronization free parallel accesses,
+ /// however if the PassManager's crash recovery is enabled this will incur
+ /// synchronization cost in exchange for persisting the storage.
DistinctAttrStorage *allocate(Attribute referencedAttr) {
- return new (allocatorCache.get().Allocate<DistinctAttrStorage>())
- DistinctAttrStorage(referencedAttr);
+ return std::invoke(allocatorFn.load(), *this, referencedAttr);
+ }
+
+ /// Sets the allocation mode. This is used by the MLIRContext to switch
+ /// allocation strategies, for example, during crash recovery.
+ void setAllocationMode(AllocationMode mode) {
+ switch (mode) {
+ case AllocationMode::ThreadLocal:
+ allocatorFn.store(&DistinctAttributeAllocator::allocateThreadLocalWrapper,
+ std::memory_order_release);
+ break;
+ case AllocationMode::SharedLocked:
+ allocatorFn.store(
+ &DistinctAttributeAllocator::allocateLockedSharedWrapper,
+ std::memory_order_release);
+ break;
+ }
}
private:
+ // Define some static wrappers for allocating using either the thread-local
+ // allocator or the shared locked allocator. We use these static wrappers to
+ // ensure that the width of the function pointer to these functions is the
+ // size of regular pointer on the platform (which is not always the case for
+ // pointer-to-member-function), and so should be handled by atomic
+ // instructions.
+ static DistinctAttrStorage *
+ allocateThreadLocalWrapper(DistinctAttributeAllocator &self,
+ Attribute referencedAttr) {
+ return self.allocateUsingThreadLocal(referencedAttr);
+ }
+
+ static DistinctAttrStorage *
+ allocateLockedSharedWrapper(DistinctAttributeAllocator &self,
+ Attribute referencedAttr) {
+ return self.allocateUsingLockedShared(referencedAttr);
+ }
+
+ DistinctAttrStorage *allocateUsingLockedShared(Attribute referencedAttr) {
+ std::scoped_lock<std::mutex> guard(sharedAllocatorMutex);
+ return doAllocate(sharedAllocator, referencedAttr);
+ }
+
+ DistinctAttrStorage *allocateUsingThreadLocal(Attribute referencedAttr) {
+ return doAllocate(allocatorCache.get(), referencedAttr);
+ };
+
+ DistinctAttrStorage *doAllocate(llvm::BumpPtrAllocator &allocator,
+ Attribute referencedAttr) {
+ return new (allocator.Allocate<DistinctAttrStorage>())
+ DistinctAttrStorage(referencedAttr);
+ };
+
+ /// Used to allocate Distinct Attribute storage without synchronization
ThreadLocalCache<llvm::BumpPtrAllocator> allocatorCache;
+
+ /// Used to allocate Distict Attribute storage with synchronization,
+ /// but without bounding the lifetime of the allocated memory to the
+ /// lifetime of a thread.
+ llvm::BumpPtrAllocator sharedAllocator;
+
+ /// Used to lock access to the shared allocator
+ std::mutex sharedAllocatorMutex;
+
+ using AllocatorFuncPtr =
+ DistinctAttrStorage *(*)(DistinctAttributeAllocator &, Attribute);
+
+ /// Atomic function pointer for the allocation strategy. Using a pointer to a
+ /// static member function as opposed to a non-static one should guarantee
+ /// that the function pointer fits into a standard pointer size, and so will
+ /// atomic instruction support for the corresponding width.
+ std::atomic<AllocatorFuncPtr> allocatorFn;
};
} // namespace detail
} // namespace mlir
diff --git a/mlir/lib/IR/CMakeLists.txt b/mlir/lib/IR/CMakeLists.txt
index 4cabac185171c..ac1110a6ebda4 100644
--- a/mlir/lib/IR/CMakeLists.txt
+++ b/mlir/lib/IR/CMakeLists.txt
@@ -68,5 +68,7 @@ add_mlir_library(MLIRIR
LINK_LIBS PUBLIC
MLIRSupport
+ PRIVATE
+ atomic
)
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 716d9c85a377d..dd7999201f116 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -711,6 +711,13 @@ bool MLIRContext::isOperationRegistered(StringRef name) {
return RegisteredOperationName::lookup(name, this).has_value();
}
+void MLIRContext::disableThreadLocalStorage(bool disable) {
+ const auto mode =
+ disable ? DistinctAttributeAllocator::AllocationMode::SharedLocked
+ : DistinctAttributeAllocator::AllocationMode::ThreadLocal;
+ getImpl().distinctAttributeAllocator.setAllocationMode(mode);
+}
+
void Dialect::addType(TypeID typeID, AbstractType &&typeInfo) {
auto &impl = context->getImpl();
assert(impl.multiThreadedExecutionContext == 0 &&
diff --git a/mlir/lib/Pass/PassCrashRecovery.cpp b/mlir/lib/Pass/PassCrashRecovery.cpp
index b048ff9462392..ceeef353c2f2f 100644
--- a/mlir/lib/Pass/PassCrashRecovery.cpp
+++ b/mlir/lib/Pass/PassCrashRecovery.cpp
@@ -414,14 +414,36 @@ struct FileReproducerStream : public mlir::ReproducerStream {
LogicalResult PassManager::runWithCrashRecovery(Operation *op,
AnalysisManager am) {
+ // Notify the context to disable the use of thread-local storage for
+ // allocating attribute storage while the pass manager is running in a crash
+ // recovery context thread. Re-enable the thread-local storage upon function
+ // exit. This is required to persist any attribute storage allocated in
+ // thread-local storage during passes beyond the lifetime of the recovery
+ // context thread.
+ const bool threadingEnabled = getContext()->isMultithreadingEnabled();
+ if (threadingEnabled) {
+ crashRecoveryLock.lock();
+ getContext()->disableThreadLocalStorage();
+ }
+ auto guard = llvm::make_scope_exit([this]() {
+ getContext()->enableThreadLocalStorage();
+ crashRecoveryLock.unlock();
+ });
+
crashReproGenerator->initialize(getPasses(), op, verifyPasses);
// Safely invoke the passes within a recovery context.
LogicalResult passManagerResult = failure();
llvm::CrashRecoveryContext recoveryContext;
- recoveryContext.RunSafelyOnThread(
- [&] { passManagerResult = runPasses(op, am); });
+ const auto runPassesFn = [&] { passManagerResult = runPasses(op, am); };
+ if (threadingEnabled)
+ recoveryContext.RunSafelyOnThread(runPassesFn);
+ else
+ recoveryContext.RunSafely(runPassesFn);
crashReproGenerator->finalize(op, passManagerResult);
+
+ if (!threadingEnabled)
+ guard.release();
return passManagerResult;
}
diff --git a/mlir/unittests/IR/CMakeLists.txt b/mlir/unittests/IR/CMakeLists.txt
index d22afb3003e76..a46e64718dab9 100644
--- a/mlir/unittests/IR/CMakeLists.txt
+++ b/mlir/unittests/IR/CMakeLists.txt
@@ -6,6 +6,7 @@ add_mlir_unittest(MLIRIRTests
AttrTypeReplacerTest.cpp
Diagnostic.cpp
DialectTest.cpp
+ DistinctAttributeAllocatorTest.cpp
InterfaceTest.cpp
IRMapping.cpp
InterfaceAttachmentTest.cpp
diff --git a/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp b/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp
new file mode 100644
index 0000000000000..f03b3fab87ace
--- /dev/null
+++ b/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp
@@ -0,0 +1,54 @@
+//===- DistinctAttributeAllocatorTest.cpp - DistinctAttr storage alloc test
+//-===//
+//
+// 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 "gtest/gtest.h"
+
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/MLIRContext.h"
+#include "llvm/ADT/ScopeExit.h"
+#include "llvm/Support/CrashRecoveryContext.h"
+#include <thread>
+
+using namespace mlir;
+
+//
+// Test that a DistinctAttr that is created on a separate thread does
+// not have its storage deleted when thread local storage is disabled
+// on the MLIRContext.
+//
+TEST(DistinctAttributeAllocatorTest, TestAttributeWellFormedAfterThreadJoin) {
+ MLIRContext ctx;
+ OpBuilder builder(&ctx);
+ DistinctAttr attr;
+
+ {
+ ctx.disableThreadLocalStorage();
+ auto guard =
+ llvm::make_scope_exit([&ctx]() { ctx.enableThreadLocalStorage(); });
+
+ std::thread t([&ctx, &attr]() {
+ attr = DistinctAttr::create(UnitAttr::get(&ctx));
+ ASSERT_TRUE(attr);
+ });
+ t.join();
+ }
+
+ // If the attribute storage got deleted after the thread joins (which we don't
+ // want) then trying to access it triggers an assert in Debug mode, and a
+ // crash otherwise. Run this in a CrashRecoveryContext to avoid bringing down
+ // the whole test suite if this test fails. Additionally, MSAN and/or TSAN
+ // should raise failures here if the attribute storage was deleted.
+ llvm::CrashRecoveryContext crc;
+ EXPECT_TRUE(crc.RunSafely([attr]() { (void)attr.getAbstractAttribute(); }));
+ EXPECT_TRUE(
+ crc.RunSafely([attr]() { (void)*cast<Attribute>(attr).getImpl(); }));
+
+ ASSERT_TRUE(attr);
+}
>From 0bb10c37f0a44da06dd59cc76bf49e3fc1132fdc Mon Sep 17 00:00:00 2001
From: Artemiy Bulavin <artemiyb at graphcore.ai>
Date: Mon, 14 Jul 2025 17:40:10 +0000
Subject: [PATCH 2/8] fixup: use LLVM_ATOMIC_LIB to link atomics
---
mlir/CMakeLists.txt | 1 +
mlir/lib/IR/CMakeLists.txt | 2 +-
2 files changed, 2 insertions(+), 1 deletion(-)
diff --git a/mlir/CMakeLists.txt b/mlir/CMakeLists.txt
index a1ad81f625cd6..03ad0a9c5ea22 100644
--- a/mlir/CMakeLists.txt
+++ b/mlir/CMakeLists.txt
@@ -69,6 +69,7 @@ list(INSERT CMAKE_MODULE_PATH 0
include(AddMLIR)
include(IRDLToCpp)
+include(CheckAtomic)
# -BSymbolic is incompatible with TypeID
if("${CMAKE_SHARED_LINKER_FLAGS}" MATCHES "-Bsymbolic[^-]")
diff --git a/mlir/lib/IR/CMakeLists.txt b/mlir/lib/IR/CMakeLists.txt
index ac1110a6ebda4..200f150ad4993 100644
--- a/mlir/lib/IR/CMakeLists.txt
+++ b/mlir/lib/IR/CMakeLists.txt
@@ -69,6 +69,6 @@ add_mlir_library(MLIRIR
LINK_LIBS PUBLIC
MLIRSupport
PRIVATE
- atomic
+ ${LLVM_ATOMIC_LIB}
)
>From 38d38b5a7a40608bab1b5e5ec57bb332bb6753e0 Mon Sep 17 00:00:00 2001
From: Artemiy Bulavin <artemiyb at graphcore.ai>
Date: Tue, 15 Jul 2025 07:56:56 +0000
Subject: [PATCH 3/8] fixup: revert to synchronized allocator
---
mlir/CMakeLists.txt | 1 -
mlir/include/mlir/IR/MLIRContext.h | 7 --
mlir/lib/IR/AttributeDetail.h | 108 ++----------------
mlir/lib/IR/CMakeLists.txt | 2 -
mlir/lib/IR/MLIRContext.cpp | 8 +-
mlir/lib/Pass/PassCrashRecovery.cpp | 12 --
.../IR/DistinctAttributeAllocatorTest.cpp | 17 +--
7 files changed, 17 insertions(+), 138 deletions(-)
diff --git a/mlir/CMakeLists.txt b/mlir/CMakeLists.txt
index 03ad0a9c5ea22..a1ad81f625cd6 100644
--- a/mlir/CMakeLists.txt
+++ b/mlir/CMakeLists.txt
@@ -69,7 +69,6 @@ list(INSERT CMAKE_MODULE_PATH 0
include(AddMLIR)
include(IRDLToCpp)
-include(CheckAtomic)
# -BSymbolic is incompatible with TypeID
if("${CMAKE_SHARED_LINKER_FLAGS}" MATCHES "-Bsymbolic[^-]")
diff --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h
index a4ac3ab08b53c..ef8dab87f131a 100644
--- a/mlir/include/mlir/IR/MLIRContext.h
+++ b/mlir/include/mlir/IR/MLIRContext.h
@@ -153,13 +153,6 @@ class MLIRContext {
disableMultithreading(!enable);
}
- /// Set the flag specifying if thread-local storage should be used by
- /// attribute storage allocators in this context.
- void disableThreadLocalStorage(bool disable = true);
- void enableThreadLocalStorage(bool enable = true) {
- disableThreadLocalStorage(!enable);
- }
-
/// 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/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index 820fb021d96fe..b5a332618453c 100644
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/mlir/lib/IR/AttributeDetail.h
@@ -19,12 +19,8 @@
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/IntegerSet.h"
#include "mlir/IR/MLIRContext.h"
-#include "mlir/Support/StorageUniquer.h"
-#include "mlir/Support/ThreadLocalCache.h"
#include "llvm/ADT/APFloat.h"
#include "llvm/Support/Allocator.h"
-#include <atomic>
-#include <cstdint>
#include <mutex>
namespace mlir {
@@ -398,112 +394,30 @@ class DistinctAttributeUniquer {
Attribute referencedAttr);
};
-/// An allocator for distinct attribute storage instances. It uses thread-local
-/// bump pointer allocators stored in a thread-local cache to ensure the storage
-/// is freed after the destruction of the distinct attribute allocator. The way
-/// in which storage is allocated may be changed from a thread-local allocator
-/// to a shared, locked allocator. This can be used to prevent use-after-free
-/// errors if storage is allocated on a thread that has a lifetime less than the
-/// lifetime of the storage.
+/// An allocator for distinct attribute storage instances. Uses a synchronized
+/// BumpPtrAllocator to ensure allocated storage is deleted when the
+/// DistinctAttributeAllocator is destroyed
class DistinctAttributeAllocator final {
public:
- /// The allocation strategy for DistinctAttribute storage.
- enum class AllocationMode : uint8_t {
- /// Use a thread-local allocator. Lock-free, but storage lifetime is tied to
- /// a thread, which may have shorter lifetime than the storage allocated
- /// here.
- ThreadLocal,
- /// Use a single, shared allocator protected by a mutex. Slower due to
- /// locking, but storage persists for the lifetime of the MLIR context.
- SharedLocked,
- };
-
- DistinctAttributeAllocator() {
- setAllocationMode(AllocationMode::ThreadLocal);
- };
-
+ DistinctAttributeAllocator() = default;
DistinctAttributeAllocator(DistinctAttributeAllocator &&) = delete;
DistinctAttributeAllocator(const DistinctAttributeAllocator &) = delete;
DistinctAttributeAllocator &
operator=(const DistinctAttributeAllocator &) = delete;
- /// Allocates a distinct attribute storage. In most cases this will be
- /// using a thread-local allocator for synchronization free parallel accesses,
- /// however if the PassManager's crash recovery is enabled this will incur
- /// synchronization cost in exchange for persisting the storage.
DistinctAttrStorage *allocate(Attribute referencedAttr) {
- return std::invoke(allocatorFn.load(), *this, referencedAttr);
- }
-
- /// Sets the allocation mode. This is used by the MLIRContext to switch
- /// allocation strategies, for example, during crash recovery.
- void setAllocationMode(AllocationMode mode) {
- switch (mode) {
- case AllocationMode::ThreadLocal:
- allocatorFn.store(&DistinctAttributeAllocator::allocateThreadLocalWrapper,
- std::memory_order_release);
- break;
- case AllocationMode::SharedLocked:
- allocatorFn.store(
- &DistinctAttributeAllocator::allocateLockedSharedWrapper,
- std::memory_order_release);
- break;
- }
- }
-
-private:
- // Define some static wrappers for allocating using either the thread-local
- // allocator or the shared locked allocator. We use these static wrappers to
- // ensure that the width of the function pointer to these functions is the
- // size of regular pointer on the platform (which is not always the case for
- // pointer-to-member-function), and so should be handled by atomic
- // instructions.
- static DistinctAttrStorage *
- allocateThreadLocalWrapper(DistinctAttributeAllocator &self,
- Attribute referencedAttr) {
- return self.allocateUsingThreadLocal(referencedAttr);
- }
-
- static DistinctAttrStorage *
- allocateLockedSharedWrapper(DistinctAttributeAllocator &self,
- Attribute referencedAttr) {
- return self.allocateUsingLockedShared(referencedAttr);
- }
-
- DistinctAttrStorage *allocateUsingLockedShared(Attribute referencedAttr) {
- std::scoped_lock<std::mutex> guard(sharedAllocatorMutex);
- return doAllocate(sharedAllocator, referencedAttr);
- }
-
- DistinctAttrStorage *allocateUsingThreadLocal(Attribute referencedAttr) {
- return doAllocate(allocatorCache.get(), referencedAttr);
- };
-
- DistinctAttrStorage *doAllocate(llvm::BumpPtrAllocator &allocator,
- Attribute referencedAttr) {
+ std::scoped_lock<std::mutex> guard(allocatorMutex);
return new (allocator.Allocate<DistinctAttrStorage>())
DistinctAttrStorage(referencedAttr);
};
- /// Used to allocate Distinct Attribute storage without synchronization
- ThreadLocalCache<llvm::BumpPtrAllocator> allocatorCache;
-
- /// Used to allocate Distict Attribute storage with synchronization,
- /// but without bounding the lifetime of the allocated memory to the
- /// lifetime of a thread.
- llvm::BumpPtrAllocator sharedAllocator;
-
- /// Used to lock access to the shared allocator
- std::mutex sharedAllocatorMutex;
-
- using AllocatorFuncPtr =
- DistinctAttrStorage *(*)(DistinctAttributeAllocator &, Attribute);
+private:
+ /// Used to allocate Distict Attribute storage. When this allocator destroyed
+ /// in till also dealllocate any storage instances
+ llvm::BumpPtrAllocator allocator;
- /// Atomic function pointer for the allocation strategy. Using a pointer to a
- /// static member function as opposed to a non-static one should guarantee
- /// that the function pointer fits into a standard pointer size, and so will
- /// atomic instruction support for the corresponding width.
- std::atomic<AllocatorFuncPtr> allocatorFn;
+ /// Used to lock access to the allocator
+ std::mutex allocatorMutex;
};
} // namespace detail
} // namespace mlir
diff --git a/mlir/lib/IR/CMakeLists.txt b/mlir/lib/IR/CMakeLists.txt
index 200f150ad4993..4cabac185171c 100644
--- a/mlir/lib/IR/CMakeLists.txt
+++ b/mlir/lib/IR/CMakeLists.txt
@@ -68,7 +68,5 @@ add_mlir_library(MLIRIR
LINK_LIBS PUBLIC
MLIRSupport
- PRIVATE
- ${LLVM_ATOMIC_LIB}
)
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index dd7999201f116..06ec1c85fb4d5 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -31,6 +31,7 @@
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Compiler.h"
#include "llvm/Support/Debug.h"
+#include "llvm/Support/ManagedStatic.h"
#include "llvm/Support/Mutex.h"
#include "llvm/Support/RWMutex.h"
#include "llvm/Support/ThreadPool.h"
@@ -711,13 +712,6 @@ bool MLIRContext::isOperationRegistered(StringRef name) {
return RegisteredOperationName::lookup(name, this).has_value();
}
-void MLIRContext::disableThreadLocalStorage(bool disable) {
- const auto mode =
- disable ? DistinctAttributeAllocator::AllocationMode::SharedLocked
- : DistinctAttributeAllocator::AllocationMode::ThreadLocal;
- getImpl().distinctAttributeAllocator.setAllocationMode(mode);
-}
-
void Dialect::addType(TypeID typeID, AbstractType &&typeInfo) {
auto &impl = context->getImpl();
assert(impl.multiThreadedExecutionContext == 0 &&
diff --git a/mlir/lib/Pass/PassCrashRecovery.cpp b/mlir/lib/Pass/PassCrashRecovery.cpp
index ceeef353c2f2f..4699145cf5a13 100644
--- a/mlir/lib/Pass/PassCrashRecovery.cpp
+++ b/mlir/lib/Pass/PassCrashRecovery.cpp
@@ -15,7 +15,6 @@
#include "mlir/Pass/Pass.h"
#include "mlir/Support/FileUtilities.h"
#include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/CrashRecoveryContext.h"
@@ -421,15 +420,6 @@ LogicalResult PassManager::runWithCrashRecovery(Operation *op,
// thread-local storage during passes beyond the lifetime of the recovery
// context thread.
const bool threadingEnabled = getContext()->isMultithreadingEnabled();
- if (threadingEnabled) {
- crashRecoveryLock.lock();
- getContext()->disableThreadLocalStorage();
- }
- auto guard = llvm::make_scope_exit([this]() {
- getContext()->enableThreadLocalStorage();
- crashRecoveryLock.unlock();
- });
-
crashReproGenerator->initialize(getPasses(), op, verifyPasses);
// Safely invoke the passes within a recovery context.
@@ -442,8 +432,6 @@ LogicalResult PassManager::runWithCrashRecovery(Operation *op,
recoveryContext.RunSafely(runPassesFn);
crashReproGenerator->finalize(op, passManagerResult);
- if (!threadingEnabled)
- guard.release();
return passManagerResult;
}
diff --git a/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp b/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp
index f03b3fab87ace..515f9bc93fcc7 100644
--- a/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp
+++ b/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp
@@ -12,7 +12,6 @@
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/MLIRContext.h"
-#include "llvm/ADT/ScopeExit.h"
#include "llvm/Support/CrashRecoveryContext.h"
#include <thread>
@@ -28,17 +27,11 @@ TEST(DistinctAttributeAllocatorTest, TestAttributeWellFormedAfterThreadJoin) {
OpBuilder builder(&ctx);
DistinctAttr attr;
- {
- ctx.disableThreadLocalStorage();
- auto guard =
- llvm::make_scope_exit([&ctx]() { ctx.enableThreadLocalStorage(); });
-
- std::thread t([&ctx, &attr]() {
- attr = DistinctAttr::create(UnitAttr::get(&ctx));
- ASSERT_TRUE(attr);
- });
- t.join();
- }
+ std::thread t([&ctx, &attr]() {
+ attr = DistinctAttr::create(UnitAttr::get(&ctx));
+ ASSERT_TRUE(attr);
+ });
+ t.join();
// If the attribute storage got deleted after the thread joins (which we don't
// want) then trying to access it triggers an assert in Debug mode, and a
>From 58982c2280bc58406b32b3c40c7bae375fdc83f7 Mon Sep 17 00:00:00 2001
From: Artemiy Bulavin <artemiyb at graphcore.ai>
Date: Tue, 15 Jul 2025 08:13:24 +0000
Subject: [PATCH 4/8] fixup: update comments
---
mlir/lib/Pass/PassCrashRecovery.cpp | 6 ------
mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp | 3 +--
2 files changed, 1 insertion(+), 8 deletions(-)
diff --git a/mlir/lib/Pass/PassCrashRecovery.cpp b/mlir/lib/Pass/PassCrashRecovery.cpp
index 4699145cf5a13..e16016f21bccb 100644
--- a/mlir/lib/Pass/PassCrashRecovery.cpp
+++ b/mlir/lib/Pass/PassCrashRecovery.cpp
@@ -413,12 +413,6 @@ struct FileReproducerStream : public mlir::ReproducerStream {
LogicalResult PassManager::runWithCrashRecovery(Operation *op,
AnalysisManager am) {
- // Notify the context to disable the use of thread-local storage for
- // allocating attribute storage while the pass manager is running in a crash
- // recovery context thread. Re-enable the thread-local storage upon function
- // exit. This is required to persist any attribute storage allocated in
- // thread-local storage during passes beyond the lifetime of the recovery
- // context thread.
const bool threadingEnabled = getContext()->isMultithreadingEnabled();
crashReproGenerator->initialize(getPasses(), op, verifyPasses);
diff --git a/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp b/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp
index 515f9bc93fcc7..61c7c59ae8868 100644
--- a/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp
+++ b/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp
@@ -19,8 +19,7 @@ using namespace mlir;
//
// Test that a DistinctAttr that is created on a separate thread does
-// not have its storage deleted when thread local storage is disabled
-// on the MLIRContext.
+// not have its storage deleted when the thread joins
//
TEST(DistinctAttributeAllocatorTest, TestAttributeWellFormedAfterThreadJoin) {
MLIRContext ctx;
>From 2fed089eafe5232b23e43eaf83dcdb8fe8f7f14d Mon Sep 17 00:00:00 2001
From: Artemiy Bulavin <artemiyb at graphcore.ai>
Date: Tue, 15 Jul 2025 08:13:32 +0000
Subject: [PATCH 5/8] fixup: remove mutex from pass manager
---
mlir/include/mlir/Pass/PassManager.h | 5 -----
1 file changed, 5 deletions(-)
diff --git a/mlir/include/mlir/Pass/PassManager.h b/mlir/include/mlir/Pass/PassManager.h
index e6260415a83e1..6e59b0f32ac6f 100644
--- a/mlir/include/mlir/Pass/PassManager.h
+++ b/mlir/include/mlir/Pass/PassManager.h
@@ -17,7 +17,6 @@
#include "llvm/Support/raw_ostream.h"
#include <functional>
-#include <mutex>
#include <optional>
namespace mlir {
@@ -498,10 +497,6 @@ class PassManager : public OpPassManager {
/// A flag that indicates if the IR should be verified in between passes.
bool verifyPasses : 1;
-
- /// A mutex used to serialixe thread access when running the pass manager with
- /// crash reproduction enabled.
- std::mutex crashRecoveryLock;
};
/// Register a set of useful command-line options that can be used to configure
>From be7461a0b0e31d9f94d926c1b2aaf11bcfeb1219 Mon Sep 17 00:00:00 2001
From: Artemiy Bulavin <artemiyb at graphcore.ai>
Date: Tue, 15 Jul 2025 13:21:09 +0000
Subject: [PATCH 6/8] fixup: correct typos
---
mlir/lib/IR/AttributeDetail.h | 10 +++++-----
mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp | 7 +++----
2 files changed, 8 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index b5a332618453c..202985ec466d7 100644
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/mlir/lib/IR/AttributeDetail.h
@@ -395,8 +395,8 @@ class DistinctAttributeUniquer {
};
/// An allocator for distinct attribute storage instances. Uses a synchronized
-/// BumpPtrAllocator to ensure allocated storage is deleted when the
-/// DistinctAttributeAllocator is destroyed
+/// BumpPtrAllocator to ensure thread-safety. The allocated storage is deleted
+/// when the DistinctAttributeAllocator is destroyed.
class DistinctAttributeAllocator final {
public:
DistinctAttributeAllocator() = default;
@@ -412,11 +412,11 @@ class DistinctAttributeAllocator final {
};
private:
- /// Used to allocate Distict Attribute storage. When this allocator destroyed
- /// in till also dealllocate any storage instances
+ /// Used to allocate Distict Attribute storages. The managed memory is freed
+ /// automatically when the allocator instance is destroyed.
llvm::BumpPtrAllocator allocator;
- /// Used to lock access to the allocator
+ /// Used to lock access to the allocator.
std::mutex allocatorMutex;
};
} // namespace detail
diff --git a/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp b/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp
index 61c7c59ae8868..99067d09f7bed 100644
--- a/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp
+++ b/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp
@@ -1,11 +1,10 @@
-//===- DistinctAttributeAllocatorTest.cpp - DistinctAttr storage alloc test
-//-===//
+//=== DistinctAttributeAllocatorTest.cpp - DistinctAttr storage alloc test ===//
//
// 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 "gtest/gtest.h"
@@ -19,7 +18,7 @@ using namespace mlir;
//
// Test that a DistinctAttr that is created on a separate thread does
-// not have its storage deleted when the thread joins
+// not have its storage deleted when the thread joins.
//
TEST(DistinctAttributeAllocatorTest, TestAttributeWellFormedAfterThreadJoin) {
MLIRContext ctx;
>From 00c70f824848e99a38d7e0975bef1132e3abef99 Mon Sep 17 00:00:00 2001
From: Artemiy Bulavin <artemiyb at graphcore.ai>
Date: Tue, 15 Jul 2025 13:24:17 +0000
Subject: [PATCH 7/8] fixup: hopefully final typo
---
mlir/lib/IR/AttributeDetail.h | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index 202985ec466d7..cb9d21bf3e611 100644
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/mlir/lib/IR/AttributeDetail.h
@@ -412,7 +412,7 @@ class DistinctAttributeAllocator final {
};
private:
- /// Used to allocate Distict Attribute storages. The managed memory is freed
+ /// Used to allocate distict attribute storages. The managed memory is freed
/// automatically when the allocator instance is destroyed.
llvm::BumpPtrAllocator allocator;
>From b07b3137d5224b4d5e32657f8b96e88c54ba41a6 Mon Sep 17 00:00:00 2001
From: Artemiy Bulavin <artemiyb at graphcore.ai>
Date: Tue, 15 Jul 2025 16:51:28 +0100
Subject: [PATCH 8/8] Update mlir/lib/IR/AttributeDetail.h
Co-authored-by: Tobias Gysi <tobias.gysi at nextsilicon.com>
---
mlir/lib/IR/AttributeDetail.h | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index cb9d21bf3e611..88d067bcf04b5 100644
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/mlir/lib/IR/AttributeDetail.h
@@ -412,7 +412,7 @@ class DistinctAttributeAllocator final {
};
private:
- /// Used to allocate distict attribute storages. The managed memory is freed
+ /// Used to allocate distinct attribute storages. The managed memory is freed
/// automatically when the allocator instance is destroyed.
llvm::BumpPtrAllocator allocator;
More information about the Mlir-commits
mailing list