Add support for generating reproducers on pass crash and failure.

This cl adds support for generating a .mlir file containing a reproducer for crashes and failures that happen during pass execution. The reproducer contains a comment detailing the configuration of the pass manager(e.g. the textual description of the pass pipeline that the pass manager was executing), along with the original input module.

Example Output:

// configuration: -pass-pipeline='func(cse, canonicalize), inline'
// note: verifyPasses=false

module {
  ...
}

PiperOrigin-RevId: 274088134

third_party/mlir/include/mlir/Pass/PassManager.h

@@ -19,6 +19,7 @@
 #define MLIR_PASS_PASSMANAGER_H
 
 #include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallVector.h"
 
 namespace llvm {
@@ -127,6 +128,11 @@ public:
   /// Disable support for multi-threading within the pass manager.
   void disableMultithreading(bool disable = true);
 
+  /// Enable support for the pass manager to generate a reproducer on the event
+  /// of a crash or a pass failure. `outputFile` is a .mlir filename used to
+  /// write the generated reproducer.
+  void enableCrashReproducerGeneration(StringRef outputFile);
+
   //===--------------------------------------------------------------------===//
   // Instrumentations
   //===--------------------------------------------------------------------===//
@@ -159,6 +165,9 @@ private:
 
   /// A manager for pass instrumentations.
   std::unique_ptr<PassInstrumentor> instrumentor;
+
+  /// An optional filename to use when generating a crash reproducer if valid.
+  Optional<std::string> crashReproducerFileName;
 };
 
 /// Register a set of useful command-line options that can be used to configure

third_party/mlir/lib/Pass/Pass.cpp

@@ -25,11 +25,14 @@
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Module.h"
+#include "mlir/Support/FileUtilities.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/CrashRecoveryContext.h"
 #include "llvm/Support/Mutex.h"
 #include "llvm/Support/Parallel.h"
 #include "llvm/Support/Threading.h"
+#include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
 using namespace mlir::detail;
@@ -52,6 +55,8 @@ void Pass::printAsTextualPipeline(raw_ostream &os) {
       pm.printAsTextualPipeline(os);
       os << ")";
     });
+  } else if (const PassInfo *info = lookupPassInfo()) {
+    os << info->getPassArgument();
   } else {
     os << getName();
   }
