[Mlir-commits] [mlir] 38be53a - [MLIR] Fix use-after-frees when accessing DistinctAttr storage (#148666)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 16 03:11:41 PDT 2025
Author: Artemiy Bulavin
Date: 2025-07-16T12:11:38+02:00
New Revision: 38be53aa04de8c6d494de8074328ac8907f3f631
URL: https://github.com/llvm/llvm-project/commit/38be53aa04de8c6d494de8074328ac8907f3f631
DIFF: https://github.com/llvm/llvm-project/commit/38be53aa04de8c6d494de8074328ac8907f3f631.diff
LOG: [MLIR] Fix use-after-frees when accessing DistinctAttr storage (#148666)
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.
Moreover, even without crash reproduction disabling multithreading on
the context will destroy the context's thread pool, and in turn delete
the threadlocal storage. This means a call to
`ctx->disableMulthithreading()` breaks the IR.
This PR replaces the thread local allocator with a synchronised
allocator that's shared between threads. This persists the lifetime of
allocated DistinctAttr storage instances to the lifetime of the context.
### 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.
As mentioned previously, this is also seen after creating some
`DistinctAttr`s and then calling `ctx->disableMulthithreading()`.
### Solution
`DistinctAttrStorageAllocator` uses a synchronised, shared allocator
instead of one wrapped in a `ThreadLocalCache`. The former is what
stores the allocator in transient thread_local storage.
### 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/133000. I believe I've
addressed the TSAN and race condition concerns.
Added:
mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp
Modified:
mlir/lib/IR/AttributeDetail.h
mlir/lib/IR/MLIRContext.cpp
mlir/lib/Pass/PassCrashRecovery.cpp
mlir/unittests/IR/CMakeLists.txt
Removed:
################################################################################
diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index 26d40ac3a38f6..cb9d21bf3e611 100644
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/mlir/lib/IR/AttributeDetail.h
@@ -19,11 +19,9 @@
#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/ADT/PointerIntPair.h"
-#include "llvm/Support/TrailingObjects.h"
+#include "llvm/Support/Allocator.h"
+#include <mutex>
namespace mlir {
namespace detail {
@@ -396,27 +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.
-class DistinctAttributeAllocator {
+/// An allocator for distinct attribute storage instances. Uses a synchronized
+/// BumpPtrAllocator to ensure thread-safety. The allocated storage is deleted
+/// when the DistinctAttributeAllocator is destroyed.
+class DistinctAttributeAllocator final {
public:
DistinctAttributeAllocator() = default;
-
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.
DistinctAttrStorage *allocate(Attribute referencedAttr) {
- return new (allocatorCache.get().Allocate<DistinctAttrStorage>())
+ std::scoped_lock<std::mutex> guard(allocatorMutex);
+ return new (allocator.Allocate<DistinctAttrStorage>())
DistinctAttrStorage(referencedAttr);
- }
+ };
private:
- ThreadLocalCache<llvm::BumpPtrAllocator> allocatorCache;
+ /// 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.
+ std::mutex allocatorMutex;
};
} // namespace detail
} // namespace mlir
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 716d9c85a377d..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"
diff --git a/mlir/lib/Pass/PassCrashRecovery.cpp b/mlir/lib/Pass/PassCrashRecovery.cpp
index 08f5114ae6eb2..3c9735f910094 100644
--- a/mlir/lib/Pass/PassCrashRecovery.cpp
+++ b/mlir/lib/Pass/PassCrashRecovery.cpp
@@ -411,14 +411,19 @@ struct FileReproducerStream : public mlir::ReproducerStream {
LogicalResult PassManager::runWithCrashRecovery(Operation *op,
AnalysisManager am) {
+ const bool threadingEnabled = getContext()->isMultithreadingEnabled();
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);
+
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..99067d09f7bed
--- /dev/null
+++ b/mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp
@@ -0,0 +1,45 @@
+//=== 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/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 the thread joins.
+//
+TEST(DistinctAttributeAllocatorTest, TestAttributeWellFormedAfterThreadJoin) {
+ MLIRContext ctx;
+ OpBuilder builder(&ctx);
+ DistinctAttr attr;
+
+ 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);
+}
More information about the Mlir-commits
mailing list