[Mlir-commits] [mlir] [MLIR] Fix use-after-free in DistinctAttr when using crash recovery (PR #148666)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jul 14 09:45:46 PDT 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-core
Author: Artemiy Bulavin (abulavin)
<details>
<summary>Changes</summary>
This PR fixes a use-after-free error that happens when `DistinctAttr` instances are created within a `PassManager` running with crash recovery enabled. The root cause is that `DistinctAttr` storage is allocated in a thread_local allocator, which is destroyed when the crash recovery thread joins, invalidating the storage.
This PR introduces a way to switch the allocation strategy for `DistinctAttr` storage from the default lock-free thread_local allocator to a thread-safe, shared allocator. This switch is activated when the `PassManager` runs with crash recovery, ensuring the attribute storage persists beyond the lifetime of the recovery thread.
### Problem Details:
The `DistinctAttributeAllocator` uses a `ThreadLocalCache<BumpPtrAllocator>` for lock-free allocation of `DistinctAttr` storage in a multithreaded context. The issue occurs when a `PassManager` is run with crash recovery (`runWithCrashRecovery`), the pass pipeline is executed on a temporary thread spawned by `llvm::CrashRecoveryContext`. Any `DistinctAttr`s created during this execution have their storage allocated in the thread_local cache of this temporary thread. When the thread joins, the thread_local storage is destroyed, freeing the `DistinctAttr`s' memory. If this attribute is accessed later, e.g. when printing, it results in a use-after-free.
### Solution:
The `DistinctAttributeAllocator` now has two 'allocation modes':
1. ThreadLocal (Default): The existing lock-free thread_local allocation. Keeps the hot path lock free.
2. SharedLocked: uses a single `BumpPtrAllocator` shared by all threads, protected by a `std::mutex`. This pays the price of locking in exchange for persisting the storage and ensuring correctness.
To switch between these modes an `std::atomic<AllocatorFuncPtr>` is used. This allows the allocation strategy to be changed dynamically by setting the atomic function pointer to the allocation function that corresponds to the allocation mode.
`PassManager::runWithCrashRecovery` now calls `context->disableThreadLocalStorage()` before spawning the recovery thread to set the mode to `SharedLocked` and restores the default mode upon exit. A mutex is used to serialize this change if multithreading is enabled.
### Testing:
A C++ unit test has been added to validate this fix. (I was previously reproducing this failure with `mlir-opt` but I can no longer do so and I am unsure why.)
-----
Note: This is a 2nd attempt at my previous PR https://github.com/llvm/llvm-project/pull/128566 that was reverted in https://github.com/llvm/llvm-project/pull/132935. I believe I've addressed the TSAN and race condition concerns.
---
Full diff: https://github.com/llvm/llvm-project/pull/148666.diff
8 Files Affected:
- (modified) mlir/include/mlir/IR/MLIRContext.h (+7)
- (modified) mlir/include/mlir/Pass/PassManager.h (+5)
- (modified) mlir/lib/IR/AttributeDetail.h (+98-11)
- (modified) mlir/lib/IR/CMakeLists.txt (+2)
- (modified) mlir/lib/IR/MLIRContext.cpp (+7)
- (modified) mlir/lib/Pass/PassCrashRecovery.cpp (+24-2)
- (modified) mlir/unittests/IR/CMakeLists.txt (+1)
- (added) mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp (+54)
``````````diff
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);
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/148666
More information about the Mlir-commits
mailing list