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