[mlir] Enable passing crash reproducer stream factory method

[jpienaar]

Summary

Add factory to create streams for logging the reproducer. Allows for more general logging (beyond file) and logging the configuration/module separately (logged in order, configuration before module).

[mehdi_amini]

mlir/include/mlir/Pass/PassManager.h

213

Could we reduce the API surface?
I feel like like close() could be handled by the dtor.
A single stream could be handle by one method: Optional<raw_ostream> getStream().

mlir/lib/Pass/Pass.cpp

709

Could we reduce the API surface?
I feel like

752–754

(Spell the type)

757

Making this configurable is breaking the mlir-opt option to load a reproducer: the reproducer is no longer a self-contained "known" format, do you really need this hook? Can we always print to the same stream instead?

[jpienaar]

mlir/include/mlir/Pass/PassManager.h

213

close is required by ToolOutput today to differentiate between completed and not completed dumps. It is not required if one just streams without needing any cleanup - but then one needs some checking later to verify if a failure got dumped. Or are you suggesting best effort, just always assume it will complete and other parts (e.g. if file is not saved in the end, just pass the name and assume consumers are robust enough).

So would you propose that getStream opens the stream too and return none if failed? (not sure why optional)

I considered a single stream, but there seems to be 2 pieces of information being captured here: pass pipeline and module/top-level object, we mix them by using comments but that means you need to code a convention and re-extract. Which is needed if we got to file, but if stored separately then the convention doesn't help. I don't feel super strong here, as extracting the string out is easy. An option is also to not make this pure virtual and so the common case is you are reusing the same stream.

mlir/lib/Pass/Pass.cpp

757

It is self-contained if you use a filename :) it isn't if you log it to (say) a DB, but for that case you would need to do work. But you are correct: if you now implement a different file logger using the specialization, then you'd need to duplicate the same conventions.

[rdzhabarov]

mlir/include/mlir/Pass/PassManager.h

200

StringRef?

mlir/lib/Pass/Pass.cpp

705

private this and the outputFile.

746

nit: file reproducer stream.

751–752

how error is used?

[mehdi_amini]

mlir/include/mlir/Pass/PassManager.h

213

close is required by ToolOutput today to differentiate between completed and not completed dumps.

Is there a code path today where we discard the output? It seems to me that we only invoke the factory to create the ReproducerStream after we decide we want to log in the first place, i.e. close() is called unconditionally which drove my comment here.

So would you propose that getStream opens the stream too and return none if failed? (not sure why optional)

Yes, if you want to say that this can fail. Another way would be to make getStream not able to fail and rely on the factory to return nullptr.

I considered a single stream, but there seems to be 2 pieces of information being captured here: pass pipeline and module/top-level object, we mix them by using comments but that means you need to code a convention and re-extract.

Right, and we implemented this convention at this point I think with mlir-opt able to load a reproducer.
I'd just be in favor of YAGNI and KIS here :)

[rriddle]

mlir/include/mlir/Pass/PassManager.h

213

Is there a code path today where we discard the output? It seems to me that we only invoke the factory to create the ReproducerStream after we decide we want to log in the first place, i.e. close() is called unconditionally which drove my comment here.
Yes, if you want to say that this can fail. Another way would be to make getStream not able to fail and rely on the factory to return nullptr.

+1 to both of these points. I don't see a reason for an explicit close, the destructor of the stream can just handle that.

Right, and we implemented this convention at this point I think with mlir-opt able to load a reproducer.

+1 here as well. Unless there is a concrete and compelling reason, I'd rather just have one stream.

Does this struct need anything other than just:

virtual std::string description() = 0;
virtual raw_ostream &os() = 0;

[jpienaar]

mlir/include/mlir/Pass/PassManager.h

213

Is there a code path today where we discard the output?

Yes, if the process gets interrupted. So the assumption would be that no matter what we have outputted to stream successfully. And any post processing is fine (e.g., any corruption of the stream is acceptable, atomicity is assumed).

Does this struct need anything other than just:

virtual std::string description() = 0;
virtual raw_ostream &os() = 0;

Opening the stream may be expensive, but can be done in the factory. It needs no more than that if we flatten the information using a convention and assume commiting the written output on destruction is always fine.

mlir/lib/Pass/Pass.cpp

