[Mlir-commits] [mlir] [MLIR][Pass] Full & deterministic diagnostics (PR #110311)

[Billy Zhu]

https://github.com/zyx-billy updated https://github.com/llvm/llvm-project/pull/110311

>From 2ecc322305fa13e00dd960e93261c99123cec0a7 Mon Sep 17 00:00:00 2001
From: Billy Zhu <billyzhu at modular.com>
Date: Fri, 27 Sep 2024 11:29:03 -0700
Subject: [PATCH 1/2] full diagnostics

---
 mlir/lib/Pass/Pass.cpp                          |  9 ++++++---
 .../Pass/crash-recovery-dynamic-failure.mlir    |  2 +-
 mlir/test/Pass/full_diagnostics.mlir            | 17 +++++++++++++++++
 mlir/test/lib/Pass/TestPassManager.cpp          | 12 +++++++++++-
 4 files changed, 35 insertions(+), 5 deletions(-)
 create mode 100644 mlir/test/Pass/full_diagnostics.mlir

diff --git a/mlir/lib/Pass/Pass.cpp b/mlir/lib/Pass/Pass.cpp
index 83b46e8f5680a0..90fd9a713e41d7 100644
--- a/mlir/lib/Pass/Pass.cpp
+++ b/mlir/lib/Pass/Pass.cpp
@@ -732,7 +732,7 @@ void OpToOpPassAdaptor::runOnOperationImpl(bool verifyPasses) {
         unsigned initGeneration = mgr->impl->initializationGeneration;
         if (failed(runPipeline(*mgr, &op, am.nest(&op), verifyPasses,
                                initGeneration, instrumentor, &parentInfo)))
-          return signalPassFailure();
+          signalPassFailure();
       }
     }
   }
@@ -799,6 +799,7 @@ void OpToOpPassAdaptor::runOnOperationAsyncImpl(bool verifyPasses) {
   // An atomic failure variable for the async executors.
   std::vector<std::atomic<bool>> activePMs(asyncExecutors.size());
   std::fill(activePMs.begin(), activePMs.end(), false);
+  std::atomic<bool> hasFailure = false;
   auto processFn = [&](OpPMInfo &opInfo) {
     // Find an executor for this operation.
     auto it = llvm::find_if(activePMs, [](std::atomic<bool> &isActive) {
@@ -812,14 +813,16 @@ void OpToOpPassAdaptor::runOnOperationAsyncImpl(bool verifyPasses) {
     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);
-    return pipelineResult;
   };
+  parallelForEach(context, opInfos, processFn);
 
   // Signal a failure if any of the executors failed.
-  if (failed(failableParallelForEach(context, opInfos, processFn)))
+  if (hasFailure)
     signalPassFailure();
 }
 
diff --git a/mlir/test/Pass/crash-recovery-dynamic-failure.mlir b/mlir/test/Pass/crash-recovery-dynamic-failure.mlir
index cea4f6a79963bd..6d2bd696391086 100644
--- a/mlir/test/Pass/crash-recovery-dynamic-failure.mlir
+++ b/mlir/test/Pass/crash-recovery-dynamic-failure.mlir
@@ -15,4 +15,4 @@ module @inner_mod1 {
 // REPRO_LOCAL_DYNAMIC_FAILURE: module @inner_mod1
 // REPRO_LOCAL_DYNAMIC_FAILURE: module @foo {
 
-// REPRO_LOCAL_DYNAMIC_FAILURE: pipeline: "builtin.module(builtin.module(test-pass-failure))"
+// REPRO_LOCAL_DYNAMIC_FAILURE: pipeline: "builtin.module(builtin.module(test-pass-failure{gen-diagnostics=false}))"
diff --git a/mlir/test/Pass/full_diagnostics.mlir b/mlir/test/Pass/full_diagnostics.mlir
new file mode 100644
index 00000000000000..881260ce60d321
--- /dev/null
+++ b/mlir/test/Pass/full_diagnostics.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-pass-failure{gen-diagnostics}))' -verify-diagnostics
+
+// Test that all errors are reported.
+// expected-error at below {{illegal operation}}
+func.func @TestAlwaysIllegalOperationPass1() {
+  return
+}
+
+// expected-error at below {{illegal operation}}
+func.func @TestAlwaysIllegalOperationPass2() {
+  return
+}
+
+// expected-error at below {{illegal operation}}
+func.func @TestAlwaysIllegalOperationPass3() {
+  return
+}
diff --git a/mlir/test/lib/Pass/TestPassManager.cpp b/mlir/test/lib/Pass/TestPassManager.cpp
index ee32bec0c79bd4..7afe2109f04db3 100644
--- a/mlir/test/lib/Pass/TestPassManager.cpp
+++ b/mlir/test/lib/Pass/TestPassManager.cpp
@@ -151,11 +151,21 @@ struct TestCrashRecoveryPass
 struct TestFailurePass : public PassWrapper<TestFailurePass, OperationPass<>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestFailurePass)
 
-  void runOnOperation() final { signalPassFailure(); }
+  TestFailurePass() = default;
+  TestFailurePass(const TestFailurePass &other) : PassWrapper(other) {}
+
+  void runOnOperation() final {
+    signalPassFailure();
+    if (genDiagnostics)
+      mlir::emitError(getOperation()->getLoc(), "illegal operation");
+  }
   StringRef getArgument() const final { return "test-pass-failure"; }
   StringRef getDescription() const final {
     return "Test a pass in the pass manager that always fails";
   }
+
+  Option<bool> genDiagnostics{*this, "gen-diagnostics",
+                              llvm::cl::desc("Generate a diagnostic message")};
 };
 
 /// A test pass that creates an invalid operation in a function body.

>From e3f62cf117453e1acd9a1d68b0b8b78cf6d94c57 Mon Sep 17 00:00:00 2001
From: Billy Zhu <billyzhu at modular.com>
Date: Tue, 1 Oct 2024 10:08:14 -0700
Subject: [PATCH 2/2] inline func

---
 mlir/lib/Pass/Pass.cpp | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Pass/Pass.cpp b/mlir/lib/Pass/Pass.cpp
index 90fd9a713e41d7..6fd51c1e3cb538 100644
--- a/mlir/lib/Pass/Pass.cpp
+++ b/mlir/lib/Pass/Pass.cpp
@@ -800,7 +800,7 @@ void OpToOpPassAdaptor::runOnOperationAsyncImpl(bool verifyPasses) {
   std::vector<std::atomic<bool>> activePMs(asyncExecutors.size());
   std::fill(activePMs.begin(), activePMs.end(), false);
   std::atomic<bool> hasFailure = false;
-  auto processFn = [&](OpPMInfo &opInfo) {
+  parallelForEach(context, opInfos, [&](OpPMInfo &opInfo) {
     // Find an executor for this operation.
     auto it = llvm::find_if(activePMs, [](std::atomic<bool> &isActive) {
       bool expectedInactive = false;
@@ -818,8 +818,7 @@ void OpToOpPassAdaptor::runOnOperationAsyncImpl(bool verifyPasses) {
 
     // Reset the active bit for this pass manager.
     activePMs[pmIndex].store(false);
-  };
-  parallelForEach(context, opInfos, processFn);
+  });
 
   // Signal a failure if any of the executors failed.
   if (hasFailure)