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