[mlir] Refactor pass crash reproducer generation to be an assembly resource

We currently generate reproducer configurations using a comment placed at
the top of the generated .mlir file. This is kind of hacky given that comments
have no semantic context in the source file and can easily be dropped. This
strategy also wouldn't work if/when we have a bitcode format. This commit
switches to using an external assembly resource, which is verifiable/can
work with a hypothetical bitcode naturally/and removes the awkward processing
from mlir-opt for splicing comments and re-applying command line options. With
the removal of command line munging, this opens up new possibilities for
executing reproducers in memory.

Differential Revision: https://reviews.llvm.org/D126447

mlir/include/mlir/Pass/PassRegistry.h

@@ -21,7 +21,9 @@
 
 namespace mlir {
 class OpPassManager;
+class ParserConfig;
 class Pass;
+class PassManager;
 
 namespace detail {
 class PassOptions;
@@ -272,6 +274,18 @@ private:
   std::unique_ptr<detail::PassPipelineCLParserImpl> impl;
 };
 
+//===----------------------------------------------------------------------===//
+// Pass Reproducer
+//===----------------------------------------------------------------------===//
+
+/// Attach an assembly resource parser that handles MLIR reproducer
+/// configurations. Any found reproducer information will be attached to the
+/// given pass manager, e.g. the reproducer pipeline, verification flags, etc.
+// FIXME: Remove the `enableThreading` flag when possible. Some tools, e.g.
+//        mlir-opt, force disable threading during parsing.
+void attachPassReproducerAsmResource(ParserConfig &config, PassManager &pm,
+                                     bool &enableThreading);
+
 } // namespace mlir
 
 #endif // MLIR_PASS_PASSREGISTRY_H_

mlir/lib/Pass/PassCrashRecovery.cpp

@@ -11,6 +11,7 @@
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/IR/Verifier.h"
+#include "mlir/Parser/Parser.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/FileUtilities.h"
 #include "llvm/ADT/STLExtras.h"
@@ -117,17 +118,16 @@ void RecoveryReproducerContext::generate(std::string &description) {
   }
   descOS << "reproducer generated at `" << stream->description() << "`";
 
-  // Output the current pass manager configuration to the crash stream.
-  auto &os = stream->os();
-  os << "// configuration: -pass-pipeline='" << pipeline << "'";
-  if (disableThreads)
-    os << " -mlir-disable-threading";
-  if (verifyPasses)
-    os << " -verify-each";
-  os << '\n';
+  AsmState state(preCrashOperation);
+  state.attachResourcePrinter(
+      "mlir_reproducer", [&](Operation *op, AsmResourceBuilder &builder) {
+        builder.buildString("pipeline", pipeline);
+        builder.buildBool("disable_threading", disableThreads);
+        builder.buildBool("verify_each", verifyPasses);
+      });
 
   // Output the .mlir module.
-  preCrashOperation->print(os);
+  preCrashOperation->print(stream->os(), state);
 }
 
 void RecoveryReproducerContext::disable() {
@@ -438,3 +438,38 @@ void PassManager::enableCrashReproducerGeneration(
   addInstrumentation(
       std::make_unique<CrashReproducerInstrumentation>(*crashReproGenerator));
 }
+
+//===----------------------------------------------------------------------===//
+// Asm Resource
+//===----------------------------------------------------------------------===//
+
+void mlir::attachPassReproducerAsmResource(ParserConfig &config,
+                                           PassManager &pm,
+                                           bool &enableThreading) {
+  auto parseFn = [&](AsmParsedResourceEntry &entry) -> LogicalResult {
+    if (entry.getKey() == "pipeline") {
+      FailureOr<std::string> pipeline = entry.parseAsString();
+      if (failed(pipeline))
+        return failure();
+      return parsePassPipeline(*pipeline, pm);
+    }
+    if (entry.getKey() == "disable_threading") {
+      FailureOr<bool> value = entry.parseAsBool();
+
+      // FIXME: We should just update the context directly, but some places
+      // force disable threading during parsing.
+      if (succeeded(value))
+        enableThreading = !(*value);
+      return value;
+    }
+    if (entry.getKey() == "verify_each") {
+      FailureOr<bool> value = entry.parseAsBool();
+      if (succeeded(value))
+        pm.enableVerifier(*value);
+      return value;
+    }
+    return entry.emitError() << "unknown 'mlir_reproducer' resource key '"
+                             << entry.getKey() << "'";
+  };
+  config.attachResourceParser("mlir_reproducer", parseFn);
+}

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

@@ -57,20 +57,25 @@ static LogicalResult performActions(raw_ostream &os, bool verifyDiagnostics,
   bool wasThreadingEnabled = context->isMultithreadingEnabled();
   context->disableMultithreading();
 
-  // Parse the input file and reset the context threading state.
-  TimingScope parserTiming = timing.nest("Parser");
-  OwningOpRef<ModuleOp> module(parseSourceFile<ModuleOp>(sourceMgr, context));
-  context->enableMultithreading(wasThreadingEnabled);
-  if (!module)
-    return failure();
-  parserTiming.stop();
-
-  // Apply any pass manager command line options.
+  // Prepare the pass manager and apply any command line options.
   PassManager pm(context, OpPassManager::Nesting::Implicit);
   pm.enableVerifier(verifyPasses);
   applyPassManagerCLOptions(pm);
   pm.enableTiming(timing);
 
+  // Prepare the parser config, and attach any useful/necessary resource
+  // handlers.
+  ParserConfig config(context);
+  attachPassReproducerAsmResource(config, pm, wasThreadingEnabled);
+
+  // Parse the input file and reset the context threading state.
+  TimingScope parserTiming = timing.nest("Parser");
+  OwningOpRef<ModuleOp> module(parseSourceFile<ModuleOp>(sourceMgr, config));
+  context->enableMultithreading(wasThreadingEnabled);
+  if (!module)
+    return failure();
+  parserTiming.stop();
+
   // Callback to build the pipeline.
   if (failed(passManagerSetupFn(pm)))
     return failure();
@@ -219,11 +224,6 @@ LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
       "show-dialects", cl::desc("Print the list of registered dialects"),
       cl::init(false));
 
-  static cl::opt<bool> runRepro(
-      "run-reproducer",
-      cl::desc("Append the command line options of the reproducer"),
-      cl::init(false));
-
   InitLLVM y(argc, argv);
 
   // Register any command line options.
@@ -260,23 +260,6 @@ LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
     return failure();
   }
 
