From 1cf6159cb472d1f4de21928d949a1de61fb30988 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Thu, 30 Jan 2025 14:32:03 +0000 Subject: [PATCH 1/7] 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 +#include #include 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 +#include +#include 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(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 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(referencedAttr); + }; + + /// Used to allocate Distinct Attribute storage without synchronization ThreadLocalCache 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 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 08f5114ae6eb2..e327bcc2d87da 100644 --- a/mlir/lib/Pass/PassCrashRecovery.cpp +++ b/mlir/lib/Pass/PassCrashRecovery.cpp @@ -411,14 +411,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 + +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(attr).getImpl(); })); + + ASSERT_TRUE(attr); +} From b0bfc02400605a8f9c3b4484aeecf1cf580ecad0 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Mon, 14 Jul 2025 17:40:10 +0000 Subject: [PATCH 2/7] 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 1d9a104b187cb2496e9fb9f0388649f464a216b1 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Tue, 15 Jul 2025 07:56:56 +0000 Subject: [PATCH 3/7] 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 | 11 -- .../IR/DistinctAttributeAllocatorTest.cpp | 17 +-- 7 files changed, 17 insertions(+), 137 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 -#include #include 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 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 guard(allocatorMutex); return new (allocator.Allocate()) DistinctAttrStorage(referencedAttr); }; - /// Used to allocate Distinct Attribute storage without synchronization - ThreadLocalCache 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 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 e327bcc2d87da..087ced92b90f3 100644 --- a/mlir/lib/Pass/PassCrashRecovery.cpp +++ b/mlir/lib/Pass/PassCrashRecovery.cpp @@ -418,15 +418,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. @@ -439,8 +430,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 @@ -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 5facc37e6d6ba05075756bc0789bdf1053460c63 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Tue, 15 Jul 2025 08:13:24 +0000 Subject: [PATCH 4/7] 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 087ced92b90f3..3c9735f910094 100644 --- a/mlir/lib/Pass/PassCrashRecovery.cpp +++ b/mlir/lib/Pass/PassCrashRecovery.cpp @@ -411,12 +411,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 c02a1b8f09c906027bb936a217f854970386f7e5 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Tue, 15 Jul 2025 08:13:32 +0000 Subject: [PATCH 5/7] 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 -#include #include 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 74494535994dd1ffcb21e5d00b1de2f4cd008f50 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Tue, 15 Jul 2025 13:21:09 +0000 Subject: [PATCH 6/7] 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 40fca8dbe9be948439374f6cb78c2e8cc411c330 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Tue, 15 Jul 2025 13:24:17 +0000 Subject: [PATCH 7/7] 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;