[Mlir-commits] [mlir] [mlir] Use non thread-local allocator for DistinctAttr when threading is disabled (PR #128566)
Artemiy Bulavin
llvmlistbot at llvm.org
Wed Mar 5 09:38:59 PST 2025
https://github.com/abulavin updated https://github.com/llvm/llvm-project/pull/128566
>From b4a909e27ecd6970611c55ebf7208ba62c2e24b1 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/8] 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<DistinctAttrStorage>())
+ return new (getAllocatorInUse().Allocate<DistinctAttrStorage>())
DistinctAttrStorage(referencedAttr);
}
+ void disableMultithreading(bool disable = true) {
+ useThreadLocalCache = !disable;
+ };
+
private:
+ llvm::BumpPtrAllocator &getAllocatorInUse() {
+ if (useThreadLocalCache) {
+ return allocatorCache.get();
+ }
+ return allocator;
+ }
+
ThreadLocalCache<llvm::BumpPtrAllocator> 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<"<unknown>" in "">
+// CHECK: #di_subprogram = #llvm.di_subprogram<id = distinct[{{.*}}]<>, 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 <artemiyb at graphcore.ai>
Date: Fri, 28 Feb 2025 15:28:22 +0000
Subject: [PATCH 2/8] 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 <mutex>
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>())
- DistinctAttrStorage(referencedAttr);
+ if (useSynchronizedAllocator && !useThreadLocalAllocator) {
+ std::scoped_lock<std::mutex> 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>())
+ 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<llvm::BumpPtrAllocator> 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 <artemiyb at graphcore.ai>
Date: Fri, 28 Feb 2025 15:28:41 +0000
Subject: [PATCH 3/8] 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<"<unknown>" in "">
-// CHECK: #di_subprogram = #llvm.di_subprogram<id = distinct[{{.*}}]<>, 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<id = distinct[{{.*}}]<>,
+// CHECK: #di_subprogram = #llvm.di_subprogram<id = distinct[{{.*}}]<>
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 <artemiyb at graphcore.ai>
Date: Mon, 3 Mar 2025 18:06:43 +0000
Subject: [PATCH 4/8] 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 <artemiyb at graphcore.ai>
Date: Mon, 3 Mar 2025 18:07:16 +0000
Subject: [PATCH 5/8] 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 <artemiyb at graphcore.ai>
Date: Tue, 4 Mar 2025 23:39:24 +0000
Subject: [PATCH 6/8] 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 <artemiyb at graphcore.ai>
Date: Tue, 4 Mar 2025 23:45:03 +0000
Subject: [PATCH 7/8] 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<std::mutex> 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<llvm::DefaultThreadPool>();
threadPool = ownedThreadPool.get();
>From e7fd52231ef1c59385d137a567793f12273d1ff8 Mon Sep 17 00:00:00 2001
From: Artemiy Bulavin <artemiyb at graphcore.ai>
Date: Wed, 5 Mar 2025 17:38:40 +0000
Subject: [PATCH 8/8] 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;
More information about the Mlir-commits
mailing list