[mlir][Pass] Register a signal handler when generating crash reproducers.

[rriddle]

Summary

The current implementation uses CrashRecoveryContext, but this only supports recovering in a certain number of cases. This revision adds a signal handler to support even more situations.

This revision was able to properly generate a reproducer for a segfault in the Inliner, that the current recovery couldn't.

Depends On D78314

[mehdi_amini]

Can you clarify a bit in the description what kind of situation this allows to handle that weren't possible before?

Also it isn't clear to me why do we need a stack of RecoveryReproducerContext? (maybe you can describe this part in the code comments)

[rriddle]

Resolve comments

[rriddle]

In D78315#2005841, @mehdi_amini wrote:

Can you clarify a bit in the description what kind of situation this allows to handle that weren't possible before?

Added an example. This allowed for me to generate a reproducer during an Inliner segfault that I couldn't otherwise.

Also it isn't clear to me why do we need a stack of RecoveryReproducerContext? (maybe you can describe this part in the code comments)

There could be multiple pass managers running simultaneously, with different MLIRContexts. Unless we guarantee that there is only ever one context at the same time(something I would be against), we need a stack to handle this case. There is also another situation right now, caused by the lack of dynamic pass pipelines ;), where a pass manager could get run within another. Though regardless of when we fix the second problem, the first one will still exist.

[mehdi_amini]

In D78315#2008460, @rriddle wrote:

In D78315#2005841, @mehdi_amini wrote:

There could be multiple pass managers running simultaneously, with different MLIRContexts.

Make sense, but then I have another question (see inline)

mlir/lib/Pass/Pass.cpp

617

So this seems to enforce that we always delete the stack of RecoveryReproducerContext in the order in which they were created, how does it happen? If multiple threads from the application are running different PassManagers for example?

[rriddle]

Resolve comments

[rriddle]

mlir/lib/Pass/Pass.cpp

617

Ooops, fixed.

[rriddle]

Add missing signals include

mlir/lib/Pass/Pass.cpp

 #include "mlir/Pass/Pass.h"
 #include "PassDetail.h"
 #include "mlir/Analysis/Verifier.h"
 #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/ADT/SetVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/CrashRecoveryContext.h"
 #include "llvm/Support/Mutex.h"
 #include "llvm/Support/Parallel.h"
+#include "llvm/Support/Signals.h"
 #include "llvm/Support/Threading.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 //===----------------------------------------------------------------------===//
 // Pass
void OpToOpPassAdaptor::runOnOperationAsyncImpl() {
   if (passFailed)
     signalPassFailure();
 }
 
 //===----------------------------------------------------------------------===//
 // PassCrashReproducer
 //===----------------------------------------------------------------------===//
 
