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/IR/Diagnostics.h" |
15 | #include "toy/AST.h" |
16 | #include "toy/Dialect.h" |
17 | #include "toy/Lexer.h" |
18 | #include "toy/MLIRGen.h" |
19 | #include "toy/Parser.h" |
20 | #include "toy/Passes.h" |
21 | |
22 | #include "mlir/Dialect/Affine/Passes.h" |
23 | #include "mlir/IR/AsmState.h" |
24 | #include "mlir/IR/BuiltinOps.h" |
25 | #include "mlir/IR/MLIRContext.h" |
26 | #include "mlir/IR/Verifier.h" |
27 | #include "mlir/InitAllDialects.h" |
28 | #include "mlir/Parser/Parser.h" |
29 | #include "mlir/Pass/PassManager.h" |
30 | #include "mlir/Transforms/Passes.h" |
31 | |
32 | #include "llvm/ADT/StringRef.h" |
33 | #include "llvm/Support/CommandLine.h" |
34 | #include "llvm/Support/ErrorOr.h" |
35 | #include "llvm/Support/MemoryBuffer.h" |
36 | #include "llvm/Support/SourceMgr.h" |
37 | #include "llvm/Support/raw_ostream.h" |
38 | #include <memory> |
39 | #include <string> |
40 | #include <system_error> |
41 | #include <utility> |
42 | |
43 | using namespace toy; |
44 | namespace cl = llvm::cl; |
45 | |
46 | static cl::opt<std::string> inputFilename(cl::Positional, |
47 | cl::desc("<input toy file>" ), |
48 | cl::init(Val: "-" ), |
49 | cl::value_desc("filename" )); |
50 | |
51 | namespace { |
52 | enum InputType { Toy, MLIR }; |
53 | } // namespace |
54 | static cl::opt<enum InputType> inputType( |
55 | "x" , cl::init(Val: Toy), cl::desc("Decided the kind of output desired" ), |
56 | cl::values(clEnumValN(Toy, "toy" , "load the input file as a Toy source." )), |
57 | cl::values(clEnumValN(MLIR, "mlir" , |
58 | "load the input file as an MLIR file" ))); |
59 | |
60 | namespace { |
61 | enum Action { None, DumpAST, DumpMLIR, DumpMLIRAffine }; |
62 | } // namespace |
63 | static cl::opt<enum Action> emitAction( |
64 | "emit" , cl::desc("Select the kind of output desired" ), |
65 | cl::values(clEnumValN(DumpAST, "ast" , "output the AST dump" )), |
66 | cl::values(clEnumValN(DumpMLIR, "mlir" , "output the MLIR dump" )), |
67 | cl::values(clEnumValN(DumpMLIRAffine, "mlir-affine" , |
68 | "output the MLIR dump after affine lowering" ))); |
69 | |
70 | static cl::opt<bool> enableOpt("opt" , cl::desc("Enable optimizations" )); |
71 | |
72 | /// Returns a Toy AST resulting from parsing the file or a nullptr on error. |
73 | std::unique_ptr<toy::ModuleAST> parseInputFile(llvm::StringRef filename) { |
74 | llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> fileOrErr = |
75 | llvm::MemoryBuffer::getFileOrSTDIN(Filename: filename); |
76 | if (std::error_code ec = fileOrErr.getError()) { |
77 | llvm::errs() << "Could not open input file: " << ec.message() << "\n" ; |
78 | return nullptr; |
79 | } |
80 | auto buffer = fileOrErr.get()->getBuffer(); |
81 | LexerBuffer lexer(buffer.begin(), buffer.end(), std::string(filename)); |
82 | Parser parser(lexer); |
83 | return parser.parseModule(); |
84 | } |
85 | |
86 | int loadMLIR(llvm::SourceMgr &sourceMgr, mlir::MLIRContext &context, |
87 | mlir::OwningOpRef<mlir::ModuleOp> &module) { |
88 | // Handle '.toy' input to the compiler. |
89 | if (inputType != InputType::MLIR && |
90 | !llvm::StringRef(inputFilename).ends_with(Suffix: ".mlir" )) { |
91 | auto moduleAST = parseInputFile(filename: inputFilename); |
92 | if (!moduleAST) |
93 | return 6; |
94 | module = mlirGen(context, moduleAST&: *moduleAST); |
95 | return !module ? 1 : 0; |
96 | } |
97 | |
98 | // Otherwise, the input is '.mlir'. |
99 | llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> fileOrErr = |
100 | llvm::MemoryBuffer::getFileOrSTDIN(Filename: inputFilename); |
101 | if (std::error_code ec = fileOrErr.getError()) { |
102 | llvm::errs() << "Could not open input file: " << ec.message() << "\n" ; |
103 | return -1; |
104 | } |
105 | |
106 | // Parse the input mlir. |
107 | sourceMgr.AddNewSourceBuffer(F: std::move(*fileOrErr), IncludeLoc: llvm::SMLoc()); |
108 | module = mlir::parseSourceFile<mlir::ModuleOp>(sourceMgr, config: &context); |
109 | if (!module) { |
110 | llvm::errs() << "Error can't load file " << inputFilename << "\n" ; |
111 | return 3; |
112 | } |
113 | return 0; |
114 | } |
115 | |
116 | int dumpMLIR() { |
117 | mlir::DialectRegistry registry; |
118 | mlir::func::registerAllExtensions(registry); |
119 | |
120 | mlir::MLIRContext context(registry); |
121 | // Load our Dialect in this MLIR Context. |
122 | context.getOrLoadDialect<mlir::toy::ToyDialect>(); |
123 | |
124 | mlir::OwningOpRef<mlir::ModuleOp> module; |
125 | llvm::SourceMgr sourceMgr; |
126 | mlir::SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context); |
127 | if (int error = loadMLIR(sourceMgr, context, module)) |
128 | return error; |
129 | |
130 | mlir::PassManager pm(module.get()->getName()); |
131 | // Apply any generic pass manager command line options and run the pipeline. |
132 | if (mlir::failed(Result: mlir::applyPassManagerCLOptions(pm))) |
133 | return 4; |
134 | |
135 | // Check to see what granularity of MLIR we are compiling to. |
136 | bool isLoweringToAffine = emitAction >= Action::DumpMLIRAffine; |
137 | |
138 | if (enableOpt || isLoweringToAffine) { |
139 | // Inline all functions into main and then delete them. |
140 | pm.addPass(mlir::pass: createInlinerPass()); |
141 | |
142 | // Now that there is only one function, we can infer the shapes of each of |
143 | // the operations. |
144 | mlir::OpPassManager &optPM = pm.nest<mlir::toy::FuncOp>(); |
145 | optPM.addPass(pass: mlir::toy::createShapeInferencePass()); |
146 | optPM.addPass(pass: mlir::createCanonicalizerPass()); |
147 | optPM.addPass(pass: mlir::createCSEPass()); |
148 | } |
149 | |
150 | if (isLoweringToAffine) { |
151 | // Partially lower the toy dialect. |
152 | pm.addPass(pass: mlir::toy::createLowerToAffinePass()); |
153 | |
154 | // Add a few cleanups post lowering. |
155 | mlir::OpPassManager &optPM = pm.nest<mlir::func::FuncOp>(); |
156 | optPM.addPass(pass: mlir::createCanonicalizerPass()); |
157 | optPM.addPass(pass: mlir::createCSEPass()); |
158 | |
159 | // Add optimizations if enabled. |
160 | if (enableOpt) { |
161 | optPM.addPass(mlir::affine::pass: createLoopFusionPass()); |
162 | optPM.addPass(mlir::affine::pass: createAffineScalarReplacementPass()); |
163 | } |
164 | } |
165 | |
166 | if (mlir::failed(Result: pm.run(op: *module))) |
167 | return 4; |
168 | |
169 | module->dump(); |
170 | return 0; |
171 | } |
172 | |
173 | int dumpAST() { |
174 | if (inputType == InputType::MLIR) { |
175 | llvm::errs() << "Can't dump a Toy AST when the input is MLIR\n" ; |
176 | return 5; |
177 | } |
178 | |
179 | auto moduleAST = parseInputFile(filename: inputFilename); |
180 | if (!moduleAST) |
181 | return 1; |
182 | |
183 | dump(*moduleAST); |
184 | return 0; |
185 | } |
186 | |
187 | int main(int argc, char **argv) { |
188 | // Register any command line options. |
189 | mlir::registerAsmPrinterCLOptions(); |
190 | mlir::registerMLIRContextCLOptions(); |
191 | mlir::registerPassManagerCLOptions(); |
192 | |
193 | cl::ParseCommandLineOptions(argc, argv, Overview: "toy compiler\n" ); |
194 | |
195 | switch (emitAction) { |
196 | case Action::DumpAST: |
197 | return dumpAST(); |
198 | case Action::DumpMLIR: |
199 | case Action::DumpMLIRAffine: |
200 | return dumpMLIR(); |
201 | default: |
202 | llvm::errs() << "No action specified (parsing only?), use -emit=<action>\n" ; |
203 | } |
204 | |
205 | return 0; |
206 | } |
207 | |