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

[rriddle]

Summary

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.

mlir/include/mlir/Pass/PassRegistry.h

 
 #include "mlir/Pass/PassOptions.h"
 #include "mlir/Support/TypeID.h"
 #include <functional>
 #include <utility>
 
 namespace mlir {
 class OpPassManager;
+class AsmParserConfig;
 class Pass;
+class PassManager;
 
 namespace detail {
 class PassOptions;
 } // namespace detail
 
 /// A registry function that adds passes to the given pass manager. This should
 /// also parse options and return success() if parsing succeeded.
 /// `errorHandler` is a functor used to emit errors during parsing.
public:
   /// Returns true if the given pass registry entry was registered at the
   /// top-level of the parser, i.e. not within an explicit textual pipeline.
   bool contains(const PassRegistryEntry *entry) const;
 
 private:
   std::unique_ptr<detail::PassPipelineCLParserImpl> impl;
 };
 
+//===----------------------------------------------------------------------===//
+// Pass Reproducer
+//===----------------------------------------------------------------------===//
+
+/// Attach an assembly configuration 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 attachPassReproducerAsmConfig(AsmParserConfig &config, PassManager &pm,
+                                   bool &enableThreading);
+
 } // namespace mlir
 
 #endif // MLIR_PASS_PASSREGISTRY_H_

mlir/lib/Pass/PassCrashRecovery.cpp

 //===- PassCrashRecovery.cpp - Pass Crash Recovery Implementation ---------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 #include "mlir/IR/Diagnostics.h"
 #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"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/CrashRecoveryContext.h"
 #include "llvm/Support/Mutex.h"
void RecoveryReproducerContext::generate(std::string &description) {
   std::string error;
   std::unique_ptr<PassManager::ReproducerStream> stream = streamFactory(error);
   if (!stream) {
     descOS << "failed to create output stream: " << error;
     return;
   }
   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.attachConfigPrinter(
+      "mlir_reproducer", [&](Operation *op, AsmConfigBuilder &provider) {
+        provider.buildString("pipeline", pipeline);
+        provider.buildBool("disable_threading", disableThreads);
+        provider.buildBool("verify_each", verifyPasses);
+      });
 
   // Output the .mlir module.
-  preCrashOperation->print(os);
+  preCrashOperation->print(stream->os(), state);
 }
 
 void RecoveryReproducerContext::disable() {
   llvm::sys::SmartScopedLock<true> lock(*reproducerMutex);
   reproducerSet->remove(this);
   if (reproducerSet->empty())
     llvm::CrashRecoveryContext::Disable();
 }
llvm::report_fatal_error(
         "Local crash reproduction can't be setup on a "
         "pass-manager without disabling multi-threading first.");
 
   crashReproGenerator = std::make_unique<PassCrashReproducerGenerator>(
       factory, genLocalReproducer);
   addInstrumentation(
       std::make_unique<CrashReproducerInstrumentation>(*crashReproGenerator));
 }
+
+//===----------------------------------------------------------------------===//
+// Asm Configuration
+//===----------------------------------------------------------------------===//
+
+void mlir::attachPassReproducerAsmConfig(AsmParserConfig &config,
+                                         PassManager &pm,
+                                         bool &enableThreading) {
+  auto parseFn = [&](AsmConfigEntry &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' configuration key '"
+                             << entry.getKey() << "'";
+  };
+  config.attachConfigParser("mlir_reproducer", parseFn);
+}

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