751–752

That is the error being reported.

[rdzhabarov]

mlir/lib/Pass/Pass.cpp

751–752

I do not see it used anywhere in the method =)

[mehdi_amini]

mlir/include/mlir/Pass/PassManager.h

213

Yes, if the process gets interrupted.

You mean if the process crash while printing? during the first line here:

// Output the .mlir module.
preCrashOperation->print(crashStream->operationStream());
crashStream->close();

In this case we won't call the destructor either, so it should be equivalent isn't it?
You can slightly change the client code to make sure the crashStream is destroyed right at the point you call close here instead, if you really want to be strict about it.

[jpienaar]

Exposing simpler API

[jpienaar]

Fix error return string not piped through.

[jpienaar]

Spell out type & return error rather than report locally

mlir/lib/Pass/Pass.cpp

751–752

Ah, good catch thanks

[rdzhabarov]

nice!

mlir/lib/Pass/Pass.cpp

689

const std::string&

747

comment seems outdated.

[mehdi_amini]

LGTM with a minor tweak on FileReproducerStream

mlir/lib/Pass/Pass.cpp

702

I'd rather patch ToolOutputFile to get access to this (it is already stored there!):

diff --git a/llvm/include/llvm/Support/ToolOutputFile.h b/llvm/include/llvm/Support/ToolOutputFile.h
index cf01b9ecefc5..441aa5881fb8 100644
--- a/llvm/include/llvm/Support/ToolOutputFile.h
+++ b/llvm/include/llvm/Support/ToolOutputFile.h
@@ -34,7 +34,7 @@ class ToolOutputFile {
   public:
     /// The flag which indicates whether we should not delete the file.
     bool Keep;
-
+    StringRef GetFilename() { return Filename; }
     explicit CleanupInstaller(StringRef Filename);
     ~CleanupInstaller();
   } Installer;
@@ -57,6 +57,9 @@ public:
   /// Return the contained raw_fd_ostream.
   raw_fd_ostream &os() { return *OS; }
 
+  /// Return the filename it was initialized with.
+  StringRef GetFilename() { return Installer.GetFilename(); }
+
   /// Indicate that the tool's job wrt this output file has been successful and
   /// the file should not be deleted.
   void keep() { Installer.Keep = true; }

[rriddle]

LGTM with Mehdi's comments addressed.

Please update the documentation on the website detailing on how users can hook into this.

[jpienaar]

Adding short usage information and updating ToolOutputFile to query name

[jpienaar]

Update make site and config prefix

llvm/include/llvm/Support/ToolOutputFile.h

class ToolOutputFile {
   /// uninstalls them in its destructor.
   class CleanupInstaller {
     /// The name of the file.
     std::string Filename;
   public:
     /// The flag which indicates whether we should not delete the file.
     bool Keep;
 
+    StringRef getFilename() { return Filename; }
     explicit CleanupInstaller(StringRef Filename);
     ~CleanupInstaller();
   } Installer;
 
   /// Storage for the stream, if we're owning our own stream. This is
   /// intentionally declared after Installer.
   Optional<raw_fd_ostream> OSHolder;
 
   /// The actual stream to use.
   raw_fd_ostream *OS;
 
 public:
   /// This constructor's arguments are passed to raw_fd_ostream's
   /// constructor.
   ToolOutputFile(StringRef Filename, std::error_code &EC,
                  sys::fs::OpenFlags Flags);
 
   ToolOutputFile(StringRef Filename, int FD);
 
   /// Return the contained raw_fd_ostream.
   raw_fd_ostream &os() { return *OS; }
 
+  /// Return the filename initialized with.
+  StringRef getFilename() { return Installer.getFilename(); }
+
   /// Indicate that the tool's job wrt this output file has been successful and
   /// the file should not be deleted.
   void keep() { Installer.Keep = true; }
 };
 
 } // end llvm namespace
 
 #endif

