Refactor mlir-opt setup in a new helper function (NFC)

[mehdi_amini]

Summary

This will help refactoring some of the tools to prepare for the explicit registration of
Dialects.

[mehdi_amini]

Minor update

[mehdi_amini]

Update standalone-opt.cpp

[jpienaar]

LGTM if towards removing static registration

mlir/include/mlir/Support/MlirOptMain.h

29

passPipeline as name for arg that is a CL pass pipeline parser, seems confusing. Could passes be passed in rather than parsed inside? (And so be a passed in pass pipeline). Else why not call it a parser?

mlir/lib/Support/MlirOptMain.cpp

181

Mmm, this almost feels like something that should always be displayed in help message rather than with separate option (like what targets are linked in effectively). Passes effectively are. Up to you.

[mehdi_amini]

List dialects in --help

mlir/include/mlir/Support/MlirOptMain.h

29

The naming is weird but PassPipelineCLParser is the object that contains the result of the parsing.
It is added to the PM with LogicalResult addToPipeline(OpPassManager &pm) const;
I agree it can be a bit confusing, that said the pm has an MLIR context while the PassPipelineCLParser does not.

I'm open to suggestion on how to change this, but right now I'm not changing this API, just documenting it ; so I'll leave it for a subsequent revision

mlir/lib/Support/MlirOptMain.cpp

181

Added the list in the output of --help

I'll keep this for now as the --help is very long and one may just want the list and nothing else. We can remove later if needed.

mlir/include/mlir/Support/MlirOptMain.h

 //===- MlirOptMain.h - MLIR Optimizer Driver main ---------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Main entry function for mlir-opt for when built as standalone binary.
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/StringRef.h"
+
 #include <memory>
-#include <vector>
 
 namespace llvm {
 class raw_ostream;
 class MemoryBuffer;
 } // end namespace llvm
 
 namespace mlir {
-struct LogicalResult;
 class PassPipelineCLParser;
 
-LogicalResult MlirOptMain(llvm::raw_ostream &os,
+/// 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.
+/// - 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.
+LogicalResult MlirOptMain(llvm::raw_ostream &outputStream,
                           std::unique_ptr<llvm::MemoryBuffer> buffer,
                           const PassPipelineCLParser &passPipeline,
                           bool splitInputFile, bool verifyDiagnostics,
                           bool verifyPasses, bool allowUnregisteredDialects);
 
+/// Implementation for tools like `mlir-opt`.
+LogicalResult MlirOptMain(int argc, char **argv, llvm::StringRef toolName);
+
 } // end namespace mlir

mlir/lib/Support/MlirOptMain.cpp

 //===- MlirOptMain.cpp - MLIR Optimizer Driver ----------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This is a utility that runs an optimization pass and prints the result back
 // out. It is designed to support unit testing.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Support/MlirOptMain.h"
+#include "mlir/IR/AsmState.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Dialect.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Module.h"
 #include "mlir/Parser.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
+#include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/ToolUtilities.h"
 #include "mlir/Transforms/Passes.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileUtilities.h"
+#include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
 using namespace llvm;
 using llvm::SMLoc;
 
 /// Perform the actions on the input file indicated by the command line flags
 /// within the specified context.
 ///
static LogicalResult processBuffer(raw_ostream &os,
   performActions(os, verifyDiagnostics, verifyPasses, sourceMgr, &context,
                  passPipeline);
 
   // Verify the diagnostic handler to make sure that each of the diagnostics
   // matched.
   return sourceMgrHandler.verify();
 }
 
-LogicalResult mlir::MlirOptMain(raw_ostream &os,
+LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
                                 std::unique_ptr<MemoryBuffer> buffer,
                                 const PassPipelineCLParser &passPipeline,
                                 bool splitInputFile, bool verifyDiagnostics,
                                 bool verifyPasses,
                                 bool allowUnregisteredDialects) {
   // The split-input-file mode is a very specific mode that slices the file
   // up into small pieces and checks each independently.
   if (splitInputFile)
     return splitAndProcessBuffer(
         std::move(buffer),
         [&](std::unique_ptr<MemoryBuffer> chunkBuffer, raw_ostream &os) {
           return processBuffer(os, std::move(chunkBuffer), verifyDiagnostics,
                                verifyPasses, allowUnregisteredDialects,
                                passPipeline);
         },
-        os);
+        outputStream);
 
-  return processBuffer(os, std::move(buffer), verifyDiagnostics, verifyPasses,
-                       allowUnregisteredDialects, passPipeline);
+  return processBuffer(outputStream, std::move(buffer), verifyDiagnostics,
+                       verifyPasses, allowUnregisteredDialects, passPipeline);
+}
+
+LogicalResult mlir::MlirOptMain(int argc, char **argv, StringRef toolName) {
+  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("-"));
+
+  static cl::opt<bool> splitInputFile(
+      "split-input-file",
+      cl::desc("Split the input file into pieces and process each "
+               "chunk independently"),
+      cl::init(false));
+
+  static cl::opt<bool> verifyDiagnostics(
+      "verify-diagnostics",
+      cl::desc("Check that emitted diagnostics match "
+               "expected-* lines on the corresponding line"),
+      cl::init(false));
+
+  static cl::opt<bool> verifyPasses(
+      "verify-each",
+      cl::desc("Run the verifier after each transformation pass"),
+      cl::init(true));
+
+  static cl::opt<bool> allowUnregisteredDialects(
+      "allow-unregistered-dialect",
+      cl::desc("Allow operation with no registered dialects"), cl::init(false));
+
+  static cl::opt<bool> showDialects(
+      "show-dialects", cl::desc("Print the list of registered dialects"),
+      cl::init(false));
+
+  InitLLVM y(argc, argv);
+
+  // Register any command line options.
+  registerAsmPrinterCLOptions();
+  registerMLIRContextCLOptions();
+  registerPassManagerCLOptions();
+  PassPipelineCLParser passPipeline("", "Compiler passes to run");
+
+  // Parse pass names in main to ensure static initialization completed.
+  cl::ParseCommandLineOptions(argc, argv, (toolName + "\n").str());
+
+  if (showDialects) {
+    llvm::outs() << "Registered Dialects:\n";
+    MLIRContext context;
+    interleave(
+        context.getRegisteredDialects(), llvm::outs(),
+        [](Dialect *dialect) { llvm::outs() << dialect->getNamespace(); },
+        "\n");
+    return success();
+  }
+
+  // Set up the input file.
+  std::string errorMessage;
+  auto file = openInputFile(inputFilename, &errorMessage);
+  if (!file) {
+    llvm::errs() << errorMessage << "\n";
+    return failure();
+  }
+
+  auto output = openOutputFile(outputFilename, &errorMessage);
+  if (!output) {
+    llvm::errs() << errorMessage << "\n";
+    return failure();
+  }
+
+  if (failed(MlirOptMain(output->os(), std::move(file), passPipeline,
+                         splitInputFile, verifyDiagnostics, verifyPasses,
+                         allowUnregisteredDialects)))
+    return failure();
+
+  // Keep the output file if the invocation of MlirOptMain was successful.
+  output->keep();
+  return success();
 }

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestRecursiveTypesPass();
 void registerTestReducer();
 void registerTestSpirvEntryPointABIPass();
 void registerTestSCFUtilsPass();
 void registerTestVectorConversions();
 void registerVectorizerTestPass();
 } // namespace mlir
 
