[Mlir-commits] [mlir] [mlir][Diagnostics] Set parallel diagnostic order ids without a lock (PR #232520)
Bruno Cardoso Lopes
llvmlistbot at llvm.org
Thu Oct 15 18:08:22 PDT 2026
https://github.com/bcardosolopes updated https://github.com/llvm/llvm-project/pull/232520
>From 6db5bf8314db297fc78742c2460cd86d35e9fdbf Mon Sep 17 00:00:00 2001
From: Bruno Cardoso Lopes <bruno.cardoso at gmail.com>
Date: Wed, 14 Oct 2026 20:40:10 -0700
Subject: [PATCH 1/3] [mlir] Set parallel diagnostic order ids without a lock
failableParallelForEach tells its ParallelDiagnosticHandler which
element each thread is processing before every element and clears it
after, so diagnostics come out in element order. Both went through the
handler mutex plus a get_threadid syscall. Over many small elements
that pair was the most contended lock in the loop, mostly kernel spin.
The order id now lives in a thread-local slot naming its handler by a
monotonic id (no raw pointers, so a handler destroyed with a stale slot
naming it is harmless). A thread handles one element at a time per
handler, so the slot covers every thread whose loop is not nested in
another loop on the same thread; a nested loop falls back to the map,
and lookups check the slot first. As a side effect, a nested flush is
now attributed to the outer element instead of bypassing it.
200K single-op func.func through per-function cse, 96 threads: sys time
22.9s to 3.8s, wall 3.4s to 3.2s. Unittest pins the nested fallback and
outer-id preservation.
Assisted-by: Muse Code (Muse)
---
mlir/lib/IR/Diagnostics.cpp | 52 +++++++++++++++++++++++++++-----
mlir/unittests/IR/Diagnostic.cpp | 40 ++++++++++++++++++++++++
2 files changed, 84 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/IR/Diagnostics.cpp b/mlir/lib/IR/Diagnostics.cpp
index a3ddb76a66313c..46733a5323f89e 100644
--- a/mlir/lib/IR/Diagnostics.cpp
+++ b/mlir/lib/IR/Diagnostics.cpp
@@ -22,6 +22,7 @@
#include "llvm/Support/Signals.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/raw_ostream.h"
+#include <atomic>
#include <optional>
using namespace mlir;
@@ -991,6 +992,23 @@ void SourceMgrDiagnosticVerifierHandler::process(LocationAttr loc,
namespace mlir {
namespace detail {
+/// Unique ids for parallel diagnostic handlers, so the thread-local slot
+/// below can name a handler without holding a raw pointer to it: a handler
+/// may be destroyed while a thread that skipped its erase still names it.
+static std::atomic<uint64_t> nextSlotID{1};
+
+/// The handler, if any, whose order id for the current thread is held in
+/// `threadOrderID` rather than in the handler's `threadToOrderID` map. A
+/// parallel loop sets and erases a thread's order id once per element, and
+/// doing that under the handler's mutex made the mutex the most contended lock
+/// of a loop over many small elements. A thread processes one element at a
+/// time for each handler, so one slot covers every thread whose loop is not
+/// nested in another loop on the same thread; a nested loop's handler falls
+/// back to the map. Reentrant use of the same handler on one thread remains
+/// unsupported, as before.
+static LLVM_THREAD_LOCAL uint64_t threadSlotID = 0;
+static LLVM_THREAD_LOCAL size_t threadOrderID = 0;
+
struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
struct ThreadDiagnostic {
ThreadDiagnostic(size_t id, Diagnostic diag)
@@ -1008,16 +1026,21 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
ParallelDiagnosticHandlerImpl(MLIRContext *ctx) : context(ctx) {
handlerID = ctx->getDiagEngine().registerHandler([this](Diagnostic &diag) {
+ // Fast path: this thread's order id is in the slot, no lock needed.
+ if (threadSlotID == slotID) {
+ llvm::sys::SmartScopedLock<true> lock(mutex);
+ diagnostics.emplace_back(threadOrderID, std::move(diag));
+ return success();
+ }
+ // Slow path: a nested loop on this thread, consult the map. If this
+ // thread is not tracked, return failure to let another handler
+ // process this diagnostic.
uint64_t tid = llvm::get_threadid();
llvm::sys::SmartScopedLock<true> lock(mutex);
-
- // If this thread is not tracked, then return failure to let another
- // handler process this diagnostic.
- if (!threadToOrderID.count(tid))
+ auto it = threadToOrderID.find(tid);
+ if (it == threadToOrderID.end())
return failure();
-
- // Append a new diagnostic.
- diagnostics.emplace_back(threadToOrderID[tid], std::move(diag));
+ diagnostics.emplace_back(it->second, std::move(diag));
return success();
});
}
@@ -1050,6 +1073,11 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
/// Set the order id for the current thread.
void setOrderIDForThread(size_t orderID) {
+ if (!threadSlotID || threadSlotID == slotID) {
+ threadSlotID = slotID;
+ threadOrderID = orderID;
+ return;
+ }
uint64_t tid = llvm::get_threadid();
llvm::sys::SmartScopedLock<true> lock(mutex);
threadToOrderID[tid] = orderID;
@@ -1057,6 +1085,10 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
/// Remove the order id for the current thread.
void eraseOrderIDForThread() {
+ if (threadSlotID == slotID) {
+ threadSlotID = 0;
+ return;
+ }
uint64_t tid = llvm::get_threadid();
llvm::sys::SmartScopedLock<true> lock(mutex);
threadToOrderID.erase(tid);
@@ -1097,7 +1129,8 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
/// A smart mutex to lock access to the internal state.
llvm::sys::SmartMutex<true> mutex;
- /// A mapping between the thread id and the current order id.
+ /// A mapping between the thread id and the current order id, for the threads
+ /// whose order id is not in `threadOrderID`.
DenseMap<uint64_t, size_t> threadToOrderID;
/// An unordered list of diagnostics that were emitted.
@@ -1106,6 +1139,9 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
/// The unique id for the parallel handler.
DiagnosticEngine::HandlerID handlerID = 0;
+ /// Unique id naming this handler in the thread-local order-id slot.
+ const uint64_t slotID = nextSlotID.fetch_add(1, std::memory_order_relaxed);
+
/// The context to emit the diagnostics to.
MLIRContext *context;
};
diff --git a/mlir/unittests/IR/Diagnostic.cpp b/mlir/unittests/IR/Diagnostic.cpp
index 6390cdfe85acca..5a840bb9a2cacf 100644
--- a/mlir/unittests/IR/Diagnostic.cpp
+++ b/mlir/unittests/IR/Diagnostic.cpp
@@ -13,6 +13,8 @@
#include "gtest/gtest.h"
#include <memory>
+#include <string>
+#include <vector>
using namespace mlir;
using namespace mlir::detail;
@@ -104,6 +106,44 @@ TEST(SourceMgrDiagnosticVerifierHandler, ErasesOwnRegistration) {
EXPECT_EQ(seen, 1u);
}
+TEST(ParallelDiagnosticHandler, NestedFallbackPreservesOrder) {
+ MLIRContext ctx;
+ std::vector<std::string> seen;
+ ScopedDiagnosticHandler capture(&ctx, [&](Diagnostic &diag) {
+ std::string msg;
+ llvm::raw_string_ostream os(msg);
+ diag.print(os);
+ seen.push_back(msg);
+ return success();
+ });
+
+ {
+ ParallelDiagnosticHandler outer(&ctx);
+ outer.setOrderIDForThread(1);
+ {
+ ParallelDiagnosticHandler inner(&ctx);
+ // The slot is taken by outer, so these go through the map fallback.
+ inner.setOrderIDForThread(2);
+ emitError(UnknownLoc::get(&ctx), "inner-2");
+ inner.setOrderIDForThread(1);
+ emitError(UnknownLoc::get(&ctx), "inner-1");
+ inner.eraseOrderIDForThread();
+ } // inner emits sorted: inner-1, inner-2.
+ // The slot still names outer after the nested use: the inner flush was
+ // attributed to outer's element, and this is held with id 1, so the
+ // id-0 diagnostic below sorts first.
+ emitError(UnknownLoc::get(&ctx), "outer-A");
+ outer.setOrderIDForThread(0);
+ emitError(UnknownLoc::get(&ctx), "outer-B");
+ outer.eraseOrderIDForThread();
+ } // outer emits sorted: outer-B, then the id-1 group in arrival order.
+ ASSERT_EQ(seen.size(), 4u);
+ EXPECT_NE(seen[0].find("outer-B"), std::string::npos);
+ EXPECT_NE(seen[1].find("inner-1"), std::string::npos);
+ EXPECT_NE(seen[2].find("inner-2"), std::string::npos);
+ EXPECT_NE(seen[3].find("outer-A"), std::string::npos);
+}
+
TEST(SourceMgrDiagnosticVerifierHandler, OutlivesRegisteredContext) {
// Match mlir-opt's per-buffer context lifetime.
MLIRContext context;
>From 20bff9b0d7d4b8e3c6fe519e662073f68e783d8e Mon Sep 17 00:00:00 2001
From: Bruno Cardoso Lopes <bruno.cardoso at gmail.com>
Date: Thu, 15 Oct 2026 17:55:47 -0700
Subject: [PATCH 2/3] Address review: handlerUniqueID names, docs,
multithreaded test
---
mlir/include/mlir/IR/Diagnostics.h | 7 ++++-
mlir/lib/IR/Diagnostics.cpp | 28 +++++++++++-------
mlir/unittests/IR/Diagnostic.cpp | 47 +++++++++++++++++++++++++-----
3 files changed, 63 insertions(+), 19 deletions(-)
diff --git a/mlir/include/mlir/IR/Diagnostics.h b/mlir/include/mlir/IR/Diagnostics.h
index 43a80e0ee1c91f..62b15ec8a31790 100644
--- a/mlir/include/mlir/IR/Diagnostics.h
+++ b/mlir/include/mlir/IR/Diagnostics.h
@@ -737,10 +737,15 @@ class ParallelDiagnosticHandler {
/// '1'; and so on and so forth. This provides a way for the handler to
/// deterministically order the diagnostics that it receives given the thread
/// that it is receiving on.
+ ///
+ /// Each set must be paired with an erase on the same thread before the
+ /// handler is destroyed. Nesting distinct handlers on one thread is
+ /// supported; reentrant use of one handler on one thread is not.
void setOrderIDForThread(size_t orderID);
/// Remove the order id for the current thread. This removes the thread from
- /// diagnostics tracking.
+ /// diagnostics tracking. Must be called on the same thread that set it,
+ /// before the handler is destroyed.
void eraseOrderIDForThread();
private:
diff --git a/mlir/lib/IR/Diagnostics.cpp b/mlir/lib/IR/Diagnostics.cpp
index 46733a5323f89e..4968aa5f6cd1c9 100644
--- a/mlir/lib/IR/Diagnostics.cpp
+++ b/mlir/lib/IR/Diagnostics.cpp
@@ -995,7 +995,7 @@ namespace detail {
/// Unique ids for parallel diagnostic handlers, so the thread-local slot
/// below can name a handler without holding a raw pointer to it: a handler
/// may be destroyed while a thread that skipped its erase still names it.
-static std::atomic<uint64_t> nextSlotID{1};
+static std::atomic<uint64_t> nextHandlerUniqueID{1};
/// The handler, if any, whose order id for the current thread is held in
/// `threadOrderID` rather than in the handler's `threadToOrderID` map. A
@@ -1006,7 +1006,7 @@ static std::atomic<uint64_t> nextSlotID{1};
/// nested in another loop on the same thread; a nested loop's handler falls
/// back to the map. Reentrant use of the same handler on one thread remains
/// unsupported, as before.
-static LLVM_THREAD_LOCAL uint64_t threadSlotID = 0;
+static LLVM_THREAD_LOCAL uint64_t threadHandlerUniqueID = 0;
static LLVM_THREAD_LOCAL size_t threadOrderID = 0;
struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
@@ -1026,8 +1026,9 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
ParallelDiagnosticHandlerImpl(MLIRContext *ctx) : context(ctx) {
handlerID = ctx->getDiagEngine().registerHandler([this](Diagnostic &diag) {
- // Fast path: this thread's order id is in the slot, no lock needed.
- if (threadSlotID == slotID) {
+ // Fast path: this thread's order id is in the slot, so looking it up
+ // needs no lock. Appending below still locks.
+ if (threadHandlerUniqueID == handlerUniqueID) {
llvm::sys::SmartScopedLock<true> lock(mutex);
diagnostics.emplace_back(threadOrderID, std::move(diag));
return success();
@@ -1073,8 +1074,8 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
/// Set the order id for the current thread.
void setOrderIDForThread(size_t orderID) {
- if (!threadSlotID || threadSlotID == slotID) {
- threadSlotID = slotID;
+ if (!threadHandlerUniqueID || threadHandlerUniqueID == handlerUniqueID) {
+ threadHandlerUniqueID = handlerUniqueID;
threadOrderID = orderID;
return;
}
@@ -1085,8 +1086,8 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
/// Remove the order id for the current thread.
void eraseOrderIDForThread() {
- if (threadSlotID == slotID) {
- threadSlotID = 0;
+ if (threadHandlerUniqueID == handlerUniqueID) {
+ threadHandlerUniqueID = 0;
return;
}
uint64_t tid = llvm::get_threadid();
@@ -1129,8 +1130,12 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
/// A smart mutex to lock access to the internal state.
llvm::sys::SmartMutex<true> mutex;
- /// A mapping between the thread id and the current order id, for the threads
- /// whose order id is not in `threadOrderID`.
+ /// Order ids, keyed by thread id, for threads whose thread-local slot is
+ /// owned by another handler: a loop nested inside another parallel loop on
+ /// the same thread. Idle workers running nested elements use the slot, so
+ /// this map only holds threads that are also inside an enclosing loop.
+ /// Guarded by `mutex`. A thread in neither the slot nor this map is not
+ /// tracked, and its diagnostics pass on to the next handler.
DenseMap<uint64_t, size_t> threadToOrderID;
/// An unordered list of diagnostics that were emitted.
@@ -1140,7 +1145,8 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
DiagnosticEngine::HandlerID handlerID = 0;
/// Unique id naming this handler in the thread-local order-id slot.
- const uint64_t slotID = nextSlotID.fetch_add(1, std::memory_order_relaxed);
+ const uint64_t handlerUniqueID =
+ nextHandlerUniqueID.fetch_add(1, std::memory_order_relaxed);
/// The context to emit the diagnostics to.
MLIRContext *context;
diff --git a/mlir/unittests/IR/Diagnostic.cpp b/mlir/unittests/IR/Diagnostic.cpp
index 5a840bb9a2cacf..370f56fe15b9b3 100644
--- a/mlir/unittests/IR/Diagnostic.cpp
+++ b/mlir/unittests/IR/Diagnostic.cpp
@@ -7,12 +7,14 @@
//===----------------------------------------------------------------------===//
#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Threading.h"
#include "mlir/Support/TypeID.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/raw_ostream.h"
#include "gtest/gtest.h"
#include <memory>
+#include <numeric>
#include <string>
#include <vector>
@@ -123,12 +125,13 @@ TEST(ParallelDiagnosticHandler, NestedFallbackPreservesOrder) {
{
ParallelDiagnosticHandler inner(&ctx);
// The slot is taken by outer, so these go through the map fallback.
- inner.setOrderIDForThread(2);
- emitError(UnknownLoc::get(&ctx), "inner-2");
- inner.setOrderIDForThread(1);
- emitError(UnknownLoc::get(&ctx), "inner-1");
+ // Inner ids are deliberately disjoint from outer's 0/1.
+ inner.setOrderIDForThread(12);
+ emitError(UnknownLoc::get(&ctx), "inner-12");
+ inner.setOrderIDForThread(11);
+ emitError(UnknownLoc::get(&ctx), "inner-11");
inner.eraseOrderIDForThread();
- } // inner emits sorted: inner-1, inner-2.
+ } // inner emits sorted: inner-11, inner-12.
// The slot still names outer after the nested use: the inner flush was
// attributed to outer's element, and this is held with id 1, so the
// id-0 diagnostic below sorts first.
@@ -139,11 +142,41 @@ TEST(ParallelDiagnosticHandler, NestedFallbackPreservesOrder) {
} // outer emits sorted: outer-B, then the id-1 group in arrival order.
ASSERT_EQ(seen.size(), 4u);
EXPECT_NE(seen[0].find("outer-B"), std::string::npos);
- EXPECT_NE(seen[1].find("inner-1"), std::string::npos);
- EXPECT_NE(seen[2].find("inner-2"), std::string::npos);
+ EXPECT_NE(seen[1].find("inner-11"), std::string::npos);
+ EXPECT_NE(seen[2].find("inner-12"), std::string::npos);
EXPECT_NE(seen[3].find("outer-A"), std::string::npos);
}
+TEST(ParallelDiagnosticHandler, MultithreadedOrderingHonored) {
+ MLIRContext ctx;
+ // Element i emits "diag-i"; whatever the thread interleaving, the flush
+ // must come out in index order.
+ constexpr unsigned numElements = 1000;
+ std::vector<std::string> seen;
+ ScopedDiagnosticHandler capture(&ctx, [&](Diagnostic &diag) {
+ std::string msg;
+ llvm::raw_string_ostream os(msg);
+ diag.print(os);
+ seen.push_back(msg);
+ return success();
+ });
+
+ std::vector<unsigned> indices(numElements);
+ std::iota(indices.begin(), indices.end(), 0);
+ ASSERT_TRUE(succeeded(
+ failableParallelForEach(&ctx, indices, [&](unsigned i) -> LogicalResult {
+ emitError(UnknownLoc::get(&ctx), "diag-" + Twine(i));
+ return success();
+ })));
+
+ ASSERT_EQ(seen.size(), numElements);
+ for (size_t i = 0; i < seen.size(); ++i) {
+ size_t pos = seen[i].find("diag-");
+ ASSERT_NE(pos, std::string::npos);
+ EXPECT_EQ(std::stoul(seen[i].substr(pos + 5)), i);
+ }
+}
+
TEST(SourceMgrDiagnosticVerifierHandler, OutlivesRegisteredContext) {
// Match mlir-opt's per-buffer context lifetime.
MLIRContext context;
>From 92486292ab6c86abdaea47f56387c5934f4a1657 Mon Sep 17 00:00:00 2001
From: Bruno Cardoso Lopes <bruno.cardoso at gmail.com>
Date: Thu, 15 Oct 2026 18:07:58 -0700
Subject: [PATCH 3/3] Address review: thread-local nesting stack with map
overflow
---
mlir/lib/IR/Diagnostics.cpp | 79 ++++++++++++++++++++------------
mlir/unittests/IR/Diagnostic.cpp | 74 ++++++++++++++++++++++++++++--
2 files changed, 120 insertions(+), 33 deletions(-)
diff --git a/mlir/lib/IR/Diagnostics.cpp b/mlir/lib/IR/Diagnostics.cpp
index 4968aa5f6cd1c9..8d1cdb4df12bef 100644
--- a/mlir/lib/IR/Diagnostics.cpp
+++ b/mlir/lib/IR/Diagnostics.cpp
@@ -997,17 +997,29 @@ namespace detail {
/// may be destroyed while a thread that skipped its erase still names it.
static std::atomic<uint64_t> nextHandlerUniqueID{1};
-/// The handler, if any, whose order id for the current thread is held in
-/// `threadOrderID` rather than in the handler's `threadToOrderID` map. A
-/// parallel loop sets and erases a thread's order id once per element, and
-/// doing that under the handler's mutex made the mutex the most contended lock
-/// of a loop over many small elements. A thread processes one element at a
-/// time for each handler, so one slot covers every thread whose loop is not
-/// nested in another loop on the same thread; a nested loop's handler falls
-/// back to the map. Reentrant use of the same handler on one thread remains
-/// unsupported, as before.
-static LLVM_THREAD_LOCAL uint64_t threadHandlerUniqueID = 0;
-static LLVM_THREAD_LOCAL size_t threadOrderID = 0;
+/// One entry of the thread-local nesting stack: the handler an element
+/// belongs to and its order id. Plain data so the thread-locals below stay
+/// trivially destructible.
+struct ThreadOrderEntry {
+ uint64_t handlerID;
+ size_t orderID;
+};
+
+/// How many nested parallel loops one thread tracks without a lock. Past
+/// this the handler's map is the overflow.
+constexpr unsigned kMaxThreadOrderEntries = 4;
+
+/// The order ids for the current thread, one entry per enclosing parallel
+/// loop, innermost last. A parallel loop sets and erases a thread's order id
+/// once per element, and doing that under the handler's mutex made the mutex
+/// the most contended lock of a loop over many small elements. A thread
+/// processes one element at a time for each handler, so a small stack covers
+/// every realistic nesting depth. Erases are LIFO, matching loop nesting.
+/// Reentrant use of the same handler on one thread remains unsupported, as
+/// before.
+static LLVM_THREAD_LOCAL ThreadOrderEntry
+ threadOrderStack[kMaxThreadOrderEntries];
+static LLVM_THREAD_LOCAL unsigned threadOrderDepth = 0;
struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
struct ThreadDiagnostic {
@@ -1026,14 +1038,17 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
ParallelDiagnosticHandlerImpl(MLIRContext *ctx) : context(ctx) {
handlerID = ctx->getDiagEngine().registerHandler([this](Diagnostic &diag) {
- // Fast path: this thread's order id is in the slot, so looking it up
- // needs no lock. Appending below still locks.
- if (threadHandlerUniqueID == handlerUniqueID) {
- llvm::sys::SmartScopedLock<true> lock(mutex);
- diagnostics.emplace_back(threadOrderID, std::move(diag));
- return success();
+ // Fast path: scan the thread-local stack for this handler, no lock.
+ // Appending below still locks.
+ for (unsigned i = threadOrderDepth; i > 0; --i) {
+ if (threadOrderStack[i - 1].handlerID == handlerUniqueID) {
+ llvm::sys::SmartScopedLock<true> lock(mutex);
+ diagnostics.emplace_back(threadOrderStack[i - 1].orderID,
+ std::move(diag));
+ return success();
+ }
}
- // Slow path: a nested loop on this thread, consult the map. If this
+ // Slow path: nested past the stack, consult the map. If this
// thread is not tracked, return failure to let another handler
// process this diagnostic.
uint64_t tid = llvm::get_threadid();
@@ -1074,9 +1089,13 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
/// Set the order id for the current thread.
void setOrderIDForThread(size_t orderID) {
- if (!threadHandlerUniqueID || threadHandlerUniqueID == handlerUniqueID) {
- threadHandlerUniqueID = handlerUniqueID;
- threadOrderID = orderID;
+ if (threadOrderDepth > 0 &&
+ threadOrderStack[threadOrderDepth - 1].handlerID == handlerUniqueID) {
+ threadOrderStack[threadOrderDepth - 1].orderID = orderID;
+ return;
+ }
+ if (threadOrderDepth < kMaxThreadOrderEntries) {
+ threadOrderStack[threadOrderDepth++] = {handlerUniqueID, orderID};
return;
}
uint64_t tid = llvm::get_threadid();
@@ -1086,8 +1105,9 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
/// Remove the order id for the current thread.
void eraseOrderIDForThread() {
- if (threadHandlerUniqueID == handlerUniqueID) {
- threadHandlerUniqueID = 0;
+ if (threadOrderDepth > 0 &&
+ threadOrderStack[threadOrderDepth - 1].handlerID == handlerUniqueID) {
+ --threadOrderDepth;
return;
}
uint64_t tid = llvm::get_threadid();
@@ -1130,12 +1150,11 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
/// A smart mutex to lock access to the internal state.
llvm::sys::SmartMutex<true> mutex;
- /// Order ids, keyed by thread id, for threads whose thread-local slot is
- /// owned by another handler: a loop nested inside another parallel loop on
- /// the same thread. Idle workers running nested elements use the slot, so
- /// this map only holds threads that are also inside an enclosing loop.
- /// Guarded by `mutex`. A thread in neither the slot nor this map is not
- /// tracked, and its diagnostics pass on to the next handler.
+ /// Order ids, keyed by thread id, for threads nested past what the
+ /// thread-local stack holds: more enclosing parallel loops on one thread
+ /// than entries. Guarded by `mutex`. A thread in neither the stack nor
+ /// this map is not tracked, and its diagnostics pass on to the next
+ /// handler.
DenseMap<uint64_t, size_t> threadToOrderID;
/// An unordered list of diagnostics that were emitted.
@@ -1144,7 +1163,7 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
/// The unique id for the parallel handler.
DiagnosticEngine::HandlerID handlerID = 0;
- /// Unique id naming this handler in the thread-local order-id slot.
+ /// Unique id naming this handler in the thread-local order-id stack.
const uint64_t handlerUniqueID =
nextHandlerUniqueID.fetch_add(1, std::memory_order_relaxed);
diff --git a/mlir/unittests/IR/Diagnostic.cpp b/mlir/unittests/IR/Diagnostic.cpp
index 370f56fe15b9b3..5df046dd9c3589 100644
--- a/mlir/unittests/IR/Diagnostic.cpp
+++ b/mlir/unittests/IR/Diagnostic.cpp
@@ -108,7 +108,7 @@ TEST(SourceMgrDiagnosticVerifierHandler, ErasesOwnRegistration) {
EXPECT_EQ(seen, 1u);
}
-TEST(ParallelDiagnosticHandler, NestedFallbackPreservesOrder) {
+TEST(ParallelDiagnosticHandler, NestedStackPreservesOrder) {
MLIRContext ctx;
std::vector<std::string> seen;
ScopedDiagnosticHandler capture(&ctx, [&](Diagnostic &diag) {
@@ -124,7 +124,7 @@ TEST(ParallelDiagnosticHandler, NestedFallbackPreservesOrder) {
outer.setOrderIDForThread(1);
{
ParallelDiagnosticHandler inner(&ctx);
- // The slot is taken by outer, so these go through the map fallback.
+ // Outer sits at the bottom of the stack, so inner pushes above it.
// Inner ids are deliberately disjoint from outer's 0/1.
inner.setOrderIDForThread(12);
emitError(UnknownLoc::get(&ctx), "inner-12");
@@ -132,7 +132,7 @@ TEST(ParallelDiagnosticHandler, NestedFallbackPreservesOrder) {
emitError(UnknownLoc::get(&ctx), "inner-11");
inner.eraseOrderIDForThread();
} // inner emits sorted: inner-11, inner-12.
- // The slot still names outer after the nested use: the inner flush was
+ // Inner popped, so outer is on top again: the inner flush was
// attributed to outer's element, and this is held with id 1, so the
// id-0 diagnostic below sorts first.
emitError(UnknownLoc::get(&ctx), "outer-A");
@@ -147,6 +147,74 @@ TEST(ParallelDiagnosticHandler, NestedFallbackPreservesOrder) {
EXPECT_NE(seen[3].find("outer-A"), std::string::npos);
}
+TEST(ParallelDiagnosticHandler, DeepNestingOverflowPreservesOrder) {
+ MLIRContext ctx;
+ std::vector<std::string> seen;
+ ScopedDiagnosticHandler capture(&ctx, [&](Diagnostic &diag) {
+ std::string msg;
+ llvm::raw_string_ostream os(msg);
+ diag.print(os);
+ seen.push_back(msg);
+ return success();
+ });
+
+ // Five nested handlers: the thread-local stack holds four, so the
+ // innermost falls back to the map.
+ constexpr unsigned kDepth = 5;
+ std::vector<std::unique_ptr<ParallelDiagnosticHandler>> handlers;
+ for (unsigned i = 0; i < kDepth; ++i) {
+ handlers.push_back(std::make_unique<ParallelDiagnosticHandler>(&ctx));
+ handlers.back()->setOrderIDForThread(0);
+ }
+ handlers.back()->setOrderIDForThread(2);
+ emitError(UnknownLoc::get(&ctx), "x-2");
+ handlers.back()->setOrderIDForThread(1);
+ emitError(UnknownLoc::get(&ctx), "x-1");
+ for (auto it = handlers.rbegin(); it != handlers.rend(); ++it)
+ (*it)->eraseOrderIDForThread();
+ // Destroy innermost first: the handler impl is a pretty-stack-trace entry,
+ // which requires LIFO destruction.
+ while (!handlers.empty())
+ handlers.pop_back();
+
+ ASSERT_EQ(seen.size(), 2u);
+ EXPECT_NE(seen[0].find("x-1"), std::string::npos);
+ EXPECT_NE(seen[1].find("x-2"), std::string::npos);
+}
+
+TEST(ParallelDiagnosticHandler, InterleavedNestingFindsDeeperEntry) {
+ MLIRContext ctx;
+ std::vector<std::string> seen;
+ ScopedDiagnosticHandler capture(&ctx, [&](Diagnostic &diag) {
+ std::string msg;
+ llvm::raw_string_ostream os(msg);
+ diag.print(os);
+ seen.push_back(msg);
+ return success();
+ });
+
+ {
+ ParallelDiagnosticHandler a(&ctx); // Registered first, fires last.
+ ParallelDiagnosticHandler b(&ctx); // Registered later, fires first.
+ // b's loop is outer on this thread, a's nested inside it: the stack
+ // holds b below a, so b's callback must scan past the top to find
+ // itself.
+ b.setOrderIDForThread(0);
+ a.setOrderIDForThread(7);
+ emitError(UnknownLoc::get(&ctx), "d1");
+ a.eraseOrderIDForThread();
+ b.setOrderIDForThread(1);
+ emitError(UnknownLoc::get(&ctx), "d2");
+ b.eraseOrderIDForThread();
+ } // b flushes both, sorted.
+
+ // b holds both, sorted: d1, d2. If b's lookup only saw the top of the
+ // stack, d1 would land in a and flush after d2.
+ ASSERT_EQ(seen.size(), 2u);
+ EXPECT_NE(seen[0].find("d1"), std::string::npos);
+ EXPECT_NE(seen[1].find("d2"), std::string::npos);
+}
+
TEST(ParallelDiagnosticHandler, MultithreadedOrderingHonored) {
MLIRContext ctx;
// Element i emits "diag-i"; whatever the thread interleaving, the flush
More information about the Mlir-commits
mailing list