[mlir][NFC] Split registerAndParseCLIOptions() in mlir-opt

[andrey-golubev]

mlir-opt's registerAndParseCLIOptions() forces users to both register default MLIR options and parse the command line string. Custom mlir-opt implementations, however, may need to provide own options or own parsing. It seems that separating the two functions makes it easier to achieve necessary customizations.

For example, one can register "default" options, then register custom options (not available in standard mlir-opt), then parse all of them. Other cases include two-stage parsing where some additional options become available based on parsed information (e.g. compilation target can allow additional options to be present).

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Andrei Golubev (andrey-golubev)

Changes

mlir-opt's registerAndParseCLIOptions() forces users to both register default MLIR options and parse the command line string. Custom mlir-opt implementations, however, may need to provide own options or own parsing. It seems that separating the two functions makes it easier to achieve necessary customizations.

For example, one can register "default" options, then register custom options (not available in standard mlir-opt), then parse all of them. Other cases include two-stage parsing where some additional options become available based on parsed information (e.g. compilation target can allow additional options to be present).

---

Full diff: https://github.com/llvm/llvm-project/pull/166538.diff

2 Files Affected:

• (modified) mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h (+14)
• (modified) mlir/lib/Tools/mlir-opt/MlirOptMain.cpp (+21-12)

diff --git a/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h b/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
index b7394387b0f9a..c972578d84620 100644
--- a/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
+++ b/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
@@ -359,6 +359,20 @@ class MlirOptMainConfig {
 /// the loaded IR.
 using PassPipelineFn = llvm::function_ref<LogicalResult(PassManager &pm)>;
 
+/// Register basic command line options.
+/// - toolName is used for the header displayed by `--help`.
+/// - registry should contain all the dialects that can be parsed in the source.
+/// - return std::string for help header
+std::string registerCLIOptions(llvm::StringRef toolName,
+                               DialectRegistry &registry);
+
+/// Parse command line options.
+/// - helpHeader is used for the header displayed by `--help`.
+/// - return std::pair<std::string, std::string> for
+///   inputFilename and outputFilename command line option values.
+std::pair<std::string, std::string> parseCLIOptions(int argc, char **argv,
+                                                    llvm::StringRef helpHeader);
+
 /// Register and parse command line options.
 /// - toolName is used for the header displayed by `--help`.
 /// - registry should contain all the dialects that can be parsed in the source.
diff --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
index 9ef405dad5a70..018a188d09109 100644
--- a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
+++ b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
@@ -681,17 +681,8 @@ processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
   return success();
 }
 
-std::pair<std::string, std::string>
-mlir::registerAndParseCLIOptions(int argc, char **argv,
-                                 llvm::StringRef toolName,
-                                 DialectRegistry &registry) {
-  static cl::opt<std::string> inputFilename(
-      cl::Positional, cl::desc("<input file>"), cl::init("-"));
-
-  static cl::opt<std::string> outputFilename("o", cl::desc("Output filename"),
-                                             cl::value_desc("filename"),
-                                             cl::init("-"));
-  // Register any command line options.
+std::string mlir::registerCLIOptions(llvm::StringRef toolName,
+                                     DialectRegistry &registry) {
   MlirOptMainConfig::registerCLOptions(registry);
   registerAsmPrinterCLOptions();
   registerMLIRContextCLOptions();
@@ -706,11 +697,29 @@ mlir::registerAndParseCLIOptions(int argc, char **argv,
     interleaveComma(registry.getDialectNames(), os,
                     [&](auto name) { os << name; });
   }
-  // Parse pass names in main to ensure static initialization completed.
+  return helpHeader;
+}
+
+std::pair<std::string, std::string>
+mlir::parseCLIOptions(int argc, char **argv, llvm::StringRef helpHeader) {
+  static cl::opt<std::string> inputFilename(
+      cl::Positional, cl::desc("<input file>"), cl::init("-"));
+
+  static cl::opt<std::string> outputFilename("o", cl::desc("Output filename"),
+                                             cl::value_desc("filename"),
+                                             cl::init("-"));
   cl::ParseCommandLineOptions(argc, argv, helpHeader);
   return std::make_pair(inputFilename.getValue(), outputFilename.getValue());
 }
 
+std::pair<std::string, std::string>
+mlir::registerAndParseCLIOptions(int argc, char **argv,
+                                 llvm::StringRef toolName,
+                                 DialectRegistry &registry) {
+  auto helpHeader = registerCLIOptions(toolName, registry);
+  return parseCLIOptions(argc, argv, helpHeader);
+}
+
 static LogicalResult printRegisteredDialects(DialectRegistry &registry) {
   llvm::outs() << "Available Dialects: ";
   interleave(registry.getDialectNames(), llvm::outs(), ",");

[weiweichen]

Should these two stay in mlir::registerCLIOptions since they are both registering the options rather than parsing?

[joker-eph]

I'm not against this, but wonder: can't you register your option before calling this function?

[andrey-golubev]

I'm not against this, but wonder: can't you register your option before calling this function?

It's complicated: we actually do have to "pre parse" the CLI in order to add extra options. E.g. we have different passes depending on HW target, and in order to know which passes to register we need to parse only some options, then register more, then parse fully (this is a mess to be honest).
From MLIR's perspective, seems like it's enough to just separate "native" registration from parsing, since we can intervene in the middle whilst not having to reimplement the complete MlirOptMain.

[andrey-golubev]

@weiweichen i need the values here when returning from the "parse" method. if I move them into registration "as is", i'll lose the local variable access.

potentially, these could be moved to MlirOptMainConfigCLOptions though?

[andrey-golubev]

gentle ping. if this is fine, i'd merge it "as is". worst case I can make a follow-up.

[weiweichen]

I see, this looks fine then.