[Mlir-commits] [mlir] [mlir] Reduce contention in parallel pass execution and timing (PR #232948)

[Markus Böck]

================
@@ -218,6 +225,197 @@ 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 {
+      unsigned previous = calls;
+      std::this_thread::yield();
+      calls = previous + 1;
----------------
zero9178 wrote:

Can we leave a comment here stating the intention and what the failure mode here would be if instrumentations were not serialized?

My understanding here is that the yield is used to effectively increase the chances that the "UB" from the data race has an observable effect by causing a lost update. If the instrumentations were not serialized this would also be noticed in a TSAN build. Is that correct?

https://github.com/llvm/llvm-project/pull/232948