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