diff --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h index ef8dab87f131a..ad89d631c8a5f 100644 --- a/mlir/include/mlir/IR/MLIRContext.h +++ b/mlir/include/mlir/IR/MLIRContext.h @@ -153,6 +153,14 @@ class MLIRContext { disableMultithreading(!enable); } + /// Set the flag specifying if thread-local storage should be used by storage + /// allocators in this context. Note that disabling mutlithreading implies + /// thread-local storage is also disabled. + 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 26d40ac3a38f6..8fed18140433c 100644 --- a/mlir/lib/IR/AttributeDetail.h +++ b/mlir/lib/IR/AttributeDetail.h @@ -24,6 +24,7 @@ #include "llvm/ADT/APFloat.h" #include "llvm/ADT/PointerIntPair.h" #include "llvm/Support/TrailingObjects.h" +#include namespace mlir { namespace detail { @@ -401,7 +402,8 @@ class DistinctAttributeUniquer { /// is freed after the destruction of the distinct attribute allocator. class DistinctAttributeAllocator { public: - DistinctAttributeAllocator() = default; + DistinctAttributeAllocator(bool threadingIsEnabled) + : threadingIsEnabled(threadingIsEnabled), useThreadLocalAllocator(true) {}; DistinctAttributeAllocator(DistinctAttributeAllocator &&) = delete; DistinctAttributeAllocator(const DistinctAttributeAllocator &) = delete; @@ -411,12 +413,49 @@ class DistinctAttributeAllocator { /// 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(referencedAttr); + if (!useThreadLocalAllocator && threadingIsEnabled) { + std::scoped_lock lock(allocatorMutex); + return allocateImpl(referencedAttr); + } + return allocateImpl(referencedAttr); + } + + /// Sets a flag that stores if multithreading is enabled. The flag is used to + /// decide if locking is needed when using a non thread-safe allocator. + void disableMultiThreading(bool disable = true) { + threadingIsEnabled = !disable; + } + + /// Sets a flag to disable using thread local bump pointer allocators and use + /// a single thread-safe allocator. Use this to persist allocated storage + /// beyond the lifetime of a child thread calling this function while ensuring + /// thread-safe allocation. + void disableThreadLocalStorage(bool disable = true) { + useThreadLocalAllocator = !disable; } private: + DistinctAttrStorage *allocateImpl(Attribute referencedAttr) { + return new (getAllocatorInUse().Allocate()) + DistinctAttrStorage(referencedAttr); + } + + /// If threading is disabled on the owning MLIR context, a normal non + /// thread-local, non-thread safe bump pointer allocator is used instead to + /// prevent use-after-free errors whenever attribute storage created on a + /// crash recover thread is accessed after the thread joins. + llvm::BumpPtrAllocator &getAllocatorInUse() { + if (useThreadLocalAllocator) + return allocatorCache.get(); + return allocator; + } + ThreadLocalCache allocatorCache; + llvm::BumpPtrAllocator allocator; + std::mutex allocatorMutex; + + bool threadingIsEnabled : 1; + bool useThreadLocalAllocator : 1; }; } // namespace detail } // namespace mlir diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp index 87782e84dd6e4..ab6a5ac8b76e8 100644 --- a/mlir/lib/IR/MLIRContext.cpp +++ b/mlir/lib/IR/MLIRContext.cpp @@ -268,7 +268,8 @@ class MLIRContextImpl { public: MLIRContextImpl(bool threadingIsEnabled) - : threadingIsEnabled(threadingIsEnabled) { + : threadingIsEnabled(threadingIsEnabled), + distinctAttributeAllocator(threadingIsEnabled) { if (threadingIsEnabled) { ownedThreadPool = std::make_unique(); threadPool = ownedThreadPool.get(); @@ -596,6 +597,7 @@ void MLIRContext::disableMultithreading(bool disable) { // Update the threading mode for each of the uniquers. impl->affineUniquer.disableMultithreading(disable); impl->attributeUniquer.disableMultithreading(disable); + impl->distinctAttributeAllocator.disableMultiThreading(disable); impl->typeUniquer.disableMultithreading(disable); // Destroy thread pool (stop all threads) if it is no longer needed, or create @@ -717,6 +719,10 @@ bool MLIRContext::isOperationRegistered(StringRef name) { return RegisteredOperationName::lookup(name, this).has_value(); } +void MLIRContext::disableThreadLocalStorage(bool disable) { + getImpl().distinctAttributeAllocator.disableThreadLocalStorage(disable); +} + 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 8c6d865cb31dd..7ea2a57693217 100644 --- a/mlir/lib/Pass/PassCrashRecovery.cpp +++ b/mlir/lib/Pass/PassCrashRecovery.cpp @@ -414,6 +414,15 @@ struct FileReproducerStream : public mlir::ReproducerStream { LogicalResult PassManager::runWithCrashRecovery(Operation *op, AnalysisManager am) { + // Notify the context to disable the use of thread-local 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 during passes beyond the lifetime of the + // recovery context thread. + MLIRContext *ctx = getContext(); + ctx->disableThreadLocalStorage(); + auto guard = + llvm::make_scope_exit([ctx]() { ctx->enableThreadLocalStorage(); }); crashReproGenerator->initialize(getPasses(), op, verifyPasses); // Safely invoke the passes within a recovery context. diff --git a/mlir/test/Dialect/LLVMIR/add-debuginfo-func-scope-with-crash-reproduction.mlir b/mlir/test/Dialect/LLVMIR/add-debuginfo-func-scope-with-crash-reproduction.mlir new file mode 100644 index 0000000000000..b36ed367adb76 --- /dev/null +++ b/mlir/test/Dialect/LLVMIR/add-debuginfo-func-scope-with-crash-reproduction.mlir @@ -0,0 +1,22 @@ +// Test that the enable-debug-info-scope-on-llvm-func pass can create its +// distinct attributes when running in the crash reproducer thread. + +// RUN: mlir-opt --mlir-disable-threading --mlir-pass-pipeline-crash-reproducer=. \ +// RUN: --pass-pipeline="builtin.module(ensure-debug-info-scope-on-llvm-func)" \ +// RUN: --mlir-print-debuginfo %s | FileCheck %s + +// RUN: mlir-opt --mlir-pass-pipeline-crash-reproducer=. \ +// RUN: --pass-pipeline="builtin.module(ensure-debug-info-scope-on-llvm-func)" \ +// RUN: --mlir-print-debuginfo %s | FileCheck %s + +module { + llvm.func @func_no_debug() { + llvm.return loc(unknown) + } loc(unknown) +} loc(unknown) + +// CHECK-LABEL: llvm.func @func_no_debug() +// CHECK: llvm.return loc(#loc +// CHECK: loc(#loc[[LOC:[0-9]+]]) +// CHECK: #di_compile_unit = #llvm.di_compile_unit, +// CHECK: #di_subprogram = #llvm.di_subprogram diff --git a/mlir/test/IR/test-builtin-distinct-attrs-with-crash-reproduction.mlir b/mlir/test/IR/test-builtin-distinct-attrs-with-crash-reproduction.mlir new file mode 100644 index 0000000000000..af6dead31e263 --- /dev/null +++ b/mlir/test/IR/test-builtin-distinct-attrs-with-crash-reproduction.mlir @@ -0,0 +1,18 @@ +// This test verifies that when running with crash reproduction enabled, distinct +// attribute storage is not allocated in thread-local storage. Since crash +// reproduction runs the pass manager in a separate thread, using thread-local +// storage for distinct attributes causes use-after-free errors once the thread +// that runs the pass manager joins. + +// RUN: mlir-opt --mlir-disable-threading --mlir-pass-pipeline-crash-reproducer=. %s -test-distinct-attrs | FileCheck %s +// RUN: mlir-opt --mlir-pass-pipeline-crash-reproducer=. %s -test-distinct-attrs | FileCheck %s + +// CHECK: #[[DIST0:.*]] = distinct[0]<42 : i32> +// CHECK: #[[DIST1:.*]] = distinct[1]<42 : i32> +#distinct = distinct[0]<42 : i32> + +// CHECK: @foo_1 +func.func @foo_1() { + // CHECK: "test.op"() {distinct.input = #[[DIST0]], distinct.output = #[[DIST1]]} + "test.op"() {distinct.input = #distinct} : () -> () +}