1//===- toyc.cpp - The Toy Compiler ----------------------------------------===//
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 file implements the entry point for the Toy compiler.
10//
11//===----------------------------------------------------------------------===//
12
13#include "mlir/Dialect/Func/Extensions/AllExtensions.h"
14#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
15#include "mlir/Support/LogicalResult.h"
16#include "toy/AST.h"
17#include "toy/Dialect.h"
18#include "toy/Lexer.h"
19#include "toy/MLIRGen.h"
20#include "toy/Parser.h"
21#include "toy/Passes.h"
22
23#include "mlir/Dialect/Affine/Passes.h"
24#include "mlir/Dialect/LLVMIR/Transforms/Passes.h"
25#include "mlir/ExecutionEngine/ExecutionEngine.h"
26#include "mlir/ExecutionEngine/OptUtils.h"
27#include "mlir/IR/AsmState.h"
28#include "mlir/IR/BuiltinOps.h"
29#include "mlir/IR/MLIRContext.h"
30#include "mlir/IR/Verifier.h"
31#include "mlir/InitAllDialects.h"
32#include "mlir/Parser/Parser.h"
33#include "mlir/Pass/PassManager.h"
34#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
35#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
36#include "mlir/Target/LLVMIR/Export.h"
37#include "mlir/Transforms/Passes.h"
38
39#include "llvm/ADT/StringRef.h"
40#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
41#include "llvm/IR/Module.h"
42#include "llvm/Support/CommandLine.h"
43#include "llvm/Support/ErrorOr.h"
44#include "llvm/Support/MemoryBuffer.h"
45#include "llvm/Support/SourceMgr.h"
46#include "llvm/Support/TargetSelect.h"
47#include "llvm/Support/raw_ostream.h"
48#include <cassert>
49#include <memory>
50#include <string>
51#include <system_error>
52#include <utility>
53
54using namespace toy;
55namespace cl = llvm::cl;
56
57static cl::opt<std::string> inputFilename(cl::Positional,
58 cl::desc("<input toy file>"),
59 cl::init(Val: "-"),
60 cl::value_desc("filename"));
61
62namespace {
63enum InputType { Toy, MLIR };
64} // namespace
65static cl::opt<enum InputType> inputType(
66 "x", cl::init(Val: Toy), cl::desc("Decided the kind of output desired"),
67 cl::values(clEnumValN(Toy, "toy", "load the input file as a Toy source.")),
68 cl::values(clEnumValN(MLIR, "mlir",
69 "load the input file as an MLIR file")));
70
71namespace {
72enum Action {
73 None,
74 DumpAST,
75 DumpMLIR,
76 DumpMLIRAffine,
77 DumpMLIRLLVM,
78 DumpLLVMIR,
79 RunJIT
80};
81} // namespace
82static cl::opt<enum Action> emitAction(
83 "emit", cl::desc("Select the kind of output desired"),
84 cl::values(clEnumValN(DumpAST, "ast", "output the AST dump")),
85 cl::values(clEnumValN(DumpMLIR, "mlir", "output the MLIR dump")),
86 cl::values(clEnumValN(DumpMLIRAffine, "mlir-affine",
87 "output the MLIR dump after affine lowering")),
88 cl::values(clEnumValN(DumpMLIRLLVM, "mlir-llvm",
89 "output the MLIR dump after llvm lowering")),
90 cl::values(clEnumValN(DumpLLVMIR, "llvm", "output the LLVM IR dump")),
91 cl::values(
92 clEnumValN(RunJIT, "jit",
93 "JIT the code and run it by invoking the main function")));
94
95static cl::opt<bool> enableOpt("opt", cl::desc("Enable optimizations"));
96
97/// Returns a Toy AST resulting from parsing the file or a nullptr on error.
98std::unique_ptr<toy::ModuleAST> parseInputFile(llvm::StringRef filename) {
99 llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> fileOrErr =
100 llvm::MemoryBuffer::getFileOrSTDIN(Filename: filename);
101 if (std::error_code ec = fileOrErr.getError()) {
102 llvm::errs() << "Could not open input file: " << ec.message() << "\n";
103 return nullptr;
104 }
105 auto buffer = fileOrErr.get()->getBuffer();
106 LexerBuffer lexer(buffer.begin(), buffer.end(), std::string(filename));
107 Parser parser(lexer);
108 return parser.parseModule();
109}
110
111int loadMLIR(mlir::MLIRContext &context,
112 mlir::OwningOpRef<mlir::ModuleOp> &module) {
113 // Handle '.toy' input to the compiler.
114 if (inputType != InputType::MLIR &&
115 !llvm::StringRef(inputFilename).ends_with(Suffix: ".mlir")) {
116 auto moduleAST = parseInputFile(filename: inputFilename);
117 if (!moduleAST)
118 return 6;
119 module = mlirGen(context, moduleAST&: *moduleAST);
120 return !module ? 1 : 0;
121 }
122
123 // Otherwise, the input is '.mlir'.
124 llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> fileOrErr =
125 llvm::MemoryBuffer::getFileOrSTDIN(Filename: inputFilename);
126 if (std::error_code ec = fileOrErr.getError()) {
127 llvm::errs() << "Could not open input file: " << ec.message() << "\n";
128 return -1;
129 }
130
131 // Parse the input mlir.
132 llvm::SourceMgr sourceMgr;
133 sourceMgr.AddNewSourceBuffer(F: std::move(*fileOrErr), IncludeLoc: llvm::SMLoc());
134 module = mlir::parseSourceFile<mlir::ModuleOp>(sourceMgr, config: &context);
135 if (!module) {
136 llvm::errs() << "Error can't load file " << inputFilename << "\n";
137 return 3;
138 }
139 return 0;
140}
141
142int loadAndProcessMLIR(mlir::MLIRContext &context,
143 mlir::OwningOpRef<mlir::ModuleOp> &module) {
144 if (int error = loadMLIR(context, module))
145 return error;
146
147 mlir::PassManager pm(module.get()->getName());
148 // Apply any generic pass manager command line options and run the pipeline.
149 if (mlir::failed(result: mlir::applyPassManagerCLOptions(pm)))
150 return 4;
151
152 // Check to see what granularity of MLIR we are compiling to.
153 bool isLoweringToAffine = emitAction >= Action::DumpMLIRAffine;
154 bool isLoweringToLLVM = emitAction >= Action::DumpMLIRLLVM;
155
156 if (enableOpt || isLoweringToAffine) {
157 // Inline all functions into main and then delete them.
158 pm.addPass(mlir::pass: createInlinerPass());
159
160 // Now that there is only one function, we can infer the shapes of each of
161 // the operations.
162 mlir::OpPassManager &optPM = pm.nest<mlir::toy::FuncOp>();
163 optPM.addPass(pass: mlir::createCanonicalizerPass());
164 optPM.addPass(pass: mlir::toy::createShapeInferencePass());
165 optPM.addPass(pass: mlir::createCanonicalizerPass());
166 optPM.addPass(pass: mlir::createCSEPass());
167 }
168
169 if (isLoweringToAffine) {
170 // Partially lower the toy dialect.
171 pm.addPass(pass: mlir::toy::createLowerToAffinePass());
172
173 // Add a few cleanups post lowering.
174 mlir::OpPassManager &optPM = pm.nest<mlir::func::FuncOp>();
175 optPM.addPass(pass: mlir::createCanonicalizerPass());
176 optPM.addPass(pass: mlir::createCSEPass());
177
178 // Add optimizations if enabled.
179 if (enableOpt) {
180 optPM.addPass(mlir::affine::pass: createLoopFusionPass());
181 optPM.addPass(mlir::affine::pass: createAffineScalarReplacementPass());
182 }
183 }
184
185 if (isLoweringToLLVM) {
186 // Finish lowering the toy IR to the LLVM dialect.
187 pm.addPass(pass: mlir::toy::createLowerToLLVMPass());
188 // This is necessary to have line tables emitted and basic
189 // debugger working. In the future we will add proper debug information
190 // emission directly from our frontend.
191 pm.addPass(pass: mlir::LLVM::createDIScopeForLLVMFuncOpPass());
192 }
193
194 if (mlir::failed(result: pm.run(op: *module)))
195 return 4;
196 return 0;
197}
198
199int dumpAST() {
200 if (inputType == InputType::MLIR) {
201 llvm::errs() << "Can't dump a Toy AST when the input is MLIR\n";
202 return 5;
203 }
204
205 auto moduleAST = parseInputFile(filename: inputFilename);
206 if (!moduleAST)
207 return 1;
208
209 dump(*moduleAST);
210 return 0;
211}
212
213int dumpLLVMIR(mlir::ModuleOp module) {
214 // Register the translation to LLVM IR with the MLIR context.
215 mlir::registerBuiltinDialectTranslation(*module->getContext());
216 mlir::registerLLVMDialectTranslation(*module->getContext());
217
218 // Convert the module to LLVM IR in a new LLVM IR context.
219 llvm::LLVMContext llvmContext;
220 auto llvmModule = mlir::translateModuleToLLVMIR(module: module, llvmContext);
221 if (!llvmModule) {
222 llvm::errs() << "Failed to emit LLVM IR\n";
223 return -1;
224 }
225
226 // Initialize LLVM targets.
227 llvm::InitializeNativeTarget();
228 llvm::InitializeNativeTargetAsmPrinter();
229
230 // Create target machine and configure the LLVM Module
231 auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
232 if (!tmBuilderOrError) {
233 llvm::errs() << "Could not create JITTargetMachineBuilder\n";
234 return -1;
235 }
236
237 auto tmOrError = tmBuilderOrError->createTargetMachine();
238 if (!tmOrError) {
239 llvm::errs() << "Could not create TargetMachine\n";
240 return -1;
241 }
242 mlir::ExecutionEngine::setupTargetTripleAndDataLayout(llvmModule: llvmModule.get(),
243 tm: tmOrError.get().get());
244
245 /// Optionally run an optimization pipeline over the llvm module.
246 auto optPipeline = mlir::makeOptimizingTransformer(
247 /*optLevel=*/enableOpt ? 3 : 0, /*sizeLevel=*/0,
248 /*targetMachine=*/nullptr);
249 if (auto err = optPipeline(llvmModule.get())) {
250 llvm::errs() << "Failed to optimize LLVM IR " << err << "\n";
251 return -1;
252 }
253 llvm::errs() << *llvmModule << "\n";
254 return 0;
255}
256
257int runJit(mlir::ModuleOp module) {
258 // Initialize LLVM targets.
259 llvm::InitializeNativeTarget();
260 llvm::InitializeNativeTargetAsmPrinter();
261
262 // Register the translation from MLIR to LLVM IR, which must happen before we
263 // can JIT-compile.
264 mlir::registerBuiltinDialectTranslation(*module->getContext());
265 mlir::registerLLVMDialectTranslation(*module->getContext());
266
267 // An optimization pipeline to use within the execution engine.
268 auto optPipeline = mlir::makeOptimizingTransformer(
269 /*optLevel=*/enableOpt ? 3 : 0, /*sizeLevel=*/0,
270 /*targetMachine=*/nullptr);
271
272 // Create an MLIR execution engine. The execution engine eagerly JIT-compiles
273 // the module.
274 mlir::ExecutionEngineOptions engineOptions;
275 engineOptions.transformer = optPipeline;
276 auto maybeEngine = mlir::ExecutionEngine::create(op: module, options: engineOptions);
277 assert(maybeEngine && "failed to construct an execution engine");
278 auto &engine = maybeEngine.get();
279
280 // Invoke the JIT-compiled function.
281 auto invocationResult = engine->invokePacked("main");
282 if (invocationResult) {
283 llvm::errs() << "JIT invocation failed\n";
284 return -1;
285 }
286
287 return 0;
288}
289
290int main(int argc, char **argv) {
291 // Register any command line options.
292 mlir::registerAsmPrinterCLOptions();
293 mlir::registerMLIRContextCLOptions();
294 mlir::registerPassManagerCLOptions();
295
296 cl::ParseCommandLineOptions(argc, argv, Overview: "toy compiler\n");
297
298 if (emitAction == Action::DumpAST)
299 return dumpAST();
300
301 // If we aren't dumping the AST, then we are compiling with/to MLIR.
302 mlir::DialectRegistry registry;
303 mlir::func::registerAllExtensions(registry);
304
305 mlir::MLIRContext context(registry);
306 // Load our Dialect in this MLIR Context.
307 context.getOrLoadDialect<mlir::toy::ToyDialect>();
308
309 mlir::OwningOpRef<mlir::ModuleOp> module;
310 if (int error = loadAndProcessMLIR(context, module))
311 return error;
312
313 // If we aren't exporting to non-mlir, then we are done.
314 bool isOutputingMLIR = emitAction <= Action::DumpMLIRLLVM;
315 if (isOutputingMLIR) {
316 module->dump();
317 return 0;
318 }
319
320 // Check to see if we are compiling to LLVM IR.
321 if (emitAction == Action::DumpLLVMIR)
322 return dumpLLVMIR(*module);
323
324 // Otherwise, we must be running the jit.
325 if (emitAction == Action::RunJIT)
326 return runJit(*module);
327
328 llvm::errs() << "No action specified (parsing only?), use -emit=<action>\n";
329 return -1;
330}
331

source code of mlir/examples/toy/Ch7/toyc.cpp