Refactor a MlirOptMainConfig class to hold the configuration of MlirOptMain (NFC)

[mehdi_amini]

Summary

The list of boolean flags and others is becoming unresonnably long.

[jpienaar]

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

38

Should these have default true ? E..g.,

config.setSplitInputFile();

Reads naturally as enabling to me (and I believe we follow that style with printer).

118

Should we add the [[deprecated]] marker so that we get warnings?

146

Not quite sure i follow the ordering here (at least among the bools)

[mehdi_amini]

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

38

Ah right, I set I shouldn't use set for the fluent style API but use enable or use depending on the flag: config.enableSplitInputFile()

118

We should, but can't hear, we have in-tree users. I actually plan to kill this soon (I'll put it on the new deprecations page on the website :)).

146

I think I took the order the cl::opt are defined in the source (it's a bad reason, it's just that I had split screen editor and went one by one!).

Should I put them in alphabetical order?

[rriddle]

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

98

Drop the llvm:: here?

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

329

I think the style prefers constructor calls IIRC.

[rriddle]

The windows build failures look real

[rkayaith]

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

59–60

the docstrings on a few of these are wrong, seems the setVerifyPasses comment got copied

[mehdi_amini]

• sort option implementation alphabetically
• address other comments

[mehdi_amini]

C:\ws\w4\llvm-project\premerge-checks\mlir\include\mlir/Tools/mlir-opt/MlirOptMain.h(65): error C2039: 'function': is not a member of 'std'

(Will fix windows, it is still on my radar)

[mehdi_amini]

Fix windows build

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

 
 #ifndef MLIR_TOOLS_MLIROPT_MLIROPTMAIN_H
 #define MLIR_TOOLS_MLIROPT_MLIROPTMAIN_H
 
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/StringRef.h"
 
 #include <cstdlib>
