[Mlir-commits] [mlir] b93b033 - [mlir] Reduce contention in parallel pass execution and timing (#232948)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Oct 21 08:54:33 PDT 2026
Author: Mehdi Amini
Date: 2026-10-21T15:41:30Z
New Revision: b93b033ab8a5fd92f09c4e6f793e16608ca25e6c
URL: https://github.com/llvm/llvm-project/commit/b93b033ab8a5fd92f09c4e6f793e16608ca25e6c
DIFF: https://github.com/llvm/llvm-project/commit/b93b033ab8a5fd92f09c4e6f793e16608ca25e6c.diff
LOG: [mlir] Reduce contention in parallel pass execution and timing (#232948)
Replace per-operation executor scans with stable per-task worker
indices. Allow default timing callbacks to run concurrently using
thread-local stacks and synchronized parent-scope publication. Existing
instrumentations and custom timing managers retain serialized callbacks
unless they opt in.
Benchmark: 200,000 trivial functions with nested CSE, Release with
assertions disabled. Medians of seven alternating samples against main,
which includes #232520. With timing enabled:
| Metric | Before | After |
|---|---:|---:|
| Pipeline wall | 1.0945 s | 0.1699 s |
| mlir-opt wall | 1.7564 s | 0.8329 s |
| mlir-opt system CPU | 15.4775 s | 2.0195 s |
This is a 6.44x pipeline speedup, 2.11x end-to-end speedup.
Assisted-by: Codex
Added:
Modified:
mlir/docs/PassManagement.md
mlir/include/mlir/IR/Threading.h
mlir/include/mlir/Pass/PassInstrumentation.h
mlir/include/mlir/Support/Timing.h
mlir/lib/Pass/Pass.cpp
mlir/lib/Pass/PassTiming.cpp
mlir/unittests/Pass/PassManagerTest.cpp
Removed:
################################################################################
diff --git a/mlir/docs/PassManagement.md b/mlir/docs/PassManagement.md
index 8d20b496cd3a61..70f055f789d039 100644
--- a/mlir/docs/PassManagement.md
+++ b/mlir/docs/PassManagement.md
@@ -1063,9 +1063,17 @@ PassInstrumentation instances may be registered directly with a
[PassManager](#pass-manager) instance via the `addInstrumentation` method.
Instrumentations added to the PassManager are run in a stack like fashion, i.e.
the last instrumentation to execute a `runBefore*` hook will be the first to
-execute the respective `runAfter*` hook. The hooks of a `PassInstrumentation`
-class are guaranteed to be executed in a thread-safe fashion, so additional
-synchronization is not necessary. Below in an example instrumentation that
+execute the respective `runAfter*` hook. Corresponding `runBefore*` and
+`runAfter*` callbacks are guaranteed to execute on the same thread. The hooks of
+a `PassInstrumentation` class are serialized by default, so additional
+synchronization is not necessary.
+An instrumentation may override `isThreadSafe()` to return true if its callbacks
+support concurrent invocation on
diff erent threads. The instrumentor allows
+concurrent dispatch only when all registered instrumentations opt in; otherwise
+it preserves serialization of the entire callback sequence. Instrumentations
+must be registered before running the pass manager.
+
+Below is an example instrumentation that
counts the number of times the `DominanceInfo` analysis is computed:
```c++
@@ -1110,6 +1118,11 @@ passes are taking the most time to execute, as well as how much of an effect a
pass has on the total execution time of the pipeline. Users can enable this
instrumentation directly on the PassManager via `enableTiming`. This
instrumentation is also made available in mlir-opt via the `-mlir-timing` flag.
+With the default timing manager, timing callbacks may run concurrently on
+
diff erent threads. A custom `TimingManager` retains serialized callbacks unless
+it overrides `isThreadSafe()` to return true. Registering an instrumentation
+that requires serialization also serializes timing callbacks.
+
The PassTiming instrumentation provides several
diff erent display modes for the
timing results, each of which is described below:
diff --git a/mlir/include/mlir/IR/Threading.h b/mlir/include/mlir/IR/Threading.h
index 3ceab6b3e883a5..ced44984805384 100644
--- a/mlir/include/mlir/IR/Threading.h
+++ b/mlir/include/mlir/IR/Threading.h
@@ -23,18 +23,25 @@
#include <atomic>
namespace mlir {
-
/// Invoke the given function on the elements between [begin, end)
-/// asynchronously. If the given function returns a failure when processing any
-/// of the elements, execution is stopped and a failure is returned from this
-/// function. This means that in the case of failure, not all elements of the
-/// range will be processed. Diagnostics emitted during processing are ordered
-/// relative to the element's position within [begin, end). If the provided
-/// context does not have multi-threading enabled, this function always
-/// processes elements sequentially.
+/// asynchronously, passing each element and its worker index. If the given
+/// function returns a failure when processing any of the elements, execution is
+/// stopped and a failure is returned from this function. This means that in the
+/// case of failure, not all elements of the range will be processed.
+/// Diagnostics emitted during processing are ordered relative to the element's
+/// position within [begin, end). If the provided context does not have
+/// multi-threading enabled, this function always processes elements
+/// sequentially.
+///
+/// Worker indices are in [0, min(range size, maximum concurrency)), or zero
+/// when processing sequentially. Each index is used by only one task for the
+/// entire call, including while that task waits for nested parallel work. This
+/// allows callers to keep per-worker state without synchronization or
+/// thread-local storage. Indices are local to this invocation, not to a thread.
template <typename IteratorT, typename FuncT>
-LogicalResult failableParallelForEach(MLIRContext *context, IteratorT begin,
- IteratorT end, FuncT &&func) {
+LogicalResult failableParallelForEachWithWorkerId(MLIRContext *context,
+ IteratorT begin,
+ IteratorT end, FuncT &&func) {
unsigned numElements = static_cast<unsigned>(std::distance(begin, end));
if (numElements == 0)
return success();
@@ -43,7 +50,7 @@ LogicalResult failableParallelForEach(MLIRContext *context, IteratorT begin,
// process the elements directly on this thread.
if (!context->isMultithreadingEnabled() || numElements <= 1) {
for (; begin != end; ++begin)
- if (failed(func(*begin)))
+ if (failed(func(*begin, 0)))
return failure();
return success();
}
@@ -53,13 +60,13 @@ LogicalResult failableParallelForEach(MLIRContext *context, IteratorT begin,
ParallelDiagnosticHandler handler(context);
std::atomic<unsigned> curIndex(0);
std::atomic<bool> processingFailed(false);
- auto processFn = [&] {
+ auto processFn = [&](unsigned workerId) {
while (!processingFailed) {
unsigned index = curIndex++;
if (index >= numElements)
break;
handler.setOrderIDForThread(index);
- if (failed(func(*std::next(begin, index))))
+ if (failed(func(*std::next(begin, index), workerId)))
processingFailed = true;
handler.eraseOrderIDForThread();
}
@@ -70,7 +77,7 @@ LogicalResult failableParallelForEach(MLIRContext *context, IteratorT begin,
llvm::ThreadPoolTaskGroup tasksGroup(threadPool);
size_t numActions = std::min(numElements, threadPool.getMaxConcurrency());
for (unsigned i = 0; i < numActions; ++i)
- tasksGroup.async(processFn);
+ tasksGroup.async([&, i] { processFn(i); });
// If the current thread is a worker thread from the pool, then waiting for
// the task group allows the current thread to also participate in processing
// tasks from the group, which avoid any deadlock/starvation.
@@ -78,6 +85,17 @@ LogicalResult failableParallelForEach(MLIRContext *context, IteratorT begin,
return failure(processingFailed);
}
+/// Invoke the given function on each element, with ordered diagnostics and
+/// early termination on failure. See failableParallelForEachWithWorkerId.
+template <typename IteratorT, typename FuncT>
+LogicalResult failableParallelForEach(MLIRContext *context, IteratorT begin,
+ IteratorT end, FuncT &&func) {
+ return failableParallelForEachWithWorkerId(
+ context, begin, end, [&](auto &&value, unsigned) {
+ return func(std::forward<decltype(value)>(value));
+ });
+}
+
/// Invoke the given function on the elements in the provided range
/// asynchronously. If the given function returns a failure when processing any
/// of the elements, execution is stopped and a failure is returned from this
@@ -141,7 +159,6 @@ template <typename FuncT>
void parallelFor(MLIRContext *context, size_t begin, size_t end, FuncT &&func) {
parallelForEach(context, llvm::seq(begin, end), std::forward<FuncT>(func));
}
-
} // namespace mlir
#endif // MLIR_IR_THREADING_H
diff --git a/mlir/include/mlir/Pass/PassInstrumentation.h b/mlir/include/mlir/Pass/PassInstrumentation.h
index 0247db6060e2e8..f74dc732278b44 100644
--- a/mlir/include/mlir/Pass/PassInstrumentation.h
+++ b/mlir/include/mlir/Pass/PassInstrumentation.h
@@ -41,6 +41,13 @@ class PassInstrumentation {
virtual ~PassInstrumentation() = 0;
+ /// Whether callbacks may run concurrently on
diff erent threads. By default,
+ /// the pass instrumentor serializes callbacks. Opting in requires all shared
+ /// state to support concurrent use.
+ /// Callbacks on a single thread still run in registration order before a
+ /// pipeline, pass, or analysis, and in reverse order afterwards.
+ virtual bool isThreadSafe() const { return false; }
+
/// A callback to run before a pass pipeline is executed. This function takes
/// the name of the operation type being operated on, or std::nullopt if the
/// pipeline is op-agnostic, and information related to the parent that
diff --git a/mlir/include/mlir/Support/Timing.h b/mlir/include/mlir/Support/Timing.h
index 50ae84735dbf04..1944f5019fa184 100644
--- a/mlir/include/mlir/Support/Timing.h
+++ b/mlir/include/mlir/Support/Timing.h
@@ -115,6 +115,12 @@ class TimingManager {
explicit TimingManager();
virtual ~TimingManager();
+ /// Whether timer operations may run concurrently on
diff erent threads. Each
+ /// timer must still be started and stopped on one thread, and parent timers
+ /// must outlive their children. Configuration and reporting are not covered
+ /// by this guarantee.
+ virtual bool isThreadSafe() const { return false; }
+
/// Get the root timer of this timing manager. The returned timer must be
/// started and stopped manually. Execution time can be measured by nesting
/// timers within this root timer and starting/stopping them as appropriate.
@@ -205,6 +211,9 @@ class Timer {
/// performance low.
explicit operator bool() const { return tm != nullptr; }
+ /// Whether the timing manager supports concurrent timer operations.
+ bool isThreadSafe() const { return !tm || tm->isThreadSafe(); }
+
/// Start the timer. This must be accompanied by a corresponding call to
/// `stop()` at a later point.
void start() {
@@ -292,6 +301,9 @@ class TimingScope {
/// Check if the timing scope actually contains a valid timer.
explicit operator bool() const { return bool(timer); }
+ /// Whether the timing manager supports concurrent timer operations.
+ bool isThreadSafe() const { return timer.isThreadSafe(); }
+
// Disable copying of the `TimingScope`.
TimingScope(const TimingScope &) = delete;
TimingScope &operator=(const TimingScope &) = delete;
@@ -411,6 +423,8 @@ class DefaultTimingManager : public TimingManager {
DefaultTimingManager(DefaultTimingManager &&rhs);
~DefaultTimingManager() override;
+ bool isThreadSafe() const override { return true; }
+
// Disable copying of the `DefaultTimingManager`.
DefaultTimingManager(const DefaultTimingManager &rhs) = delete;
DefaultTimingManager &operator=(const DefaultTimingManager &rhs) = delete;
diff --git a/mlir/lib/Pass/Pass.cpp b/mlir/lib/Pass/Pass.cpp
index e3b39034fcb387..656b8044587737 100644
--- a/mlir/lib/Pass/Pass.cpp
+++ b/mlir/lib/Pass/Pass.cpp
@@ -25,6 +25,7 @@
#include "llvm/Support/Mutex.h"
#include "llvm/Support/Signals.h"
#include "llvm/Support/Threading.h"
+#include <mutex>
#include <optional>
#define DEBUG_TYPE "pass-manager"
@@ -984,29 +985,22 @@ void OpToOpPassAdaptor::runOnOperationAsyncImpl(bool verifyPasses) {
this};
auto *instrumentor = am.getPassInstrumentor();
- // An atomic failure variable for the async executors.
- std::vector<std::atomic<bool>> activePMs(asyncExecutors.size());
- llvm::fill(activePMs, false);
+ // Each parallel task owns an executor for its entire lifetime. A worker
+ // index belongs to the task, rather than the thread: nested waits may run
+ // other tasks on the same thread while this task's executor is still in use.
std::atomic<bool> hasFailure = false;
- parallelForEach(context, opInfos, [&](OpPMInfo &opInfo) {
- // Find an executor for this operation.
- auto it = llvm::find_if(activePMs, [](std::atomic<bool> &isActive) {
- bool expectedInactive = false;
- return isActive.compare_exchange_strong(expectedInactive, true);
- });
- unsigned pmIndex = it - activePMs.begin();
-
- // Get the pass manager for this operation and execute it.
- OpPassManager &pm = asyncExecutors[pmIndex][opInfo.passManagerIdx];
- LogicalResult pipelineResult = runPipeline(
- pm, opInfo.op, opInfo.am, verifyPasses,
- pm.impl->initializationGeneration, instrumentor, &parentInfo);
- if (failed(pipelineResult))
- hasFailure.store(true);
-
- // Reset the active bit for this pass manager.
- activePMs[pmIndex].store(false);
- });
+ (void)failableParallelForEachWithWorkerId(
+ context, opInfos.begin(), opInfos.end(),
+ [&](OpPMInfo &opInfo, unsigned pmIndex) {
+ OpPassManager &pm = asyncExecutors[pmIndex][opInfo.passManagerIdx];
+ LogicalResult pipelineResult = runPipeline(
+ pm, opInfo.op, opInfo.am, verifyPasses,
+ pm.impl->initializationGeneration, instrumentor, &parentInfo);
+ if (failed(pipelineResult))
+ hasFailure.store(true);
+ // Continue processing independent operations after a pipeline failure.
+ return success();
+ });
// Signal a failure if any of the executors failed.
if (hasFailure)
@@ -1223,6 +1217,17 @@ struct PassInstrumentorImpl {
/// Mutex to keep instrumentation access thread-safe.
llvm::sys::SmartMutex<true> mutex;
+ /// Dispatch may run concurrently only if every instrumentation opts in.
+ /// Registration finishes before the pipeline starts running.
+ bool threadSafe = true;
+
+ auto lockInstrumentations() {
+ std::unique_lock<llvm::sys::SmartMutex<true>> lock(mutex, std::defer_lock);
+ if (!threadSafe)
+ lock.lock();
+ return lock;
+ }
+
/// Set of registered instrumentations.
std::vector<std::unique_ptr<PassInstrumentation>> instrumentations;
};
@@ -1236,7 +1241,7 @@ PassInstrumentor::~PassInstrumentor() = default;
void PassInstrumentor::runBeforePipeline(
std::optional<OperationName> name,
const PassInstrumentation::PipelineParentInfo &parentInfo) {
- llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
+ auto instrumentationLock = impl->lockInstrumentations();
for (auto &instr : impl->instrumentations)
instr->runBeforePipeline(name, parentInfo);
}
@@ -1245,28 +1250,28 @@ void PassInstrumentor::runBeforePipeline(
void PassInstrumentor::runAfterPipeline(
std::optional<OperationName> name,
const PassInstrumentation::PipelineParentInfo &parentInfo) {
- llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
+ auto instrumentationLock = impl->lockInstrumentations();
for (auto &instr : llvm::reverse(impl->instrumentations))
instr->runAfterPipeline(name, parentInfo);
}
/// See PassInstrumentation::runBeforePass for details.
void PassInstrumentor::runBeforePass(Pass *pass, Operation *op) {
- llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
+ auto instrumentationLock = impl->lockInstrumentations();
for (auto &instr : impl->instrumentations)
instr->runBeforePass(pass, op);
}
/// See PassInstrumentation::runAfterPass for details.
void PassInstrumentor::runAfterPass(Pass *pass, Operation *op) {
- llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
+ auto instrumentationLock = impl->lockInstrumentations();
for (auto &instr : llvm::reverse(impl->instrumentations))
instr->runAfterPass(pass, op);
}
/// See PassInstrumentation::runAfterPassFailed for details.
void PassInstrumentor::runAfterPassFailed(Pass *pass, Operation *op) {
- llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
+ auto instrumentationLock = impl->lockInstrumentations();
for (auto &instr : llvm::reverse(impl->instrumentations))
instr->runAfterPassFailed(pass, op);
}
@@ -1274,7 +1279,7 @@ void PassInstrumentor::runAfterPassFailed(Pass *pass, Operation *op) {
/// See PassInstrumentation::runBeforeAnalysis for details.
void PassInstrumentor::runBeforeAnalysis(StringRef name, TypeID id,
Operation *op) {
- llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
+ auto instrumentationLock = impl->lockInstrumentations();
for (auto &instr : impl->instrumentations)
instr->runBeforeAnalysis(name, id, op);
}
@@ -1282,7 +1287,7 @@ void PassInstrumentor::runBeforeAnalysis(StringRef name, TypeID id,
/// See PassInstrumentation::runAfterAnalysis for details.
void PassInstrumentor::runAfterAnalysis(StringRef name, TypeID id,
Operation *op) {
- llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
+ auto instrumentationLock = impl->lockInstrumentations();
for (auto &instr : llvm::reverse(impl->instrumentations))
instr->runAfterAnalysis(name, id, op);
}
@@ -1291,5 +1296,6 @@ void PassInstrumentor::runAfterAnalysis(StringRef name, TypeID id,
void PassInstrumentor::addInstrumentation(
std::unique_ptr<PassInstrumentation> pi) {
llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
+ impl->threadSafe &= pi->isThreadSafe();
impl->instrumentations.emplace_back(std::move(pi));
}
diff --git a/mlir/lib/Pass/PassTiming.cpp b/mlir/lib/Pass/PassTiming.cpp
index dab56f09a72ebe..27d847c2951376 100644
--- a/mlir/lib/Pass/PassTiming.cpp
+++ b/mlir/lib/Pass/PassTiming.cpp
@@ -8,9 +8,11 @@
#include "PassDetail.h"
#include "mlir/Pass/PassManager.h"
-#include "llvm/ADT/SmallVector.h"
+#include "mlir/Support/ThreadLocalCache.h"
+#include "llvm/Support/RWMutex.h"
#include "llvm/Support/Threading.h"
+#include <deque>
#include <optional>
using namespace mlir;
@@ -28,11 +30,13 @@ struct PassTiming : public PassInstrumentation {
ownedTimingScope(ownedTimingManager->getRootScope()),
rootScope(ownedTimingScope) {}
- /// If a pass can spawn additional work on other threads, it records the
- /// index to its currently active timer here. Passes that run on a
- /// newly-forked thread will check this list to find the active timer of the
- /// parent thread into which the new thread should be nested.
- DenseMap<PipelineParentInfo, unsigned> parentTimerIndices;
+ bool isThreadSafe() const override { return rootScope.isThreadSafe(); }
+
+ /// Published adaptor scopes remain alive until their nested pipelines have
+ /// joined. The deque below keeps their addresses stable when a worker helps
+ /// execute another task while waiting for nested parallel work.
+ DenseMap<PipelineParentInfo, TimingScope *> parentScopes;
+ llvm::sys::SmartRWMutex<true> parentScopesMutex;
/// The timing manager owned by this instrumentation (in case timing was
/// enabled by the user on the pass manager without providing an external
@@ -43,7 +47,7 @@ struct PassTiming : public PassInstrumentation {
TimingScope ownedTimingScope;
/// A stack of the currently active timing scopes per thread.
- DenseMap<uint64_t, SmallVector<TimingScope, 4>> activeThreadTimers;
+ ThreadLocalCache<std::deque<TimingScope>> activeThreadTimers;
/// The root timing scope into which timing is reported.
TimingScope &rootScope;
@@ -54,17 +58,17 @@ struct PassTiming : public PassInstrumentation {
void runBeforePipeline(std::optional<OperationName> name,
const PipelineParentInfo &parentInfo) override {
- auto tid = llvm::get_threadid();
- auto &activeTimers = activeThreadTimers[tid];
+ auto &activeTimers = *activeThreadTimers;
// Find the parent scope, either using the parent info or the root scope
// (e.g. in the case of the top-level pipeline).
- TimingScope *parentScope;
- auto it = parentTimerIndices.find(parentInfo);
- if (it != parentTimerIndices.end())
- parentScope = &activeThreadTimers[parentInfo.parentThreadID][it->second];
- else
- parentScope = &rootScope;
+ TimingScope *parentScope = &rootScope;
+ {
+ llvm::sys::SmartScopedReader<true> lock(parentScopesMutex);
+ auto it = parentScopes.find(parentInfo);
+ if (it != parentScopes.end())
+ parentScope = it->second;
+ }
// Use nullptr to anchor op-agnostic pipelines, otherwise use the name of
// the operation.
@@ -76,7 +80,7 @@ struct PassTiming : public PassInstrumentation {
void runAfterPipeline(std::optional<OperationName>,
const PipelineParentInfo &) override {
- auto &activeTimers = activeThreadTimers[llvm::get_threadid()];
+ auto &activeTimers = *activeThreadTimers;
assert(!activeTimers.empty() && "expected active timer");
activeTimers.pop_back();
}
@@ -86,18 +90,18 @@ struct PassTiming : public PassInstrumentation {
//===--------------------------------------------------------------------===//
void runBeforePass(Pass *pass, Operation *) override {
- auto tid = llvm::get_threadid();
- auto &activeTimers = activeThreadTimers[tid];
+ auto &activeTimers = *activeThreadTimers;
auto &parentScope = activeTimers.empty() ? rootScope : activeTimers.back();
if (auto *adaptor = dyn_cast<OpToOpPassAdaptor>(pass)) {
- parentTimerIndices[{tid, pass}] = activeTimers.size();
auto scope =
parentScope.nest(pass->getThreadingSiblingOrThis(),
[adaptor]() { return adaptor->getAdaptorName(); });
if (adaptor->getPassManagers().size() <= 1)
scope.hide();
activeTimers.push_back(std::move(scope));
+ llvm::sys::SmartScopedWriter<true> lock(parentScopesMutex);
+ parentScopes[{llvm::get_threadid(), pass}] = &activeTimers.back();
} else {
activeTimers.push_back(
parentScope.nest(pass->getThreadingSiblingOrThis(),
@@ -106,10 +110,11 @@ struct PassTiming : public PassInstrumentation {
}
void runAfterPass(Pass *pass, Operation *) override {
- auto tid = llvm::get_threadid();
- if (isa<OpToOpPassAdaptor>(pass))
- parentTimerIndices.erase({tid, pass});
- auto &activeTimers = activeThreadTimers[tid];
+ if (isa<OpToOpPassAdaptor>(pass)) {
+ llvm::sys::SmartScopedWriter<true> lock(parentScopesMutex);
+ parentScopes.erase({llvm::get_threadid(), pass});
+ }
+ auto &activeTimers = *activeThreadTimers;
assert(!activeTimers.empty() && "expected active timer");
activeTimers.pop_back();
}
@@ -123,15 +128,14 @@ struct PassTiming : public PassInstrumentation {
//===--------------------------------------------------------------------===//
void runBeforeAnalysis(StringRef name, TypeID id, Operation *) override {
- auto tid = llvm::get_threadid();
- auto &activeTimers = activeThreadTimers[tid];
+ auto &activeTimers = *activeThreadTimers;
auto &parentScope = activeTimers.empty() ? rootScope : activeTimers.back();
activeTimers.push_back(parentScope.nest(
id.getAsOpaquePointer(), [name] { return "(A) " + name.str(); }));
}
void runAfterAnalysis(StringRef, TypeID, Operation *) override {
- auto &activeTimers = activeThreadTimers[llvm::get_threadid()];
+ auto &activeTimers = *activeThreadTimers;
assert(!activeTimers.empty() && "expected active timer");
activeTimers.pop_back();
}
diff --git a/mlir/unittests/Pass/PassManagerTest.cpp b/mlir/unittests/Pass/PassManagerTest.cpp
index c0b7dcbb7eaed7..e873a68d0b3ecd 100644
--- a/mlir/unittests/Pass/PassManagerTest.cpp
+++ b/mlir/unittests/Pass/PassManagerTest.cpp
@@ -13,11 +13,18 @@
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Threading.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassInstrumentation.h"
#include "gtest/gtest.h"
+#include <array>
+#include <atomic>
+#include <chrono>
+#include <condition_variable>
#include <memory>
+#include <mutex>
+#include <thread>
using namespace mlir;
using namespace mlir::detail;
@@ -218,6 +225,201 @@ TEST(PassManagerTest, PassInstrumentation) {
}
}
+// A worker keeps exclusive ownership of its state even while nested waits let
+// the same thread execute other tasks from the outer loop.
+TEST(PassManagerTest, ParallelWorkerOwnership) {
+ llvm::StdThreadPool pool(llvm::hardware_concurrency(4));
+ MLIRContext context(MLIRContext::Threading::DISABLED);
+ context.setThreadPool(pool);
+ std::vector<std::atomic<bool>> active(pool.getMaxConcurrency());
+ for (auto &flag : active)
+ flag = false;
+ std::atomic<unsigned> visits = 0;
+ auto indices = llvm::seq<unsigned>(0, 64);
+ EXPECT_TRUE(succeeded(failableParallelForEachWithWorkerId(
+ &context, indices.begin(), indices.end(), [&](unsigned, unsigned worker) {
+ EXPECT_LT(worker, active.size());
+ EXPECT_FALSE(active[worker].exchange(true));
+ auto inner = llvm::seq<unsigned>(0, 8);
+ EXPECT_TRUE(
+ succeeded(failableParallelForEach(&context, inner, [&](unsigned) {
+ ++visits;
+ return success();
+ })));
+ EXPECT_TRUE(active[worker].exchange(false));
+ return success();
+ })));
+ EXPECT_EQ(visits.load(), 64u * 8);
+ EXPECT_TRUE(failed(failableParallelForEachWithWorkerId(
+ &context, indices.begin(), indices.end(),
+ [&](unsigned, unsigned) { return failure(); })));
+ EXPECT_TRUE(succeeded(failableParallelForEachWithWorkerId(
+ &context, indices.begin(), indices.begin(), [&](unsigned, unsigned) {
+ ADD_FAILURE() << "empty range must not invoke the callback";
+ return failure();
+ })));
+
+ context.disableMultithreading();
+ unsigned serialVisits = 0;
+ EXPECT_TRUE(succeeded(failableParallelForEachWithWorkerId(
+ &context, indices.begin(), indices.end(), [&](unsigned, unsigned worker) {
+ EXPECT_EQ(worker, 0u);
+ ++serialVisits;
+ return success();
+ })));
+ EXPECT_EQ(serialVisits, 64u);
+ EXPECT_TRUE(failed(failableParallelForEachWithWorkerId(
+ &context, indices.begin(), indices.end(),
+ [&](unsigned, unsigned) { return failure(); })));
+}
+
+struct ConcurrentInstrumentation : PassInstrumentation {
+ bool isThreadSafe() const override { return true; }
+ std::mutex mutex;
+ std::condition_variable cv;
+ unsigned entered = 0;
+ unsigned overlapped = 0;
+
+ void runBeforePass(Pass *, Operation *) override {
+ std::unique_lock<std::mutex> lock(mutex);
+ ++entered;
+ cv.notify_all();
+ overlapped += cv.wait_for(lock, std::chrono::seconds(10),
+ [&] { return entered == 2; });
+ }
+};
+
+TEST(PassManagerTest, ConcurrentInstrumentation) {
+ if (!llvm::llvm_is_multithreaded())
+ GTEST_SKIP();
+ PassInstrumentor instrumentor;
+ auto instrumentation = std::make_unique<ConcurrentInstrumentation>();
+ auto *ptr = instrumentation.get();
+ instrumentor.addInstrumentation(std::move(instrumentation));
+ std::thread first([&] { instrumentor.runBeforePass(nullptr, nullptr); });
+ std::thread second([&] { instrumentor.runBeforePass(nullptr, nullptr); });
+ first.join();
+ second.join();
+ EXPECT_EQ(ptr->entered, 2u);
+ EXPECT_EQ(ptr->overlapped, 2u);
+}
+
+// Non-atomic callback state remains protected, including when another
+// instrumentation opts into concurrent dispatch.
+TEST(PassManagerTest, SerializedInstrumentation) {
+ if (!llvm::llvm_is_multithreaded())
+ GTEST_SKIP();
+ struct Counter : PassInstrumentation {
+ Counter(bool threadSafe) : threadSafe(threadSafe) {}
+ bool threadSafe;
+ unsigned calls = 0;
+ bool isThreadSafe() const override { return threadSafe; }
+ void runBeforePass(Pass *, Operation *) override {
+ // Without serialization, this non-atomic read-modify-write would race.
+ // Yield between the read and write to increase the chance of a lost
+ // update, making the final count too small. ThreadSanitizer can also
+ // detect the race even if the final count happens to be correct.
+ unsigned previous = calls;
+ std::this_thread::yield();
+ calls = previous + 1;
+ }
+ };
+ PassInstrumentor instrumentor;
+ auto counter = std::make_unique<Counter>(false);
+ auto concurrent = std::make_unique<Counter>(true);
+ auto *counterPtr = counter.get();
+ auto *concurrentPtr = concurrent.get();
+ instrumentor.addInstrumentation(std::move(counter));
+ instrumentor.addInstrumentation(std::move(concurrent));
+ std::array<std::thread, 4> threads;
+ for (auto &thread : threads)
+ thread = std::thread([&] {
+ for (unsigned i = 0; i < 128; ++i)
+ instrumentor.runBeforePass(nullptr, nullptr);
+ });
+ for (auto &thread : threads)
+ thread.join();
+ EXPECT_EQ(counterPtr->calls, 512u);
+ EXPECT_EQ(concurrentPtr->calls, 512u);
+}
+
+TEST(PassManagerTest, CustomTimingManagerSerialization) {
+ struct CustomTimingManager : TimingManager {
+ std::atomic<bool> inCallback = false;
+ std::atomic<bool> overlap = false;
+
+ void observeCallback() {
+ if (inCallback.exchange(true))
+ overlap = true;
+ std::this_thread::yield();
+ inCallback = false;
+ }
+ std::optional<void *> rootTimer() override { return this; }
+ void startTimer(void *) override { observeCallback(); }
+ void stopTimer(void *) override { observeCallback(); }
+ void *nestTimer(void *, const void *,
+ function_ref<std::string()>) override {
+ observeCallback();
+ return this;
+ }
+ };
+
+ llvm::StdThreadPool pool(llvm::hardware_concurrency(4));
+ MLIRContext context(MLIRContext::Threading::DISABLED);
+ context.setThreadPool(pool);
+ context.loadDialect<func::FuncDialect>();
+ Builder builder(&context);
+ OwningOpRef<ModuleOp> module(ModuleOp::create(builder.getUnknownLoc()));
+ for (unsigned i = 0; i < 64; ++i) {
+ auto func =
+ func::FuncOp::create(builder.getUnknownLoc(), "f" + std::to_string(i),
+ builder.getFunctionType({}, {}));
+ func.setPrivate();
+ module->push_back(func);
+ }
+ auto timing = std::make_unique<CustomTimingManager>();
+ auto *timingPtr = timing.get();
+ auto pm = PassManager::on<ModuleOp>(&context);
+ pm.addNestedPass<func::FuncOp>(std::make_unique<AnnotateFunctionPass>());
+ pm.enableTiming(std::move(timing));
+ EXPECT_TRUE(succeeded(pm.run(module.get())));
+ EXPECT_FALSE(timingPtr->overlap.load());
+}
+
+TEST(PassManagerTest, NestedParallelTiming) {
+ llvm::StdThreadPool pool(llvm::hardware_concurrency(4));
+ MLIRContext context(MLIRContext::Threading::DISABLED);
+ context.setThreadPool(pool);
+ context.loadDialect<func::FuncDialect>();
+ Builder builder(&context);
+ auto module = ModuleOp::create(builder.getUnknownLoc());
+ OwningOpRef<ModuleOp> owner(module);
+ for (unsigned i = 0; i < 8; ++i) {
+ auto nested = ModuleOp::create(builder.getUnknownLoc());
+ module.push_back(nested);
+ for (unsigned j = 0; j < 8; ++j) {
+ auto func =
+ func::FuncOp::create(builder.getUnknownLoc(), "f" + std::to_string(j),
+ builder.getFunctionType({}, {}));
+ func.setPrivate();
+ nested.push_back(func);
+ }
+ }
+ auto pm = PassManager::on<ModuleOp>(&context);
+ pm.nest<ModuleOp>().addNestedPass<func::FuncOp>(
+ std::make_unique<AnnotateFunctionPass>());
+ pm.enableTiming();
+ for (unsigned run = 0; run < 2; ++run) {
+ ASSERT_TRUE(succeeded(pm.run(module)));
+ unsigned processed = 0;
+ module.walk([&](func::FuncOp func) {
+ EXPECT_TRUE(func->hasDiscardableAttr("isFunc"));
+ ++processed;
+ });
+ EXPECT_EQ(processed, 64u);
+ }
+}
+
TEST(PassManagerTest, ExecutionAction) {
MLIRContext context;
context.loadDialect<func::FuncDialect>();
More information about the Mlir-commits
mailing list