@@ -272,7 +277,12 @@ const OperationName &OpPassManager::getOpName() const { return impl->name; }
 /// Prints out the passes of the pass mangager as the textual representation
 /// of pipelines.
 void OpPassManager::printAsTextualPipeline(raw_ostream &os) {
-  interleaveComma(impl->passes, os, [&](const std::unique_ptr<Pass> &pass) {
+  // Filter out passes that are not part of the public pipeline.
+  auto filteredPasses = llvm::make_filter_range(
+      impl->passes, [](const std::unique_ptr<Pass> &pass) {
+        return !isa<VerifierPass>(pass);
+      });
+  interleaveComma(filteredPasses, os, [&](const std::unique_ptr<Pass> &pass) {
     pass->printAsTextualPipeline(os);
   });
 }
@@ -464,6 +474,70 @@ OpToOpPassAdaptorBase *mlir::detail::getAdaptorPassBase(Pass *pass) {
   return nullptr;
 }
 
+//===----------------------------------------------------------------------===//
+// PassCrashReproducer
+//===----------------------------------------------------------------------===//
+
+/// Safely run the pass manager over the given module, creating a reproducible
+/// on failure or crash.
+static LogicalResult runWithCrashRecovery(OpPassManager &pm,
+                                          ModuleAnalysisManager &am,
+                                          ModuleOp module,
+                                          StringRef crashReproducerFileName) {
+  /// Enable crash recovery.
+  llvm::CrashRecoveryContext::Enable();
+
+  // Grab the textual pipeline executing within the pass manager first, just in
+  // case the pass manager becomes compromised.
+  std::string pipeline;
+  {
+    llvm::raw_string_ostream pipelineOS(pipeline);
+    pm.printAsTextualPipeline(pipelineOS);
+  }
+
+  // Clone the initial module before running it through the pass pipeline.
+  OwningModuleRef reproducerModule = module.clone();
+
+  // Safely invoke the pass manager within a recovery context.
+  LogicalResult passManagerResult = failure();
+  llvm::CrashRecoveryContext recoveryContext;
+  recoveryContext.RunSafelyOnThread(
+      [&] { passManagerResult = pm.run(module, am); });
+
+  /// Disable crash recovery.
+  llvm::CrashRecoveryContext::Disable();
+  if (succeeded(passManagerResult))
+    return success();
+
+  // The conversion failed, so generate a reproducible.
+  std::string error;
+  std::unique_ptr<llvm::ToolOutputFile> outputFile =
+      mlir::openOutputFile(crashReproducerFileName, &error);
+  if (!outputFile)
+    return emitError(UnknownLoc::get(pm.getContext()),
+                     "<MLIR-PassManager-Crash-Reproducer>: ")
+           << error;
+  auto &outputOS = outputFile->os();
+
+  // Output the current pass manager configuration.
+  outputOS << "// configuration: -pass-pipeline='" << pipeline << "'";
+  if (pm.getImpl().disableThreads)
+    outputOS << " -disable-pass-threading";
+
+  // TODO(riverriddle) Should this also be configured with a pass manager flag?
+  outputOS << "\n// note: verifyPasses="
+           << (pm.getImpl().verifyPasses ? "true" : "false") << "\n";
+
+  // Output the .mlir module.
+  reproducerModule->print(outputOS);
+  outputFile->keep();
+
+  return reproducerModule->emitError()
+         << "A crash has been detected while processing the MLIR module, a "
+            "reproducer has been generated in '"
+         << crashReproducerFileName << "'";
+}
+
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//
@@ -481,8 +555,13 @@ LogicalResult PassManager::run(ModuleOp module) {
   // pipeline.
   getImpl().coalesceAdjacentAdaptorPasses();
 
-  // Construct an analysis manager for the pipeline and run it.
+  // Construct an analysis manager for the pipeline.
   ModuleAnalysisManager am(module, instrumentor.get());
+
+  // If reproducer generation is enabled, run the pass manager with crash
+  // handling enabled.
+  if (crashReproducerFileName)
+    return runWithCrashRecovery(*this, am, module, *crashReproducerFileName);
   return OpPassManager::run(module, am);
 }
 
@@ -491,6 +570,13 @@ void PassManager::disableMultithreading(bool disable) {
   getImpl().disableThreads = disable;
 }
 
+/// Enable support for the pass manager to generate a reproducer on the event
+/// of a crash or a pass failure. `outputFile` is a .mlir filename used to write
+/// the generated reproducer.
+void PassManager::enableCrashReproducerGeneration(StringRef outputFile) {
+  crashReproducerFileName = outputFile;
+}
+
 /// Add the provided instrumentation to the pass manager.
 void PassManager::addInstrumentation(std::unique_ptr<PassInstrumentation> pi) {
   if (!instrumentor)

third_party/mlir/lib/Pass/PassManagerOptions.cpp

@@ -25,6 +25,14 @@ using namespace mlir;
 
 namespace {
 struct PassManagerOptions {
+  //===--------------------------------------------------------------------===//
+  // Crash Reproducer Generator
+  //===--------------------------------------------------------------------===//
+  llvm::cl::opt<std::string> reproducerFile{
+      "pass-pipeline-crash-reproducer",
+      llvm::cl::desc("Generate a .mlir reproducer file at the given output path"
+                     " if the pass manager crashes or fails")};
+
   //===--------------------------------------------------------------------===//
   // Multi-threading
   //===--------------------------------------------------------------------===//
@@ -130,6 +138,10 @@ void mlir::registerPassManagerCLOptions() {
 }
 
 void mlir::applyPassManagerCLOptions(PassManager &pm) {
+  // Generate a reproducer on crash/failure.
+  if ((*options)->reproducerFile.getNumOccurrences())
+    pm.enableCrashReproducerGeneration((*options)->reproducerFile);
+
   // Disable multi-threading.
   if ((*options)->disableThreads)
     pm.disableMultithreading();

third_party/mlir/test/lib/Pass/TestPassManager.cpp

@@ -28,7 +28,6 @@ struct TestModulePass : public ModulePass<TestModulePass> {
 struct TestFunctionPass : public FunctionPass<TestFunctionPass> {
   void runOnFunction() final {}
 };
-
 class TestOptionsPass : public FunctionPass<TestOptionsPass> {
 public:
   struct Options : public PassOptions<Options> {
@@ -69,7 +68,13 @@ public:
   SmallVector<std::string, 4> stringListOption;
   std::string stringOption;
 };
-} // namespace
+
+/// A test pass that always aborts to enable testing the crash recovery
+/// mechanism of the pass manager.
+class TestCrashRecoveryPass : public OperationPass<TestCrashRecoveryPass> {
+  void runOnOperation() final { abort(); }
+};
+} // end anonymous namespace
 
 static void testNestedPipeline(OpPassManager &pm) {
   // Nest a module pipeline that contains:
@@ -97,6 +102,10 @@ static PassRegistration<TestModulePass>
 static PassRegistration<TestFunctionPass>
     unusedFP("test-function-pass", "Test a function pass in the pass manager");
 
+static PassRegistration<TestCrashRecoveryPass>
+    unusedCrashP("test-pass-crash",
+                 "Test a pass in the pass manager that always crashes");
+
 static PassPipelineRegistration<>
     unused("test-pm-nested-pipeline",
            "Test a nested pipeline in the pass manager", testNestedPipeline);