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 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()) + std::scoped_lock guard(allocatorMutex); + return new (allocator.Allocate()) DistinctAttrStorage(referencedAttr); - } + }; private: - ThreadLocalCache 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 + +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(attr).getImpl(); })); + + ASSERT_TRUE(attr); +}