From b4a909e27ecd6970611c55ebf7208ba62c2e24b1 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Thu, 30 Jan 2025 14:32:03 +0000 Subject: [PATCH 01/10] Use non thread-local allocator for DistinctAttr when threading is disabled --- mlir/lib/IR/AttributeDetail.h | 22 +++++++++++++++++-- mlir/lib/IR/MLIRContext.cpp | 1 + ...fo-func-scope-with-crash-reproduction.mlir | 19 ++++++++++++++++ ...istinct-attrs-with-crash-reproduction.mlir | 17 ++++++++++++++ 4 files changed, 57 insertions(+), 2 deletions(-) create mode 100644 mlir/test/Dialect/LLVMIR/add-debuginfo-func-scope-with-crash-reproduction.mlir create mode 100644 mlir/test/IR/test-builtin-distinct-attrs-with-crash-reproduction.mlir diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h index 26d40ac3a38f6..6787dc6d3e698 100644 --- a/mlir/lib/IR/AttributeDetail.h +++ b/mlir/lib/IR/AttributeDetail.h @@ -23,6 +23,7 @@ #include "mlir/Support/ThreadLocalCache.h" #include "llvm/ADT/APFloat.h" #include "llvm/ADT/PointerIntPair.h" +#include "llvm/Support/Allocator.h" #include "llvm/Support/TrailingObjects.h" namespace mlir { @@ -409,14 +410,31 @@ class DistinctAttributeAllocator { operator=(const DistinctAttributeAllocator &) = delete; /// Allocates a distinct attribute storage using a thread local bump pointer - /// allocator to enable synchronization free parallel allocations. + /// allocator to enable synchronization free parallel allocations. If + /// threading is disabled on the owning MLIR context, a normal non + /// thread-local 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. DistinctAttrStorage *allocate(Attribute referencedAttr) { - return new (allocatorCache.get().Allocate()) + return new (getAllocatorInUse().Allocate()) DistinctAttrStorage(referencedAttr); } + void disableMultithreading(bool disable = true) { + useThreadLocalCache = !disable; + }; + private: + llvm::BumpPtrAllocator &getAllocatorInUse() { + if (useThreadLocalCache) { + return allocatorCache.get(); + } + return allocator; + } + ThreadLocalCache allocatorCache; + llvm::BumpPtrAllocator allocator; + bool useThreadLocalCache : 1; }; } // namespace detail } // namespace mlir diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp index 87782e84dd6e4..04f2cf8f7b1ec 100644 --- a/mlir/lib/IR/MLIRContext.cpp +++ b/mlir/lib/IR/MLIRContext.cpp @@ -596,6 +596,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 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..5f49df419bccb --- /dev/null +++ b/mlir/test/Dialect/LLVMIR/add-debuginfo-func-scope-with-crash-reproduction.mlir @@ -0,0 +1,19 @@ +// Test that the enable-debug-info-scope-on-llvm-func pass can create its +// DI 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 + +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_file = #llvm.di_file<"" in ""> +// CHECK: #di_subprogram = #llvm.di_subprogram, compileUnit = #di_compile_unit, scope = #di_file, name = "func_no_debug", linkageName = "func_no_debug", file = #di_file, line = 1, scopeLine = 1, subprogramFlags = "Definition|Optimized", type = #di_subroutine_type> +// CHECK: #loc[[LOC]] = loc(fused<#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..5c3b78bc6e154 --- /dev/null +++ b/mlir/test/IR/test-builtin-distinct-attrs-with-crash-reproduction.mlir @@ -0,0 +1,17 @@ +// This is a regression test that 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 + +// 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} : () -> () +} From 55a7a38d3fd83addeb86e419aedc424efe8d593d Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Fri, 28 Feb 2025 15:28:22 +0000 Subject: [PATCH 02/10] fixup: Disabel threadlocal storage to handle multithreaded crash reproducer generation --- mlir/include/mlir/IR/MLIRContext.h | 7 ++++ mlir/lib/IR/AttributeDetail.h | 59 +++++++++++++++++++++-------- mlir/lib/IR/MLIRContext.cpp | 6 ++- mlir/lib/Pass/PassCrashRecovery.cpp | 13 +++++++ 4 files changed, 69 insertions(+), 16 deletions(-) diff --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h index ef8dab87f131a..221429e65c6ab 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 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/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h index 6787dc6d3e698..91aa88a5ca5cb 100644 --- a/mlir/lib/IR/AttributeDetail.h +++ b/mlir/lib/IR/AttributeDetail.h @@ -23,8 +23,9 @@ #include "mlir/Support/ThreadLocalCache.h" #include "llvm/ADT/APFloat.h" #include "llvm/ADT/PointerIntPair.h" -#include "llvm/Support/Allocator.h" +#include "llvm/Support/ErrorHandling.h" #include "llvm/Support/TrailingObjects.h" +#include namespace mlir { namespace detail { @@ -402,7 +403,8 @@ class DistinctAttributeUniquer { /// is freed after the destruction of the distinct attribute allocator. class DistinctAttributeAllocator { public: - DistinctAttributeAllocator() = default; + DistinctAttributeAllocator() + : useThreadLocalAllocator(true), useSynchronizedAllocator(false) {}; DistinctAttributeAllocator(DistinctAttributeAllocator &&) = delete; DistinctAttributeAllocator(const DistinctAttributeAllocator &) = delete; @@ -410,31 +412,58 @@ class DistinctAttributeAllocator { operator=(const DistinctAttributeAllocator &) = delete; /// Allocates a distinct attribute storage using a thread local bump pointer - /// allocator to enable synchronization free parallel allocations. If - /// threading is disabled on the owning MLIR context, a normal non - /// thread-local 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. + /// allocator to enable synchronization free parallel allocations. DistinctAttrStorage *allocate(Attribute referencedAttr) { - return new (getAllocatorInUse().Allocate()) - DistinctAttrStorage(referencedAttr); + if (useSynchronizedAllocator && !useThreadLocalAllocator) { + std::scoped_lock lock(allocatorMutex); + return allocateImpl(referencedAttr); + } + if (!useSynchronizedAllocator) + return allocateImpl(referencedAttr); + llvm_unreachable( + "Cannot use both a synchronised and thread_local allocator!"); } - void disableMultithreading(bool disable = true) { - useThreadLocalCache = !disable; - }; + /// Sets flags to use thread local bump pointer allocators or a single + /// non-thread safe bump pointer allocator depending on if multi-threading is + /// enabled. Use this to disable the use of thread local storage and bypass + /// thread safety synchronization overhead. + void disableMultiThreading(bool disable = true) { + disableThreadLocalStorage(disable); + useSynchronizedAllocator = !disable; + } + + /// Sets flags 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; + useSynchronizedAllocator = 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 (useThreadLocalCache) { + if (useThreadLocalAllocator) return allocatorCache.get(); - } return allocator; } ThreadLocalCache allocatorCache; llvm::BumpPtrAllocator allocator; - bool useThreadLocalCache : 1; + std::mutex allocatorMutex; + + bool useThreadLocalAllocator : 1; + bool useSynchronizedAllocator : 1; }; } // namespace detail } // namespace mlir diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp index 04f2cf8f7b1ec..c18bf6dc006a3 100644 --- a/mlir/lib/IR/MLIRContext.cpp +++ b/mlir/lib/IR/MLIRContext.cpp @@ -596,7 +596,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->distinctAttributeAllocator.disableThreadLocalStorage(disable); impl->typeUniquer.disableMultithreading(disable); // Destroy thread pool (stop all threads) if it is no longer needed, or create @@ -718,6 +718,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..19f2a936299ed 100644 --- a/mlir/lib/Pass/PassCrashRecovery.cpp +++ b/mlir/lib/Pass/PassCrashRecovery.cpp @@ -414,6 +414,19 @@ 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. This RAII guard + // will re-enable thread local storage use upon function exit. + class ScopedThreadLocalStorageDisable { + MLIRContext *const ctx; + + public: + ScopedThreadLocalStorageDisable(MLIRContext *ctx) : ctx(ctx) { + ctx->disableThreadLocalStorage(); + } + ~ScopedThreadLocalStorageDisable() { ctx->enableThreadLocalStorage(); } + }; + ScopedThreadLocalStorageDisable raii(this->getContext()); crashReproGenerator->initialize(getPasses(), op, verifyPasses); // Safely invoke the passes within a recovery context. From 8ace71b19269127ca7ddfb0fed933de2c541f360 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Fri, 28 Feb 2025 15:28:41 +0000 Subject: [PATCH 03/10] fixup: add tests with multithreading enabled --- ...-debuginfo-func-scope-with-crash-reproduction.mlir | 11 +++++++---- ...uiltin-distinct-attrs-with-crash-reproduction.mlir | 11 ++++++----- 2 files changed, 13 insertions(+), 9 deletions(-) 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 index 5f49df419bccb..b36ed367adb76 100644 --- 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 @@ -1,10 +1,14 @@ // Test that the enable-debug-info-scope-on-llvm-func pass can create its -// DI attributes when running in the crash reproducer thread, +// 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) @@ -14,6 +18,5 @@ module { // CHECK-LABEL: llvm.func @func_no_debug() // CHECK: llvm.return loc(#loc // CHECK: loc(#loc[[LOC:[0-9]+]]) -// CHECK: #di_file = #llvm.di_file<"" in ""> -// CHECK: #di_subprogram = #llvm.di_subprogram, compileUnit = #di_compile_unit, scope = #di_file, name = "func_no_debug", linkageName = "func_no_debug", file = #di_file, line = 1, scopeLine = 1, subprogramFlags = "Definition|Optimized", type = #di_subroutine_type> -// CHECK: #loc[[LOC]] = loc(fused<#di_subprogram> +// 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 index 5c3b78bc6e154..af6dead31e263 100644 --- a/mlir/test/IR/test-builtin-distinct-attrs-with-crash-reproduction.mlir +++ b/mlir/test/IR/test-builtin-distinct-attrs-with-crash-reproduction.mlir @@ -1,10 +1,11 @@ -// This is a regression test that 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. +// 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> From d611dba5104f832195d13dee11770e85d1665ef7 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Mon, 3 Mar 2025 18:06:43 +0000 Subject: [PATCH 04/10] fixup: use disableMultiThreading in MLIRContext::disableMultiThreading --- mlir/lib/IR/MLIRContext.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp index c18bf6dc006a3..7bfed8c298dd4 100644 --- a/mlir/lib/IR/MLIRContext.cpp +++ b/mlir/lib/IR/MLIRContext.cpp @@ -596,7 +596,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.disableThreadLocalStorage(disable); + impl->distinctAttributeAllocator.disableMultiThreading(disable); impl->typeUniquer.disableMultithreading(disable); // Destroy thread pool (stop all threads) if it is no longer needed, or create From 3792201bb7c5b1b69ac721037b50e11d01b26107 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Mon, 3 Mar 2025 18:07:16 +0000 Subject: [PATCH 05/10] fixup: use llvm::make_scope_exit instead of custom RAII guard --- mlir/lib/Pass/PassCrashRecovery.cpp | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/mlir/lib/Pass/PassCrashRecovery.cpp b/mlir/lib/Pass/PassCrashRecovery.cpp index 19f2a936299ed..f357e4fb6ab76 100644 --- a/mlir/lib/Pass/PassCrashRecovery.cpp +++ b/mlir/lib/Pass/PassCrashRecovery.cpp @@ -415,18 +415,14 @@ 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. This RAII guard - // will re-enable thread local storage use upon function exit. - class ScopedThreadLocalStorageDisable { - MLIRContext *const ctx; - - public: - ScopedThreadLocalStorageDisable(MLIRContext *ctx) : ctx(ctx) { - ctx->disableThreadLocalStorage(); - } - ~ScopedThreadLocalStorageDisable() { ctx->enableThreadLocalStorage(); } - }; - ScopedThreadLocalStorageDisable raii(this->getContext()); + // 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. + auto *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. From 68843d56ce726471ebafcc5f1282d482050847f5 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Tue, 4 Mar 2025 23:39:24 +0000 Subject: [PATCH 06/10] fixup: use MLIRContext* instead of auto* --- mlir/lib/Pass/PassCrashRecovery.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Pass/PassCrashRecovery.cpp b/mlir/lib/Pass/PassCrashRecovery.cpp index f357e4fb6ab76..7ea2a57693217 100644 --- a/mlir/lib/Pass/PassCrashRecovery.cpp +++ b/mlir/lib/Pass/PassCrashRecovery.cpp @@ -419,7 +419,7 @@ LogicalResult PassManager::runWithCrashRecovery(Operation *op, // 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. - auto *ctx = getContext(); + MLIRContext *ctx = getContext(); ctx->disableThreadLocalStorage(); auto guard = llvm::make_scope_exit([ctx]() { ctx->enableThreadLocalStorage(); }); From 8baecc9a466bc8998b1bca77df289442a5ebb073 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Tue, 4 Mar 2025 23:45:03 +0000 Subject: [PATCH 07/10] fixup: replace unreachable with clearer logic in allocate method --- mlir/lib/IR/AttributeDetail.h | 17 ++++++----------- mlir/lib/IR/MLIRContext.cpp | 3 ++- 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h index 91aa88a5ca5cb..21d7d98eddf33 100644 --- a/mlir/lib/IR/AttributeDetail.h +++ b/mlir/lib/IR/AttributeDetail.h @@ -403,8 +403,8 @@ class DistinctAttributeUniquer { /// is freed after the destruction of the distinct attribute allocator. class DistinctAttributeAllocator { public: - DistinctAttributeAllocator() - : useThreadLocalAllocator(true), useSynchronizedAllocator(false) {}; + DistinctAttributeAllocator(bool threadingIsEnabled) + : threadingIsEnabled(threadingIsEnabled), useThreadLocalAllocator(true) {}; DistinctAttributeAllocator(DistinctAttributeAllocator &&) = delete; DistinctAttributeAllocator(const DistinctAttributeAllocator &) = delete; @@ -414,14 +414,11 @@ class DistinctAttributeAllocator { /// Allocates a distinct attribute storage using a thread local bump pointer /// allocator to enable synchronization free parallel allocations. DistinctAttrStorage *allocate(Attribute referencedAttr) { - if (useSynchronizedAllocator && !useThreadLocalAllocator) { + if (!useThreadLocalAllocator && threadingIsEnabled) { std::scoped_lock lock(allocatorMutex); return allocateImpl(referencedAttr); } - if (!useSynchronizedAllocator) - return allocateImpl(referencedAttr); - llvm_unreachable( - "Cannot use both a synchronised and thread_local allocator!"); + return allocateImpl(referencedAttr); } /// Sets flags to use thread local bump pointer allocators or a single @@ -429,8 +426,7 @@ class DistinctAttributeAllocator { /// enabled. Use this to disable the use of thread local storage and bypass /// thread safety synchronization overhead. void disableMultiThreading(bool disable = true) { - disableThreadLocalStorage(disable); - useSynchronizedAllocator = !disable; + threadingIsEnabled = !disable; } /// Sets flags to disable using thread local bump pointer allocators and use a @@ -439,7 +435,6 @@ class DistinctAttributeAllocator { /// thread-safe allocation. void disableThreadLocalStorage(bool disable = true) { useThreadLocalAllocator = !disable; - useSynchronizedAllocator = disable; } private: @@ -462,8 +457,8 @@ class DistinctAttributeAllocator { llvm::BumpPtrAllocator allocator; std::mutex allocatorMutex; + bool threadingIsEnabled : 1; bool useThreadLocalAllocator : 1; - bool useSynchronizedAllocator : 1; }; } // namespace detail } // namespace mlir diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp index 7bfed8c298dd4..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(); From e7fd52231ef1c59385d137a567793f12273d1ff8 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Wed, 5 Mar 2025 17:38:40 +0000 Subject: [PATCH 08/10] fixup: update outdated docs comments --- mlir/lib/IR/AttributeDetail.h | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h index 21d7d98eddf33..9d309f3d03d5f 100644 --- a/mlir/lib/IR/AttributeDetail.h +++ b/mlir/lib/IR/AttributeDetail.h @@ -421,17 +421,15 @@ class DistinctAttributeAllocator { return allocateImpl(referencedAttr); } - /// Sets flags to use thread local bump pointer allocators or a single - /// non-thread safe bump pointer allocator depending on if multi-threading is - /// enabled. Use this to disable the use of thread local storage and bypass - /// thread safety synchronization overhead. + /// 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 flags 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 + /// 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; From 7fffb590e69b33726e8501669d9c5756a3c3dc53 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Thu, 6 Mar 2025 16:50:45 +0000 Subject: [PATCH 09/10] fixup: note that disabling threading implicitly disables thread-local storage --- mlir/include/mlir/IR/MLIRContext.h | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h index 221429e65c6ab..ad89d631c8a5f 100644 --- a/mlir/include/mlir/IR/MLIRContext.h +++ b/mlir/include/mlir/IR/MLIRContext.h @@ -153,8 +153,9 @@ class MLIRContext { disableMultithreading(!enable); } - /// Set the flag specifying if thread-local storage should be used - /// by storage allocators in this context. + /// 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); From da4e176d7ee13257bd945d9e4706e2f2836e7029 Mon Sep 17 00:00:00 2001 From: Artemiy Bulavin Date: Thu, 6 Mar 2025 16:51:06 +0000 Subject: [PATCH 10/10] fixup: remove unused error handling header --- mlir/lib/IR/AttributeDetail.h | 1 - 1 file changed, 1 deletion(-) diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h index 9d309f3d03d5f..8fed18140433c 100644 --- a/mlir/lib/IR/AttributeDetail.h +++ b/mlir/lib/IR/AttributeDetail.h @@ -23,7 +23,6 @@ #include "mlir/Support/ThreadLocalCache.h" #include "llvm/ADT/APFloat.h" #include "llvm/ADT/PointerIntPair.h" -#include "llvm/Support/ErrorHandling.h" #include "llvm/Support/TrailingObjects.h" #include