[MLIR][mlir-opt] Add `registerationAndParseCLIOptions` for `MlirOptMain`. #70581

weiweichen · 2023-10-29T04:21:12Z

Seprate registeration and CLI parsing from MlirOptMain to mlir::registrationAndParseCLIOptions and a new MlirOptMain so that external tools to drive mlir-opt can call these two new APIs to get CLI option values before running MlirOptMain.

llvmbot · 2023-10-29T04:23:03Z

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: weiwei chen (weiweichen)

Changes

Seprate registeration and CLI parsing from MlirOptMain to mlir::registrationAndParseCLIOptions and a new MlirOptMain so that external tools to drive mlir-opt can call these two new APIs to get CLI option values before running MlirOptMain.

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

2 Files Affected:

(modified) mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h (+17)
(modified) mlir/lib/Tools/mlir-opt/MlirOptMain.cpp (+50-27)

diff --git a/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h b/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
index 222a51e8db77eac..61c90ab978cf7ab 100644
--- a/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
+++ b/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
@@ -229,6 +229,13 @@ class MlirOptMainConfig {
 /// the loaded IR.
 using PassPipelineFn = llvm::function_ref<LogicalResult(PassManager &pm)>;
 
+/// Register all dialects 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.
+std::pair<std::string, std::string>
+registrationAndParseCLIOptions(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 +252,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:
diff --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
index 644113058bdc1cc..be466a4316612f2 100644
--- a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
+++ b/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::registrationAndParseCLIOptions(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 dialects and parse command line options.
+  std::string inputFilename, outputFilename;
+  std::tie(inputFilename, outputFilename) =
+      registrationAndParseCLIOptions(argc, argv, toolName, registry);
+
+  return MlirOptMain(argc, argv, inputFilename, outputFilename, registry);
+}

bzcheeseman · 2023-10-29T05:51:22Z

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

+/// - toolName is used for the header displayed by `--help`.
+/// - registry should contain all the dialects that can be parsed in the source.
+std::pair<std::string, std::string>
+registrationAndParseCLIOptions(int argc, char **argv, llvm::StringRef toolName,


spelling: registerAndParseCLIOptions?

Also, could you comment what it returns? It's not clear by looking at the header.

spelling: registerAndParseCLIOptions?

Yes!

Also, could you comment what it returns? It's not clear by looking at the header.

Sure, updated.

bzcheeseman · 2023-10-29T05:51:55Z

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

@@ -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()


whitespace changes? Did the file need reformatting?

whitespace changes? Did the file need reformatting?

Weird, I think my clang-format is acting out somehow here. Wait, the reformatting looks like needed since the original code has 88 characters in the line where llvm coding style says it should be "within 80 columns" 🤔.

joker-eph · 2023-10-29T06:22:06Z

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

@@ -229,6 +229,13 @@ class MlirOptMainConfig {
 /// the loaded IR.
 using PassPipelineFn = llvm::function_ref<LogicalResult(PassManager &pm)>;

+/// Register all dialects and parse command line options.


What does "Register all dialects" refer to here?

What does "Register all dialects" refer to here?

Oops, wrong comment, should be Register and parse command line options here.

(not fixed yet)

Yep, should be fixed now 🙏

Seprate registeration and CLI parsing from `MlirOptMain` to `mlir::registrationAndParseCLIOptions` and a new `MlirOptMain` so that external tools to drive mlir-opt can call these two new APIs to get CLI option values before running MlirOptMain.

weiweichen marked this pull request as ready for review October 29, 2023 04:22

llvmbot added mlir:core MLIR Core Infrastructure mlir labels Oct 29, 2023

weiweichen requested review from Mogball, River707 and bzcheeseman and removed request for bzcheeseman October 29, 2023 04:36

bzcheeseman reviewed Oct 29, 2023

View reviewed changes

joker-eph reviewed Oct 29, 2023

View reviewed changes

weiweichen added 2 commits October 29, 2023 22:26

Address review comments.

07cb28f

weiweichen force-pushed the weiweic/mliroptmain-api-change branch from 632dc8b to 07cb28f Compare October 30, 2023 02:27

weiweichen added 2 commits October 29, 2023 22:28

Update comment.

7ec2e5a

Update missed comment fix.

0df5fe9

joker-eph approved these changes Oct 30, 2023

View reviewed changes

weiweichen merged commit 77c6339 into llvm:main Oct 30, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[MLIR][mlir-opt] Add `registerationAndParseCLIOptions` for `MlirOptMain`. #70581

[MLIR][mlir-opt] Add `registerationAndParseCLIOptions` for `MlirOptMain`. #70581

Uh oh!

weiweichen commented Oct 29, 2023

Uh oh!

llvmbot commented Oct 29, 2023 •

edited

Loading

Uh oh!

bzcheeseman Oct 29, 2023

Uh oh!

bzcheeseman Oct 29, 2023

Uh oh!

weiweichen Oct 30, 2023

Uh oh!

weiweichen Oct 30, 2023

Uh oh!

bzcheeseman Oct 29, 2023

Uh oh!

weiweichen Oct 30, 2023

Uh oh!

joker-eph Oct 29, 2023

Uh oh!

weiweichen Oct 30, 2023

Uh oh!

joker-eph Oct 30, 2023

Uh oh!

weiweichen Oct 30, 2023

Uh oh!

Uh oh!

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

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

Uh oh!

Conversation

weiweichen commented Oct 29, 2023

Uh oh!

llvmbot commented Oct 29, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

[MLIR][mlir-opt] Add `registerationAndParseCLIOptions` for `MlirOptMain`. #70581

[MLIR][mlir-opt] Add `registerationAndParseCLIOptions` for `MlirOptMain`. #70581

llvmbot commented Oct 29, 2023 •

edited

Loading