+namespace {
+/// This class contains all of the context for generating a recovery reproducer.
+/// Each recovery context is registered globally to allow for generating
+/// reproducers when a signal is raised, such as a segfault.
+struct RecoveryReproducerContext {
+  RecoveryReproducerContext(MutableArrayRef<std::unique_ptr<Pass>> passes,
+                            ModuleOp module, StringRef filename,
+                            bool disableThreads, bool verifyPasses);
+  ~RecoveryReproducerContext();
+
+  /// Generate a reproducer with the current context.
+  LogicalResult generate(std::string &error);
+
+private:
+  /// This function is invoked in the event of a crash.
+  static void crashHandler(void *);
+
+  /// Register a signal handler to run in the event of a crash.
+  static void registerSignalHandler();
+
+  /// The textual description of the currently executing pipeline.
+  std::string pipeline;
+
+  /// The MLIR module representing the IR before the crash.
+  OwningModuleRef module;
+
+  /// The filename to use when generating the reproducer.
+  StringRef filename;
+
+  /// Various pass manager flags.
+  bool disableThreads;
+  bool verifyPasses;
+
+  /// The current set of active reproducer contexts. This is used in the event
+  /// of a crash. This is not thread_local as the pass manager may produce any
+  /// number of child threads. This uses a set to allow for multiple MLIR pass
+  /// managers to be running at the same time.
+  static llvm::ManagedStatic<llvm::sys::SmartMutex<true>> reproducerMutex;
+  static llvm::ManagedStatic<
+      llvm::SmallSetVector<RecoveryReproducerContext *, 1>>
+      reproducerSet;
+};
+} // end anonymous namespace
+
+llvm::ManagedStatic<llvm::sys::SmartMutex<true>>
+    RecoveryReproducerContext::reproducerMutex;
+llvm::ManagedStatic<llvm::SmallSetVector<RecoveryReproducerContext *, 1>>
+    RecoveryReproducerContext::reproducerSet;
+
+RecoveryReproducerContext::RecoveryReproducerContext(
+    MutableArrayRef<std::unique_ptr<Pass>> passes, ModuleOp module,
+    StringRef filename, bool disableThreads, bool verifyPasses)
+    : module(module.clone()), filename(filename),
+      disableThreads(disableThreads), verifyPasses(verifyPasses) {
+  // Grab the textual pipeline being executed..
+  {
+    llvm::raw_string_ostream pipelineOS(pipeline);
+    ::printAsTextualPipeline(passes, pipelineOS);
+  }
+
+  // Make sure that the handler is registered, and update the current context.
+  llvm::sys::SmartScopedLock<true> producerLock(*reproducerMutex);
+  registerSignalHandler();
+  reproducerSet->insert(this);
+}
+
+RecoveryReproducerContext::~RecoveryReproducerContext() {
+  llvm::sys::SmartScopedLock<true> producerLock(*reproducerMutex);
+  reproducerSet->remove(this);
+}
+
+LogicalResult RecoveryReproducerContext::generate(std::string &error) {
+  std::unique_ptr<llvm::ToolOutputFile> outputFile =
+      mlir::openOutputFile(filename, &error);
+  if (!outputFile)
+    return failure();
+  auto &outputOS = outputFile->os();
+
+  // Output the current pass manager configuration.
+  outputOS << "// configuration: -pass-pipeline='" << pipeline << "'";
+  if (disableThreads)
+    outputOS << " -disable-pass-threading";
+
+  // TODO: Should this also be configured with a pass manager flag?
+  outputOS << "\n// note: verifyPasses=" << (verifyPasses ? "true" : "false")
+           << "\n";
+
+  // Output the .mlir module.
+  module->print(outputOS);
+  outputFile->keep();
+  return success();
+}
+
+void RecoveryReproducerContext::crashHandler(void *) {
+  // Walk the current stack of contexts and generate a reproducer for each one.
+  // We can't know for certain which one was the cause, so we need to generate
+  // a reproducer for all of them.
+  std::string ignored;
+  for (RecoveryReproducerContext *context : *reproducerSet)
+    context->generate(ignored);
+}
+
+void RecoveryReproducerContext::registerSignalHandler() {
+  // Ensure that the handler is only registered once.
+  static bool registered =
+      (llvm::sys::AddSignalHandler(crashHandler, nullptr), false);
+  (void)registered;
+}
+
 /// 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
LogicalResult PassManager::runWithCrashRecovery(ModuleOp module,
   }
   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 being executed first, just in case the passes
-  // become compromised.
-  std::string pipeline;
-  {
-    llvm::raw_string_ostream pipelineOS(pipeline);
-    ::printAsTextualPipeline(passes, pipelineOS);
-  }
-
-  // Clone the initial module before running it through the pass pipeline.
-  OwningModuleRef reproducerModule = module.clone();
+  RecoveryReproducerContext context(passes, module, *crashReproducerFileName,
+                                    impl->disableThreads, impl->verifyPasses);
 
   // Safely invoke the passes within a recovery context.
+  llvm::CrashRecoveryContext::Enable();
   LogicalResult passManagerResult = failure();
   llvm::CrashRecoveryContext recoveryContext;
   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();
 
   std::string error;
-  std::unique_ptr<llvm::ToolOutputFile> outputFile =
-      mlir::openOutputFile(*crashReproducerFileName, &error);
-  if (!outputFile)
+  if (failed(context.generate(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 (impl->disableThreads)
-    outputOS << " -disable-pass-threading";
-
-  // TODO: Should this also be configured with a pass manager flag?
-  outputOS << "\n// note: verifyPasses="
-           << (impl->verifyPasses ? "true" : "false") << "\n";
-
-  // Output the .mlir module.
-  reproducerModule->print(outputOS);
-  outputFile->keep();
-
-  return reproducerModule->emitError()
+  return module.emitError()
          << "A failure has been detected while processing the MLIR module, a "
             "reproducer has been generated in '"
          << *crashReproducerFileName << "'";
 }
 
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//