[Mlir-commits] [mlir] [mlir][Pass] Enable the option for reproducer generation without crashing (PR #75421)

Puyan Lotfi llvmlistbot at llvm.org
Thu Dec 21 01:12:44 PST 2023


https://github.com/plotfi updated https://github.com/llvm/llvm-project/pull/75421

>From 15b3128ddc858565042158b182d4edecaa2695e7 Mon Sep 17 00:00:00 2001
From: Puyan Lotfi <puyan at puyan.org>
Date: Thu, 21 Dec 2023 01:09:16 -0800
Subject: [PATCH] [mlir][Pass] Enable the option for reproducer generation
 without crashing

Adding API `enableReproducerGeneration` (with an optional
alwaysGenerateReproducer parameter) and cl::opt flag
`-mlir-generate-reproducer` in order to allow for mlir
reproducer dumps even when the pipeline doesn't crash.

This will be useful for frameworks and runtimes that use MLIR where it is
needed to reproduce the pipeline behavior for reasons outside of diagnosing
crashes. An example is for diagnosing performance issue using offline tools,
where being able to dump the reproducer from a runtime compiler would be
helpful.
---
 mlir/include/mlir/Pass/PassManager.h     | 44 ++++++++-----
 mlir/lib/Pass/PassCrashRecovery.cpp      | 83 +++++++++++++++++-------
 mlir/lib/Pass/PassDetail.h               |  2 +-
 mlir/lib/Tools/mlir-opt/MlirOptMain.cpp  | 16 +++++
 mlir/test/Pass/crashless-reproducer.mlir | 13 ++++
 5 files changed, 117 insertions(+), 41 deletions(-)
 create mode 100644 mlir/test/Pass/crashless-reproducer.mlir

diff --git a/mlir/include/mlir/Pass/PassManager.h b/mlir/include/mlir/Pass/PassManager.h
index d5f1ea0fe0350d..7d843f34200515 100644
--- a/mlir/include/mlir/Pass/PassManager.h
+++ b/mlir/include/mlir/Pass/PassManager.h
@@ -207,6 +207,35 @@ enum class PassDisplayMode {
   Pipeline,
 };
 
+/// Streams on which to output crash reproducer.
+struct ReproducerStream {
+  virtual ~ReproducerStream() = default;
+
+  /// Description of the reproducer stream.
+  virtual StringRef description() = 0;
+
+  /// Stream on which to output reproducer.
+  virtual raw_ostream &os() = 0;
+};
+
+/// Method type for constructing ReproducerStream.
+using ReproducerStreamFactory =
+    std::function<std::unique_ptr<ReproducerStream>(std::string &error)>;
+
+ReproducerStreamFactory makeReproducerStreamFactory(StringRef outputFile);
+std::string produceTextualPipelineWithScope(Pass *pass, Operation *&op);
+void appendReproducer(std::string &description,
+                      Operation *op,
+                      const mlir::ReproducerStreamFactory &factory,
+                      const std::string &pipelineElements,
+                      bool disableThreads,
+                      bool verifyPasses);
+std::string makeReproducer(Pass *pass,
+                           Operation *op,
+                           StringRef outputFile,
+                           bool disableThreads = false,
+                           bool verifyPasses = false);
+
 /// The main pass manager and pipeline builder.
 class PassManager : public OpPassManager {
 public:
@@ -243,21 +272,6 @@ class PassManager : public OpPassManager {
   void enableCrashReproducerGeneration(StringRef outputFile,
                                        bool genLocalReproducer = false);
 
-  /// Streams on which to output crash reproducer.
-  struct ReproducerStream {
-    virtual ~ReproducerStream() = default;
-
-    /// Description of the reproducer stream.
-    virtual StringRef description() = 0;
-
-    /// Stream on which to output reproducer.
-    virtual raw_ostream &os() = 0;
-  };
-
-  /// Method type for constructing ReproducerStream.
-  using ReproducerStreamFactory =
-      std::function<std::unique_ptr<ReproducerStream>(std::string &error)>;
-
   /// Enable support for the pass manager to generate a reproducer on the event
   /// of a crash or a pass failure. `factory` is used to construct the streams
   /// to write the generated reproducer to. If `genLocalReproducer` is true, the
diff --git a/mlir/lib/Pass/PassCrashRecovery.cpp b/mlir/lib/Pass/PassCrashRecovery.cpp
index df1a0762ae34a5..29fa7b638fc261 100644
--- a/mlir/lib/Pass/PassCrashRecovery.cpp
+++ b/mlir/lib/Pass/PassCrashRecovery.cpp
@@ -38,7 +38,7 @@ namespace detail {
 /// reproducers when a signal is raised, such as a segfault.
 struct RecoveryReproducerContext {
   RecoveryReproducerContext(std::string passPipelineStr, Operation *op,
-                            PassManager::ReproducerStreamFactory &streamFactory,
+                            mlir::ReproducerStreamFactory &streamFactory,
                             bool verifyPasses);
   ~RecoveryReproducerContext();
 
@@ -67,7 +67,7 @@ struct RecoveryReproducerContext {
 
   /// The factory for the reproducer output stream to use when generating the
   /// reproducer.
-  PassManager::ReproducerStreamFactory &streamFactory;
+  mlir::ReproducerStreamFactory &streamFactory;
 
   /// Various pass manager and context flags.
   bool disableThreads;
@@ -92,7 +92,7 @@ llvm::ManagedStatic<llvm::SmallSetVector<RecoveryReproducerContext *, 1>>
 
 RecoveryReproducerContext::RecoveryReproducerContext(
     std::string passPipelineStr, Operation *op,
-    PassManager::ReproducerStreamFactory &streamFactory, bool verifyPasses)
+    mlir::ReproducerStreamFactory &streamFactory, bool verifyPasses)
     : pipelineElements(std::move(passPipelineStr)),
       preCrashOperation(op->clone()), streamFactory(streamFactory),
       disableThreads(!op->getContext()->isMultithreadingEnabled()),
@@ -107,21 +107,31 @@ RecoveryReproducerContext::~RecoveryReproducerContext() {
 }
 
 void RecoveryReproducerContext::generate(std::string &description) {
+  appendReproducer(description, preCrashOperation, streamFactory,
+                   pipelineElements, disableThreads, verifyPasses);
+}
+
+void mlir::appendReproducer(std::string &description,
+                            Operation *op,
+                            const mlir::ReproducerStreamFactory &factory,
+                            const std::string &pipelineElements,
+                            bool disableThreads,
+                            bool verifyPasses) {
   llvm::raw_string_ostream descOS(description);
 
   // Try to create a new output stream for this crash reproducer.
   std::string error;
-  std::unique_ptr<PassManager::ReproducerStream> stream = streamFactory(error);
+  std::unique_ptr<mlir::ReproducerStream> stream = factory(error);
   if (!stream) {
     descOS << "failed to create output stream: " << error;
     return;
   }
   descOS << "reproducer generated at `" << stream->description() << "`";
 
-  std::string pipeline = (preCrashOperation->getName().getStringRef() + "(" +
+  std::string pipeline = (op->getName().getStringRef() + "(" +
                           pipelineElements + ")")
                              .str();
-  AsmState state(preCrashOperation);
+  AsmState state(op);
   state.attachResourcePrinter(
       "mlir_reproducer", [&](Operation *op, AsmResourceBuilder &builder) {
         builder.buildString("pipeline", pipeline);
@@ -130,7 +140,17 @@ void RecoveryReproducerContext::generate(std::string &description) {
       });
 
   // Output the .mlir module.
-  preCrashOperation->print(stream->os(), state);
+  op->print(stream->os(), state);
+}
+std::string mlir::makeReproducer(Pass *pass, Operation *op,
+                                 StringRef outputFile,
+                                 bool disableThreads, bool verifyPasses) {
+  std::string description;
+  std::string pipelineStr = produceTextualPipelineWithScope(pass, op);
+  appendReproducer(description, op,
+                   makeReproducerStreamFactory(outputFile),
+                   pipelineStr, disableThreads, verifyPasses);
+  return description;
 }
 
 void RecoveryReproducerContext::disable() {
@@ -175,12 +195,12 @@ void RecoveryReproducerContext::registerSignalHandler() {
 //===----------------------------------------------------------------------===//
 
 struct PassCrashReproducerGenerator::Impl {
-  Impl(PassManager::ReproducerStreamFactory &streamFactory,
+  Impl(mlir::ReproducerStreamFactory &streamFactory,
        bool localReproducer)
       : streamFactory(streamFactory), localReproducer(localReproducer) {}
 
   /// The factory to use when generating a crash reproducer.
-  PassManager::ReproducerStreamFactory streamFactory;
+  mlir::ReproducerStreamFactory streamFactory;
 
   /// Flag indicating if reproducer generation should be localized to the
   /// failing pass.
@@ -198,7 +218,7 @@ struct PassCrashReproducerGenerator::Impl {
 };
 
 PassCrashReproducerGenerator::PassCrashReproducerGenerator(
-    PassManager::ReproducerStreamFactory &streamFactory, bool localReproducer)
+    mlir::ReproducerStreamFactory &streamFactory, bool localReproducer)
     : impl(std::make_unique<Impl>(streamFactory, localReproducer)) {}
 PassCrashReproducerGenerator::~PassCrashReproducerGenerator() = default;
 
@@ -296,6 +316,13 @@ void PassCrashReproducerGenerator::prepareReproducerFor(Pass *pass,
   if (!impl->activeContexts.empty())
     impl->activeContexts.back()->disable();
 
+  std::string pipelineStr = produceTextualPipelineWithScope(pass, op);
+  impl->activeContexts.push_back(std::make_unique<RecoveryReproducerContext>(
+      pipelineStr, op, impl->streamFactory, impl->pmFlagVerifyPasses));
+}
+
+namespace  mlir {
+std::string produceTextualPipelineWithScope(Pass *pass, Operation *&op) {
   // Collect all of the parent scopes of this operation.
   SmallVector<OperationName> scopes;
   while (Operation *parentOp = op->getParentOp()) {
@@ -313,8 +340,8 @@ void PassCrashReproducerGenerator::prepareReproducerFor(Pass *pass,
   for (unsigned i = 0, e = scopes.size(); i < e; ++i)
     passOS << ")";
 
-  impl->activeContexts.push_back(std::make_unique<RecoveryReproducerContext>(
-      passOS.str(), op, impl->streamFactory, impl->pmFlagVerifyPasses));
+  return passOS.str();
+}
 }
 void PassCrashReproducerGenerator::prepareReproducerFor(
     iterator_range<PassManager::pass_iterator> passes, Operation *op) {
@@ -382,9 +409,9 @@ struct CrashReproducerInstrumentation : public PassInstrumentation {
 //===----------------------------------------------------------------------===//
 
 namespace {
-/// This class represents a default instance of PassManager::ReproducerStream
+/// This class represents a default instance of mlir::ReproducerStream
 /// that is backed by a file.
-struct FileReproducerStream : public PassManager::ReproducerStream {
+struct FileReproducerStream : public mlir::ReproducerStream {
   FileReproducerStream(std::unique_ptr<llvm::ToolOutputFile> outputFile)
       : outputFile(std::move(outputFile)) {}
   ~FileReproducerStream() override { outputFile->keep(); }
@@ -420,20 +447,26 @@ LogicalResult PassManager::runWithCrashRecovery(Operation *op,
 
 void PassManager::enableCrashReproducerGeneration(StringRef outputFile,
                                                   bool genLocalReproducer) {
+  enableCrashReproducerGeneration(
+      makeReproducerStreamFactory(outputFile),
+      genLocalReproducer);
+}
+namespace mlir {
+ReproducerStreamFactory makeReproducerStreamFactory(StringRef outputFile) {
   // Capture the filename by value in case outputFile is out of scope when
   // invoked.
   std::string filename = outputFile.str();
-  enableCrashReproducerGeneration(
-      [filename](std::string &error) -> std::unique_ptr<ReproducerStream> {
-        std::unique_ptr<llvm::ToolOutputFile> outputFile =
-            mlir::openOutputFile(filename, &error);
-        if (!outputFile) {
-          error = "Failed to create reproducer stream: " + error;
-          return nullptr;
-        }
-        return std::make_unique<FileReproducerStream>(std::move(outputFile));
-      },
-      genLocalReproducer);
+  return
+    [filename](std::string &error) -> std::unique_ptr<ReproducerStream> {
+      std::unique_ptr<llvm::ToolOutputFile> outputFile =
+          mlir::openOutputFile(filename, &error);
+      if (!outputFile) {
+        error = "Failed to create reproducer stream: " + error;
+        return nullptr;
+      }
+      return std::make_unique<FileReproducerStream>(std::move(outputFile));
+    };
+}
 }
 
 void PassManager::enableCrashReproducerGeneration(
diff --git a/mlir/lib/Pass/PassDetail.h b/mlir/lib/Pass/PassDetail.h
index 0e964b6d6d36bc..830997079c8993 100644
--- a/mlir/lib/Pass/PassDetail.h
+++ b/mlir/lib/Pass/PassDetail.h
@@ -99,7 +99,7 @@ class OpToOpPassAdaptor
 class PassCrashReproducerGenerator {
 public:
   PassCrashReproducerGenerator(
-      PassManager::ReproducerStreamFactory &streamFactory,
+      mlir::ReproducerStreamFactory &streamFactory,
       bool localReproducer);
   ~PassCrashReproducerGenerator();
 
diff --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
index d7d47619ef4ac9..afd644fa53c049 100644
--- a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
+++ b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
@@ -324,6 +324,15 @@ static LogicalResult doVerifyRoundTrip(Operation *op,
   return doVerifyRoundTrip(op, config, /*useBytecode=*/true);
 }
 
+llvm::cl::opt<std::string> reproducerNonCrashFile{
+     "mlir-reproducer-filename",
+     llvm::cl::desc("Generate a .mlir reproducer file at the given output path"
+                    " if the pass manager crashes or fails")};
+llvm::cl::opt<bool> generateNonCrashReproducer{
+     "mlir-generate-reproducer",
+     llvm::cl::desc("Generating a reproducer even if a crash did not occur "),
+     llvm::cl::init(false)};
+
 /// Perform the actions on the input file indicated by the command line flags
 /// within the specified context.
 ///
@@ -384,6 +393,13 @@ performActions(raw_ostream &os,
   if (failed(pm.run(*op)))
     return failure();
 
+  // Generate reproducers if requested
+  for (auto &pass : pm.getPasses()) {
+    if (!generateNonCrashReproducer)
+      break;
+    makeReproducer(&pass, op.get(), reproducerNonCrashFile);
+  }
+
   // Print the output.
   TimingScope outputTiming = timing.nest("Output");
   if (config.shouldEmitBytecode()) {
diff --git a/mlir/test/Pass/crashless-reproducer.mlir b/mlir/test/Pass/crashless-reproducer.mlir
new file mode 100644
index 00000000000000..e34efaeace89ea
--- /dev/null
+++ b/mlir/test/Pass/crashless-reproducer.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-opt %s -pass-pipeline='builtin.module(builtin.module(test-module-pass))' \
+// RUN:   -mlir-pass-pipeline-crash-reproducer=%t \
+// RUN:   -mlir-pass-pipeline-always-generate-reproducer=true -verify-diagnostics
+
+// RUN: cat %t | FileCheck -check-prefix=REPRO %s
+
+module @inner_mod1 {
+  module @foo {}
+}
+
+// REPRO: module @inner_mod1
+// REPRO: module @foo {
+// REPRO: pipeline: "builtin.module(builtin.module(test-module-pass))"



More information about the Mlir-commits mailing list