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

Puyan Lotfi llvmlistbot at llvm.org
Fri Dec 29 13:30:54 PST 2023


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

>From 4066275cd052824a15a42c2188028835807051f7 Mon Sep 17 00:00:00 2001
From: Puyan Lotfi <puyan at puyan.org>
Date: Fri, 22 Dec 2023 13:08:50 -0800
Subject: [PATCH 1/2] [NFC][mlir][Pass] Reorganize MLIR Crash Recovery
 Reproducer implementation

This patch decouples the code that handles generation of an MLIR reproducer
from the crash recovery portion. The purpose is to allow for generating
reproducers outside of the context of a compiler crash.

(cherry picked from commit bfb0682f2a318a84a05f91466757a03d68ddfd1f)
---
 mlir/include/mlir/Pass/PassManager.h |  30 ++++----
 mlir/lib/Pass/Pass.cpp               |  19 +++--
 mlir/lib/Pass/PassCrashRecovery.cpp  | 101 +++++++++++++++------------
 mlir/lib/Pass/PassDetail.h           |   5 +-
 4 files changed, 88 insertions(+), 67 deletions(-)

diff --git a/mlir/include/mlir/Pass/PassManager.h b/mlir/include/mlir/Pass/PassManager.h
index d5f1ea0fe0350d..6aaa518a5cab17 100644
--- a/mlir/include/mlir/Pass/PassManager.h
+++ b/mlir/include/mlir/Pass/PassManager.h
@@ -207,6 +207,21 @@ 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)>;
+
 /// The main pass manager and pipeline builder.
 class PassManager : public OpPassManager {
 public:
@@ -243,21 +258,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/Pass.cpp b/mlir/lib/Pass/Pass.cpp
index 810d6a357d52c0..9fc9f4714b2653 100644
--- a/mlir/lib/Pass/Pass.cpp
+++ b/mlir/lib/Pass/Pass.cpp
@@ -382,16 +382,23 @@ StringRef OpPassManager::getOpAnchorName() const {
 
 /// Prints out the passes of the pass manager as the textual representation
 /// of pipelines.
-void OpPassManager::printAsTextualPipeline(raw_ostream &os) const {
-  os << getOpAnchorName() << "(";
+void printAsTextualPipeline(
+    raw_ostream &os, StringRef anchorName,
+    const llvm::iterator_range<OpPassManager::pass_iterator> &passes) {
+  os << anchorName << "(";
   llvm::interleave(
-      impl->passes,
-      [&](const std::unique_ptr<Pass> &pass) {
-        pass->printAsTextualPipeline(os);
-      },
+      passes, [&](mlir::Pass &pass) { pass.printAsTextualPipeline(os); },
       [&]() { os << ","; });
   os << ")";
 }
+void OpPassManager::printAsTextualPipeline(raw_ostream &os) const {
+  StringRef anchorName = getOpAnchorName();
+  auto *begin = MutableArrayRef<std::unique_ptr<Pass>>{impl->passes}.begin();
+  auto *end = MutableArrayRef<std::unique_ptr<Pass>>{impl->passes}.end();
+  const llvm::iterator_range<OpPassManager::pass_iterator> passes = {begin,
+                                                                     end};
+  ::printAsTextualPipeline(os, anchorName, passes);
+}
 
 void OpPassManager::dump() {
   llvm::errs() << "Pass Manager with " << impl->passes.size() << " passes:\n";
diff --git a/mlir/lib/Pass/PassCrashRecovery.cpp b/mlir/lib/Pass/PassCrashRecovery.cpp
index df1a0762ae34a5..a26be2dabaa6d9 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()),
@@ -106,22 +106,24 @@ RecoveryReproducerContext::~RecoveryReproducerContext() {
   disable();
 }
 
-void RecoveryReproducerContext::generate(std::string &description) {
+static void 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() + "(" +
-                          pipelineElements + ")")
-                             .str();
-  AsmState state(preCrashOperation);
+  std::string pipeline =
+      (op->getName().getStringRef() + "(" + pipelineElements + ")").str();
+  AsmState state(op);
   state.attachResourcePrinter(
       "mlir_reproducer", [&](Operation *op, AsmResourceBuilder &builder) {
         builder.buildString("pipeline", pipeline);
@@ -130,7 +132,12 @@ void RecoveryReproducerContext::generate(std::string &description) {
       });
 
   // Output the .mlir module.
-  preCrashOperation->print(stream->os(), state);
+  op->print(stream->os(), state);
+}
+
+void RecoveryReproducerContext::generate(std::string &description) {
+  appendReproducer(description, preCrashOperation, streamFactory,
+                   pipelineElements, disableThreads, verifyPasses);
 }
 
 void RecoveryReproducerContext::disable() {
@@ -175,12 +182,11 @@ void RecoveryReproducerContext::registerSignalHandler() {
 //===----------------------------------------------------------------------===//
 
 struct PassCrashReproducerGenerator::Impl {
-  Impl(PassManager::ReproducerStreamFactory &streamFactory,
-       bool localReproducer)
+  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 +204,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;
 
@@ -283,19 +289,7 @@ void PassCrashReproducerGenerator::finalize(Operation *rootOp,
   impl->runningPasses.clear();
 }
 
-void PassCrashReproducerGenerator::prepareReproducerFor(Pass *pass,
-                                                        Operation *op) {
-  // If not tracking local reproducers, we simply remember that this pass is
-  // running.
-  impl->runningPasses.insert(std::make_pair(pass, op));
-  if (!impl->localReproducer)
-    return;
-
-  // Disable the current pass recovery context, if there is one. This may happen
-  // in the case of dynamic pass pipelines.
-  if (!impl->activeContexts.empty())
-    impl->activeContexts.back()->disable();
-
+static 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 +307,25 @@ void PassCrashReproducerGenerator::prepareReproducerFor(Pass *pass,
   for (unsigned i = 0, e = scopes.size(); i < e; ++i)
     passOS << ")";
 
+  return passOS.str();
+}
+
+void PassCrashReproducerGenerator::prepareReproducerFor(Pass *pass,
+                                                        Operation *op) {
+  // If not tracking local reproducers, we simply remember that this pass is
+  // running.
+  impl->runningPasses.insert(std::make_pair(pass, op));
+  if (!impl->localReproducer)
+    return;
+
+  // Disable the current pass recovery context, if there is one. This may happen
+  // in the case of dynamic pass pipelines.
+  if (!impl->activeContexts.empty())
+    impl->activeContexts.back()->disable();
+
+  std::string pipelineStr = produceTextualPipelineWithScope(pass, op);
   impl->activeContexts.push_back(std::make_unique<RecoveryReproducerContext>(
-      passOS.str(), op, impl->streamFactory, impl->pmFlagVerifyPasses));
+      pipelineStr, op, impl->streamFactory, impl->pmFlagVerifyPasses));
 }
 void PassCrashReproducerGenerator::prepareReproducerFor(
     iterator_range<PassManager::pass_iterator> passes, Operation *op) {
@@ -382,9 +393,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(); }
@@ -418,22 +429,26 @@ LogicalResult PassManager::runWithCrashRecovery(Operation *op,
   return passManagerResult;
 }
 
-void PassManager::enableCrashReproducerGeneration(StringRef outputFile,
-                                                  bool genLocalReproducer) {
+static 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(StringRef outputFile,
+                                                  bool genLocalReproducer) {
+  enableCrashReproducerGeneration(makeReproducerStreamFactory(outputFile),
+                                  genLocalReproducer);
 }
 
 void PassManager::enableCrashReproducerGeneration(
diff --git a/mlir/lib/Pass/PassDetail.h b/mlir/lib/Pass/PassDetail.h
index 0e964b6d6d36bc..164955a1e32b1b 100644
--- a/mlir/lib/Pass/PassDetail.h
+++ b/mlir/lib/Pass/PassDetail.h
@@ -98,9 +98,8 @@ class OpToOpPassAdaptor
 
 class PassCrashReproducerGenerator {
 public:
-  PassCrashReproducerGenerator(
-      PassManager::ReproducerStreamFactory &streamFactory,
-      bool localReproducer);
+  PassCrashReproducerGenerator(mlir::ReproducerStreamFactory &streamFactory,
+                               bool localReproducer);
   ~PassCrashReproducerGenerator();
 
   /// Initialize the generator in preparation for reproducer generation. The

>From 67130c57606c3fd601e7ae74274ef1c71f59d0b0 Mon Sep 17 00:00:00 2001
From: Puyan Lotfi <puyan at puyan.org>
Date: Fri, 22 Dec 2023 14:04:03 -0800
Subject: [PATCH 2/2] [mlir][Pass] Enable the option for reproducer generation
 without crashing

This patch adds API `makeReproducer` and cl::opt flags
--mlir-reproducer-filename and --mlir-generate-reproducer=1 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 issues using offline tools,
where being able to dump the reproducer from a runtime compiler would be
helpful.

(cherry picked from commit fe0f91258a195388bfa150cea9a45e588e929173)
---
 mlir/include/mlir/Pass/PassManager.h          |  6 +++++
 .../include/mlir/Tools/mlir-opt/MlirOptMain.h | 10 +++++++++
 mlir/lib/Pass/PassCrashRecovery.cpp           | 19 ++++++++++++++++
 mlir/lib/Tools/mlir-opt/MlirOptMain.cpp       | 22 +++++++++++++++++++
 mlir/test/Pass/crashless-reproducer.mlir      | 13 +++++++++++
 5 files changed, 70 insertions(+)
 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 6aaa518a5cab17..1b2e6a3bc82bb4 100644
--- a/mlir/include/mlir/Pass/PassManager.h
+++ b/mlir/include/mlir/Pass/PassManager.h
@@ -222,6 +222,12 @@ struct ReproducerStream {
 using ReproducerStreamFactory =
     std::function<std::unique_ptr<ReproducerStream>(std::string &error)>;
 
+std::string
+makeReproducer(StringRef anchorName,
+               const llvm::iterator_range<OpPassManager::pass_iterator> &passes,
+               Operation *op, StringRef outputFile, bool disableThreads = false,
+               bool verifyPasses = false);
+
 /// The main pass manager and pipeline builder.
 class PassManager : public OpPassManager {
 public:
diff --git a/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h b/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
index e255d9fa70b659..1be9c46740915e 100644
--- a/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
+++ b/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
@@ -173,6 +173,10 @@ class MlirOptMainConfig {
   }
   bool shouldVerifyRoundtrip() const { return verifyRoundtripFlag; }
 
+  /// Reproducer file generation (no crash required).
+  StringRef getReproduerFilename() const { return reproduerFilenameFlag; }
+  bool shouldGenerateReproducer() const { return generateReproducerFlag; }
+
 protected:
   /// Allow operation with no registered dialects.
   /// This option is for convenience during testing only and discouraged in
@@ -228,6 +232,12 @@ class MlirOptMainConfig {
 
   /// Verify that the input IR round-trips perfectly.
   bool verifyRoundtripFlag = false;
+
+  /// The reproducer output filename (no crash required).
+  std::string reproduerFilenameFlag = "";
+
+  /// Generate a reproducer file at reproduerFilename (no crash required).
+  bool generateReproducerFlag = false;
 };
 
 /// This defines the function type used to setup the pass manager. This can be
diff --git a/mlir/lib/Pass/PassCrashRecovery.cpp b/mlir/lib/Pass/PassCrashRecovery.cpp
index a26be2dabaa6d9..6950ef513c2195 100644
--- a/mlir/lib/Pass/PassCrashRecovery.cpp
+++ b/mlir/lib/Pass/PassCrashRecovery.cpp
@@ -445,6 +445,25 @@ makeReproducerStreamFactory(StringRef outputFile) {
   };
 }
 
+void printAsTextualPipeline(
+    raw_ostream &os, StringRef anchorName,
+    const llvm::iterator_range<OpPassManager::pass_iterator> &passes);
+
+std::string mlir::makeReproducer(
+    StringRef anchorName,
+    const llvm::iterator_range<OpPassManager::pass_iterator> &passes,
+    Operation *op, StringRef outputFile, bool disableThreads,
+    bool verifyPasses) {
+
+  std::string description;
+  std::string pipelineStr;
+  llvm::raw_string_ostream passOS(pipelineStr);
+  ::printAsTextualPipeline(passOS, anchorName, passes);
+  appendReproducer(description, op, makeReproducerStreamFactory(outputFile),
+                   pipelineStr, disableThreads, verifyPasses);
+  return description;
+}
+
 void PassManager::enableCrashReproducerGeneration(StringRef outputFile,
                                                   bool genLocalReproducer) {
   enableCrashReproducerGeneration(makeReproducerStreamFactory(outputFile),
diff --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
index d7d47619ef4ac9..2d688f9294ee6a 100644
--- a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
+++ b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
@@ -151,6 +151,21 @@ struct MlirOptMainConfigCLOptions : public MlirOptMainConfig {
 
     static cl::list<std::string> passPlugins(
         "load-pass-plugin", cl::desc("Load passes from plugin library"));
+
+    static cl::opt<std::string, /*ExternalStorage=*/true> reproduerFilename(
+        "mlir-reproducer-filename",
+        llvm::cl::desc("Location of the .mlir reproducer output file path"
+                       " the --mlir-generate-reproducer=1"
+                       " (no crash required)"),
+        cl::location(reproduerFilenameFlag), cl::init(""),
+        cl::value_desc("filename"));
+    static cl::opt<bool, /*ExternalStorage=*/true> generateReproducer(
+        "mlir-generate-reproducer",
+        llvm::cl::desc("Generate a reproducer at"
+                       " --mlir-reproducer-filename=<filename> "
+                       " (no crash required)"),
+        cl::location(generateReproducerFlag), cl::init(false));
+
     /// Set the callback to load a pass plugin.
     passPlugins.setCallback([&](const std::string &pluginPath) {
       auto plugin = PassPlugin::load(pluginPath);
@@ -384,6 +399,13 @@ performActions(raw_ostream &os,
   if (failed(pm.run(*op)))
     return failure();
 
+  // Generate reproducers if requested
+  if (config.shouldGenerateReproducer()) {
+    StringRef anchorName = pm.getAnyOpAnchorName();
+    const auto &passes = pm.getPasses();
+    makeReproducer(anchorName, passes, op.get(), config.getReproduerFilename());
+  }
+
   // 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..34092ca94afa05
--- /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-reproducer-filename=%t \
+// RUN:   --mlir-generate-reproducer=1 -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(any(builtin.module(test-module-pass)))"



More information about the Mlir-commits mailing list