From b8ffcb12e2ed52e080532586ae1ed4e6f9e47b70 Mon Sep 17 00:00:00 2001
From: River Riddle <riddleriver@gmail.com>
Date: Tue, 5 Oct 2021 00:24:24 +0000
Subject: [PATCH] [mlir:Pass] Generate a reproducer as early as possible

This avoids keeping references to passes that may be freed by
the time that the pass manager has finished executing (in the
non-crash case).

Fixes PR#52069

Differential Revision: https://reviews.llvm.org/D111106
---
 mlir/lib/Pass/PassCrashRecovery.cpp             | 12 ++++++++----
 .../Pass/crash-recovery-dynamic-failure.mlir    | 17 +++++++++++++++++
 mlir/test/lib/Pass/TestDynamicPipeline.cpp      | 16 ++++++----------
 3 files changed, 31 insertions(+), 14 deletions(-)
 create mode 100644 mlir/test/Pass/crash-recovery-dynamic-failure.mlir

diff --git a/mlir/lib/Pass/PassCrashRecovery.cpp b/mlir/lib/Pass/PassCrashRecovery.cpp
index c9b916699231..d4657b072f7c 100644
--- a/mlir/lib/Pass/PassCrashRecovery.cpp
+++ b/mlir/lib/Pass/PassCrashRecovery.cpp
@@ -227,6 +227,10 @@ formatPassOpReproducerMessage(Diagnostic &os,
 
 void PassCrashReproducerGenerator::finalize(Operation *rootOp,
                                             LogicalResult executionResult) {
+  // Don't generate a reproducer if we have no active contexts.
+  if (impl->activeContexts.empty())
+    return;
+
   // If the pass manager execution succeeded, we don't generate any reproducers.
   if (succeeded(executionResult))
     return impl->activeContexts.clear();
@@ -345,20 +349,20 @@ struct CrashReproducerInstrumentation : public PassInstrumentation {
       : generator(generator) {}
   ~CrashReproducerInstrumentation() override = default;
 
-  /// A callback to run before a pass is executed.
   void runBeforePass(Pass *pass, Operation *op) override {
     if (!isa<OpToOpPassAdaptor>(pass))
       generator.prepareReproducerFor(pass, op);
   }
 
-  /// A callback to run after a pass is successfully executed. This function
-  /// takes a pointer to the pass to be executed, as well as the current
-  /// operation being operated on.
   void runAfterPass(Pass *pass, Operation *op) override {
     if (!isa<OpToOpPassAdaptor>(pass))
       generator.removeLastReproducerFor(pass, op);
   }
 
+  void runAfterPassFailed(Pass *pass, Operation *op) override {
+    generator.finalize(op, /*executionResult=*/failure());
+  }
+
 private:
   /// The generator used to create crash reproducers.
   PassCrashReproducerGenerator &generator;
diff --git a/mlir/test/Pass/crash-recovery-dynamic-failure.mlir b/mlir/test/Pass/crash-recovery-dynamic-failure.mlir
new file mode 100644
index 000000000000..22fc706cfb25
--- /dev/null
+++ b/mlir/test/Pass/crash-recovery-dynamic-failure.mlir
@@ -0,0 +1,17 @@
+// Check that local reproducers will also traverse dynamic pass pipelines.
+// RUN: mlir-opt %s -pass-pipeline='test-module-pass,test-dynamic-pipeline{op-name=inner_mod1 run-on-nested-operations=1 dynamic-pipeline=test-pass-failure}' -pass-pipeline-crash-reproducer=%t -verify-diagnostics -pass-pipeline-local-reproducer --mlir-disable-threading
+// RUN: cat %t | FileCheck -check-prefix=REPRO_LOCAL_DYNAMIC_FAILURE %s
+
+// The crash recovery mechanism will leak memory allocated in the crashing thread.
+// UNSUPPORTED: asan
+
+module @inner_mod1 {
+  // expected-error@below {{Failures have been detected while processing an MLIR pass pipeline}}
+  // expected-note@below {{Pipeline failed while executing}}
+  module @foo {}
+}
+
+// REPRO_LOCAL_DYNAMIC_FAILURE: configuration: -pass-pipeline='builtin.module(test-pass-failure)'
+
+// REPRO_LOCAL_DYNAMIC_FAILURE: module @inner_mod1
+// REPRO_LOCAL_DYNAMIC_FAILURE: module @foo {
diff --git a/mlir/test/lib/Pass/TestDynamicPipeline.cpp b/mlir/test/lib/Pass/TestDynamicPipeline.cpp
index 275043dcf3d2..d10dd4a200fd 100644
--- a/mlir/test/lib/Pass/TestDynamicPipeline.cpp
+++ b/mlir/test/lib/Pass/TestDynamicPipeline.cpp
@@ -56,16 +56,14 @@ public:
       llvm::errs() << "dynamic-pipeline skip op name: " << opName << "\n";
       return;
     }
-    if (!pm) {
-      pm = std::make_unique<OpPassManager>(currentOp->getName().getIdentifier(),
-                                           OpPassManager::Nesting::Implicit);
-      (void)parsePassPipeline(pipeline, *pm, llvm::errs());
-    }
+    OpPassManager pm(currentOp->getName().getIdentifier(),
+                     OpPassManager::Nesting::Implicit);
+    (void)parsePassPipeline(pipeline, pm, llvm::errs());
 
     // Check that running on the parent operation always immediately fails.
     if (runOnParent) {
       if (currentOp->getParentOp())
-        if (!failed(runPipeline(*pm, currentOp->getParentOp())))
+        if (!failed(runPipeline(pm, currentOp->getParentOp())))
           signalPassFailure();
       return;
     }
@@ -78,18 +76,16 @@ public:
           return;
         llvm::errs() << "Run on " << *op << "\n";
         // Run on the current operation
-        if (failed(runPipeline(*pm, op)))
+        if (failed(runPipeline(pm, op)))
           signalPassFailure();
       });
     } else {
       // Run on the current operation
-      if (failed(runPipeline(*pm, currentOp)))
+      if (failed(runPipeline(pm, currentOp)))
         signalPassFailure();
     }
   }
 
-  std::unique_ptr<OpPassManager> pm;
-
   Option<bool> runOnNestedOp{
       *this, "run-on-nested-operations",
       llvm::cl::desc("This will apply the pipeline on nested operations under "