840 lines · cpp
1//===- MlirOptMain.cpp - MLIR Optimizer Driver ----------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This is a utility that runs an optimization pass and prints the result back10// out. It is designed to support unit testing.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Tools/mlir-opt/MlirOptMain.h"15#include "mlir/Bytecode/BytecodeWriter.h"16#include "mlir/Debug/CLOptionsSetup.h"17#include "mlir/Debug/Counter.h"18#include "mlir/Dialect/IRDL/IR/IRDL.h"19#include "mlir/Dialect/IRDL/IRDLLoading.h"20#include "mlir/IR/AsmState.h"21#include "mlir/IR/Attributes.h"22#include "mlir/IR/BuiltinOps.h"23#include "mlir/IR/Diagnostics.h"24#include "mlir/IR/Location.h"25#include "mlir/IR/MLIRContext.h"26#include "mlir/IR/Remarks.h"27#include "mlir/Parser/Parser.h"28#include "mlir/Pass/PassManager.h"29#include "mlir/Pass/PassRegistry.h"30#include "mlir/Remark/RemarkStreamer.h"31#include "mlir/Support/FileUtilities.h"32#include "mlir/Support/Timing.h"33#include "mlir/Support/ToolUtilities.h"34#include "mlir/Tools/ParseUtilities.h"35#include "mlir/Tools/Plugins/DialectPlugin.h"36#include "mlir/Tools/Plugins/PassPlugin.h"37#include "llvm/ADT/StringRef.h"38#include "llvm/Remarks/RemarkFormat.h"39#include "llvm/Support/CommandLine.h"40#include "llvm/Support/Debug.h"41#include "llvm/Support/InitLLVM.h"42#include "llvm/Support/LogicalResult.h"43#include "llvm/Support/ManagedStatic.h"44#include "llvm/Support/Process.h"45#include "llvm/Support/Regex.h"46#include "llvm/Support/SourceMgr.h"47#include "llvm/Support/ThreadPool.h"48#include "llvm/Support/ToolOutputFile.h"49 50using namespace mlir;51using namespace llvm;52 53namespace {54class BytecodeVersionParser : public cl::parser<std::optional<int64_t>> {55public:56 BytecodeVersionParser(cl::Option &o)57 : cl::parser<std::optional<int64_t>>(o) {}58 59 bool parse(cl::Option &o, StringRef /*argName*/, StringRef arg,60 std::optional<int64_t> &v) {61 long long w;62 if (getAsSignedInteger(arg, 10, w))63 return o.error("Invalid argument '" + arg +64 "', only integer is supported.");65 v = w;66 return false;67 }68};69 70/// This class is intended to manage the handling of command line options for71/// creating a *-opt config. This is a singleton.72struct MlirOptMainConfigCLOptions : public MlirOptMainConfig {73 MlirOptMainConfigCLOptions() {74 // These options are static but all uses ExternalStorage to initialize the75 // members of the parent class. This is unusual but since this class is a76 // singleton it basically attaches command line option to the singleton77 // members.78 79 static cl::opt<bool, /*ExternalStorage=*/true> allowUnregisteredDialects(80 "allow-unregistered-dialect",81 cl::desc("Allow operation with no registered dialects"),82 cl::location(allowUnregisteredDialectsFlag), cl::init(false));83 84 static cl::opt<bool, /*ExternalStorage=*/true> dumpPassPipeline(85 "dump-pass-pipeline", cl::desc("Print the pipeline that will be run"),86 cl::location(dumpPassPipelineFlag), cl::init(false));87 88 static cl::opt<bool, /*ExternalStorage=*/true> emitBytecode(89 "emit-bytecode", cl::desc("Emit bytecode when generating output"),90 cl::location(emitBytecodeFlag), cl::init(false));91 92 static cl::opt<bool, /*ExternalStorage=*/true> elideResourcesFromBytecode(93 "elide-resource-data-from-bytecode",94 cl::desc("Elide resources when generating bytecode"),95 cl::location(elideResourceDataFromBytecodeFlag), cl::init(false));96 97 static cl::opt<std::optional<int64_t>, /*ExternalStorage=*/true,98 BytecodeVersionParser>99 bytecodeVersion(100 "emit-bytecode-version",101 cl::desc("Use specified bytecode when generating output"),102 cl::location(emitBytecodeVersion), cl::init(std::nullopt));103 104 static cl::opt<std::string, /*ExternalStorage=*/true> irdlFile(105 "irdl-file",106 cl::desc("IRDL file to register before processing the input"),107 cl::location(irdlFileFlag), cl::init(""), cl::value_desc("filename"));108 109 static cl::opt<VerbosityLevel, /*ExternalStorage=*/true>110 diagnosticVerbosityLevel(111 "mlir-diagnostic-verbosity-level",112 cl::desc("Choose level of diagnostic information"),113 cl::location(diagnosticVerbosityLevelFlag),114 cl::init(VerbosityLevel::ErrorsWarningsAndRemarks),115 cl::values(116 clEnumValN(VerbosityLevel::ErrorsOnly, "errors", "Errors only"),117 clEnumValN(VerbosityLevel::ErrorsAndWarnings, "warnings",118 "Errors and warnings"),119 clEnumValN(VerbosityLevel::ErrorsWarningsAndRemarks, "remarks",120 "Errors, warnings and remarks")));121 122 static cl::opt<bool, /*ExternalStorage=*/true> disableDiagnosticNotes(123 "mlir-disable-diagnostic-notes", cl::desc("Disable diagnostic notes."),124 cl::location(disableDiagnosticNotesFlag), cl::init(false));125 126 static cl::opt<bool, /*ExternalStorage=*/true> explicitModule(127 "no-implicit-module",128 cl::desc("Disable implicit addition of a top-level module op during "129 "parsing"),130 cl::location(useExplicitModuleFlag), cl::init(false));131 132 static cl::opt<bool, /*ExternalStorage=*/true> listPasses(133 "list-passes", cl::desc("Print the list of registered passes and exit"),134 cl::location(listPassesFlag), cl::init(false));135 136 static cl::opt<bool, /*ExternalStorage=*/true> runReproducer(137 "run-reproducer", cl::desc("Run the pipeline stored in the reproducer"),138 cl::location(runReproducerFlag), cl::init(false));139 140 static cl::opt<bool, /*ExternalStorage=*/true> showDialects(141 "show-dialects",142 cl::desc("Print the list of registered dialects and exit"),143 cl::location(showDialectsFlag), cl::init(false));144 145 static cl::opt<std::string, /*ExternalStorage=*/true> splitInputFile{146 "split-input-file",147 llvm::cl::ValueOptional,148 cl::callback([&](const std::string &str) {149 // Implicit value: use default marker if flag was used without value.150 if (str.empty())151 splitInputFile.setValue(kDefaultSplitMarker);152 }),153 cl::desc("Split the input file into chunks using the given or "154 "default marker and process each chunk independently"),155 cl::location(splitInputFileFlag),156 cl::init("")};157 158 static cl::opt<std::string, /*ExternalStorage=*/true> outputSplitMarker(159 "output-split-marker",160 cl::desc("Split marker to use for merging the ouput"),161 cl::location(outputSplitMarkerFlag), cl::init(kDefaultSplitMarker));162 163 static cl::opt<SourceMgrDiagnosticVerifierHandler::Level,164 /*ExternalStorage=*/true>165 verifyDiagnostics{166 "verify-diagnostics", llvm::cl::ValueOptional,167 cl::desc("Check that emitted diagnostics match expected-* lines on "168 "the corresponding line"),169 cl::location(verifyDiagnosticsFlag),170 cl::values(171 clEnumValN(SourceMgrDiagnosticVerifierHandler::Level::All,172 "all",173 "Check all diagnostics (expected, unexpected, "174 "near-misses)"),175 // Implicit value: when passed with no arguments, e.g.176 // `--verify-diagnostics` or `--verify-diagnostics=`.177 clEnumValN(SourceMgrDiagnosticVerifierHandler::Level::All, "",178 "Check all diagnostics (expected, unexpected, "179 "near-misses)"),180 clEnumValN(181 SourceMgrDiagnosticVerifierHandler::Level::OnlyExpected,182 "only-expected", "Check only expected diagnostics"))};183 184 static cl::opt<bool, /*ExternalStorage=*/true> verifyPasses(185 "verify-each",186 cl::desc("Run the verifier after each transformation pass"),187 cl::location(verifyPassesFlag), cl::init(true));188 189 static cl::opt<bool, /*ExternalStorage=*/true> disableVerifyOnParsing(190 "mlir-very-unsafe-disable-verifier-on-parsing",191 cl::desc("Disable the verifier on parsing (very unsafe)"),192 cl::location(disableVerifierOnParsingFlag), cl::init(false));193 194 static cl::opt<bool, /*ExternalStorage=*/true> verifyRoundtrip(195 "verify-roundtrip",196 cl::desc("Round-trip the IR after parsing and ensure it succeeds"),197 cl::location(verifyRoundtripFlag), cl::init(false));198 199 static cl::list<std::string> passPlugins(200 "load-pass-plugin", cl::desc("Load passes from plugin library"));201 202 static cl::opt<std::string, /*ExternalStorage=*/true>203 generateReproducerFile(204 "mlir-generate-reproducer",205 llvm::cl::desc(206 "Generate an mlir reproducer at the provided filename"207 " (no crash required)"),208 cl::location(generateReproducerFileFlag), cl::init(""),209 cl::value_desc("filename"));210 211 static cl::OptionCategory remarkCategory(212 "Remark Options",213 "Filter remarks by regular expression (llvm::Regex syntax).");214 215 static llvm::cl::opt<RemarkFormat, /*ExternalStorage=*/true> remarkFormat{216 "remark-format",217 llvm::cl::desc("Specify the format for remark output."),218 cl::location(remarkFormatFlag),219 llvm::cl::value_desc("format"),220 llvm::cl::init(RemarkFormat::REMARK_FORMAT_STDOUT),221 llvm::cl::values(clEnumValN(RemarkFormat::REMARK_FORMAT_STDOUT,222 "emitRemark",223 "Print as emitRemark to command-line"),224 clEnumValN(RemarkFormat::REMARK_FORMAT_YAML, "yaml",225 "Print yaml file"),226 clEnumValN(RemarkFormat::REMARK_FORMAT_BITSTREAM,227 "bitstream", "Print bitstream file")),228 llvm::cl::cat(remarkCategory)};229 230 static llvm::cl::opt<RemarkPolicy, /*ExternalStorage=*/true> remarkPolicy{231 "remark-policy",232 llvm::cl::desc("Specify the policy for remark output."),233 cl::location(remarkPolicyFlag),234 llvm::cl::value_desc("format"),235 llvm::cl::init(RemarkPolicy::REMARK_POLICY_ALL),236 llvm::cl::values(clEnumValN(RemarkPolicy::REMARK_POLICY_ALL, "all",237 "Print all remarks"),238 clEnumValN(RemarkPolicy::REMARK_POLICY_FINAL, "final",239 "Print final remarks")),240 llvm::cl::cat(remarkCategory)};241 242 static cl::opt<std::string, /*ExternalStorage=*/true> remarksAll(243 "remarks-filter",244 cl::desc("Show all remarks: passed, missed, failed, analysis"),245 cl::location(remarksAllFilterFlag), cl::init(""),246 cl::cat(remarkCategory));247 248 static cl::opt<std::string, /*ExternalStorage=*/true> remarksFile(249 "remarks-output-file",250 cl::desc(251 "Output file for yaml and bitstream remark formats. Default is "252 "mlir-remarks.yaml or mlir-remarks.bitstream"),253 cl::location(remarksOutputFileFlag), cl::init(""),254 cl::cat(remarkCategory));255 256 static cl::opt<std::string, /*ExternalStorage=*/true> remarksPassed(257 "remarks-filter-passed", cl::desc("Show passed remarks"),258 cl::location(remarksPassedFilterFlag), cl::init(""),259 cl::cat(remarkCategory));260 261 static cl::opt<std::string, /*ExternalStorage=*/true> remarksFailed(262 "remarks-filter-failed", cl::desc("Show failed remarks"),263 cl::location(remarksFailedFilterFlag), cl::init(""),264 cl::cat(remarkCategory));265 266 static cl::opt<std::string, /*ExternalStorage=*/true> remarksMissed(267 "remarks-filter-missed", cl::desc("Show missed remarks"),268 cl::location(remarksMissedFilterFlag), cl::init(""),269 cl::cat(remarkCategory));270 271 static cl::opt<std::string, /*ExternalStorage=*/true> remarksAnalyse(272 "remarks-filter-analyse", cl::desc("Show analysis remarks"),273 cl::location(remarksAnalyseFilterFlag), cl::init(""),274 cl::cat(remarkCategory));275 276 /// Set the callback to load a pass plugin.277 passPlugins.setCallback([&](const std::string &pluginPath) {278 auto plugin = PassPlugin::load(pluginPath);279 if (!plugin) {280 errs() << "Failed to load passes from '" << pluginPath281 << "'. Request ignored.\n";282 return;283 }284 plugin.get().registerPassRegistryCallbacks();285 });286 287 static cl::list<std::string> dialectPlugins(288 "load-dialect-plugin", cl::desc("Load dialects from plugin library"));289 this->dialectPlugins = std::addressof(dialectPlugins);290 291 static PassPipelineCLParser passPipeline("", "Compiler passes to run", "p");292 setPassPipelineParser(passPipeline);293 }294 295 /// Set the callback to load a dialect plugin.296 void setDialectPluginsCallback(DialectRegistry ®istry);297 298 /// Pointer to static dialectPlugins variable in constructor, needed by299 /// setDialectPluginsCallback(DialectRegistry&).300 cl::list<std::string> *dialectPlugins = nullptr;301};302 303/// A scoped diagnostic handler that suppresses certain diagnostics based on304/// the verbosity level and whether the diagnostic is a note.305class DiagnosticFilter : public ScopedDiagnosticHandler {306public:307 DiagnosticFilter(MLIRContext *ctx, VerbosityLevel verbosityLevel,308 bool showNotes = true)309 : ScopedDiagnosticHandler(ctx) {310 setHandler([verbosityLevel, showNotes](Diagnostic &diag) {311 auto severity = diag.getSeverity();312 switch (severity) {313 case mlir::DiagnosticSeverity::Error:314 // failure indicates that the error is not handled by the filter and315 // goes through to the default handler. Therefore, the error can be316 // successfully printed.317 return failure();318 case mlir::DiagnosticSeverity::Warning:319 if (verbosityLevel == VerbosityLevel::ErrorsOnly)320 return success();321 else322 return failure();323 case mlir::DiagnosticSeverity::Remark:324 if (verbosityLevel == VerbosityLevel::ErrorsOnly ||325 verbosityLevel == VerbosityLevel::ErrorsAndWarnings)326 return success();327 else328 return failure();329 case mlir::DiagnosticSeverity::Note:330 if (showNotes)331 return failure();332 else333 return success();334 }335 llvm_unreachable("Unknown diagnostic severity");336 });337 }338};339} // namespace340 341ManagedStatic<MlirOptMainConfigCLOptions> clOptionsConfig;342 343void MlirOptMainConfig::registerCLOptions(DialectRegistry ®istry) {344 clOptionsConfig->setDialectPluginsCallback(registry);345 tracing::DebugConfig::registerCLOptions();346}347 348MlirOptMainConfig MlirOptMainConfig::createFromCLOptions() {349 clOptionsConfig->setDebugConfig(tracing::DebugConfig::createFromCLOptions());350 return *clOptionsConfig;351}352 353MlirOptMainConfig &MlirOptMainConfig::setPassPipelineParser(354 const PassPipelineCLParser &passPipeline) {355 passPipelineCallback = [&](PassManager &pm) {356 auto errorHandler = [&](const Twine &msg) {357 emitError(UnknownLoc::get(pm.getContext())) << msg;358 return failure();359 };360 if (failed(passPipeline.addToPipeline(pm, errorHandler)))361 return failure();362 if (this->shouldDumpPassPipeline()) {363 364 pm.dump();365 llvm::errs() << "\n";366 }367 return success();368 };369 return *this;370}371 372void MlirOptMainConfigCLOptions::setDialectPluginsCallback(373 DialectRegistry ®istry) {374 dialectPlugins->setCallback([&](const std::string &pluginPath) {375 auto plugin = DialectPlugin::load(pluginPath);376 if (!plugin) {377 errs() << "Failed to load dialect plugin from '" << pluginPath378 << "'. Request ignored.\n";379 return;380 };381 plugin.get().registerDialectRegistryCallbacks(registry);382 });383}384 385LogicalResult loadIRDLDialects(StringRef irdlFile, MLIRContext &ctx) {386 DialectRegistry registry;387 registry.insert<irdl::IRDLDialect>();388 ctx.appendDialectRegistry(registry);389 390 // Set up the input file.391 std::string errorMessage;392 std::unique_ptr<MemoryBuffer> file = openInputFile(irdlFile, &errorMessage);393 if (!file) {394 emitError(UnknownLoc::get(&ctx)) << errorMessage;395 return failure();396 }397 398 // Give the buffer to the source manager.399 // This will be picked up by the parser.400 SourceMgr sourceMgr;401 sourceMgr.AddNewSourceBuffer(std::move(file), SMLoc());402 403 SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &ctx);404 405 // Parse the input file.406 OwningOpRef<ModuleOp> module(parseSourceFile<ModuleOp>(sourceMgr, &ctx));407 if (!module)408 return failure();409 410 // Load IRDL dialects.411 return irdl::loadDialects(module.get());412}413 414// Return success if the module can correctly round-trip. This intended to test415// that the custom printers/parsers are complete.416static LogicalResult doVerifyRoundTrip(Operation *op,417 const MlirOptMainConfig &config,418 bool useBytecode) {419 // We use a new context to avoid resource handle renaming issue in the diff.420 MLIRContext roundtripContext;421 OwningOpRef<Operation *> roundtripModule;422 roundtripContext.appendDialectRegistry(423 op->getContext()->getDialectRegistry());424 if (op->getContext()->allowsUnregisteredDialects())425 roundtripContext.allowUnregisteredDialects();426 StringRef irdlFile = config.getIrdlFile();427 if (!irdlFile.empty() && failed(loadIRDLDialects(irdlFile, roundtripContext)))428 return failure();429 430 std::string testType = (useBytecode) ? "bytecode" : "textual";431 // Print a first time with custom format (or bytecode) and parse it back to432 // the roundtripModule.433 {434 std::string buffer;435 llvm::raw_string_ostream ostream(buffer);436 if (useBytecode) {437 if (failed(writeBytecodeToFile(op, ostream))) {438 op->emitOpError()439 << "failed to write bytecode, cannot verify round-trip.\n";440 return failure();441 }442 } else {443 op->print(ostream,444 OpPrintingFlags().printGenericOpForm().enableDebugInfo());445 }446 FallbackAsmResourceMap fallbackResourceMap;447 ParserConfig parseConfig(&roundtripContext, config.shouldVerifyOnParsing(),448 &fallbackResourceMap);449 roundtripModule = parseSourceString<Operation *>(buffer, parseConfig);450 if (!roundtripModule) {451 op->emitOpError() << "failed to parse " << testType452 << " content back, cannot verify round-trip.\n";453 return failure();454 }455 }456 457 // Print in the generic form for the reference module and the round-tripped458 // one and compare the outputs.459 std::string reference, roundtrip;460 {461 llvm::raw_string_ostream ostreamref(reference);462 op->print(ostreamref,463 OpPrintingFlags().printGenericOpForm().enableDebugInfo());464 llvm::raw_string_ostream ostreamrndtrip(roundtrip);465 roundtripModule.get()->print(466 ostreamrndtrip,467 OpPrintingFlags().printGenericOpForm().enableDebugInfo());468 }469 if (reference != roundtrip) {470 // TODO implement a diff.471 return op->emitOpError()472 << testType473 << " roundTrip testing roundtripped module differs "474 "from reference:\n<<<<<<Reference\n"475 << reference << "\n=====\n"476 << roundtrip << "\n>>>>>roundtripped\n";477 }478 479 return success();480}481 482static LogicalResult doVerifyRoundTrip(Operation *op,483 const MlirOptMainConfig &config) {484 auto txtStatus = doVerifyRoundTrip(op, config, /*useBytecode=*/false);485 auto bcStatus = doVerifyRoundTrip(op, config, /*useBytecode=*/true);486 return success(succeeded(txtStatus) && succeeded(bcStatus));487}488 489/// Perform the actions on the input file indicated by the command line flags490/// within the specified context.491///492/// This typically parses the main source file, runs zero or more optimization493/// passes, then prints the output.494///495static LogicalResult496performActions(raw_ostream &os,497 const std::shared_ptr<llvm::SourceMgr> &sourceMgr,498 MLIRContext *context, const MlirOptMainConfig &config) {499 DefaultTimingManager tm;500 applyDefaultTimingManagerCLOptions(tm);501 TimingScope timing = tm.getRootScope();502 503 // Disable multi-threading when parsing the input file. This removes the504 // unnecessary/costly context synchronization when parsing.505 bool wasThreadingEnabled = context->isMultithreadingEnabled();506 context->disableMultithreading();507 508 // Prepare the parser config, and attach any useful/necessary resource509 // handlers. Unhandled external resources are treated as passthrough, i.e.510 // they are not processed and will be emitted directly to the output511 // untouched.512 PassReproducerOptions reproOptions;513 FallbackAsmResourceMap fallbackResourceMap;514 ParserConfig parseConfig(context, config.shouldVerifyOnParsing(),515 &fallbackResourceMap);516 if (config.shouldRunReproducer())517 reproOptions.attachResourceParser(parseConfig);518 519 // Parse the input file and reset the context threading state.520 TimingScope parserTiming = timing.nest("Parser");521 OwningOpRef<Operation *> op = parseSourceFileForTool(522 sourceMgr, parseConfig, !config.shouldUseExplicitModule());523 parserTiming.stop();524 if (!op)525 return failure();526 527 // Perform round-trip verification if requested528 if (config.shouldVerifyRoundtrip() &&529 failed(doVerifyRoundTrip(op.get(), config)))530 return failure();531 532 context->enableMultithreading(wasThreadingEnabled);533 // Set the remark categories and policy.534 remark::RemarkCategories cats{535 config.getRemarksAllFilter(), config.getRemarksPassedFilter(),536 config.getRemarksMissedFilter(), config.getRemarksAnalyseFilter(),537 config.getRemarksFailedFilter()};538 539 mlir::MLIRContext &ctx = *context;540 // Helper to create the appropriate policy based on configuration541 auto createPolicy = [&config]()542 -> std::unique_ptr<mlir::remark::detail::RemarkEmittingPolicyBase> {543 if (config.getRemarkPolicy() == RemarkPolicy::REMARK_POLICY_ALL)544 return std::make_unique<mlir::remark::RemarkEmittingPolicyAll>();545 if (config.getRemarkPolicy() == RemarkPolicy::REMARK_POLICY_FINAL)546 return std::make_unique<mlir::remark::RemarkEmittingPolicyFinal>();547 548 llvm_unreachable("Invalid remark policy");549 };550 551 switch (config.getRemarkFormat()) {552 case RemarkFormat::REMARK_FORMAT_STDOUT:553 if (failed(mlir::remark::enableOptimizationRemarks(554 ctx, nullptr, createPolicy(), cats, true /*printAsEmitRemarks*/)))555 return failure();556 break;557 558 case RemarkFormat::REMARK_FORMAT_YAML: {559 std::string file = config.getRemarksOutputFile().empty()560 ? "mlir-remarks.yaml"561 : config.getRemarksOutputFile();562 if (failed(mlir::remark::enableOptimizationRemarksWithLLVMStreamer(563 ctx, file, llvm::remarks::Format::YAML, createPolicy(), cats)))564 return failure();565 break;566 }567 568 case RemarkFormat::REMARK_FORMAT_BITSTREAM: {569 std::string file = config.getRemarksOutputFile().empty()570 ? "mlir-remarks.bitstream"571 : config.getRemarksOutputFile();572 if (failed(mlir::remark::enableOptimizationRemarksWithLLVMStreamer(573 ctx, file, llvm::remarks::Format::Bitstream, createPolicy(), cats)))574 return failure();575 break;576 }577 }578 579 // Prepare the pass manager, applying command-line and reproducer options.580 PassManager pm(op.get()->getName(), PassManager::Nesting::Implicit);581 pm.enableVerifier(config.shouldVerifyPasses());582 if (failed(applyPassManagerCLOptions(pm)))583 return failure();584 pm.enableTiming(timing);585 if (config.shouldRunReproducer() && failed(reproOptions.apply(pm)))586 return failure();587 if (failed(config.setupPassPipeline(pm)))588 return failure();589 590 // Run the pipeline.591 if (failed(pm.run(*op)))592 return failure();593 594 // Generate reproducers if requested595 if (!config.getReproducerFilename().empty()) {596 StringRef anchorName = pm.getAnyOpAnchorName();597 const auto &passes = pm.getPasses();598 makeReproducer(anchorName, passes, op.get(),599 config.getReproducerFilename());600 }601 602 // Print the output.603 TimingScope outputTiming = timing.nest("Output");604 if (config.shouldEmitBytecode()) {605 BytecodeWriterConfig writerConfig(fallbackResourceMap);606 if (auto v = config.bytecodeVersionToEmit())607 writerConfig.setDesiredBytecodeVersion(*v);608 if (config.shouldElideResourceDataFromBytecode())609 writerConfig.setElideResourceDataFlag();610 return writeBytecodeToFile(op.get(), os, writerConfig);611 }612 613 if (config.bytecodeVersionToEmit().has_value())614 return emitError(UnknownLoc::get(pm.getContext()))615 << "bytecode version while not emitting bytecode";616 AsmState asmState(op.get(), OpPrintingFlags(), /*locationMap=*/nullptr,617 &fallbackResourceMap);618 os << OpWithState(op.get(), asmState) << '\n';619 620 // This is required if the remark policy is final. Otherwise, the remarks are621 // not emitted.622 if (remark::detail::RemarkEngine *engine = ctx.getRemarkEngine())623 engine->getRemarkEmittingPolicy()->finalize();624 625 return success();626}627 628/// Parses the memory buffer. If successfully, run a series of passes against629/// it and print the result.630static LogicalResult631processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,632 llvm::MemoryBufferRef sourceBuffer,633 const MlirOptMainConfig &config, DialectRegistry ®istry,634 SourceMgrDiagnosticVerifierHandler *verifyHandler,635 llvm::ThreadPoolInterface *threadPool) {636 // Tell sourceMgr about this buffer, which is what the parser will pick up.637 auto sourceMgr = std::make_shared<SourceMgr>();638 // Add the original buffer to the source manager to use for determining639 // locations.640 sourceMgr->AddNewSourceBuffer(641 llvm::MemoryBuffer::getMemBuffer(sourceBuffer,642 /*RequiresNullTerminator=*/false),643 SMLoc());644 sourceMgr->AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());645 646 // Create a context just for the current buffer. Disable threading on647 // creation since we'll inject the thread-pool separately.648 MLIRContext context(registry, MLIRContext::Threading::DISABLED);649 if (threadPool)650 context.setThreadPool(*threadPool);651 if (verifyHandler)652 verifyHandler->registerInContext(&context);653 654 StringRef irdlFile = config.getIrdlFile();655 if (!irdlFile.empty() && failed(loadIRDLDialects(irdlFile, context)))656 return failure();657 658 // Parse the input file.659 context.allowUnregisteredDialects(config.shouldAllowUnregisteredDialects());660 if (config.shouldVerifyDiagnostics())661 context.printOpOnDiagnostic(false);662 663 tracing::InstallDebugHandler installDebugHandler(context,664 config.getDebugConfig());665 666 // If we are in verify diagnostics mode then we have a lot of work to do,667 // otherwise just perform the actions without worrying about it.668 if (!config.shouldVerifyDiagnostics()) {669 SourceMgrDiagnosticHandler sourceMgrHandler(*sourceMgr, &context);670 DiagnosticFilter diagnosticFilter(&context,671 config.getDiagnosticVerbosityLevel(),672 config.shouldShowNotes());673 return performActions(os, sourceMgr, &context, config);674 }675 676 // Do any processing requested by command line flags. We don't care whether677 // these actions succeed or fail, we only care what diagnostics they produce678 // and whether they match our expectations.679 (void)performActions(os, sourceMgr, &context, config);680 681 return success();682}683 684std::string mlir::registerCLIOptions(llvm::StringRef toolName,685 DialectRegistry ®istry) {686 MlirOptMainConfig::registerCLOptions(registry);687 registerAsmPrinterCLOptions();688 registerMLIRContextCLOptions();689 registerPassManagerCLOptions();690 registerDefaultTimingManagerCLOptions();691 tracing::DebugCounter::registerCLOptions();692 693 // Build the list of dialects as a header for the --help message.694 std::string helpHeader = (toolName + "\nAvailable Dialects: ").str();695 {696 llvm::raw_string_ostream os(helpHeader);697 interleaveComma(registry.getDialectNames(), os,698 [&](auto name) { os << name; });699 }700 return helpHeader;701}702 703std::pair<std::string, std::string>704mlir::parseCLIOptions(int argc, char **argv, llvm::StringRef helpHeader) {705 static cl::opt<std::string> inputFilename(706 cl::Positional, cl::desc("<input file>"), cl::init("-"));707 708 static cl::opt<std::string> outputFilename("o", cl::desc("Output filename"),709 cl::value_desc("filename"),710 cl::init("-"));711 cl::ParseCommandLineOptions(argc, argv, helpHeader);712 return std::make_pair(inputFilename.getValue(), outputFilename.getValue());713}714 715std::pair<std::string, std::string>716mlir::registerAndParseCLIOptions(int argc, char **argv,717 llvm::StringRef toolName,718 DialectRegistry ®istry) {719 auto helpHeader = registerCLIOptions(toolName, registry);720 return parseCLIOptions(argc, argv, helpHeader);721}722 723static LogicalResult printRegisteredDialects(DialectRegistry ®istry) {724 llvm::outs() << "Available Dialects: ";725 interleave(registry.getDialectNames(), llvm::outs(), ",");726 llvm::outs() << "\n";727 return success();728}729 730static LogicalResult printRegisteredPassesAndReturn() {731 mlir::printRegisteredPasses();732 return success();733}734 735LogicalResult mlir::MlirOptMain(llvm::raw_ostream &outputStream,736 std::unique_ptr<llvm::MemoryBuffer> buffer,737 DialectRegistry ®istry,738 const MlirOptMainConfig &config) {739 if (config.shouldShowDialects())740 return printRegisteredDialects(registry);741 742 if (config.shouldListPasses())743 return printRegisteredPassesAndReturn();744 745 // The split-input-file mode is a very specific mode that slices the file746 // up into small pieces and checks each independently.747 // We use an explicit threadpool to avoid creating and joining/destroying748 // threads for each of the split.749 ThreadPoolInterface *threadPool = nullptr;750 751 // Create a temporary context for the sake of checking if752 // --mlir-disable-threading was passed on the command line.753 // We use the thread-pool this context is creating, and avoid754 // creating any thread when disabled.755 MLIRContext threadPoolCtx;756 if (threadPoolCtx.isMultithreadingEnabled())757 threadPool = &threadPoolCtx.getThreadPool();758 759 SourceMgr sourceMgr;760 sourceMgr.AddNewSourceBuffer(761 llvm::MemoryBuffer::getMemBuffer(buffer->getMemBufferRef(),762 /*RequiresNullTerminator=*/false),763 SMLoc());764 // Note: this creates a verifier handler independent of the the flag set, as765 // internally if the flag is not set, a new scoped diagnostic handler is766 // created which would intercept the diagnostics and verify them.767 SourceMgrDiagnosticVerifierHandler sourceMgrHandler(768 sourceMgr, &threadPoolCtx, config.verifyDiagnosticsLevel());769 auto chunkFn = [&](std::unique_ptr<MemoryBuffer> chunkBuffer,770 llvm::MemoryBufferRef sourceBuffer, raw_ostream &os) {771 return processBuffer(772 os, std::move(chunkBuffer), sourceBuffer, config, registry,773 config.shouldVerifyDiagnostics() ? &sourceMgrHandler : nullptr,774 threadPool);775 };776 LogicalResult status = splitAndProcessBuffer(777 llvm::MemoryBuffer::getMemBuffer(buffer->getMemBufferRef(),778 /*RequiresNullTerminator=*/false),779 chunkFn, outputStream, config.inputSplitMarker(),780 config.outputSplitMarker());781 if (config.shouldVerifyDiagnostics() && failed(sourceMgrHandler.verify()))782 status = failure();783 return status;784}785 786LogicalResult mlir::MlirOptMain(int argc, char **argv,787 llvm::StringRef inputFilename,788 llvm::StringRef outputFilename,789 DialectRegistry ®istry) {790 791 InitLLVM y(argc, argv);792 793 MlirOptMainConfig config = MlirOptMainConfig::createFromCLOptions();794 795 if (config.shouldShowDialects())796 return printRegisteredDialects(registry);797 798 if (config.shouldListPasses())799 return printRegisteredPassesAndReturn();800 801 // When reading from stdin and the input is a tty, it is often a user802 // mistake and the process "appears to be stuck". Print a message to let the803 // user know about it!804 if (inputFilename == "-" &&805 sys::Process::FileDescriptorIsDisplayed(fileno(stdin)))806 llvm::errs() << "(processing input from stdin now, hit ctrl-c/ctrl-d to "807 "interrupt)\n";808 809 // Set up the input file.810 std::string errorMessage;811 auto file = openInputFile(inputFilename, &errorMessage);812 if (!file) {813 llvm::errs() << errorMessage << "\n";814 return failure();815 }816 817 auto output = openOutputFile(outputFilename, &errorMessage);818 if (!output) {819 llvm::errs() << errorMessage << "\n";820 return failure();821 }822 if (failed(MlirOptMain(output->os(), std::move(file), registry, config)))823 return failure();824 825 // Keep the output file if the invocation of MlirOptMain was successful.826 output->keep();827 return success();828}829 830LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,831 DialectRegistry ®istry) {832 833 // Register and parse command line options.834 std::string inputFilename, outputFilename;835 std::tie(inputFilename, outputFilename) =836 registerAndParseCLIOptions(argc, argv, toolName, registry);837 838 return MlirOptMain(argc, argv, inputFilename, outputFilename, registry);839}840