static LogicalResult performActions(raw_ostream &os, bool verifyDiagnostics,
   applyDefaultTimingManagerCLOptions(tm);
   TimingScope timing = tm.getRootScope();
 
   // Disable multi-threading when parsing the input file. This removes the
   // unnecessary/costly context synchronization when parsing.
   bool wasThreadingEnabled = context->isMultithreadingEnabled();
   context->disableMultithreading();
 
+  // 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 configurations.
+  AsmParserConfig config(context);
+  attachPassReproducerAsmConfig(config, pm, wasThreadingEnabled);
+
   // Parse the input file and reset the context threading state.
   TimingScope parserTiming = timing.nest("Parser");
-  OwningOpRef<ModuleOp> module(parseSourceFile<ModuleOp>(sourceMgr, context));
+  OwningOpRef<ModuleOp> module(parseSourceFile<ModuleOp>(sourceMgr, config));
   context->enableMultithreading(wasThreadingEnabled);
   if (!module)
     return failure();
   parserTiming.stop();
 
-  // Apply any pass manager command line options.
-  PassManager pm(context, OpPassManager::Nesting::Implicit);
-  pm.enableVerifier(verifyPasses);
-  applyPassManagerCLOptions(pm);
-  pm.enableTiming(timing);
-
   // Callback to build the pipeline.
   if (failed(passManagerSetupFn(pm)))
     return failure();
 
   // Run the pipeline.
   if (failed(pm.run(*module)))
     return failure();
 
LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
   static cl::opt<bool> allowUnregisteredDialects(
       "allow-unregistered-dialect",
       cl::desc("Allow operation with no registered dialects"), cl::init(false));
 
   static cl::opt<bool> showDialects(
       "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.
   registerAsmPrinterCLOptions();
   registerMLIRContextCLOptions();
   registerPassManagerCLOptions();
   registerDefaultTimingManagerCLOptions();
   DebugCounter::registerCLOptions();
LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
   // Set up the input file.
   std::string errorMessage;
   auto file = openInputFile(inputFilename, &errorMessage);
   if (!file) {
     llvm::errs() << errorMessage << "\n";
     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";
     return failure();
   }
 
   if (failed(MlirOptMain(output->os(), std::move(file), passPipeline, registry,
                          splitInputFile, verifyDiagnostics, verifyPasses,
                          allowUnregisteredDialects, preloadDialectsInContext)))
     return failure();
 
   // Keep the output file if the invocation of MlirOptMain was successful.
   output->keep();
   return success();
 }

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

 // Check that local reproducers will also traverse dynamic pass pipelines.
 // RUN: mlir-opt %s -pass-pipeline='test-module-pass,test-dynamic-pipeline{op-name=inner_mod1 run-on-nested-operations=1 dynamic-pipeline=test-pass-failure}' -mlir-pass-pipeline-crash-reproducer=%t -verify-diagnostics -mlir-pass-pipeline-local-reproducer --mlir-disable-threading
 // RUN: cat %t | FileCheck -check-prefix=REPRO_LOCAL_DYNAMIC_FAILURE %s
 
 // The crash recovery mechanism will leak memory allocated in the crashing thread.
 // UNSUPPORTED: asan
 
 module @inner_mod1 {
   // expected-error@below {{Failures have been detected while processing an MLIR pass pipeline}}
   // expected-note@below {{Pipeline failed while executing}}
   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

 // UNSUPPORTED: asan
 
 // expected-error@below {{Failures have been detected while processing an MLIR pass pipeline}}
 // expected-note@below {{Pipeline failed while executing}}
 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

-// 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
   return
 }
 
 func.func @bar() {
   return
 }
 
+{-#
+  config: {
+    <mlir_reproducer>: {
+      pipeline: "func.func(cse,canonicalize)",
+      disable_threading: true
+    }
+  }
+#-}
+
 // BEFORE: // -----// IR Dump Before{{.*}}CSE //----- //
 // BEFORE-NEXT: func @foo()
 // BEFORE: // -----// IR Dump Before{{.*}}CSE //----- //
 // BEFORE-NEXT: func @bar()
 // BEFORE-NOT: // -----// IR Dump Before{{.*}}Canonicalizer //----- //
 // BEFORE-NOT: // -----// IR Dump After

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

public:
 
   StringRef getArgument() const final { return "test-spirv-module-combiner"; }
   StringRef getDescription() const final {
     return "Tests SPIR-V module combiner library";
   }
   TestModuleCombinerPass() = default;
   TestModuleCombinerPass(const TestModuleCombinerPass &) {}
   void runOnOperation() override;
-
-private:
-  OwningOpRef<spirv::ModuleOp> combinedModule;
 };
 } // namespace
 
 void TestModuleCombinerPass::runOnOperation() {
   auto modules = llvm::to_vector<4>(getOperation().getOps<spirv::ModuleOp>());
 
   OpBuilder combinedModuleBuilder(modules[0]);
 
   auto listener = [](spirv::ModuleOp originalModule, StringRef oldSymbol,
                      StringRef newSymbol) {
     llvm::outs() << "[" << originalModule.getName() << "] " << oldSymbol
                  << " -> " << 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 {
 void registerTestSpirvModuleCombinerPass() {
   PassRegistration<TestModuleCombinerPass>();
 }
 } // namespace mlir