[mlir-opt] Delay pass manager creation until after parsing

[rkayaith]

Summary

Currently the pass manager is created before parsing, which requires an
assumption that the top-level operation will be builtin.module.
Delaying the creation allows for using the parsed top-level operation as
the PassManager operation instead.

A followup change will allow for parsing top-level operations other than
builtin.module.

[rkayaith]

rebase

[mehdi_amini]

mlir/include/mlir/Pass/PassRegistry.h

283–295

This function is quite tricky in terms of invariant, and in particular the lifetime expectation for enableVerifier and enableThreading. Not the FIXME above, adding enableVerifier isn't going in the right direction here.

[rkayaith]

mlir/include/mlir/Pass/PassRegistry.h

283–295

Agreed it's tricky, though to me it seems like modifying the pass pipeline or context during parsing makes it even harder to follow. The change here moves more towards "do the minimum in the resource handler, and the main work afterwards". If it helps I could hide the individual options with something like:

ReproducerOptions reproOpts;
attachPassReproducerAsmResource(config, &reproOpts);
... parse ...
PassManager pm(...);
reproOpts.applyOptions(pm, context);

Another option would be to include the top-level op in the reproducer pipeline string (I was planning on this anyways) and construct the full pass manager in the resource handler, so something like:

Optional<PassManager> reproPM;
attachPassReproducerAsmResource(config, reproPM, enableThreading);
... parse ...
PassManager pm(...);
if (reproPM)
  pm = std::move(*reproPM)

though this makes it harder to support some edge case, e.g. verify_each is ordered before pipeline, or verify_each is present in the resource but pipeline isn't

[rkayaith]

create a ReproducerOptions class to hold parsed repro config

[rriddle]

The options class approach looks nice to me.

mlir/include/mlir/Pass/PassRegistry.h

283–295

I'd name this PassReproducerOptions.

284–286

Can you move the private section to the bottom of the class?

289–292

Does this need to be a friend? Why not just a non-static method that attaches a parser for the current options config?

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

84–86

[rkayaith]

mlir/include/mlir/Pass/PassRegistry.h

283–295

I updated the diff to use the first approach, let me know if it seems reasonable. I removed the FIXME's even though I didn't follow their direction; not sure if I should still leave them.

[rkayaith]

address review comments

[rkayaith]

@rriddle @mehdi_amini any other comments here?

mlir/include/mlir/Pass/PassRegistry.h

 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);
+struct PassReproducerOptions {
+  /// Attach an assembly resource parser to 'config' that collects the MLIR
+  /// reproducer configuration into this instance.
+  void attachResourceParser(ParserConfig &config);
+
+  /// Apply the reproducer options to 'pm' and its context.
+  LogicalResult apply(PassManager &pm) const;
+
+private:
+  Optional<std::string> pipeline;
+  Optional<bool> verifyEach;
+  Optional<bool> disableThreading;
+};
 
 } // namespace mlir
 
 #endif // MLIR_PASS_PASSREGISTRY_H_

mlir/lib/Pass/PassCrashRecovery.cpp

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 {
+void PassReproducerOptions::attachResourceParser(ParserConfig &config) {
+  auto parseFn = [this](AsmParsedResourceEntry &entry) -> LogicalResult {
     if (entry.getKey() == "pipeline") {
-      FailureOr<std::string> pipeline = entry.parseAsString();
-      if (failed(pipeline))
-        return failure();
-      return parsePassPipeline(*pipeline, pm);
+      FailureOr<std::string> value = entry.parseAsString();
+      if (succeeded(value))
+        this->pipeline = std::move(*value);
+      return value;
     }
     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);
+        this->disableThreading = *value;
       return value;
     }
     if (entry.getKey() == "verify_each") {
       FailureOr<bool> value = entry.parseAsBool();
       if (succeeded(value))
-        pm.enableVerifier(*value);
+        this->verifyEach = *value;
       return value;
     }
     return entry.emitError() << "unknown 'mlir_reproducer' resource key '"
                              << entry.getKey() << "'";
   };
   config.attachResourceParser("mlir_reproducer", parseFn);
 }
+
+LogicalResult PassReproducerOptions::apply(PassManager &pm) const {
+  if (pipeline.has_value())
+    if (failed(parsePassPipeline(*pipeline, pm)))
+      return failure();
+
+  if (disableThreading.has_value())
+    pm.getContext()->disableMultithreading(*disableThreading);
+
+  if (verifyEach.has_value())
+    pm.enableVerifier(*verifyEach);
+
+  return success();
+}

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 resource
   // handlers. Unhandled external resources are treated as passthrough, i.e.
   // they are not processed and will be emitted directly to the output
   // untouched.
+  PassReproducerOptions reproOptions;
   FallbackAsmResourceMap fallbackResourceMap;
   ParserConfig config(context, &fallbackResourceMap);
-  attachPassReproducerAsmResource(config, pm, wasThreadingEnabled);
+  reproOptions.attachResourceParser(config);
 
   // 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)))
+  // Prepare the pass manager, applying command-line and reproducer options.
+  PassManager pm(context, OpPassManager::Nesting::Implicit,
+                 module->getOperationName());
+  pm.enableVerifier(verifyPasses);
+  applyPassManagerCLOptions(pm);
+  pm.enableTiming(timing);
+  if (failed(reproOptions.apply(pm)) || failed(passManagerSetupFn(pm)))
     return failure();
 
   // Run the pipeline.
   if (failed(pm.run(*module)))
     return failure();
 
   // Print the output.
   TimingScope outputTiming = timing.nest("Output");
   if (emitBytecode) {