[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

290

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

290

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

290

I'd name this PassReproducerOptions.

291–293

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

296–299

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

86–88

[rkayaith]

mlir/include/mlir/Pass/PassRegistry.h

290

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.
+/// configurations.
 // FIXME: Remove the `enableThreading` flag when possible. Some tools, e.g.
 //        mlir-opt, force disable threading during parsing.
-void attachPassReproducerAsmResource(ParserConfig &config, PassManager &pm,
+void attachPassReproducerAsmResource(ParserConfig &config,
+                                     std::string &pipelineStr,
+                                     bool &enableVerifier,
                                      bool &enableThreading);
 
 } // namespace mlir
 
 #endif // MLIR_PASS_PASSREGISTRY_H_

mlir/lib/Pass/PassCrashRecovery.cpp

addInstrumentation(
       std::make_unique<CrashReproducerInstrumentation>(*crashReproGenerator));
 }
 
 //===----------------------------------------------------------------------===//
 // Asm Resource
 //===----------------------------------------------------------------------===//
 
 void mlir::attachPassReproducerAsmResource(ParserConfig &config,
-                                           PassManager &pm,
+                                           std::string &pipelineStr,
+                                           bool &enableVerifier,
                                            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);
+      FailureOr<std::string> value = entry.parseAsString();
+      if (succeeded(value))
+        pipelineStr = 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);
       return value;
     }
     if (entry.getKey() == "verify_each") {
       FailureOr<bool> value = entry.parseAsBool();
       if (succeeded(value))
-        pm.enableVerifier(*value);
+        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

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.
   FallbackAsmResourceMap fallbackResourceMap;
   ParserConfig config(context, &fallbackResourceMap);
-  attachPassReproducerAsmResource(config, pm, wasThreadingEnabled);
+  std::string reproPipeline;
+  attachPassReproducerAsmResource(config, reproPipeline, verifyPasses,
+                                  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.
+  // 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(parsePassPipeline(reproPipeline, pm)))
+    return failure();
   if (failed(passManagerSetupFn(pm)))
     return failure();
 
   // Run the pipeline.
   if (failed(pm.run(*module)))
     return failure();
 
   // Print the output.
   TimingScope outputTiming = timing.nest("Output");