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-exception |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | // |
9 | // This is a utility that runs an optimization pass and prints the result back |
10 | // 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/Debug/DebuggerExecutionContextHook.h" |
19 | #include "mlir/Debug/ExecutionContext.h" |
20 | #include "mlir/Debug/Observers/ActionLogging.h" |
21 | #include "mlir/Dialect/IRDL/IR/IRDL.h" |
22 | #include "mlir/Dialect/IRDL/IRDLLoading.h" |
23 | #include "mlir/IR/AsmState.h" |
24 | #include "mlir/IR/Attributes.h" |
25 | #include "mlir/IR/BuiltinOps.h" |
26 | #include "mlir/IR/Diagnostics.h" |
27 | #include "mlir/IR/Dialect.h" |
28 | #include "mlir/IR/Location.h" |
29 | #include "mlir/IR/MLIRContext.h" |
30 | #include "mlir/Parser/Parser.h" |
31 | #include "mlir/Pass/Pass.h" |
32 | #include "mlir/Pass/PassManager.h" |
33 | #include "mlir/Support/FileUtilities.h" |
34 | #include "mlir/Support/LogicalResult.h" |
35 | #include "mlir/Support/Timing.h" |
36 | #include "mlir/Support/ToolUtilities.h" |
37 | #include "mlir/Tools/ParseUtilities.h" |
38 | #include "mlir/Tools/Plugins/DialectPlugin.h" |
39 | #include "mlir/Tools/Plugins/PassPlugin.h" |
40 | #include "llvm/ADT/StringRef.h" |
41 | #include "llvm/Support/CommandLine.h" |
42 | #include "llvm/Support/FileUtilities.h" |
43 | #include "llvm/Support/InitLLVM.h" |
44 | #include "llvm/Support/ManagedStatic.h" |
45 | #include "llvm/Support/Process.h" |
46 | #include "llvm/Support/Regex.h" |
47 | #include "llvm/Support/SourceMgr.h" |
48 | #include "llvm/Support/StringSaver.h" |
49 | #include "llvm/Support/ThreadPool.h" |
50 | #include "llvm/Support/ToolOutputFile.h" |
51 | |
52 | using namespace mlir; |
53 | using namespace llvm; |
54 | |
55 | namespace { |
56 | class BytecodeVersionParser : public cl::parser<std::optional<int64_t>> { |
57 | public: |
58 | BytecodeVersionParser(cl::Option &o) |
59 | : cl::parser<std::optional<int64_t>>(o) {} |
60 | |
61 | bool parse(cl::Option &o, StringRef /*argName*/, StringRef arg, |
62 | std::optional<int64_t> &v) { |
63 | long long w; |
64 | if (getAsSignedInteger(Str: arg, Radix: 10, Result&: w)) |
65 | return o.error(Message: "Invalid argument '" + arg + |
66 | "', only integer is supported." ); |
67 | v = w; |
68 | return false; |
69 | } |
70 | }; |
71 | |
72 | /// This class is intended to manage the handling of command line options for |
73 | /// creating a *-opt config. This is a singleton. |
74 | struct MlirOptMainConfigCLOptions : public MlirOptMainConfig { |
75 | MlirOptMainConfigCLOptions() { |
76 | // These options are static but all uses ExternalStorage to initialize the |
77 | // members of the parent class. This is unusual but since this class is a |
78 | // singleton it basically attaches command line option to the singleton |
79 | // members. |
80 | |
81 | static cl::opt<bool, /*ExternalStorage=*/true> allowUnregisteredDialects( |
82 | "allow-unregistered-dialect" , |
83 | cl::desc("Allow operation with no registered dialects" ), |
84 | cl::location(L&: allowUnregisteredDialectsFlag), cl::init(Val: false)); |
85 | |
86 | static cl::opt<bool, /*ExternalStorage=*/true> dumpPassPipeline( |
87 | "dump-pass-pipeline" , cl::desc("Print the pipeline that will be run" ), |
88 | cl::location(L&: dumpPassPipelineFlag), cl::init(Val: false)); |
89 | |
90 | static cl::opt<bool, /*ExternalStorage=*/true> emitBytecode( |
91 | "emit-bytecode" , cl::desc("Emit bytecode when generating output" ), |
92 | cl::location(L&: emitBytecodeFlag), cl::init(Val: false)); |
93 | |
94 | static cl::opt<bool, /*ExternalStorage=*/true> elideResourcesFromBytecode( |
95 | "elide-resource-data-from-bytecode" , |
96 | cl::desc("Elide resources when generating bytecode" ), |
97 | cl::location(L&: elideResourceDataFromBytecodeFlag), cl::init(Val: false)); |
98 | |
99 | static cl::opt<std::optional<int64_t>, /*ExternalStorage=*/true, |
100 | BytecodeVersionParser> |
101 | bytecodeVersion( |
102 | "emit-bytecode-version" , |
103 | cl::desc("Use specified bytecode when generating output" ), |
104 | cl::location(L&: emitBytecodeVersion), cl::init(Val: std::nullopt)); |
105 | |
106 | static cl::opt<std::string, /*ExternalStorage=*/true> irdlFile( |
107 | "irdl-file" , |
108 | cl::desc("IRDL file to register before processing the input" ), |
109 | cl::location(L&: irdlFileFlag), cl::init(Val: "" ), cl::value_desc("filename" )); |
110 | |
111 | static cl::opt<bool, /*ExternalStorage=*/true> enableDebuggerHook( |
112 | "mlir-enable-debugger-hook" , |
113 | cl::desc("Enable Debugger hook for debugging MLIR Actions" ), |
114 | cl::location(L&: enableDebuggerActionHookFlag), cl::init(Val: false)); |
115 | |
116 | static cl::opt<bool, /*ExternalStorage=*/true> explicitModule( |
117 | "no-implicit-module" , |
118 | cl::desc("Disable implicit addition of a top-level module op during " |
119 | "parsing" ), |
120 | cl::location(L&: useExplicitModuleFlag), cl::init(Val: false)); |
121 | |
122 | static cl::opt<bool, /*ExternalStorage=*/true> runReproducer( |
123 | "run-reproducer" , cl::desc("Run the pipeline stored in the reproducer" ), |
124 | cl::location(L&: runReproducerFlag), cl::init(Val: false)); |
125 | |
126 | static cl::opt<bool, /*ExternalStorage=*/true> showDialects( |
127 | "show-dialects" , |
128 | cl::desc("Print the list of registered dialects and exit" ), |
129 | cl::location(L&: showDialectsFlag), cl::init(Val: false)); |
130 | |
131 | static cl::opt<std::string, /*ExternalStorage=*/true> splitInputFile{ |
132 | "split-input-file" , llvm::cl::ValueOptional, |
133 | cl::callback(CB: [&](const std::string &str) { |
134 | // Implicit value: use default marker if flag was used without value. |
135 | if (str.empty()) |
136 | splitInputFile.setValue(V: kDefaultSplitMarker); |
137 | }), |
138 | cl::desc("Split the input file into chunks using the given or " |
139 | "default marker and process each chunk independently" ), |
140 | cl::location(L&: splitInputFileFlag), cl::init(Val: "" )}; |
141 | |
142 | static cl::opt<std::string, /*ExternalStorage=*/true> outputSplitMarker( |
143 | "output-split-marker" , |
144 | cl::desc("Split marker to use for merging the ouput" ), |
145 | cl::location(L&: outputSplitMarkerFlag), cl::init(Val: kDefaultSplitMarker)); |
146 | |
147 | static cl::opt<bool, /*ExternalStorage=*/true> verifyDiagnostics( |
148 | "verify-diagnostics" , |
149 | cl::desc("Check that emitted diagnostics match " |
150 | "expected-* lines on the corresponding line" ), |
151 | cl::location(L&: verifyDiagnosticsFlag), cl::init(Val: false)); |
152 | |
153 | static cl::opt<bool, /*ExternalStorage=*/true> verifyPasses( |
154 | "verify-each" , |
155 | cl::desc("Run the verifier after each transformation pass" ), |
156 | cl::location(L&: verifyPassesFlag), cl::init(Val: true)); |
157 | |
158 | static cl::opt<bool, /*ExternalStorage=*/true> verifyRoundtrip( |
159 | "verify-roundtrip" , |
160 | cl::desc("Round-trip the IR after parsing and ensure it succeeds" ), |
161 | cl::location(L&: verifyRoundtripFlag), cl::init(Val: false)); |
162 | |
163 | static cl::list<std::string> passPlugins( |
164 | "load-pass-plugin" , cl::desc("Load passes from plugin library" )); |
165 | |
166 | static cl::opt<std::string, /*ExternalStorage=*/true> |
167 | generateReproducerFile( |
168 | "mlir-generate-reproducer" , |
169 | llvm::cl::desc( |
170 | "Generate an mlir reproducer at the provided filename" |
171 | " (no crash required)" ), |
172 | cl::location(L&: generateReproducerFileFlag), cl::init(Val: "" ), |
173 | cl::value_desc("filename" )); |
174 | |
175 | /// Set the callback to load a pass plugin. |
176 | passPlugins.setCallback([&](const std::string &pluginPath) { |
177 | auto plugin = PassPlugin::load(filename: pluginPath); |
178 | if (!plugin) { |
179 | errs() << "Failed to load passes from '" << pluginPath |
180 | << "'. Request ignored.\n" ; |
181 | return; |
182 | } |
183 | plugin.get().registerPassRegistryCallbacks(); |
184 | }); |
185 | |
186 | static cl::list<std::string> dialectPlugins( |
187 | "load-dialect-plugin" , cl::desc("Load dialects from plugin library" )); |
188 | this->dialectPlugins = std::addressof(r&: dialectPlugins); |
189 | |
190 | static PassPipelineCLParser passPipeline("" , "Compiler passes to run" , "p" ); |
191 | setPassPipelineParser(passPipeline); |
192 | } |
193 | |
194 | /// Set the callback to load a dialect plugin. |
195 | void setDialectPluginsCallback(DialectRegistry ®istry); |
196 | |
197 | /// Pointer to static dialectPlugins variable in constructor, needed by |
198 | /// setDialectPluginsCallback(DialectRegistry&). |
199 | cl::list<std::string> *dialectPlugins = nullptr; |
200 | }; |
201 | } // namespace |
202 | |
203 | ManagedStatic<MlirOptMainConfigCLOptions> clOptionsConfig; |
204 | |
205 | void MlirOptMainConfig::registerCLOptions(DialectRegistry ®istry) { |
206 | clOptionsConfig->setDialectPluginsCallback(registry); |
207 | tracing::DebugConfig::registerCLOptions(); |
208 | } |
209 | |
210 | MlirOptMainConfig MlirOptMainConfig::createFromCLOptions() { |
211 | clOptionsConfig->setDebugConfig(tracing::DebugConfig::createFromCLOptions()); |
212 | return *clOptionsConfig; |
213 | } |
214 | |
215 | MlirOptMainConfig &MlirOptMainConfig::setPassPipelineParser( |
216 | const PassPipelineCLParser &passPipeline) { |
217 | passPipelineCallback = [&](PassManager &pm) { |
218 | auto errorHandler = [&](const Twine &msg) { |
219 | emitError(UnknownLoc::get(pm.getContext())) << msg; |
220 | return failure(); |
221 | }; |
222 | if (failed(result: passPipeline.addToPipeline(pm, errorHandler))) |
223 | return failure(); |
224 | if (this->shouldDumpPassPipeline()) { |
225 | |
226 | pm.dump(); |
227 | llvm::errs() << "\n" ; |
228 | } |
229 | return success(); |
230 | }; |
231 | return *this; |
232 | } |
233 | |
234 | void MlirOptMainConfigCLOptions::setDialectPluginsCallback( |
235 | DialectRegistry ®istry) { |
236 | dialectPlugins->setCallback([&](const std::string &pluginPath) { |
237 | auto plugin = DialectPlugin::load(filename: pluginPath); |
238 | if (!plugin) { |
239 | errs() << "Failed to load dialect plugin from '" << pluginPath |
240 | << "'. Request ignored.\n" ; |
241 | return; |
242 | }; |
243 | plugin.get().registerDialectRegistryCallbacks(dialectRegistry&: registry); |
244 | }); |
245 | } |
246 | |
247 | LogicalResult loadIRDLDialects(StringRef irdlFile, MLIRContext &ctx) { |
248 | DialectRegistry registry; |
249 | registry.insert<irdl::IRDLDialect>(); |
250 | ctx.appendDialectRegistry(registry); |
251 | |
252 | // Set up the input file. |
253 | std::string errorMessage; |
254 | std::unique_ptr<MemoryBuffer> file = openInputFile(inputFilename: irdlFile, errorMessage: &errorMessage); |
255 | if (!file) { |
256 | emitError(UnknownLoc::get(&ctx)) << errorMessage; |
257 | return failure(); |
258 | } |
259 | |
260 | // Give the buffer to the source manager. |
261 | // This will be picked up by the parser. |
262 | SourceMgr sourceMgr; |
263 | sourceMgr.AddNewSourceBuffer(F: std::move(file), IncludeLoc: SMLoc()); |
264 | |
265 | SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &ctx); |
266 | |
267 | // Parse the input file. |
268 | OwningOpRef<ModuleOp> module(parseSourceFile<ModuleOp>(sourceMgr, &ctx)); |
269 | |
270 | // Load IRDL dialects. |
271 | return irdl::loadDialects(module.get()); |
272 | } |
273 | |
274 | // Return success if the module can correctly round-trip. This intended to test |
275 | // that the custom printers/parsers are complete. |
276 | static LogicalResult doVerifyRoundTrip(Operation *op, |
277 | const MlirOptMainConfig &config, |
278 | bool useBytecode) { |
279 | // We use a new context to avoid resource handle renaming issue in the diff. |
280 | MLIRContext roundtripContext; |
281 | OwningOpRef<Operation *> roundtripModule; |
282 | roundtripContext.appendDialectRegistry( |
283 | registry: op->getContext()->getDialectRegistry()); |
284 | if (op->getContext()->allowsUnregisteredDialects()) |
285 | roundtripContext.allowUnregisteredDialects(); |
286 | StringRef irdlFile = config.getIrdlFile(); |
287 | if (!irdlFile.empty() && failed(result: loadIRDLDialects(irdlFile, ctx&: roundtripContext))) |
288 | return failure(); |
289 | |
290 | std::string testType = (useBytecode) ? "bytecode" : "textual" ; |
291 | // Print a first time with custom format (or bytecode) and parse it back to |
292 | // the roundtripModule. |
293 | { |
294 | std::string buffer; |
295 | llvm::raw_string_ostream ostream(buffer); |
296 | if (useBytecode) { |
297 | if (failed(result: writeBytecodeToFile(op, os&: ostream))) { |
298 | op->emitOpError() |
299 | << "failed to write bytecode, cannot verify round-trip.\n" ; |
300 | return failure(); |
301 | } |
302 | } else { |
303 | op->print(os&: ostream, |
304 | flags: OpPrintingFlags().printGenericOpForm().enableDebugInfo()); |
305 | } |
306 | FallbackAsmResourceMap fallbackResourceMap; |
307 | ParserConfig parseConfig(&roundtripContext, /*verifyAfterParse=*/true, |
308 | &fallbackResourceMap); |
309 | roundtripModule = |
310 | parseSourceString<Operation *>(sourceStr: ostream.str(), config: parseConfig); |
311 | if (!roundtripModule) { |
312 | op->emitOpError() << "failed to parse " << testType |
313 | << " content back, cannot verify round-trip.\n" ; |
314 | return failure(); |
315 | } |
316 | } |
317 | |
318 | // Print in the generic form for the reference module and the round-tripped |
319 | // one and compare the outputs. |
320 | std::string reference, roundtrip; |
321 | { |
322 | llvm::raw_string_ostream ostreamref(reference); |
323 | op->print(os&: ostreamref, |
324 | flags: OpPrintingFlags().printGenericOpForm().enableDebugInfo()); |
325 | llvm::raw_string_ostream ostreamrndtrip(roundtrip); |
326 | roundtripModule.get()->print( |
327 | os&: ostreamrndtrip, |
328 | flags: OpPrintingFlags().printGenericOpForm().enableDebugInfo()); |
329 | } |
330 | if (reference != roundtrip) { |
331 | // TODO implement a diff. |
332 | return op->emitOpError() |
333 | << testType |
334 | << " roundTrip testing roundtripped module differs " |
335 | "from reference:\n<<<<<<Reference\n" |
336 | << reference << "\n=====\n" |
337 | << roundtrip << "\n>>>>>roundtripped\n" ; |
338 | } |
339 | |
340 | return success(); |
341 | } |
342 | |
343 | static LogicalResult doVerifyRoundTrip(Operation *op, |
344 | const MlirOptMainConfig &config) { |
345 | auto txtStatus = doVerifyRoundTrip(op, config, /*useBytecode=*/false); |
346 | auto bcStatus = doVerifyRoundTrip(op, config, /*useBytecode=*/true); |
347 | return success(isSuccess: succeeded(result: txtStatus) && succeeded(result: bcStatus)); |
348 | } |
349 | |
350 | /// Perform the actions on the input file indicated by the command line flags |
351 | /// within the specified context. |
352 | /// |
353 | /// This typically parses the main source file, runs zero or more optimization |
354 | /// passes, then prints the output. |
355 | /// |
356 | static LogicalResult |
357 | performActions(raw_ostream &os, |
358 | const std::shared_ptr<llvm::SourceMgr> &sourceMgr, |
359 | MLIRContext *context, const MlirOptMainConfig &config) { |
360 | DefaultTimingManager tm; |
361 | applyDefaultTimingManagerCLOptions(tm); |
362 | TimingScope timing = tm.getRootScope(); |
363 | |
364 | // Disable multi-threading when parsing the input file. This removes the |
365 | // unnecessary/costly context synchronization when parsing. |
366 | bool wasThreadingEnabled = context->isMultithreadingEnabled(); |
367 | context->disableMultithreading(); |
368 | |
369 | // Prepare the parser config, and attach any useful/necessary resource |
370 | // handlers. Unhandled external resources are treated as passthrough, i.e. |
371 | // they are not processed and will be emitted directly to the output |
372 | // untouched. |
373 | PassReproducerOptions reproOptions; |
374 | FallbackAsmResourceMap fallbackResourceMap; |
375 | ParserConfig parseConfig(context, /*verifyAfterParse=*/true, |
376 | &fallbackResourceMap); |
377 | if (config.shouldRunReproducer()) |
378 | reproOptions.attachResourceParser(config&: parseConfig); |
379 | |
380 | // Parse the input file and reset the context threading state. |
381 | TimingScope parserTiming = timing.nest(args: "Parser" ); |
382 | OwningOpRef<Operation *> op = parseSourceFileForTool( |
383 | sourceMgr, config: parseConfig, insertImplicitModule: !config.shouldUseExplicitModule()); |
384 | parserTiming.stop(); |
385 | if (!op) |
386 | return failure(); |
387 | |
388 | // Perform round-trip verification if requested |
389 | if (config.shouldVerifyRoundtrip() && |
390 | failed(result: doVerifyRoundTrip(op: op.get(), config))) |
391 | return failure(); |
392 | |
393 | context->enableMultithreading(enable: wasThreadingEnabled); |
394 | |
395 | // Prepare the pass manager, applying command-line and reproducer options. |
396 | PassManager pm(op.get()->getName(), PassManager::Nesting::Implicit); |
397 | pm.enableVerifier(enabled: config.shouldVerifyPasses()); |
398 | if (failed(result: applyPassManagerCLOptions(pm))) |
399 | return failure(); |
400 | pm.enableTiming(timingScope&: timing); |
401 | if (config.shouldRunReproducer() && failed(result: reproOptions.apply(pm))) |
402 | return failure(); |
403 | if (failed(result: config.setupPassPipeline(pm))) |
404 | return failure(); |
405 | |
406 | // Run the pipeline. |
407 | if (failed(result: pm.run(op: *op))) |
408 | return failure(); |
409 | |
410 | // Generate reproducers if requested |
411 | if (!config.getReproducerFilename().empty()) { |
412 | StringRef anchorName = pm.getAnyOpAnchorName(); |
413 | const auto &passes = pm.getPasses(); |
414 | makeReproducer(anchorName, passes, op: op.get(), |
415 | outputFile: config.getReproducerFilename()); |
416 | } |
417 | |
418 | // Print the output. |
419 | TimingScope outputTiming = timing.nest(args: "Output" ); |
420 | if (config.shouldEmitBytecode()) { |
421 | BytecodeWriterConfig writerConfig(fallbackResourceMap); |
422 | if (auto v = config.bytecodeVersionToEmit()) |
423 | writerConfig.setDesiredBytecodeVersion(*v); |
424 | if (config.shouldElideResourceDataFromBytecode()) |
425 | writerConfig.setElideResourceDataFlag(); |
426 | return writeBytecodeToFile(op: op.get(), os, config: writerConfig); |
427 | } |
428 | |
429 | if (config.bytecodeVersionToEmit().has_value()) |
430 | return emitError(UnknownLoc::get(pm.getContext())) |
431 | << "bytecode version while not emitting bytecode" ; |
432 | AsmState asmState(op.get(), OpPrintingFlags(), /*locationMap=*/nullptr, |
433 | &fallbackResourceMap); |
434 | op.get()->print(os, state&: asmState); |
435 | os << '\n'; |
436 | return success(); |
437 | } |
438 | |
439 | /// Parses the memory buffer. If successfully, run a series of passes against |
440 | /// it and print the result. |
441 | static LogicalResult processBuffer(raw_ostream &os, |
442 | std::unique_ptr<MemoryBuffer> ownedBuffer, |
443 | const MlirOptMainConfig &config, |
444 | DialectRegistry ®istry, |
445 | llvm::ThreadPoolInterface *threadPool) { |
446 | // Tell sourceMgr about this buffer, which is what the parser will pick up. |
447 | auto sourceMgr = std::make_shared<SourceMgr>(); |
448 | sourceMgr->AddNewSourceBuffer(F: std::move(ownedBuffer), IncludeLoc: SMLoc()); |
449 | |
450 | // Create a context just for the current buffer. Disable threading on creation |
451 | // since we'll inject the thread-pool separately. |
452 | MLIRContext context(registry, MLIRContext::Threading::DISABLED); |
453 | if (threadPool) |
454 | context.setThreadPool(*threadPool); |
455 | |
456 | StringRef irdlFile = config.getIrdlFile(); |
457 | if (!irdlFile.empty() && failed(result: loadIRDLDialects(irdlFile, ctx&: context))) |
458 | return failure(); |
459 | |
460 | // Parse the input file. |
461 | context.allowUnregisteredDialects(allow: config.shouldAllowUnregisteredDialects()); |
462 | if (config.shouldVerifyDiagnostics()) |
463 | context.printOpOnDiagnostic(enable: false); |
464 | |
465 | tracing::InstallDebugHandler installDebugHandler(context, |
466 | config.getDebugConfig()); |
467 | |
468 | // If we are in verify diagnostics mode then we have a lot of work to do, |
469 | // otherwise just perform the actions without worrying about it. |
470 | if (!config.shouldVerifyDiagnostics()) { |
471 | SourceMgrDiagnosticHandler sourceMgrHandler(*sourceMgr, &context); |
472 | return performActions(os, sourceMgr, context: &context, config); |
473 | } |
474 | |
475 | SourceMgrDiagnosticVerifierHandler sourceMgrHandler(*sourceMgr, &context); |
476 | |
477 | // Do any processing requested by command line flags. We don't care whether |
478 | // these actions succeed or fail, we only care what diagnostics they produce |
479 | // and whether they match our expectations. |
480 | (void)performActions(os, sourceMgr, context: &context, config); |
481 | |
482 | // Verify the diagnostic handler to make sure that each of the diagnostics |
483 | // matched. |
484 | return sourceMgrHandler.verify(); |
485 | } |
486 | |
487 | std::pair<std::string, std::string> |
488 | mlir::registerAndParseCLIOptions(int argc, char **argv, |
489 | llvm::StringRef toolName, |
490 | DialectRegistry ®istry) { |
491 | static cl::opt<std::string> inputFilename( |
492 | cl::Positional, cl::desc("<input file>" ), cl::init(Val: "-" )); |
493 | |
494 | static cl::opt<std::string> outputFilename("o" , cl::desc("Output filename" ), |
495 | cl::value_desc("filename" ), |
496 | cl::init(Val: "-" )); |
497 | // Register any command line options. |
498 | MlirOptMainConfig::registerCLOptions(registry); |
499 | registerAsmPrinterCLOptions(); |
500 | registerMLIRContextCLOptions(); |
501 | registerPassManagerCLOptions(); |
502 | registerDefaultTimingManagerCLOptions(); |
503 | tracing::DebugCounter::registerCLOptions(); |
504 | |
505 | // Build the list of dialects as a header for the --help message. |
506 | std::string = (toolName + "\nAvailable Dialects: " ).str(); |
507 | { |
508 | llvm::raw_string_ostream os(helpHeader); |
509 | interleaveComma(c: registry.getDialectNames(), os, |
510 | each_fn: [&](auto name) { os << name; }); |
511 | } |
512 | // Parse pass names in main to ensure static initialization completed. |
513 | cl::ParseCommandLineOptions(argc, argv, Overview: helpHeader); |
514 | return std::make_pair(x&: inputFilename.getValue(), y&: outputFilename.getValue()); |
515 | } |
516 | |
517 | static LogicalResult printRegisteredDialects(DialectRegistry ®istry) { |
518 | llvm::outs() << "Available Dialects: " ; |
519 | interleave(c: registry.getDialectNames(), os&: llvm::outs(), separator: "," ); |
520 | llvm::outs() << "\n" ; |
521 | return success(); |
522 | } |
523 | |
524 | LogicalResult mlir::MlirOptMain(llvm::raw_ostream &outputStream, |
525 | std::unique_ptr<llvm::MemoryBuffer> buffer, |
526 | DialectRegistry ®istry, |
527 | const MlirOptMainConfig &config) { |
528 | if (config.shouldShowDialects()) |
529 | return printRegisteredDialects(registry); |
530 | |
531 | // The split-input-file mode is a very specific mode that slices the file |
532 | // up into small pieces and checks each independently. |
533 | // We use an explicit threadpool to avoid creating and joining/destroying |
534 | // threads for each of the split. |
535 | ThreadPoolInterface *threadPool = nullptr; |
536 | |
537 | // Create a temporary context for the sake of checking if |
538 | // --mlir-disable-threading was passed on the command line. |
539 | // We use the thread-pool this context is creating, and avoid |
540 | // creating any thread when disabled. |
541 | MLIRContext threadPoolCtx; |
542 | if (threadPoolCtx.isMultithreadingEnabled()) |
543 | threadPool = &threadPoolCtx.getThreadPool(); |
544 | |
545 | auto chunkFn = [&](std::unique_ptr<MemoryBuffer> chunkBuffer, |
546 | raw_ostream &os) { |
547 | return processBuffer(os, ownedBuffer: std::move(chunkBuffer), config, registry, |
548 | threadPool); |
549 | }; |
550 | return splitAndProcessBuffer(originalBuffer: std::move(buffer), processChunkBuffer: chunkFn, os&: outputStream, |
551 | inputSplitMarker: config.inputSplitMarker(), |
552 | outputSplitMarker: config.outputSplitMarker()); |
553 | } |
554 | |
555 | LogicalResult mlir::MlirOptMain(int argc, char **argv, |
556 | llvm::StringRef inputFilename, |
557 | llvm::StringRef outputFilename, |
558 | DialectRegistry ®istry) { |
559 | |
560 | InitLLVM y(argc, argv); |
561 | |
562 | MlirOptMainConfig config = MlirOptMainConfig::createFromCLOptions(); |
563 | |
564 | if (config.shouldShowDialects()) |
565 | return printRegisteredDialects(registry); |
566 | |
567 | // When reading from stdin and the input is a tty, it is often a user mistake |
568 | // and the process "appears to be stuck". Print a message to let the user know |
569 | // about it! |
570 | if (inputFilename == "-" && |
571 | sys::Process::FileDescriptorIsDisplayed(fd: fileno(stdin))) |
572 | llvm::errs() << "(processing input from stdin now, hit ctrl-c/ctrl-d to " |
573 | "interrupt)\n" ; |
574 | |
575 | // Set up the input file. |
576 | std::string errorMessage; |
577 | auto file = openInputFile(inputFilename, errorMessage: &errorMessage); |
578 | if (!file) { |
579 | llvm::errs() << errorMessage << "\n" ; |
580 | return failure(); |
581 | } |
582 | |
583 | auto output = openOutputFile(outputFilename, errorMessage: &errorMessage); |
584 | if (!output) { |
585 | llvm::errs() << errorMessage << "\n" ; |
586 | return failure(); |
587 | } |
588 | if (failed(result: MlirOptMain(outputStream&: output->os(), buffer: std::move(file), registry, config))) |
589 | return failure(); |
590 | |
591 | // Keep the output file if the invocation of MlirOptMain was successful. |
592 | output->keep(); |
593 | return success(); |
594 | } |
595 | |
596 | LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName, |
597 | DialectRegistry ®istry) { |
598 | |
599 | // Register and parse command line options. |
600 | std::string inputFilename, outputFilename; |
601 | std::tie(args&: inputFilename, args&: outputFilename) = |
602 | registerAndParseCLIOptions(argc, argv, toolName, registry); |
603 | |
604 | return MlirOptMain(argc, argv, inputFilename, outputFilename, registry); |
605 | } |
606 | |