mlir/docs/PassManagement.md

 `PassManager::enableCrashReproducerGeneration` or via the command line flag
 `pass-pipeline-crash-reproducer`. In either case, an argument is provided that
 corresponds to the output `.mlir` file name that the reproducible should be
 written to. The reproducible contains the configuration of the pass manager that
 was executing, as well as the initial IR before any passes were run. A potential
 reproducible may have the form:
 
 ```mlir
-// configuration: -pass-pipeline='func(cse,canonicalize),inline'
-// note: verifyPasses=false
+// configuration: -pass-pipeline='func(cse,canonicalize),inline' -verify-each
 
 module {
   func @foo() {
     ...
   }
 }
 ```
 
 The configuration dumped can be passed to `mlir-opt` by specifying
 `-run-reproducer` flag. This will result in parsing the first line configuration
 of the reproducer and adding those to the command line options.
 
+Beyond specifying a filename, one can also register a `ReproducerStreamFactory`
+function that would be invoked in the case of a crash and the reproducer written
+to its stream.
+
 ### 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

mlir/include/mlir/Pass/PassManager.h

public:
   /// 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. 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);
 
+  /// 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
+  /// pass manager will attempt to generate a local reproducer that contains the
+  /// smallest pipeline.
+  void enableCrashReproducerGeneration(ReproducerStreamFactory factory,
+                                       bool genLocalReproducer = false);
+
   /// Runs the verifier after each individual pass.
   void enableVerifier(bool enabled = true);
 
   //===--------------------------------------------------------------------===//
   // Instrumentations
   //===--------------------------------------------------------------------===//
 
   /// Add the provided instrumentation to the pass manager.
private:
   MLIRContext *context;
 
   /// 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;
+  /// An optional factory to use when generating a crash reproducer if valid.
+  ReproducerStreamFactory crashReproducerStreamFactory;
 
   /// 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;
 
   /// A flag that indicates if the IR should be verified in between passes.

mlir/lib/Pass/Pass.cpp

 //===----------------------------------------------------------------------===//
 
 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,
-                            Operation *op, StringRef filename,
+                            Operation *op,
+                            PassManager::ReproducerStreamFactory &crashStream,
                             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 operation representing the IR before the crash.
   Operation *preCrashOperation;
 
-  /// The filename to use when generating the reproducer.
-  StringRef filename;
+  /// The factory for the reproducer output stream to use when generating the
+  /// reproducer.
+  PassManager::ReproducerStreamFactory &crashStreamFactory;
 
   /// Various pass manager and context 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;
 };
+
+/// Instance of ReproducerStream backed by file.
+struct FileReproducerStream : public PassManager::ReproducerStream {
+  FileReproducerStream(std::unique_ptr<llvm::ToolOutputFile> outputFile)
+      : outputFile(std::move(outputFile)) {}
+  ~FileReproducerStream() override;
+
+  /// Description of the reproducer stream.
+  StringRef description() override;
+
+  /// Stream on which to output reprooducer.
+  raw_ostream &os() override;
+
+private:
+  /// ToolOutputFile corresponding to opened `filename`.
+  std::unique_ptr<llvm::ToolOutputFile> outputFile = nullptr;
+};
+
 } // 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, Operation *op,
-    StringRef filename, bool disableThreads, bool verifyPasses)
-    : preCrashOperation(op->clone()), filename(filename),
+    PassManager::ReproducerStreamFactory &crashStreamFactory,
+    bool disableThreads, bool verifyPasses)
+    : preCrashOperation(op->clone()), crashStreamFactory(crashStreamFactory),
       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.
RecoveryReproducerContext::~RecoveryReproducerContext() {
   preCrashOperation->erase();
 
   llvm::sys::SmartScopedLock<true> producerLock(*reproducerMutex);
   reproducerSet->remove(this);
   if (reproducerSet->empty())
     llvm::CrashRecoveryContext::Disable();
 }
 