-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("-"));
-
-static cl::opt<bool>
-    splitInputFile("split-input-file",
-                   cl::desc("Split the input file into pieces and process each "
-                            "chunk independently"),
-                   cl::init(false));
-
-static cl::opt<bool>
-    verifyDiagnostics("verify-diagnostics",
-                      cl::desc("Check that emitted diagnostics match "
-                               "expected-* lines on the corresponding line"),
-                      cl::init(false));
-
-static cl::opt<bool>
-    verifyPasses("verify-each",
-                 cl::desc("Run the verifier after each transformation pass"),
-                 cl::init(true));
-
-static cl::opt<bool> allowUnregisteredDialects(
-    "allow-unregistered-dialect",
-    cl::desc("Allow operation with no registered dialects"), cl::init(false));
-
 #ifdef MLIR_INCLUDE_TESTS
 void registerTestPasses() {
   registerConvertToTargetEnvPass();
   registerInliner();
   registerMemRefBoundCheck();
   registerPassManagerTestPass();
   registerPatternsTestPass();
   registerPrintOpAvailabilityPass();
#endif
   registerTestGpuParallelLoopMappingPass();
   registerTestSpirvEntryPointABIPass();
   registerTestSCFUtilsPass();
   registerTestVectorConversions();
   registerVectorizerTestPass();
 }
 #endif
 
-static cl::opt<bool>
-    showDialects("show-dialects",
-                 cl::desc("Print the list of registered dialects"),
-                 cl::init(false));
-
 int main(int argc, char **argv) {
   registerAllDialects();
   registerAllPasses();
 #ifdef MLIR_INCLUDE_TESTS
   registerTestPasses();
 #endif
-  InitLLVM y(argc, argv);
-
-  // Register any command line options.
-  registerAsmPrinterCLOptions();
-  registerMLIRContextCLOptions();
-  registerPassManagerCLOptions();
-  PassPipelineCLParser passPipeline("", "Compiler passes to run");
-
-  // Parse pass names in main to ensure static initialization completed.
-  cl::ParseCommandLineOptions(argc, argv, "MLIR modular optimizer driver\n");
-
-  if(showDialects) {
-    llvm::outs() << "Registered Dialects:\n";
-    MLIRContext context;
-    for(Dialect *dialect : context.getRegisteredDialects()) {
-      llvm::outs() << dialect->getNamespace() << "\n";
-    }
-    return 0;
-  }
-
-  // Set up the input file.
-  std::string errorMessage;
-  auto file = openInputFile(inputFilename, &errorMessage);
-  if (!file) {
-    llvm::errs() << errorMessage << "\n";
-    return 1;
-  }
-
-  auto output = openOutputFile(outputFilename, &errorMessage);
-  if (!output) {
-    llvm::errs() << errorMessage << "\n";
-    exit(1);
-  }
-
-  if (failed(MlirOptMain(output->os(), std::move(file), passPipeline,
-                         splitInputFile, verifyDiagnostics, verifyPasses,
-                         allowUnregisteredDialects))) {
-    return 1;
-  }
-  // Keep the output file if the invocation of MlirOptMain was successful.
-  output->keep();
-  return 0;
+  return failed(MlirOptMain(argc, argv, "MLIR modular optimizer driver\n"));
 }