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