[MLIR][mlir-opt] Add registerationAndParseCLIOptions for MlirOptMain.

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

@@ -229,6 +229,15 @@ class MlirOptMainConfig {
 /// the loaded IR.
 using PassPipelineFn = llvm::function_ref<LogicalResult(PassManager &pm)>;
 
+/// 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.
+/// - return std::pair<std::string, std::string> for
+///   inputFilename and outputFilename command line option values.
+std::pair<std::string, std::string>
+registerAndParseCLIOptions(int argc, char **argv, llvm::StringRef toolName,
+                           DialectRegistry &registry);
+
 /// 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.
@@ -245,6 +254,16 @@ LogicalResult MlirOptMain(llvm::raw_ostream &outputStream,
 LogicalResult MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
                           DialectRegistry &registry);
 
+/// Implementation for tools like `mlir-opt`.
+/// This function can be used with registrationAndParseCLIOptions so that
+/// CLI options can be accessed before running MlirOptMain.
+/// - inputFilename is the name of the input mlir file.
+/// - outputFilename is the name of the output file.
+/// - registry should contain all the dialects that can be parsed in the source.
+LogicalResult MlirOptMain(int argc, char **argv, llvm::StringRef inputFilename,
+                          llvm::StringRef outputFilename,
+                          DialectRegistry &registry);
+
 /// Helper wrapper to return the result of MlirOptMain directly from main.
 ///
 /// Example:

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

@@ -268,7 +268,8 @@ static LogicalResult doVerifyRoundTrip(Operation *op,
     llvm::raw_string_ostream ostream(buffer);
     if (useBytecode) {
       if (failed(writeBytecodeToFile(op, ostream))) {
-        op->emitOpError() << "failed to write bytecode, cannot verify round-trip.\n";
+        op->emitOpError()
+            << "failed to write bytecode, cannot verify round-trip.\n";
         return failure();
       }
     } else {
@@ -281,7 +282,8 @@ static LogicalResult doVerifyRoundTrip(Operation *op,
     roundtripModule =
         parseSourceString<Operation *>(ostream.str(), parseConfig);
     if (!roundtripModule) {
-      op->emitOpError() << "failed to parse bytecode back, cannot verify round-trip.\n";
+      op->emitOpError()
+          << "failed to parse bytecode back, cannot verify round-trip.\n";
       return failure();
     }
   }
@@ -300,7 +302,8 @@ static LogicalResult doVerifyRoundTrip(Operation *op,
   }
   if (reference != roundtrip) {
     // TODO implement a diff.
-    return op->emitOpError() << "roundTrip testing roundtripped module differs from reference:\n<<<<<<Reference\n"
+    return op->emitOpError() << "roundTrip testing roundtripped module differs "
+                                "from reference:\n<<<<<<Reference\n"
                              << reference << "\n=====\n"
                              << roundtrip << "\n>>>>>roundtripped\n";
   }
@@ -443,6 +446,36 @@ static LogicalResult processBuffer(raw_ostream &os,
   return sourceMgrHandler.verify();
 }
 
+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.
+  MlirOptMainConfig::registerCLOptions(registry);
+  registerAsmPrinterCLOptions();
+  registerMLIRContextCLOptions();
+  registerPassManagerCLOptions();
+  registerDefaultTimingManagerCLOptions();
+  tracing::DebugCounter::registerCLOptions();
+
+  // Build the list of dialects as a header for the --help message.
+  std::string helpHeader = (toolName + "\nAvailable Dialects: ").str();
+  {
+    llvm::raw_string_ostream os(helpHeader);
+    interleaveComma(registry.getDialectNames(), os,
+                    [&](auto name) { os << name; });
+  }
+  // Parse pass names in main to ensure static initialization completed.
+  cl::ParseCommandLineOptions(argc, argv, helpHeader);
+  return std::make_pair(inputFilename.getValue(), outputFilename.getValue());
+}
+
 LogicalResult mlir::MlirOptMain(llvm::raw_ostream &outputStream,
                                 std::unique_ptr<llvm::MemoryBuffer> buffer,
                                 DialectRegistry &registry,
@@ -477,34 +510,13 @@ LogicalResult mlir::MlirOptMain(llvm::raw_ostream &outputStream,
                                /*insertMarkerInOutput=*/true);
 }
 
-LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
+LogicalResult mlir::MlirOptMain(int argc, char **argv,
+                                llvm::StringRef inputFilename,
+                                llvm::StringRef outputFilename,
                                 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("-"));
 
   InitLLVM y(argc, argv);
 
-  // Register any command line options.
-  MlirOptMainConfig::registerCLOptions(registry);
-  registerAsmPrinterCLOptions();
-  registerMLIRContextCLOptions();
-  registerPassManagerCLOptions();
-  registerDefaultTimingManagerCLOptions();
-  tracing::DebugCounter::registerCLOptions();
-
-  // Build the list of dialects as a header for the --help message.
-  std::string helpHeader = (toolName + "\nAvailable Dialects: ").str();
-  {
-    llvm::raw_string_ostream os(helpHeader);
-    interleaveComma(registry.getDialectNames(), os,
-                    [&](auto name) { os << name; });
-  }
-  // Parse pass names in main to ensure static initialization completed.
-  cl::ParseCommandLineOptions(argc, argv, helpHeader);
   MlirOptMainConfig config = MlirOptMainConfig::createFromCLOptions();
 
   // When reading from stdin and the input is a tty, it is often a user mistake
@@ -535,3 +547,14 @@ LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
   output->keep();
   return success();
 }
+
+LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
+                                DialectRegistry &registry) {
+
+  // Register and parse command line options.
+  std::string inputFilename, outputFilename;
+  std::tie(inputFilename, outputFilename) =
+      registerAndParseCLIOptions(argc, argv, toolName, registry);
+
+  return MlirOptMain(argc, argv, inputFilename, outputFilename, registry);
+}