[Mlir-commits] [mlir] [MLIR] Fix use-after-free in DistinctAttr when using crash recovery (PR #148666)
Artemiy Bulavin
llvmlistbot at llvm.org
Mon Jul 14 10:40:32 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/2] 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/2] 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}
)
More information about the Mlir-commits
mailing list