+#include <functional>
 #include <memory>
 
 namespace llvm {
 class raw_ostream;
 class MemoryBuffer;
 } // namespace llvm
 
 namespace mlir {
 class DialectRegistry;
 class PassPipelineCLParser;
 class PassManager;
 
+/// Configuration options for the mlir-opt tool.
+/// This is intended to help building tools like mlir-opt by collecting the
+/// supported options.
+/// The API is fluent, and the options are sorted in alphabetical order below.
+class MlirOptMainConfig {
+public:
+  /// Allow operation with no registered dialects.
+  /// This option is for convenience during testing only and discouraged in
+  /// general.
+  MlirOptMainConfig &allowUnregisteredDialects(bool allow) {
+    allowUnregisteredDialectsFlag = allow;
+    return *this;
+  }
+  bool shouldAllowUnregisteredDialects() const {
+    return allowUnregisteredDialectsFlag;
+  }
+
+  /// Print the pass-pipeline as text before executing.
+  MlirOptMainConfig &dumpPassPipeline(bool dump) {
+    dumpPassPipelineFlag = dump;
+    return *this;
+  }
+  bool shouldDumpPassPipeline() const { return dumpPassPipelineFlag; }
+
+  /// Set the output format to bytecode instead of textual IR.
+  MlirOptMainConfig &emitBytecode(bool emit) {
+    emitBytecodeFlag = emit;
+    return *this;
+  }
+  bool shouldEmitBytecode() const { return emitBytecodeFlag; }
+
+  /// Set the callback to populate the pass manager.
+  MlirOptMainConfig &
+  setPassPipelineSetupFn(std::function<LogicalResult(PassManager &)> callback) {
+    passPipelineCallback = std::move(callback);
+    return *this;
+  }
+
+  /// Set the parser to use to populate the pass manager.
+  MlirOptMainConfig &setPassPipelineParser(const PassPipelineCLParser &parser);
+
+  /// Populate the passmanager, if any callback was set.
+  LogicalResult setupPassPipeline(PassManager &pm) const {
+    if (passPipelineCallback)
+      return passPipelineCallback(pm);
+    return success();
+  }
+
+  // Deprecated.
+  MlirOptMainConfig &preloadDialectsInContext(bool preload) {
+    preloadDialectsInContextFlag = preload;
+    return *this;
+  }
+  bool shouldPreloadDialectsInContext() const {
+    return preloadDialectsInContextFlag;
+  }
+
+  /// Show the registered dialects before trying to load the input file.
+  MlirOptMainConfig &showDialects(bool show) {
+    showDialectsFlag = show;
+    return *this;
+  }
+  bool shouldShowDialects() const { return showDialectsFlag; }
+
+  /// Set whether to split the input file based on the `// -----` marker into
+  /// pieces and process each chunk independently.
+  MlirOptMainConfig &splitInputFile(bool split = true) {
+    splitInputFileFlag = split;
+    return *this;
+  }
+  bool shouldSplitInputFile() const { return splitInputFileFlag; }
+
+  /// Disable implicit addition of a top-level module op during parsing.
+  MlirOptMainConfig &useImplicitModule(bool useImplicitModule) {
+    useImplicitModuleFlag = useImplicitModule;
+    return *this;
+  }
+  bool shouldUseImplicitModule() const { return useImplicitModuleFlag; }
+
+  /// Set whether to check that emitted diagnostics match `expected-*` lines on
+  /// the corresponding line. This is meant for implementing diagnostic tests.
+  MlirOptMainConfig &verifyDiagnostics(bool verify) {
+    verifyDiagnosticsFlag = verify;
+    return *this;
+  }
+  bool shouldVerifyDiagnostics() const { return verifyDiagnosticsFlag; }
+
+  /// Set whether to run the verifier after each transformation pass.
+  MlirOptMainConfig &verifyPasses(bool verify) {
+    verifyPassesFlag = verify;
+    return *this;
+  }
+  bool shouldVerifyPasses() const { return verifyPassesFlag; }
+
+private:
+  /// Allow operation with no registered dialects.
+  /// This option is for convenience during testing only and discouraged in
+  /// general.
+  bool allowUnregisteredDialectsFlag = false;
+
+  /// Print the pipeline that will be run.
+  bool dumpPassPipelineFlag = false;
+
+  /// Emit bytecode instead of textual assembly when generating output.
+  bool emitBytecodeFlag = false;
+
+  /// The callback to populate the pass manager.
+  std::function<LogicalResult(PassManager &)> passPipelineCallback;
+
+  /// Deprecated.
+  bool preloadDialectsInContextFlag = false;
+
+  /// Show the registered dialects before trying to load the input file.
+  bool showDialectsFlag = false;
+
+  /// Split the input file based on the `// -----` marker into pieces and
+  /// process each chunk independently.
+  bool splitInputFileFlag = false;
+
+  /// Use an implicit top-level module op during parsing.
+  bool useImplicitModuleFlag = true;
+
+  /// Set whether to check that emitted diagnostics match `expected-*` lines on
+  /// the corresponding line. This is meant for implementing diagnostic tests.
+  bool verifyDiagnosticsFlag = false;
+
+  /// Run the verifier after each transformation pass.
+  bool verifyPassesFlag = true;
+};
+
 /// This defines the function type used to setup the pass manager. This can be
 /// used to pass in a callback to setup a default pass pipeline to be applied on
 /// the loaded IR.
 using PassPipelineFn = llvm::function_ref<LogicalResult(PassManager &pm)>;
 
-/// Perform the core processing behind `mlir-opt`:
+/// Perform the core processing behind `mlir-opt`.
 /// - outputStream is the stream where the resulting IR is printed.
 /// - buffer is the in-memory file to parser and process.
-/// - passPipeline is the specification of the pipeline that will be applied.
 /// - registry should contain all the dialects that can be parsed in the source.
-/// - splitInputFile will look for a "-----" marker in the input file, and load
-/// each chunk in an individual ModuleOp processed separately.
-/// - verifyDiagnostics enables a verification mode where comments starting with
-/// "expected-(error|note|remark|warning)" are parsed in the input and matched
-/// against emitted diagnostics.
-/// - verifyPasses enables the IR verifier in-between each pass in the pipeline.
-/// - allowUnregisteredDialects allows to parse and create operation without
-/// registering the Dialect in the MLIRContext.
-/// - preloadDialectsInContext will trigger the upfront loading of all
-///   dialects from the global registry in the MLIRContext. This option is
-///   deprecated and will be removed soon.
-/// - emitBytecode will generate bytecode output instead of text.
-/// - implicitModule will enable implicit addition of a top-level
-/// 'builtin.module' if one doesn't already exist.
-/// - dumpPassPipeline will dump the pipeline being run to stderr
+/// - config contains the configuration options for the tool.
+LogicalResult MlirOptMain(llvm::raw_ostream &outputStream,
+                          std::unique_ptr<llvm::MemoryBuffer> buffer,
+                          DialectRegistry &registry,
+                          const MlirOptMainConfig &config);
+
+/// Perform the core processing behind `mlir-opt`.
+/// This API is deprecated, use the MlirOptMainConfig version above instead.
 LogicalResult
 MlirOptMain(llvm::raw_ostream &outputStream,
             std::unique_ptr<llvm::MemoryBuffer> buffer,
             const PassPipelineCLParser &passPipeline, DialectRegistry &registry,
             bool splitInputFile, bool verifyDiagnostics, bool verifyPasses,
             bool allowUnregisteredDialects,
             bool preloadDialectsInContext = false, bool emitBytecode = false,
             bool implicitModule = false, bool dumpPassPipeline = false);
 
-/// Support a callback to setup the pass manager.
-/// - passManagerSetupFn is the callback invoked to setup the pass manager to
-///   apply on the loaded IR.
+/// Perform the core processing behind `mlir-opt`.
+/// This API is deprecated, use the MlirOptMainConfig version above instead.
 LogicalResult MlirOptMain(
     llvm::raw_ostream &outputStream, std::unique_ptr<llvm::MemoryBuffer> buffer,
     PassPipelineFn passManagerSetupFn, DialectRegistry &registry,
     bool splitInputFile, bool verifyDiagnostics, bool verifyPasses,
     bool allowUnregisteredDialects, bool preloadDialectsInContext = false,
     bool emitBytecode = false, bool implicitModule = false);
 
 /// Implementation for tools like `mlir-opt`.

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

 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/ThreadPool.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
 using namespace llvm;
 
+MlirOptMainConfig &MlirOptMainConfig::setPassPipelineParser(
+    const PassPipelineCLParser &passPipeline) {
+  passPipelineCallback = [&](PassManager &pm) {
+    auto errorHandler = [&](const Twine &msg) {
+      emitError(UnknownLoc::get(pm.getContext())) << msg;
+      return failure();
+    };
+    if (failed(passPipeline.addToPipeline(pm, errorHandler)))
+      return failure();
+    if (this->shouldDumpPassPipeline()) {
+
+      pm.dump();
+      llvm::errs() << "\n";
+    }
+    return success();
+  };
+  return *this;
+}
+
 /// Perform the actions on the input file indicated by the command line flags
 /// within the specified context.
 ///
 /// This typically parses the main source file, runs zero or more optimization
 /// passes, then prints the output.
 ///
 static LogicalResult
-performActions(raw_ostream &os, bool verifyDiagnostics, bool verifyPasses,
+performActions(raw_ostream &os,
                const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
-               MLIRContext *context, PassPipelineFn passManagerSetupFn,
-               bool emitBytecode, bool implicitModule) {
+               MLIRContext *context, const MlirOptMainConfig &config) {
   DefaultTimingManager tm;
   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 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, /*verifyAfterParse=*/true, &fallbackResourceMap);
-  reproOptions.attachResourceParser(config);
+  ParserConfig parseConfig(context, /*verifyAfterParse=*/true,
+                           &fallbackResourceMap);
+  reproOptions.attachResourceParser(parseConfig);
 
   // Parse the input file and reset the context threading state.
   TimingScope parserTiming = timing.nest("Parser");
-  OwningOpRef<Operation *> op =
-      parseSourceFileForTool(sourceMgr, config, implicitModule);
+  OwningOpRef<Operation *> op = parseSourceFileForTool(
+      sourceMgr, parseConfig, config.shouldUseImplicitModule());
   context->enableMultithreading(wasThreadingEnabled);
   if (!op)
     return failure();
   parserTiming.stop();
 
   // Prepare the pass manager, applying command-line and reproducer options.
   PassManager pm(op.get()->getName(), PassManager::Nesting::Implicit);
-  pm.enableVerifier(verifyPasses);
+  pm.enableVerifier(config.shouldVerifyPasses());
   applyPassManagerCLOptions(pm);
   pm.enableTiming(timing);
-  if (failed(reproOptions.apply(pm)) || failed(passManagerSetupFn(pm)))
+  if (failed(reproOptions.apply(pm)) || failed(config.setupPassPipeline(pm)))
     return failure();
 
   // Run the pipeline.
   if (failed(pm.run(*op)))
     return failure();
 
   // Print the output.
   TimingScope outputTiming = timing.nest("Output");
-  if (emitBytecode) {
+  if (config.shouldEmitBytecode()) {
     BytecodeWriterConfig writerConfig(fallbackResourceMap);
     writeBytecodeToFile(op.get(), os, writerConfig);
   } else {
     AsmState asmState(op.get(), OpPrintingFlags(), /*locationMap=*/nullptr,
                       &fallbackResourceMap);
     op.get()->print(os, asmState);
     os << '\n';
   }
   return success();
 }
 
 /// Parses the memory buffer.  If successfully, run a series of passes against
 /// it and print the result.
-static LogicalResult
-processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
-              bool verifyDiagnostics, bool verifyPasses,
-              bool allowUnregisteredDialects, bool preloadDialectsInContext,
-              bool emitBytecode, bool implicitModule,
-              PassPipelineFn passManagerSetupFn, DialectRegistry &registry,
+static LogicalResult processBuffer(raw_ostream &os,
+                                   std::unique_ptr<MemoryBuffer> ownedBuffer,
+                                   const MlirOptMainConfig &config,
+                                   DialectRegistry &registry,
                                    llvm::ThreadPool *threadPool) {
   // Tell sourceMgr about this buffer, which is what the parser will pick up.
   auto sourceMgr = std::make_shared<SourceMgr>();
   sourceMgr->AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
 
   // Create a context just for the current buffer. Disable threading on creation
   // since we'll inject the thread-pool separately.
   MLIRContext context(registry, MLIRContext::Threading::DISABLED);
   if (threadPool)
     context.setThreadPool(*threadPool);
 
   // Parse the input file.
-  if (preloadDialectsInContext)
+  if (config.shouldPreloadDialectsInContext())
     context.loadAllAvailableDialects();
-  context.allowUnregisteredDialects(allowUnregisteredDialects);
-  if (verifyDiagnostics)
+  context.allowUnregisteredDialects(config.shouldAllowUnregisteredDialects());
+  if (config.shouldVerifyDiagnostics())
     context.printOpOnDiagnostic(false);
   context.getDebugActionManager().registerActionHandler<DebugCounter>();
 
   // If we are in verify diagnostics mode then we have a lot of work to do,
   // otherwise just perform the actions without worrying about it.
-  if (!verifyDiagnostics) {
+  if (!config.shouldVerifyDiagnostics()) {
     SourceMgrDiagnosticHandler sourceMgrHandler(*sourceMgr, &context);
-    return performActions(os, verifyDiagnostics, verifyPasses, sourceMgr,
-                          &context, passManagerSetupFn, emitBytecode,
-                          implicitModule);
+    return performActions(os, sourceMgr, &context, config);
   }
 
   SourceMgrDiagnosticVerifierHandler sourceMgrHandler(*sourceMgr, &context);
 
   // Do any processing requested by command line flags.  We don't care whether
   // these actions succeed or fail, we only care what diagnostics they produce
   // and whether they match our expectations.
-  (void)performActions(os, verifyDiagnostics, verifyPasses, sourceMgr, &context,
-                       passManagerSetupFn, emitBytecode, implicitModule);
+  (void)performActions(os, sourceMgr, &context, config);
 
   // Verify the diagnostic handler to make sure that each of the diagnostics
   // matched.
   return sourceMgrHandler.verify();
 }
 
-LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
-                                std::unique_ptr<MemoryBuffer> buffer,
-                                PassPipelineFn passManagerSetupFn,
-                                DialectRegistry &registry, bool splitInputFile,
-                                bool verifyDiagnostics, bool verifyPasses,
-                                bool allowUnregisteredDialects,
-                                bool preloadDialectsInContext,
-                                bool emitBytecode, bool implicitModule) {
+LogicalResult mlir::MlirOptMain(llvm::raw_ostream &outputStream,
+                                std::unique_ptr<llvm::MemoryBuffer> buffer,
+                                DialectRegistry &registry,
+                                const MlirOptMainConfig &config) {
   // The split-input-file mode is a very specific mode that slices the file
   // up into small pieces and checks each independently.
   // We use an explicit threadpool to avoid creating and joining/destroying
   // threads for each of the split.
   ThreadPool *threadPool = nullptr;
 
   // Create a temporary context for the sake of checking if
   // --mlir-disable-threading was passed on the command line.
   // We use the thread-pool this context is creating, and avoid
   // creating any thread when disabled.
   MLIRContext threadPoolCtx;
   if (threadPoolCtx.isMultithreadingEnabled())
     threadPool = &threadPoolCtx.getThreadPool();
 
   auto chunkFn = [&](std::unique_ptr<MemoryBuffer> chunkBuffer,
                      raw_ostream &os) {
-    return processBuffer(os, std::move(chunkBuffer), verifyDiagnostics,
-                         verifyPasses, allowUnregisteredDialects,
-                         preloadDialectsInContext, emitBytecode, implicitModule,
-                         passManagerSetupFn, registry, threadPool);
+    return processBuffer(os, std::move(chunkBuffer), config, registry,
+                         threadPool);
   };
   return splitAndProcessBuffer(std::move(buffer), chunkFn, outputStream,
-                               splitInputFile, /*insertMarkerInOutput=*/true);
+                               config.shouldSplitInputFile(),
+                               /*insertMarkerInOutput=*/true);
+}
+
+LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
+                                std::unique_ptr<MemoryBuffer> buffer,
+                                PassPipelineFn passManagerSetupFn,
+                                DialectRegistry &registry, bool splitInputFile,
+                                bool verifyDiagnostics, bool verifyPasses,
+                                bool allowUnregisteredDialects,
+                                bool preloadDialectsInContext,
+                                bool emitBytecode, bool implicitModule) {
+  return MlirOptMain(outputStream, std::move(buffer), registry,
+                     MlirOptMainConfig{}
+                         .splitInputFile(splitInputFile)
+                         .verifyDiagnostics(verifyDiagnostics)
+                         .verifyPasses(verifyPasses)
+                         .allowUnregisteredDialects(allowUnregisteredDialects)
+                         .preloadDialectsInContext(preloadDialectsInContext)
+                         .emitBytecode(emitBytecode)
+                         .useImplicitModule(implicitModule)
+                         .setPassPipelineSetupFn(passManagerSetupFn));
 }
 
 LogicalResult mlir::MlirOptMain(
     raw_ostream &outputStream, std::unique_ptr<MemoryBuffer> buffer,
     const PassPipelineCLParser &passPipeline, DialectRegistry &registry,
     bool splitInputFile, bool verifyDiagnostics, bool verifyPasses,
     bool allowUnregisteredDialects, bool preloadDialectsInContext,
     bool emitBytecode, bool implicitModule, bool dumpPassPipeline) {
-  auto passManagerSetupFn = [&](PassManager &pm) {
-    auto errorHandler = [&](const Twine &msg) {
-      emitError(UnknownLoc::get(pm.getContext())) << msg;
-      return failure();
-    };
-    if (failed(passPipeline.addToPipeline(pm, errorHandler)))
-      return failure();
-    if (dumpPassPipeline) {
-      pm.dump();
-      llvm::errs() << "\n";
-    }
-    return success();
-  };
-  return MlirOptMain(outputStream, std::move(buffer), passManagerSetupFn,
-                     registry, splitInputFile, verifyDiagnostics, verifyPasses,
-                     allowUnregisteredDialects, preloadDialectsInContext,
-                     emitBytecode, implicitModule);
+  return MlirOptMain(outputStream, std::move(buffer), registry,
+                     MlirOptMainConfig{}
+                         .splitInputFile(splitInputFile)
+                         .verifyDiagnostics(verifyDiagnostics)
+                         .verifyPasses(verifyPasses)
+                         .allowUnregisteredDialects(allowUnregisteredDialects)
+                         .preloadDialectsInContext(preloadDialectsInContext)
+                         .emitBytecode(emitBytecode)
+                         .useImplicitModule(implicitModule)
+                         .dumpPassPipeline(dumpPassPipeline)
+                         .setPassPipelineParser(passPipeline));
 }
 
 LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
                                 DialectRegistry &registry,
                                 bool preloadDialectsInContext) {
   static cl::opt<std::string> inputFilename(
       cl::Positional, cl::desc("<input file>"), cl::init("-"));
 
if (!file) {
     return failure();
   }
 
   auto output = openOutputFile(outputFilename, &errorMessage);
   if (!output) {
     llvm::errs() << errorMessage << "\n";
     return failure();
   }
+  // Setup the configuration for the main function.
+  MlirOptMainConfig config;
+  config.setPassPipelineParser(passPipeline)
+      .splitInputFile(splitInputFile)
+      .verifyDiagnostics(verifyDiagnostics)
+      .verifyPasses(verifyPasses)
+      .allowUnregisteredDialects(allowUnregisteredDialects)
+      .preloadDialectsInContext(preloadDialectsInContext)
+      .emitBytecode(emitBytecode)
+      .useImplicitModule(!noImplicitModule)
+      .dumpPassPipeline(dumpPassPipeline);
 
-  if (failed(MlirOptMain(output->os(), std::move(file), passPipeline, registry,
-                         splitInputFile, verifyDiagnostics, verifyPasses,
-                         allowUnregisteredDialects, preloadDialectsInContext,
-                         emitBytecode, /*implicitModule=*/!noImplicitModule,
-                         dumpPassPipeline)))
+  if (failed(MlirOptMain(output->os(), std::move(file), registry, config)))
     return failure();
 
   // Keep the output file if the invocation of MlirOptMain was successful.
   output->keep();
   return success();
 }