[Mlir-commits] [mlir] 2c0d9a6 - [mlir][Diagnostics] Set parallel diagnostic order ids without a lock (#232520)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 15 20:04:15 PDT 2026
Author: Bruno Cardoso Lopes
Date: 2026-10-15T20:04:07-07:00
New Revision: 2c0d9a6a800787f04e0bac5049f44c6abaed12e3
URL: https://github.com/llvm/llvm-project/commit/2c0d9a6a800787f04e0bac5049f44c6abaed12e3
DIFF: https://github.com/llvm/llvm-project/commit/2c0d9a6a800787f04e0bac5049f44c6abaed12e3.diff
LOG: [mlir][Diagnostics] Set parallel diagnostic order ids without a lock (#232520)
failableParallelForEach tells its ParallelDiagnosticHandler which
element each thread is processing before every element and clears it
after, so diagnostics come out in element order.
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.
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.
Added:
Modified:
mlir/include/mlir/IR/Diagnostics.h
mlir/lib/IR/Diagnostics.cpp
mlir/unittests/IR/Diagnostic.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/IR/Diagnostics.h b/mlir/include/mlir/IR/Diagnostics.h
index 43a80e0ee1c91..62b15ec8a3179 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 a3ddb76a66313..8d1cdb4df12be 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,35 @@ 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> nextHandlerUniqueID{1};
+
+/// 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 {
ThreadDiagnostic(size_t id, Diagnostic diag)
@@ -1008,16 +1038,25 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
ParallelDiagnosticHandlerImpl(MLIRContext *ctx) : context(ctx) {
handlerID = ctx->getDiagEngine().registerHandler([this](Diagnostic &diag) {
+ // 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: 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();
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 +1089,15 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
/// Set the order id for the current thread.
void setOrderIDForThread(size_t 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();
llvm::sys::SmartScopedLock<true> lock(mutex);
threadToOrderID[tid] = orderID;
@@ -1057,6 +1105,11 @@ struct ParallelDiagnosticHandlerImpl : public llvm::PrettyStackTraceEntry {
/// Remove the order id for the current thread.
void eraseOrderIDForThread() {
+ if (threadOrderDepth > 0 &&
+ threadOrderStack[threadOrderDepth - 1].handlerID == handlerUniqueID) {
+ --threadOrderDepth;
+ return;
+ }
uint64_t tid = llvm::get_threadid();
llvm::sys::SmartScopedLock<true> lock(mutex);
threadToOrderID.erase(tid);
@@ -1097,7 +1150,11 @@ 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.
+ /// 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.
@@ -1106,6 +1163,10 @@ 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 stack.
+ 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 6390cdfe85acc..5df046dd9c358 100644
--- a/mlir/unittests/IR/Diagnostic.cpp
+++ b/mlir/unittests/IR/Diagnostic.cpp
@@ -7,12 +7,16 @@
//===----------------------------------------------------------------------===//
#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>
using namespace mlir;
using namespace mlir::detail;
@@ -104,6 +108,143 @@ TEST(SourceMgrDiagnosticVerifierHandler, ErasesOwnRegistration) {
EXPECT_EQ(seen, 1u);
}
+TEST(ParallelDiagnosticHandler, NestedStackPreservesOrder) {
+ 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);
+ // 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");
+ inner.setOrderIDForThread(11);
+ emitError(UnknownLoc::get(&ctx), "inner-11");
+ inner.eraseOrderIDForThread();
+ } // inner emits sorted: inner-11, inner-12.
+ // 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");
+ 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-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, 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
+ // 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;
More information about the Mlir-commits
mailing list