-  // Parse reproducer options.
-  BumpPtrAllocator a;
-  StringSaver saver(a);
-  if (runRepro) {
-    auto pair = file->getBuffer().split('\n');
-    if (!pair.first.consume_front("// configuration:")) {
-      llvm::errs() << "Failed to find repro configuration, expect file to "
-                      "begin with '// configuration:'\n";
-      return failure();
-    }
-    // Tokenize & parse the first line.
-    SmallVector<const char *, 4> newArgv;
-    newArgv.push_back(argv[0]);
-    llvm::cl::TokenizeGNUCommandLine(pair.first, saver, newArgv);
-    cl::ParseCommandLineOptions(newArgv.size(), &newArgv[0], helpHeader);
-  }
-
   auto output = openOutputFile(outputFilename, &errorMessage);
   if (!output) {
     llvm::errs() << errorMessage << "\n";

mlir/test/Pass/crash-recovery-dynamic-failure.mlir

@@ -11,7 +11,8 @@ module @inner_mod1 {
   module @foo {}
 }
 
-// REPRO_LOCAL_DYNAMIC_FAILURE: configuration: -pass-pipeline='builtin.module(test-pass-failure)'
 
 // REPRO_LOCAL_DYNAMIC_FAILURE: module @inner_mod1
 // REPRO_LOCAL_DYNAMIC_FAILURE: module @foo {
+
+// REPRO_LOCAL_DYNAMIC_FAILURE: pipeline: "builtin.module(test-pass-failure)"

mlir/test/Pass/crash-recovery.mlir

@@ -20,17 +20,14 @@ module @inner_mod1 {
   module @foo {}
 }
 
-// REPRO: configuration: -pass-pipeline='builtin.module(test-module-pass,test-pass-crash)'
-
 // REPRO: module @inner_mod1
 // REPRO: module @foo {
-
-// REPRO_LOCAL: configuration: -pass-pipeline='builtin.module(test-pass-crash)'
+// REPRO: pipeline: "builtin.module(test-module-pass,test-pass-crash)"
 
 // REPRO_LOCAL: module @inner_mod1
 // REPRO_LOCAL: module @foo {
-
-// REPRO_LOCAL_DYNAMIC: configuration: -pass-pipeline='builtin.module(test-pass-crash)'
+// REPRO_LOCAL: pipeline: "builtin.module(test-pass-crash)"
 
 // REPRO_LOCAL_DYNAMIC: module @inner_mod1
 // REPRO_LOCAL_DYNAMIC: module @foo {
+// REPRO_LOCAL_DYNAMIC: pipeline: "builtin.module(test-pass-crash)"

mlir/test/Pass/run-reproducer.mlir

@@ -1,9 +1,4 @@
-// configuration: -mlir-disable-threading=true -pass-pipeline='func.func(cse,canonicalize)' -mlir-print-ir-before=cse
-
-// Test of the reproducer run option. The first line has to be the
-// configuration (matching what is produced by reproducer).
-
-// RUN: mlir-opt %s -run-reproducer 2>&1 | FileCheck -check-prefix=BEFORE %s
+// RUN: mlir-opt %s -mlir-print-ir-before=cse 2>&1 | FileCheck -check-prefix=BEFORE %s
 
 func.func @foo() {
   %0 = arith.constant 0 : i32
@@ -14,6 +9,15 @@ func.func @bar() {
   return
 }
 
+{-#
+  external_resources: {
+    mlir_reproducer: {
+      pipeline: "func.func(cse,canonicalize)",
+      disable_threading: true
+    }
+  }
+#-}
+
 // BEFORE: // -----// IR Dump Before{{.*}}CSE //----- //
 // BEFORE-NEXT: func @foo()
 // BEFORE: // -----// IR Dump Before{{.*}}CSE //----- //

mlir/test/lib/Dialect/SPIRV/TestModuleCombiner.cpp

@@ -29,9 +29,6 @@ public:
   TestModuleCombinerPass() = default;
   TestModuleCombinerPass(const TestModuleCombinerPass &) {}
   void runOnOperation() override;
-
-private:
-  OwningOpRef<spirv::ModuleOp> combinedModule;
 };
 } // namespace
 
@@ -46,10 +43,12 @@ void TestModuleCombinerPass::runOnOperation() {
                  << " -> " << newSymbol << "\n";
   };
 
-  combinedModule = spirv::combine(modules, combinedModuleBuilder, listener);
+  OwningOpRef<spirv::ModuleOp> combinedModule =
+      spirv::combine(modules, combinedModuleBuilder, listener);
 
   for (spirv::ModuleOp module : modules)
     module.erase();
+  combinedModule.release();
 }
 
 namespace mlir {