+/// Description of the reproducer stream.
+StringRef FileReproducerStream::description() {
+  return outputFile->getFilename();
+}
+
+/// Stream on which to output reproducer.
+raw_ostream &FileReproducerStream::os() { return outputFile->os(); }
+
+FileReproducerStream::~FileReproducerStream() { outputFile->keep(); }
+
 LogicalResult RecoveryReproducerContext::generate(std::string &error) {
-  std::unique_ptr<llvm::ToolOutputFile> outputFile =
-      mlir::openOutputFile(filename, &error);
-  if (!outputFile)
+  std::unique_ptr<PassManager::ReproducerStream> crashStream =
+      crashStreamFactory(error);
+  if (!crashStream)
     return failure();
-  auto &outputOS = outputFile->os();
 
   // Output the current pass manager configuration.
-  outputOS << "// configuration: -pass-pipeline='" << pipeline << "'";
+  auto &os = crashStream->os();
+  os << "// configuration: -pass-pipeline='" << pipeline << "'";
   if (disableThreads)
-    outputOS << " -mlir-disable-threading";
-
-  // TODO: Should this also be configured with a pass manager flag?
-  outputOS << "\n// note: verifyPasses=" << (verifyPasses ? "true" : "false")
-           << "\n";
+    os << " -mlir-disable-threading";
+  if (verifyPasses)
+    os << " -verify-each";
+  os << '\n';
 
   // Output the .mlir module.
-  preCrashOperation->print(outputOS);
-  outputFile->keep();
+  preCrashOperation->print(os);
+
+  bool shouldPrintOnOp =
+      preCrashOperation->getContext()->shouldPrintOpOnDiagnostic();
+  preCrashOperation->getContext()->printOpOnDiagnostic(false);
+  preCrashOperation->emitError()
+      << "A failure has been detected while processing the MLIR module, a "
+         "reproducer has been generated in '"
+      << crashStream->description() << "'";
+  preCrashOperation->getContext()->printOpOnDiagnostic(shouldPrintOnOp);
   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;
if (failed(runWithCrashRecovery(pass, op, am)))
       return failure();
   return success();
 }
 
 /// Run the given passes with crash recover enabled.
 LogicalResult
 PassManager::runWithCrashRecovery(MutableArrayRef<std::unique_ptr<Pass>> passes,
                                   Operation *op, AnalysisManager am) {
-  RecoveryReproducerContext context(passes, op, *crashReproducerFileName,
+  RecoveryReproducerContext context(passes, op, crashReproducerStreamFactory,
                                     !getContext()->isMultithreadingEnabled(),
                                     verifyPasses);
 
   // Safely invoke the passes within a recovery context.
   LogicalResult passManagerResult = failure();
   llvm::CrashRecoveryContext recoveryContext;
   recoveryContext.RunSafelyOnThread([&] {
     for (std::unique_ptr<Pass> &pass : passes)
       if (failed(OpToOpPassAdaptor::run(pass.get(), op, am, verifyPasses,
                                         impl->initializationGeneration)))
         return;
     passManagerResult = success();
   });
   if (succeeded(passManagerResult))
     return success();
 
   std::string error;
   if (failed(context.generate(error)))
     return op->emitError("<MLIR-PassManager-Crash-Reproducer>: ") << error;
-  bool shouldPrintOnOp = op->getContext()->shouldPrintOpOnDiagnostic();
-  op->getContext()->printOpOnDiagnostic(false);
-  op->emitError()
-      << "A failure has been detected while processing the MLIR module, a "
-         "reproducer has been generated in '"
-      << *crashReproducerFileName << "'";
-  op->getContext()->printOpOnDiagnostic(shouldPrintOnOp);
   return failure();
 }
 
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//
 
 PassManager::PassManager(MLIRContext *ctx, Nesting nesting,
LogicalResult PassManager::run(Operation *op) {
   ModuleAnalysisManager am(op, instrumentor.get());
 
   // Notify the context that we start running a pipeline for book keeping.
   context->enterMultiThreadedExecution();
 
   // If reproducer generation is enabled, run the pass manager with crash
   // handling enabled.
   LogicalResult result =
-      crashReproducerFileName
+      crashReproducerStreamFactory
           ? runWithCrashRecovery(op, am)
           : OpToOpPassAdaptor::runPipeline(getPasses(), op, am, verifyPasses,
                                            impl->initializationGeneration);
 
   // Notify the context that the run is done.
   context->exitMultiThreadedExecution();
 
   // Dump all of the pass statistics if necessary.
   if (passStatisticsMode)
     dumpStatistics();
   return result;
 }
 
 /// 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. 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);
+  // 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);
+}
+
+/// 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
+/// pass manager will attempt to generate a local reproducer that contains the
+/// smallest pipeline.
+void PassManager::enableCrashReproducerGeneration(
+    ReproducerStreamFactory factory, bool genLocalReproducer) {
+  crashReproducerStreamFactory = factory;
   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>();