1 | //===- MlirOptMain.h - MLIR Optimizer Driver main ---------------*- C++ -*-===// |
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 | // Main entry function for mlir-opt for when built as standalone binary. |
10 | // |
11 | //===----------------------------------------------------------------------===// |
12 | |
13 | #ifndef MLIR_TOOLS_MLIROPT_MLIROPTMAIN_H |
14 | #define MLIR_TOOLS_MLIROPT_MLIROPTMAIN_H |
15 | |
16 | #include "mlir/Debug/CLOptionsSetup.h" |
17 | #include "mlir/Support/ToolUtilities.h" |
18 | #include "llvm/ADT/StringRef.h" |
19 | |
20 | #include <cstdlib> |
21 | #include <functional> |
22 | #include <memory> |
23 | |
24 | namespace llvm { |
25 | class raw_ostream; |
26 | class MemoryBuffer; |
27 | } // namespace llvm |
28 | |
29 | namespace mlir { |
30 | class DialectRegistry; |
31 | class PassPipelineCLParser; |
32 | class PassManager; |
33 | |
34 | /// enum class to indicate the verbosity level of the diagnostic filter. |
35 | enum class VerbosityLevel { |
36 | ErrorsOnly = 0, |
37 | ErrorsAndWarnings, |
38 | ErrorsWarningsAndRemarks |
39 | }; |
40 | |
41 | /// Configuration options for the mlir-opt tool. |
42 | /// This is intended to help building tools like mlir-opt by collecting the |
43 | /// supported options. |
44 | /// The API is fluent, and the options are sorted in alphabetical order below. |
45 | /// The options can be exposed to the LLVM command line by registering them |
46 | /// with `MlirOptMainConfig::registerCLOptions(DialectRegistry &);` and creating |
47 | /// a config using `auto config = MlirOptMainConfig::createFromCLOptions();`. |
48 | class MlirOptMainConfig { |
49 | public: |
50 | /// Register the options as global LLVM command line options. |
51 | static void registerCLOptions(DialectRegistry &dialectRegistry); |
52 | |
53 | /// Create a new config with the default set from the CL options. |
54 | static MlirOptMainConfig createFromCLOptions(); |
55 | |
56 | /// |
57 | /// Options. |
58 | /// |
59 | |
60 | /// Allow operation with no registered dialects. |
61 | /// This option is for convenience during testing only and discouraged in |
62 | /// general. |
63 | MlirOptMainConfig &allowUnregisteredDialects(bool allow) { |
64 | allowUnregisteredDialectsFlag = allow; |
65 | return *this; |
66 | } |
67 | bool shouldAllowUnregisteredDialects() const { |
68 | return allowUnregisteredDialectsFlag; |
69 | } |
70 | |
71 | /// Set the debug configuration to use. |
72 | MlirOptMainConfig &setDebugConfig(tracing::DebugConfig config) { |
73 | debugConfig = std::move(config); |
74 | return *this; |
75 | } |
76 | tracing::DebugConfig &getDebugConfig() { return debugConfig; } |
77 | const tracing::DebugConfig &getDebugConfig() const { return debugConfig; } |
78 | |
79 | /// Print the pass-pipeline as text before executing. |
80 | MlirOptMainConfig &dumpPassPipeline(bool dump) { |
81 | dumpPassPipelineFlag = dump; |
82 | return *this; |
83 | } |
84 | |
85 | VerbosityLevel getDiagnosticVerbosityLevel() const { |
86 | return diagnosticVerbosityLevelFlag; |
87 | } |
88 | |
89 | bool shouldDumpPassPipeline() const { return dumpPassPipelineFlag; } |
90 | |
91 | /// Set the output format to bytecode instead of textual IR. |
92 | MlirOptMainConfig &emitBytecode(bool emit) { |
93 | emitBytecodeFlag = emit; |
94 | return *this; |
95 | } |
96 | bool shouldEmitBytecode() const { return emitBytecodeFlag; } |
97 | |
98 | bool shouldElideResourceDataFromBytecode() const { |
99 | return elideResourceDataFromBytecodeFlag; |
100 | } |
101 | |
102 | bool shouldShowNotes() const { return !disableDiagnosticNotesFlag; } |
103 | |
104 | /// Set the IRDL file to load before processing the input. |
105 | MlirOptMainConfig &setIrdlFile(StringRef file) { |
106 | irdlFileFlag = file; |
107 | return *this; |
108 | } |
109 | StringRef getIrdlFile() const { return irdlFileFlag; } |
110 | |
111 | /// Set the bytecode version to emit. |
112 | MlirOptMainConfig &setEmitBytecodeVersion(int64_t version) { |
113 | emitBytecodeVersion = version; |
114 | return *this; |
115 | } |
116 | std::optional<int64_t> bytecodeVersionToEmit() const { |
117 | return emitBytecodeVersion; |
118 | } |
119 | |
120 | /// Set the callback to populate the pass manager. |
121 | MlirOptMainConfig & |
122 | setPassPipelineSetupFn(std::function<LogicalResult(PassManager &)> callback) { |
123 | passPipelineCallback = std::move(callback); |
124 | return *this; |
125 | } |
126 | |
127 | /// Set the parser to use to populate the pass manager. |
128 | MlirOptMainConfig &setPassPipelineParser(const PassPipelineCLParser &parser); |
129 | |
130 | /// Populate the passmanager, if any callback was set. |
131 | LogicalResult setupPassPipeline(PassManager &pm) const { |
132 | if (passPipelineCallback) |
133 | return passPipelineCallback(pm); |
134 | return success(); |
135 | } |
136 | |
137 | /// List the registered passes and return. |
138 | MlirOptMainConfig &listPasses(bool list) { |
139 | listPassesFlag = list; |
140 | return *this; |
141 | } |
142 | bool shouldListPasses() const { return listPassesFlag; } |
143 | |
144 | /// Enable running the reproducer information stored in resources (if |
145 | /// present). |
146 | MlirOptMainConfig &runReproducer(bool enableReproducer) { |
147 | runReproducerFlag = enableReproducer; |
148 | return *this; |
149 | }; |
150 | |
151 | /// Return true if the reproducer should be run. |
152 | bool shouldRunReproducer() const { return runReproducerFlag; } |
153 | |
154 | /// Show the registered dialects before trying to load the input file. |
155 | MlirOptMainConfig &showDialects(bool show) { |
156 | showDialectsFlag = show; |
157 | return *this; |
158 | } |
159 | bool shouldShowDialects() const { return showDialectsFlag; } |
160 | |
161 | /// Set the marker on which to split the input into chunks and process each |
162 | /// chunk independently. Input is not split if empty. |
163 | MlirOptMainConfig & |
164 | splitInputFile(std::string splitMarker = kDefaultSplitMarker) { |
165 | splitInputFileFlag = std::move(splitMarker); |
166 | return *this; |
167 | } |
168 | StringRef inputSplitMarker() const { return splitInputFileFlag; } |
169 | |
170 | /// Set whether to merge the output chunks into one file using the given |
171 | /// marker. |
172 | MlirOptMainConfig & |
173 | outputSplitMarker(std::string splitMarker = kDefaultSplitMarker) { |
174 | outputSplitMarkerFlag = std::move(splitMarker); |
175 | return *this; |
176 | } |
177 | StringRef outputSplitMarker() const { return outputSplitMarkerFlag; } |
178 | |
179 | /// Disable implicit addition of a top-level module op during parsing. |
180 | MlirOptMainConfig &useExplicitModule(bool useExplicitModule) { |
181 | useExplicitModuleFlag = useExplicitModule; |
182 | return *this; |
183 | } |
184 | bool shouldUseExplicitModule() const { return useExplicitModuleFlag; } |
185 | |
186 | /// Set whether to check that emitted diagnostics match `expected-*` lines on |
187 | /// the corresponding line. This is meant for implementing diagnostic tests. |
188 | MlirOptMainConfig & |
189 | verifyDiagnostics(SourceMgrDiagnosticVerifierHandler::Level verify) { |
190 | verifyDiagnosticsFlag = verify; |
191 | return *this; |
192 | } |
193 | |
194 | bool shouldVerifyDiagnostics() const { |
195 | return verifyDiagnosticsFlag != |
196 | SourceMgrDiagnosticVerifierHandler::Level::None; |
197 | } |
198 | |
199 | SourceMgrDiagnosticVerifierHandler::Level verifyDiagnosticsLevel() const { |
200 | return verifyDiagnosticsFlag; |
201 | } |
202 | |
203 | /// Set whether to run the verifier after each transformation pass. |
204 | MlirOptMainConfig &verifyPasses(bool verify) { |
205 | verifyPassesFlag = verify; |
206 | return *this; |
207 | } |
208 | bool shouldVerifyPasses() const { return verifyPassesFlag; } |
209 | |
210 | /// Set whether to run the verifier on parsing. |
211 | MlirOptMainConfig &verifyOnParsing(bool verify) { |
212 | disableVerifierOnParsingFlag = !verify; |
213 | return *this; |
214 | } |
215 | bool shouldVerifyOnParsing() const { return !disableVerifierOnParsingFlag; } |
216 | |
217 | /// Set whether to run the verifier after each transformation pass. |
218 | MlirOptMainConfig &verifyRoundtrip(bool verify) { |
219 | verifyRoundtripFlag = verify; |
220 | return *this; |
221 | } |
222 | bool shouldVerifyRoundtrip() const { return verifyRoundtripFlag; } |
223 | |
224 | /// Reproducer file generation (no crash required). |
225 | StringRef getReproducerFilename() const { return generateReproducerFileFlag; } |
226 | |
227 | protected: |
228 | /// Allow operation with no registered dialects. |
229 | /// This option is for convenience during testing only and discouraged in |
230 | /// general. |
231 | bool allowUnregisteredDialectsFlag = false; |
232 | |
233 | /// Configuration for the debugging hooks. |
234 | tracing::DebugConfig debugConfig; |
235 | |
236 | /// Verbosity level of diagnostic information. 0: Errors only, |
237 | /// 1: Errors and warnings, 2: Errors, warnings and remarks. |
238 | VerbosityLevel diagnosticVerbosityLevelFlag = |
239 | VerbosityLevel::ErrorsWarningsAndRemarks; |
240 | |
241 | /// Print the pipeline that will be run. |
242 | bool dumpPassPipelineFlag = false; |
243 | |
244 | /// Emit bytecode instead of textual assembly when generating output. |
245 | bool emitBytecodeFlag = false; |
246 | |
247 | /// Elide resources when generating bytecode. |
248 | bool elideResourceDataFromBytecodeFlag = false; |
249 | |
250 | /// IRDL file to register before processing the input. |
251 | std::string irdlFileFlag = "" ; |
252 | |
253 | /// Location Breakpoints to filter the action logging. |
254 | std::vector<tracing::BreakpointManager *> logActionLocationFilter; |
255 | |
256 | /// Emit bytecode at given version. |
257 | std::optional<int64_t> emitBytecodeVersion = std::nullopt; |
258 | |
259 | /// The callback to populate the pass manager. |
260 | std::function<LogicalResult(PassManager &)> passPipelineCallback; |
261 | |
262 | /// List the registered passes and return. |
263 | bool listPassesFlag = false; |
264 | |
265 | /// Enable running the reproducer. |
266 | bool runReproducerFlag = false; |
267 | |
268 | /// Show the registered dialects before trying to load the input file. |
269 | bool showDialectsFlag = false; |
270 | |
271 | /// Show the notes in diagnostic information. Notes can be included in |
272 | /// any diagnostic information, so it is not specified in the verbosity |
273 | /// level. |
274 | bool disableDiagnosticNotesFlag = true; |
275 | |
276 | /// Split the input file based on the given marker into chunks and process |
277 | /// each chunk independently. Input is not split if empty. |
278 | std::string splitInputFileFlag = "" ; |
279 | |
280 | /// Merge output chunks into one file using the given marker. |
281 | std::string outputSplitMarkerFlag = "" ; |
282 | |
283 | /// Use an explicit top-level module op during parsing. |
284 | bool useExplicitModuleFlag = false; |
285 | |
286 | /// Set whether to check that emitted diagnostics match `expected-*` lines on |
287 | /// the corresponding line. This is meant for implementing diagnostic tests. |
288 | SourceMgrDiagnosticVerifierHandler::Level verifyDiagnosticsFlag = |
289 | SourceMgrDiagnosticVerifierHandler::Level::None; |
290 | |
291 | /// Run the verifier after each transformation pass. |
292 | bool verifyPassesFlag = true; |
293 | |
294 | /// Disable the verifier on parsing. |
295 | bool disableVerifierOnParsingFlag = false; |
296 | |
297 | /// Verify that the input IR round-trips perfectly. |
298 | bool verifyRoundtripFlag = false; |
299 | |
300 | /// The reproducer output filename (no crash required). |
301 | std::string generateReproducerFileFlag = "" ; |
302 | }; |
303 | |
304 | /// This defines the function type used to setup the pass manager. This can be |
305 | /// used to pass in a callback to setup a default pass pipeline to be applied on |
306 | /// the loaded IR. |
307 | using PassPipelineFn = llvm::function_ref<LogicalResult(PassManager &pm)>; |
308 | |
309 | /// Register and parse command line options. |
310 | /// - toolName is used for the header displayed by `--help`. |
311 | /// - registry should contain all the dialects that can be parsed in the source. |
312 | /// - return std::pair<std::string, std::string> for |
313 | /// inputFilename and outputFilename command line option values. |
314 | std::pair<std::string, std::string> |
315 | registerAndParseCLIOptions(int argc, char **argv, llvm::StringRef toolName, |
316 | DialectRegistry ®istry); |
317 | |
318 | /// Perform the core processing behind `mlir-opt`. |
319 | /// - outputStream is the stream where the resulting IR is printed. |
320 | /// - buffer is the in-memory file to parser and process. |
321 | /// - registry should contain all the dialects that can be parsed in the source. |
322 | /// - config contains the configuration options for the tool. |
323 | LogicalResult MlirOptMain(llvm::raw_ostream &outputStream, |
324 | std::unique_ptr<llvm::MemoryBuffer> buffer, |
325 | DialectRegistry ®istry, |
326 | const MlirOptMainConfig &config); |
327 | |
328 | /// Implementation for tools like `mlir-opt`. |
329 | /// - toolName is used for the header displayed by `--help`. |
330 | /// - registry should contain all the dialects that can be parsed in the source. |
331 | LogicalResult MlirOptMain(int argc, char **argv, llvm::StringRef toolName, |
332 | DialectRegistry ®istry); |
333 | |
334 | /// Implementation for tools like `mlir-opt`. |
335 | /// This function can be used with registerAndParseCLIOptions so that |
336 | /// CLI options can be accessed before running MlirOptMain. |
337 | /// - inputFilename is the name of the input mlir file. |
338 | /// - outputFilename is the name of the output file. |
339 | /// - registry should contain all the dialects that can be parsed in the source. |
340 | LogicalResult MlirOptMain(int argc, char **argv, llvm::StringRef inputFilename, |
341 | llvm::StringRef outputFilename, |
342 | DialectRegistry ®istry); |
343 | |
344 | /// Helper wrapper to return the result of MlirOptMain directly from main. |
345 | /// |
346 | /// Example: |
347 | /// |
348 | /// int main(int argc, char **argv) { |
349 | /// // ... |
350 | /// return mlir::asMainReturnCode(mlir::MlirOptMain( |
351 | /// argc, argv, /* ... */); |
352 | /// } |
353 | /// |
354 | inline int asMainReturnCode(LogicalResult r) { |
355 | return r.succeeded() ? EXIT_SUCCESS : EXIT_FAILURE; |
356 | } |
357 | |
358 | } // namespace mlir |
359 | |
360 | #endif // MLIR_TOOLS_MLIROPT_MLIROPTMAIN_H |
361 | |