[mlir][Pass] Add support for generating local crash reproducers

[rriddle]

Summary

This revision adds a mode to the crash reproducer generator to attempt to generate a more local reproducer. This will attempt to generate a reproducer right before the offending pass that fails. This is useful for the majority of failures that are specific to a single pass, and situations where some passes in the pipeline are not registered with a specific tool.

Depends On D78313

[jpienaar]

mlir/docs/PassManagement.md

1001

minimal has to me a size connotation (e.g., "minimum amount of code needed to make the bug appear when executed."), and here it is more about scope than size (e.g., you only need to run one pass, but your reproducer may be huge). Would "local" or "pass-pipeline-dump-before-first-fail" work? Alternatively, where do you see someone not wanting this behavior?

1004

Missing space before (

[rriddle]

mlir/docs/PassManagement.md

1001

I'm fine with local.

This type of behavior isn't desirable if there are multi-pass interactions. It isn't guaranteed that the problem is with the pass that it is failing.

[rriddle]

Resolve comments

[mehdi_amini]

mlir/docs/PassManagement.md

1001

Interestingly local was confusing to me since I thought of it as opposite to "remote" ;)

[rriddle]

Rebase

mlir/docs/PassManagement.md

 // note: verifyPasses=false
 
 module {
   func @foo() {
     ...
   }
 }
 ```
+
+### Local Reproducer Generation
+
+An additional flag may be passed to
+`PassManager::enableCrashReproducerGeneration`, and specified via
+`pass-pipeline-local-reproducer` on the command line, that signals that the pass
+manager should attempt to generate a "local" reproducer. This will attempt to
+generate a reproducer containing IR right before the pass that fails. This is
+useful for situations where the crash is known to be within a specific pass, or
+when the original input relies on components (like dialects or passes) that may
+not always be available.
+
+For example, if the failure in the previous example came from `canonicalize`,
+the following reproducer will be generated:
+
+```mlir
+// configuration: -pass-pipeline='func(canonicalize)'
+// note: verifyPasses=false
+
+module {
+  func @foo() {
+    ...
+  }
+}
+```

mlir/include/mlir/Pass/Pass.h

private:
   std::vector<Statistic *> statistics;
 
   /// The pass options registered to this pass instance.
   detail::PassOptions passOptions;
 
   /// Allow access to 'clone' and 'run'.
   friend class OpPassManager;
 
+  /// Allow access to 'run'.
+  friend class PassManager;
+
   /// Allow access to 'passOptions'.
   friend class PassInfo;
 };
 
 //===----------------------------------------------------------------------===//
 // Pass Model Definitions
 //===----------------------------------------------------------------------===//
 

mlir/include/mlir/Pass/PassManager.h

 private:
   OpPassManager(OperationName name, bool disableThreads, bool verifyPasses);
 
   /// A pointer to an internal implementation instance.
   std::unique_ptr<detail::OpPassManagerImpl> impl;
 
   /// Allow access to the constructor.
   friend class PassManager;
+
+  /// Allow access.
+  friend detail::OpPassManagerImpl;
 };
 
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//
 
 /// An enum describing the different display modes for the information within
 /// the pass manager.
public:
   void disableMultithreading(bool disable = true);
 
   /// Return true if the pass manager is configured with multi-threading
   /// enabled.
   bool isMultithreadingEnabled();
 
   /// 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);
+  /// write the generated reproducer. If `genLocalReproducer` is true, the pass
+  /// manager will attempt to generate a local reproducer that contains the
+  /// smallest pipeline.
+  void enableCrashReproducerGeneration(StringRef outputFile,
+                                       bool genLocalReproducer = false);
 
   //===--------------------------------------------------------------------===//
   // Instrumentations
   //===--------------------------------------------------------------------===//
 
   /// Add the provided instrumentation to the pass manager.
   void addInstrumentation(std::unique_ptr<PassInstrumentation> pi);
 
public:
   /// held passes after each call to 'run'.
   void
   enableStatistics(PassDisplayMode displayMode = PassDisplayMode::Pipeline);
 
 private:
   /// Dump the statistics of the passes within this pass manager.
   void dumpStatistics();
 
-  /// Flag that specifies if pass timing is enabled.
-  bool passTiming : 1;
+  /// Run the pass manager with crash recover enabled.
+  LogicalResult runWithCrashRecovery(ModuleOp module, AnalysisManager am);
+  /// Run the given passes with crash recover enabled.
+  LogicalResult
+  runWithCrashRecovery(MutableArrayRef<std::unique_ptr<Pass>> passes,
+                       ModuleOp module, AnalysisManager am);
 
   /// Flag that specifies if pass statistics should be dumped.
   Optional<PassDisplayMode> passStatisticsMode;
 
   /// 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;
+
+  /// Flag that specifies if pass timing is enabled.
+  bool passTiming : 1;
+
+  /// Flag that specifies if the generated crash reproducer should be local.
+  bool localReproducer : 1;
 };
 
 /// Register a set of useful command-line options that can be used to configure
 /// a pass manager. The values of these options can be applied via the
 /// 'applyPassManagerCLOptions' method below.
 void registerPassManagerCLOptions();
 
 /// Apply any values provided to the pass manager options that were registered
 /// with 'registerPassManagerOptions'.
 void applyPassManagerCLOptions(PassManager &pm);
 } // end namespace mlir
 
 #endif // MLIR_PASS_PASSMANAGER_H

mlir/lib/Pass/Pass.cpp

 namespace mlir {
 namespace detail {
 struct OpPassManagerImpl {
   OpPassManagerImpl(OperationName name, bool disableThreads, bool verifyPasses)
       : name(name), disableThreads(disableThreads), verifyPasses(verifyPasses) {
   }
 
   /// Merge the passes of this pass manager into the one provided.
-  void mergeInto(OpPassManagerImpl &rhs) {
-    assert(name == rhs.name && "merging unrelated pass managers");
-    for (auto &pass : passes)
-      rhs.passes.push_back(std::move(pass));
-    passes.clear();
+  void mergeInto(OpPassManagerImpl &rhs);
+
+  /// Nest a new operation pass manager for the given operation kind under this
+  /// pass manager.
+  OpPassManager &nest(const OperationName &nestedName);
+  OpPassManager &nest(StringRef nestedName) {
+    return nest(OperationName(nestedName, getContext()));
   }
 
+  /// Add the given pass to this pass manager. If this pass has a concrete
+  /// operation type, it must be the same type as this pass manager.
+  void addPass(std::unique_ptr<Pass> pass);
+
   /// Coalesce adjacent AdaptorPasses into one large adaptor. This runs
   /// recursively through the pipeline graph.
   void coalesceAdjacentAdaptorPasses();
 
+  /// Split all of AdaptorPasses such that each adaptor only contains one leaf
+  /// pass.
+  void splitAdaptorPasses();
+
+  /// Return an instance of the context.
+  MLIRContext *getContext() const {
+    return name.getAbstractOperation()->dialect.getContext();
+  }
+
   /// The name of the operation that passes of this pass manager operate on.
   OperationName name;
 
   /// Flag to disable multi-threading of passes.
   bool disableThreads : 1;
 
   /// Flag that specifies if the IR should be verified after each pass has run.
   bool verifyPasses : 1;
 
   /// The set of passes to run as part of this pass manager.
   std::vector<std::unique_ptr<Pass>> passes;
 };
 } // end namespace detail
 } // end namespace mlir
 
-/// Coalesce adjacent AdaptorPasses into one large adaptor. This runs
-/// recursively through the pipeline graph.
+void OpPassManagerImpl::mergeInto(OpPassManagerImpl &rhs) {
+  assert(name == rhs.name && "merging unrelated pass managers");
+  for (auto &pass : passes)
+    rhs.passes.push_back(std::move(pass));
+  passes.clear();
+}
+
+OpPassManager &OpPassManagerImpl::nest(const OperationName &nestedName) {
+  OpPassManager nested(nestedName, disableThreads, verifyPasses);
+  auto *adaptor = new OpToOpPassAdaptor(std::move(nested));
+  addPass(std::unique_ptr<Pass>(adaptor));
+  return adaptor->getPassManagers().front();
+}
+
+void OpPassManagerImpl::addPass(std::unique_ptr<Pass> pass) {
+  // If this pass runs on a different operation than this pass manager, then
+  // implicitly nest a pass manager for this operation.
+  auto passOpName = pass->getOpName();
+  if (passOpName && passOpName != name.getStringRef())
+    return nest(*passOpName).addPass(std::move(pass));
+
+  passes.emplace_back(std::move(pass));
+  if (verifyPasses)
+    passes.emplace_back(std::make_unique<VerifierPass>());
+}
+
 void OpPassManagerImpl::coalesceAdjacentAdaptorPasses() {
   // Bail out early if there are no adaptor passes.
   if (llvm::none_of(passes, [](std::unique_ptr<Pass> &pass) {
         return isa<OpToOpPassAdaptor>(pass.get());
       }))
     return;
 
   // Walk the pass list and merge adjacent adaptors.
for (auto &pm : lastAdaptor->getPassManagers())
       pm.getImpl().coalesceAdjacentAdaptorPasses();
   }
 
   // Now that the adaptors have been merged, erase the empty slot corresponding
   // to the merged adaptors that were nulled-out in the loop above.
   llvm::erase_if(passes, std::logical_not<std::unique_ptr<Pass>>());
 }
 
+void OpPassManagerImpl::splitAdaptorPasses() {
+  std::vector<std::unique_ptr<Pass>> oldPasses;
+  std::swap(passes, oldPasses);
+
+  for (std::unique_ptr<Pass> &pass : oldPasses) {
+    // If this pass isn't an adaptor, move it directly to the new pass list.
+    auto *currentAdaptor = dyn_cast<OpToOpPassAdaptor>(pass.get());
+    if (!currentAdaptor) {
+      passes.push_back(std::move(pass));
+      continue;
+    }
+
+    // Otherwise, split the adaptors of each manager within the adaptor.
+    for (OpPassManager &adaptorPM : currentAdaptor->getPassManagers()) {
+      adaptorPM.getImpl().splitAdaptorPasses();
+
+      // Add all non-verifier passes to this pass manager.
+      for (std::unique_ptr<Pass> &nestedPass : adaptorPM.getImpl().passes) {
+        if (!isa<VerifierPass>(nestedPass.get()))
+          nest(adaptorPM.getOpName()).addPass(std::move(nestedPass));
+      }
+    }
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // OpPassManager
 //===----------------------------------------------------------------------===//
 
 OpPassManager::OpPassManager(OperationName name, bool disableThreads,
                              bool verifyPasses)
     : impl(new OpPassManagerImpl(name, disableThreads, verifyPasses)) {
   assert(name.getAbstractOperation() &&
for (auto &pass : impl->passes)
     if (failed(pass->run(op, am)))
       return failure();
   return success();
 }
 
 /// Nest a new operation pass manager for the given operation kind under this
 /// pass manager.
 OpPassManager &OpPassManager::nest(const OperationName &nestedName) {
-  OpPassManager nested(nestedName, impl->disableThreads, impl->verifyPasses);
-  auto *adaptor = new OpToOpPassAdaptor(std::move(nested));
-  addPass(std::unique_ptr<Pass>(adaptor));
-  return adaptor->getPassManagers().front();
+  return impl->nest(nestedName);
 }
 OpPassManager &OpPassManager::nest(StringRef nestedName) {
-  return nest(OperationName(nestedName, getContext()));
+  return impl->nest(nestedName);
 }
 
 /// Add the given pass to this pass manager. If this pass has a concrete
 /// operation type, it must be the same type as this pass manager.
 void OpPassManager::addPass(std::unique_ptr<Pass> pass) {
-  // If this pass runs on a different operation than this pass manager, then
-  // implicitly nest a pass manager for this operation.
-  auto passOpName = pass->getOpName();
-  if (passOpName && passOpName != impl->name.getStringRef())
-    return nest(*passOpName).addPass(std::move(pass));
-
-  impl->passes.emplace_back(std::move(pass));
-  if (impl->verifyPasses)
-    impl->passes.emplace_back(std::make_unique<VerifierPass>());
+  impl->addPass(std::move(pass));
 }
 
 /// Returns the number of passes held by this manager.
 size_t OpPassManager::size() const { return impl->passes.size(); }
 
 /// Returns the internal implementation instance.
 OpPassManagerImpl &OpPassManager::getImpl() { return *impl; }
 
 /// Return an instance of the context.
-MLIRContext *OpPassManager::getContext() const {
-  return impl->name.getAbstractOperation()->dialect.getContext();
-}
+MLIRContext *OpPassManager::getContext() const { return impl->getContext(); }
 
 /// Return the operation name that this pass manager operates on.
 const OperationName &OpPassManager::getOpName() const { return impl->name; }
 
-/// Prints out the passes of the pass manager as the textual representation
-/// of pipelines.
-void OpPassManager::printAsTextualPipeline(raw_ostream &os) {
+/// Prints out the given passes as the textual representation of a pipeline.
+static void printAsTextualPipeline(ArrayRef<std::unique_ptr<Pass>> passes,
+                                   raw_ostream &os) {
   // 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) {
+  auto filteredPasses =
+      llvm::make_filter_range(passes, [](const std::unique_ptr<Pass> &pass) {
         return !isa<VerifierPass>(pass);
       });
   llvm::interleaveComma(filteredPasses, os,
                         [&](const std::unique_ptr<Pass> &pass) {
                           pass->printAsTextualPipeline(os);
                         });
 }
 
+/// Prints out the passes of the pass manager as the textual representation
+/// of pipelines.
+void OpPassManager::printAsTextualPipeline(raw_ostream &os) {
+  ::printAsTextualPipeline(impl->passes, os);
+}
+
 //===----------------------------------------------------------------------===//
 // OpToOpPassAdaptor
 //===----------------------------------------------------------------------===//
 
 /// Utility to run the given operation and analysis manager on a provided op
 /// pass manager.
 static LogicalResult runPipeline(OpPassManager &pm, Operation *op,
                                  AnalysisManager am) {
void OpToOpPassAdaptor::runOnOperationAsyncImpl() {
   if (passFailed)
     signalPassFailure();
 }
 
 //===----------------------------------------------------------------------===//
 // 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) {
+/// Run the pass manager with crash recover enabled.
+LogicalResult PassManager::runWithCrashRecovery(ModuleOp module,
+                                                AnalysisManager am) {
+  // If this isn't a local producer, run all of the passes in recovery mode.
+  if (!localReproducer)
+    return runWithCrashRecovery(impl->passes, module, am);
+
+  // Split the passes within adaptors to ensure that each pass can be run in
+  // isolation.
+  impl->splitAdaptorPasses();
+
+  // If this is a local producer, run each of the passes individually. If the
+  // verifier is enabled, each pass will have a verifier after. This is included
+  // in the recovery run.
+  unsigned stride = impl->verifyPasses ? 2 : 1;
+  MutableArrayRef<std::unique_ptr<Pass>> passes = impl->passes;
+  for (unsigned i = 0, e = passes.size(); i != e; i += stride) {
+    if (failed(runWithCrashRecovery(passes.slice(i, stride), module, am)))
+      return failure();
+  }
+  return success();
+}
+
+/// Run the given passes with crash recover enabled.
+LogicalResult
+PassManager::runWithCrashRecovery(MutableArrayRef<std::unique_ptr<Pass>> passes,
+                                  ModuleOp module, AnalysisManager am) {
   /// Enable crash recovery.
   llvm::CrashRecoveryContext::Enable();
 
-  // Grab the textual pipeline executing within the pass manager first, just in
-  // case the pass manager becomes compromised.
+  // Grab the textual pipeline being executed first, just in case the passes
+  // become compromised.
   std::string pipeline;
   {
     llvm::raw_string_ostream pipelineOS(pipeline);
-    pm.printAsTextualPipeline(pipelineOS);
+    ::printAsTextualPipeline(passes, 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.
+  // Safely invoke the passes within a recovery context.
   LogicalResult passManagerResult = failure();
   llvm::CrashRecoveryContext recoveryContext;
-  recoveryContext.RunSafelyOnThread(
-      [&] { passManagerResult = pm.run(module, am); });
-
-  /// Disable crash recovery.
+  recoveryContext.RunSafelyOnThread([&] {
+    for (std::unique_ptr<Pass> &pass : passes)
+      if (failed(pass->run(module, am)))
+        return;
+    passManagerResult = success();
+  });
   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);
+      mlir::openOutputFile(*crashReproducerFileName, &error);
   if (!outputFile)
-    return emitError(UnknownLoc::get(pm.getContext()),
-                     "<MLIR-PassManager-Crash-Reproducer>: ")
-           << error;
+    return module.emitError("<MLIR-PassManager-Crash-Reproducer>: ") << error;
   auto &outputOS = outputFile->os();
 
   // Output the current pass manager configuration.
   outputOS << "// configuration: -pass-pipeline='" << pipeline << "'";
-  if (pm.getImpl().disableThreads)
+  if (impl->disableThreads)
     outputOS << " -disable-pass-threading";
 
-  // TODO(riverriddle) Should this also be configured with a pass manager flag?
+  // TODO: Should this also be configured with a pass manager flag?
   outputOS << "\n// note: verifyPasses="
-           << (pm.getImpl().verifyPasses ? "true" : "false") << "\n";
+           << (impl->verifyPasses ? "true" : "false") << "\n";
 
   // Output the .mlir module.
   reproducerModule->print(outputOS);
   outputFile->keep();
 
   return reproducerModule->emitError()
          << "A failure has been detected while processing the MLIR module, a "
             "reproducer has been generated in '"
-         << crashReproducerFileName << "'";
+         << *crashReproducerFileName << "'";
 }
 
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//
 
 PassManager::PassManager(MLIRContext *ctx, bool verifyPasses)
     : OpPassManager(OperationName(ModuleOp::getOperationName(), ctx),
                     /*disableThreads=*/false, verifyPasses),
-      passTiming(false) {}
+      passTiming(false), localReproducer(false) {}
 
 PassManager::~PassManager() {}
 
 /// Run the passes within this manager on the provided module.
 LogicalResult PassManager::run(ModuleOp module) {
   // Before running, make sure to coalesce any adjacent pass adaptors in the
   // pipeline.
   getImpl().coalesceAdjacentAdaptorPasses();
 
   // 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.
-  LogicalResult result =
-      crashReproducerFileName
-          ? runWithCrashRecovery(*this, am, module, *crashReproducerFileName)
+  LogicalResult result = crashReproducerFileName
+                             ? runWithCrashRecovery(module, am)
                              : OpPassManager::run(module, am);
 
   // Dump all of the pass statistics if necessary.
   if (passStatisticsMode)
     dumpStatistics();
   return result;
 }
 
 /// Disable support for multi-threading within the pass manager.
 void PassManager::disableMultithreading(bool disable) {
   getImpl().disableThreads = disable;
 }
 
 bool PassManager::isMultithreadingEnabled() {
   return !getImpl().disableThreads;
 }
 
 /// 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) {
+/// the generated reproducer. If `genLocalReproducer` is true, the pass manager
+/// will attempt to generate a local reproducer that contains the smallest
+/// pipeline.
+void PassManager::enableCrashReproducerGeneration(StringRef outputFile,
+                                                  bool genLocalReproducer) {
   crashReproducerFileName = std::string(outputFile);
+  localReproducer = genLocalReproducer;
 }
 
 /// Add the provided instrumentation to the pass manager.
 void PassManager::addInstrumentation(std::unique_ptr<PassInstrumentation> pi) {
   if (!instrumentor)
     instrumentor = std::make_unique<PassInstrumentor>();
 
   instrumentor->addInstrumentation(std::move(pi));

mlir/lib/Pass/PassManagerOptions.cpp

 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")};
+  llvm::cl::opt<bool> localReproducer{
+      "pass-pipeline-local-reproducer",
+      llvm::cl::desc("When generating a crash reproducer, attempt to generated "
+                     "a reproducer with the smallest pipeline."),
+      llvm::cl::init(false)};
 
   //===--------------------------------------------------------------------===//
   // Multi-threading
   //===--------------------------------------------------------------------===//
   llvm::cl::opt<bool> disableThreads{
       "disable-pass-threading",
       llvm::cl::desc("Disable multithreading in the pass manager"),
       llvm::cl::init(false)};

 }
 
 void mlir::applyPassManagerCLOptions(PassManager &pm) {
   if (!options.isConstructed())
     return;
 
   // Generate a reproducer on crash/failure.
   if (options->reproducerFile.getNumOccurrences())
-    pm.enableCrashReproducerGeneration(options->reproducerFile);
+    pm.enableCrashReproducerGeneration(options->reproducerFile,
+                                       options->localReproducer);
 
   // Disable multi-threading.
   if (options->disableThreads)
     pm.disableMultithreading();
 
   // Enable statistics dumping.
   if (options->passStatistics)
     pm.enableStatistics(options->passStatisticsDisplayMode);

mlir/test/Pass/crash-recovery.mlir

 // RUN: mlir-opt %s -pass-pipeline='func(test-function-pass, test-pass-crash)' -pass-pipeline-crash-reproducer=%t -verify-diagnostics
 // RUN: cat %t | FileCheck -check-prefix=REPRO %s
+// RUN: mlir-opt %s -pass-pipeline='func(test-function-pass, test-pass-crash)' -pass-pipeline-crash-reproducer=%t -verify-diagnostics -pass-pipeline-local-reproducer
+// RUN: cat %t | FileCheck -check-prefix=REPRO_LOCAL %s
 
 // expected-error@+1 {{A failure has been detected while processing the MLIR module}}
 module {
   func @foo() {
     return
   }
 }
 
 // REPRO: configuration: -pass-pipeline='func(test-function-pass, test-pass-crash)'
 
 // REPRO: module
 // REPRO: func @foo() {
 // REPRO-NEXT: return
+
+// REPRO_LOCAL: configuration: -pass-pipeline='func(test-pass-crash)'
+
+// REPRO_LOCAL: module
+// REPRO_LOCAL: func @foo() {
+// REPRO_LOCAL-NEXT: return