diff --git a/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h b/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h index c003cbd95ee6..b3cdd2ea6404 100644 --- a/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h +++ b/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h @@ -294,16 +294,25 @@ class MlirOptMainConfig { std::string generateReproducerFileFlag = ""; }; -/// VPUX-specific method to get value for arch kind from command line -/// and register HW-specific passes and pipelines -using AdditionalRegistrationFn = - std::function; - /// 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; +/// 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 ®istry); + +/// Parse command line options. +/// - helpHeader is used for the header displayed by `--help`. +/// - return std::pair for +/// inputFilename and outputFilename command line option values. +std::pair 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. @@ -326,12 +335,18 @@ LogicalResult MlirOptMain(llvm::raw_ostream &outputStream, /// Implementation for tools like `mlir-opt`. /// - toolName is used for the header displayed by `--help`. /// - registry should contain all the dialects that can be parsed in the source. -/// - additionalRegistration will be called before the main command line parsing -/// to perform additional registrations. LogicalResult MlirOptMain(int argc, char **argv, llvm::StringRef toolName, - DialectRegistry ®istry, - const AdditionalRegistrationFn &additionalRegistration - = [](llvm::StringRef){}); + DialectRegistry ®istry); + +/// Implementation for tools like `mlir-opt`. +/// This function can be used with registerAndParseCLIOptions 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 ®istry); /// Helper wrapper to return the result of MlirOptMain directly from main. /// diff --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp index df8109f7ca0e..fb0a861d45e7 100644 --- a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp +++ b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp @@ -554,17 +554,8 @@ static LogicalResult processBuffer(raw_ostream &os, return sourceMgrHandler.verify(); } -std::pair -mlir::registerAndParseCLIOptions(int argc, char **argv, - llvm::StringRef toolName, - DialectRegistry ®istry) { - static cl::opt inputFilename( - cl::Positional, cl::desc(""), cl::init("-")); - - static cl::opt 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 ®istry) { MlirOptMainConfig::registerCLOptions(registry); registerAsmPrinterCLOptions(); registerMLIRContextCLOptions(); @@ -579,11 +570,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 +mlir::parseCLIOptions(int argc, char **argv, llvm::StringRef helpHeader) { + static cl::opt inputFilename( + cl::Positional, cl::desc(""), cl::init("-")); + + static cl::opt 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 +mlir::registerAndParseCLIOptions(int argc, char **argv, + llvm::StringRef toolName, + DialectRegistry ®istry) { + auto helpHeader = registerCLIOptions(toolName, registry); + return parseCLIOptions(argc, argv, helpHeader); +} + static LogicalResult printRegisteredDialects(DialectRegistry ®istry) { llvm::outs() << "Available Dialects: "; interleave(registry.getDialectNames(), llvm::outs(), ","); @@ -630,41 +639,13 @@ LogicalResult mlir::MlirOptMain(llvm::raw_ostream &outputStream, config.outputSplitMarker()); } -LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName, - DialectRegistry ®istry, - const AdditionalRegistrationFn &additionalRegistration) { - static cl::opt inputFilename( - cl::Positional, cl::desc(""), cl::init("-")); - - static cl::opt outputFilename("o", cl::desc("Output filename"), - cl::value_desc("filename"), - cl::init("-")); +LogicalResult mlir::MlirOptMain(int argc, char **argv, + llvm::StringRef inputFilename, + llvm::StringRef outputFilename, + DialectRegistry ®istry) { InitLLVM y(argc, argv); - // Register any command line options. - 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; }); - } - - // It is not possible to place a call after command line parser - // since not all options are registered at the moment - additionalRegistration(helpHeader); - - MlirOptMainConfig::registerCLOptions(registry); - - // Parse pass names in main to ensure static initialization completed. - cl::ParseCommandLineOptions(argc, argv, helpHeader); MlirOptMainConfig config = MlirOptMainConfig::createFromCLOptions(); if (config.shouldShowDialects()) @@ -701,3 +682,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 ®istry) { + + // 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); +}