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/IR/AsmState.h"
16 #include "mlir/IR/Attributes.h"
17 #include "mlir/IR/BuiltinOps.h"
18 #include "mlir/IR/Diagnostics.h"
19 #include "mlir/IR/Dialect.h"
20 #include "mlir/IR/Location.h"
21 #include "mlir/IR/MLIRContext.h"
22 #include "mlir/Parser/Parser.h"
23 #include "mlir/Pass/Pass.h"
24 #include "mlir/Pass/PassManager.h"
25 #include "mlir/Support/DebugCounter.h"
26 #include "mlir/Support/FileUtilities.h"
27 #include "mlir/Support/Timing.h"
28 #include "mlir/Support/ToolUtilities.h"
29 #include "llvm/Support/CommandLine.h"
30 #include "llvm/Support/FileUtilities.h"
31 #include "llvm/Support/InitLLVM.h"
32 #include "llvm/Support/Regex.h"
33 #include "llvm/Support/SourceMgr.h"
34 #include "llvm/Support/StringSaver.h"
35 #include "llvm/Support/ThreadPool.h"
36 #include "llvm/Support/ToolOutputFile.h"
37 
38 using namespace mlir;
39 using namespace llvm;
40 
41 /// Perform the actions on the input file indicated by the command line flags
42 /// within the specified context.
43 ///
44 /// This typically parses the main source file, runs zero or more optimization
45 /// passes, then prints the output.
46 ///
47 static LogicalResult performActions(raw_ostream &os, bool verifyDiagnostics,
48                                     bool verifyPasses, SourceMgr &sourceMgr,
49                                     MLIRContext *context,
50                                     PassPipelineFn passManagerSetupFn) {
51   DefaultTimingManager tm;
52   applyDefaultTimingManagerCLOptions(tm);
53   TimingScope timing = tm.getRootScope();
54 
55   // Disable multi-threading when parsing the input file. This removes the
56   // unnecessary/costly context synchronization when parsing.
57   bool wasThreadingEnabled = context->isMultithreadingEnabled();
58   context->disableMultithreading();
59 
60   // Parse the input file and reset the context threading state.
61   TimingScope parserTiming = timing.nest("Parser");
62   OwningOpRef<ModuleOp> module(parseSourceFile<ModuleOp>(sourceMgr, context));
63   context->enableMultithreading(wasThreadingEnabled);
64   if (!module)
65     return failure();
66   parserTiming.stop();
67 
68   // Apply any pass manager command line options.
69   PassManager pm(context, OpPassManager::Nesting::Implicit);
70   pm.enableVerifier(verifyPasses);
71   applyPassManagerCLOptions(pm);
72   pm.enableTiming(timing);
73 
74   // Callback to build the pipeline.
75   if (failed(passManagerSetupFn(pm)))
76     return failure();
77 
78   // Run the pipeline.
79   if (failed(pm.run(*module)))
80     return failure();
81 
82   // Print the output.
83   TimingScope outputTiming = timing.nest("Output");
84   module->print(os);
85   os << '\n';
86   return success();
87 }
88 
89 /// Parses the memory buffer.  If successfully, run a series of passes against
90 /// it and print the result.
91 static LogicalResult
92 processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
93               bool verifyDiagnostics, bool verifyPasses,
94               bool allowUnregisteredDialects, bool preloadDialectsInContext,
95               PassPipelineFn passManagerSetupFn, DialectRegistry &registry,
96               llvm::ThreadPool *threadPool) {
97   // Tell sourceMgr about this buffer, which is what the parser will pick up.
98   SourceMgr sourceMgr;
99   sourceMgr.AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
100 
101   // Create a context just for the current buffer. Disable threading on creation
102   // since we'll inject the thread-pool separately.
103   MLIRContext context(registry, MLIRContext::Threading::DISABLED);
104   if (threadPool)
105     context.setThreadPool(*threadPool);
106 
107   // Parse the input file.
108   if (preloadDialectsInContext)
109     context.loadAllAvailableDialects();
110   context.allowUnregisteredDialects(allowUnregisteredDialects);
111   if (verifyDiagnostics)
112     context.printOpOnDiagnostic(false);
113   context.getDebugActionManager().registerActionHandler<DebugCounter>();
114 
115   // If we are in verify diagnostics mode then we have a lot of work to do,
116   // otherwise just perform the actions without worrying about it.
117   if (!verifyDiagnostics) {
118     SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
119     return performActions(os, verifyDiagnostics, verifyPasses, sourceMgr,
120                           &context, passManagerSetupFn);
121   }
122 
123   SourceMgrDiagnosticVerifierHandler sourceMgrHandler(sourceMgr, &context);
124 
125   // Do any processing requested by command line flags.  We don't care whether
126   // these actions succeed or fail, we only care what diagnostics they produce
127   // and whether they match our expectations.
128   (void)performActions(os, verifyDiagnostics, verifyPasses, sourceMgr, &context,
129                        passManagerSetupFn);
130 
131   // Verify the diagnostic handler to make sure that each of the diagnostics
132   // matched.
133   return sourceMgrHandler.verify();
134 }
135 
136 LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
137                                 std::unique_ptr<MemoryBuffer> buffer,
138                                 PassPipelineFn passManagerSetupFn,
139                                 DialectRegistry &registry, bool splitInputFile,
140                                 bool verifyDiagnostics, bool verifyPasses,
141                                 bool allowUnregisteredDialects,
142                                 bool preloadDialectsInContext) {
143   // The split-input-file mode is a very specific mode that slices the file
144   // up into small pieces and checks each independently.
145   // We use an explicit threadpool to avoid creating and joining/destroying
146   // threads for each of the split.
147   ThreadPool *threadPool = nullptr;
148   // Create a temporary context for the sake of checking if
149   // --mlir-disable-threading was passed on the command line.
150   // We use the thread-pool this context is creating, and avoid
151   // creating any thread when disabled.
152   MLIRContext threadPoolCtx;
153   if (threadPoolCtx.isMultithreadingEnabled())
154     threadPool = &threadPoolCtx.getThreadPool();
155 
156   if (splitInputFile)
157     return splitAndProcessBuffer(
158         std::move(buffer),
159         [&](std::unique_ptr<MemoryBuffer> chunkBuffer, raw_ostream &os) {
160           LogicalResult result = processBuffer(
161               os, std::move(chunkBuffer), verifyDiagnostics, verifyPasses,
162               allowUnregisteredDialects, preloadDialectsInContext,
163               passManagerSetupFn, registry, threadPool);
164           os << "// -----\n";
165           return result;
166         },
167         outputStream);
168 
169   return processBuffer(outputStream, std::move(buffer), verifyDiagnostics,
170                        verifyPasses, allowUnregisteredDialects,
171                        preloadDialectsInContext, passManagerSetupFn, registry,
172                        threadPool);
173 }
174 
175 LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
176                                 std::unique_ptr<MemoryBuffer> buffer,
177                                 const PassPipelineCLParser &passPipeline,
178                                 DialectRegistry &registry, bool splitInputFile,
179                                 bool verifyDiagnostics, bool verifyPasses,
180                                 bool allowUnregisteredDialects,
181                                 bool preloadDialectsInContext) {
182   auto passManagerSetupFn = [&](PassManager &pm) {
183     auto errorHandler = [&](const Twine &msg) {
184       emitError(UnknownLoc::get(pm.getContext())) << msg;
185       return failure();
186     };
187     return passPipeline.addToPipeline(pm, errorHandler);
188   };
189   return MlirOptMain(outputStream, std::move(buffer), passManagerSetupFn,
190                      registry, splitInputFile, verifyDiagnostics, verifyPasses,
191                      allowUnregisteredDialects, preloadDialectsInContext);
192 }
193 
194 LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
195                                 DialectRegistry &registry,
196                                 bool preloadDialectsInContext) {
197   static cl::opt<std::string> inputFilename(
198       cl::Positional, cl::desc("<input file>"), cl::init("-"));
199 
200   static cl::opt<std::string> outputFilename("o", cl::desc("Output filename"),
201                                              cl::value_desc("filename"),
202                                              cl::init("-"));
203 
204   static cl::opt<bool> splitInputFile(
205       "split-input-file",
206       cl::desc("Split the input file into pieces and process each "
207                "chunk independently"),
208       cl::init(false));
209 
210   static cl::opt<bool> verifyDiagnostics(
211       "verify-diagnostics",
212       cl::desc("Check that emitted diagnostics match "
213                "expected-* lines on the corresponding line"),
214       cl::init(false));
215 
216   static cl::opt<bool> verifyPasses(
217       "verify-each",
218       cl::desc("Run the verifier after each transformation pass"),
219       cl::init(true));
220 
221   static cl::opt<bool> allowUnregisteredDialects(
222       "allow-unregistered-dialect",
223       cl::desc("Allow operation with no registered dialects"), cl::init(false));
224 
225   static cl::opt<bool> showDialects(
226       "show-dialects", cl::desc("Print the list of registered dialects"),
227       cl::init(false));
228 
229   static cl::opt<bool> runRepro(
230       "run-reproducer",
231       cl::desc("Append the command line options of the reproducer"),
232       cl::init(false));
233 
234   InitLLVM y(argc, argv);
235 
236   // Register any command line options.
237   registerAsmPrinterCLOptions();
238   registerMLIRContextCLOptions();
239   registerPassManagerCLOptions();
240   registerDefaultTimingManagerCLOptions();
241   DebugCounter::registerCLOptions();
242   PassPipelineCLParser passPipeline("", "Compiler passes to run");
243 
244   // Build the list of dialects as a header for the --help message.
245   std::string helpHeader = (toolName + "\nAvailable Dialects: ").str();
246   {
247     llvm::raw_string_ostream os(helpHeader);
248     interleaveComma(registry.getDialectNames(), os,
249                     [&](auto name) { os << name; });
250   }
251   // Parse pass names in main to ensure static initialization completed.
252   cl::ParseCommandLineOptions(argc, argv, helpHeader);
253 
254   if (showDialects) {
255     llvm::outs() << "Available Dialects:\n";
256     interleave(
257         registry.getDialectNames(), llvm::outs(),
258         [](auto name) { llvm::outs() << name; }, "\n");
259     return success();
260   }
261 
262   // Set up the input file.
263   std::string errorMessage;
264   auto file = openInputFile(inputFilename, &errorMessage);
265   if (!file) {
266     llvm::errs() << errorMessage << "\n";
267     return failure();
268   }
269 
270   // Parse reproducer options.
271   BumpPtrAllocator a;
272   StringSaver saver(a);
273   if (runRepro) {
274     auto pair = file->getBuffer().split('\n');
275     if (!pair.first.consume_front("// configuration:")) {
276       llvm::errs() << "Failed to find repro configuration, expect file to "
277                       "begin with '// configuration:'\n";
278       return failure();
279     }
280     // Tokenize & parse the first line.
281     SmallVector<const char *, 4> newArgv;
282     newArgv.push_back(argv[0]);
283     llvm::cl::TokenizeGNUCommandLine(pair.first, saver, newArgv);
284     cl::ParseCommandLineOptions(newArgv.size(), &newArgv[0], helpHeader);
285   }
286 
287   auto output = openOutputFile(outputFilename, &errorMessage);
288   if (!output) {
289     llvm::errs() << errorMessage << "\n";
290     return failure();
291   }
292 
293   if (failed(MlirOptMain(output->os(), std::move(file), passPipeline, registry,
294                          splitInputFile, verifyDiagnostics, verifyPasses,
295                          allowUnregisteredDialects, preloadDialectsInContext)))
296     return failure();
297 
298   // Keep the output file if the invocation of MlirOptMain was successful.
299   output->keep();
300   return success();
301 }
302