mlir/examples/toy/Ch2/toyc.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-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file implements the entry point for the Toy compiler.
  10 //
  11 //===----------------------------------------------------------------------===//
  12
  13 #include "toy/AST.h"
  14 #include "toy/Dialect.h"
  15 #include "toy/Lexer.h"
  16 #include "toy/MLIRGen.h"
  17 #include "toy/Parser.h"
  18 #include <memory>
  19 #include <string>
  20 #include <system_error>
  21 #include <utility>
  22
  23 #include "mlir/IR/AsmState.h"
  24 #include "mlir/IR/BuiltinOps.h"
  25 #include "mlir/IR/MLIRContext.h"
  26 #include "mlir/Parser/Parser.h"
  27
  28 #include "llvm/ADT/StringRef.h"
  29 #include "llvm/Support/CommandLine.h"
  30 #include "llvm/Support/ErrorOr.h"
  31 #include "llvm/Support/MemoryBuffer.h"
  32 #include "llvm/Support/SourceMgr.h"
  33 #include "llvm/Support/raw_ostream.h"
  34
  35 using namespace toy;
  36 namespace cl = llvm::cl;
  37
  38 static cl::opt<std::string> inputFilename(cl::Positional,
  39                                           cl::desc("<input toy file>"),
  40                                           cl::init("-"),
  41                                           cl::value_desc("filename"));
  42
  43 namespace {
  44 enum InputType { Toy, MLIR };
  45 } // namespace
  46 static cl::opt<enum InputType> inputType(
  47     "x", cl::init(Toy), cl::desc("Decided the kind of output desired"),
  48     cl::values(clEnumValN(Toy, "toy", "load the input file as a Toy source.")),
  49     cl::values(clEnumValN(MLIR, "mlir",
  50                           "load the input file as an MLIR file")));
  51
  52 namespace {
  53 enum Action { None, DumpAST, DumpMLIR };
  54 } // namespace
  55 static cl::opt<enum Action> emitAction(
  56     "emit", cl::desc("Select the kind of output desired"),
  57     cl::values(clEnumValN(DumpAST, "ast", "output the AST dump")),
  58     cl::values(clEnumValN(DumpMLIR, "mlir", "output the MLIR dump")));
  59
  60 /// Returns a Toy AST resulting from parsing the file or a nullptr on error.
  61 std::unique_ptr<toy::ModuleAST> parseInputFile(llvm::StringRef filename) {
  62   llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> fileOrErr =
  63       llvm::MemoryBuffer::getFileOrSTDIN(filename);
  64   if (std::error_code ec = fileOrErr.getError()) {
  65     llvm::errs() << "Could not open input file: " << ec.message() << "\n";
  66     return nullptr;
  67   }
  68   auto buffer = fileOrErr.get()->getBuffer();
  69   LexerBuffer lexer(buffer.begin(), buffer.end(), std::string(filename));
  70   Parser parser(lexer);
  71   return parser.parseModule();
  72 }
  73
  74 int dumpMLIR() {
  75   mlir::MLIRContext context;
  76   // Load our Dialect in this MLIR Context.
  77   context.getOrLoadDialect<mlir::toy::ToyDialect>();
  78
  79   // Handle '.toy' input to the compiler.
  80   if (inputType != InputType::MLIR &&
  81       !llvm::StringRef(inputFilename).ends_with(".mlir")) {
  82     auto moduleAST = parseInputFile(inputFilename);
  83     if (!moduleAST)
  84       return 6;
  85     mlir::OwningOpRef<mlir::ModuleOp> module = mlirGen(context, *moduleAST);
  86     if (!module)
  87       return 1;
  88
  89     module->dump();
  90     return 0;
  91   }
  92
  93   // Otherwise, the input is '.mlir'.
  94   llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> fileOrErr =
  95       llvm::MemoryBuffer::getFileOrSTDIN(inputFilename);
  96   if (std::error_code ec = fileOrErr.getError()) {
  97     llvm::errs() << "Could not open input file: " << ec.message() << "\n";
  98     return -1;
  99   }
 100
 101   // Parse the input mlir.
 102   llvm::SourceMgr sourceMgr;
 103   sourceMgr.AddNewSourceBuffer(std::move(*fileOrErr), llvm::SMLoc());
 104   mlir::OwningOpRef<mlir::ModuleOp> module =
 105       mlir::parseSourceFile<mlir::ModuleOp>(sourceMgr, &context);
 106   if (!module) {
 107     llvm::errs() << "Error can't load file " << inputFilename << "\n";
 108     return 3;
 109   }
 110
 111   module->dump();
 112   return 0;
 113 }
 114
 115 int dumpAST() {
 116   if (inputType == InputType::MLIR) {
 117     llvm::errs() << "Can't dump a Toy AST when the input is MLIR\n";
 118     return 5;
 119   }
 120
 121   auto moduleAST = parseInputFile(inputFilename);
 122   if (!moduleAST)
 123     return 1;
 124
 125   dump(*moduleAST);
 126   return 0;
 127 }
 128
 129 int main(int argc, char **argv) {
 130   // Register any command line options.
 131   mlir::registerAsmPrinterCLOptions();
 132   mlir::registerMLIRContextCLOptions();
 133   cl::ParseCommandLineOptions(argc, argv, "toy compiler\n");
 134
 135   switch (emitAction) {
 136   case Action::DumpAST:
 137     return dumpAST();
 138   case Action::DumpMLIR:
 139     return dumpMLIR();
 140   default:
 141     llvm::errs() << "No action specified (parsing only?), use -emit=<action>\n";
 142   }
 143
 144   return 0;
 145 }