mlir/tools/mlir-pdll/mlir-pdll.cpp

   1 //===- mlir-pdll.cpp - MLIR PDLL frontend -----------------------*- C++ -*-===//
   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 #include "mlir/IR/BuiltinOps.h"
  10 #include "mlir/Support/FileUtilities.h"
  11 #include "mlir/Support/ToolUtilities.h"
  12 #include "mlir/Tools/PDLL/AST/Context.h"
  13 #include "mlir/Tools/PDLL/AST/Nodes.h"
  14 #include "mlir/Tools/PDLL/CodeGen/CPPGen.h"
  15 #include "mlir/Tools/PDLL/CodeGen/MLIRGen.h"
  16 #include "mlir/Tools/PDLL/ODS/Context.h"
  17 #include "mlir/Tools/PDLL/Parser/Parser.h"
  18 #include "llvm/Support/CommandLine.h"
  19 #include "llvm/Support/InitLLVM.h"
  20 #include "llvm/Support/SourceMgr.h"
  21 #include "llvm/Support/ToolOutputFile.h"
  22 #include <set>
  23
  24 using namespace mlir;
  25 using namespace mlir::pdll;
  26
  27 //===----------------------------------------------------------------------===//
  28 // main
  29 //===----------------------------------------------------------------------===//
  30
  31 /// The desired output type.
  32 enum class OutputType {
  33   AST,
  34   MLIR,
  35   CPP,
  36 };
  37
  38 static LogicalResult
  39 processBuffer(raw_ostream &os, std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,
  40               OutputType outputType, std::vector<std::string> &includeDirs,
  41               bool dumpODS, std::set<std::string> *includedFiles) {
  42   llvm::SourceMgr sourceMgr;
  43   sourceMgr.setIncludeDirs(includeDirs);
  44   sourceMgr.AddNewSourceBuffer(std::move(chunkBuffer), SMLoc());
  45
  46   // If we are dumping ODS information, also enable documentation to ensure the
  47   // summary and description information is imported as well.
  48   bool enableDocumentation = dumpODS;
  49
  50   ods::Context odsContext;
  51   ast::Context astContext(odsContext);
  52   FailureOr<ast::Module *> module =
  53       parsePDLLAST(astContext, sourceMgr, enableDocumentation);
  54   if (failed(module))
  55     return failure();
  56
  57   // Add the files that were included to the set.
  58   if (includedFiles) {
  59     for (unsigned i = 1, e = sourceMgr.getNumBuffers(); i < e; ++i) {
  60       includedFiles->insert(
  61           sourceMgr.getMemoryBuffer(i + 1)->getBufferIdentifier().str());
  62     }
  63   }
  64
  65   // Print out the ODS information if requested.
  66   if (dumpODS)
  67     odsContext.print(llvm::errs());
  68
  69   // Generate the output.
  70   if (outputType == OutputType::AST) {
  71     (*module)->print(os);
  72     return success();
  73   }
  74
  75   MLIRContext mlirContext;
  76   OwningOpRef<ModuleOp> pdlModule =
  77       codegenPDLLToMLIR(&mlirContext, astContext, sourceMgr, **module);
  78   if (!pdlModule)
  79     return failure();
  80
  81   if (outputType == OutputType::MLIR) {
  82     pdlModule->print(os, OpPrintingFlags().enableDebugInfo());
  83     return success();
  84   }
  85   codegenPDLLToCPP(**module, *pdlModule, os);
  86   return success();
  87 }
  88
  89 /// Create a dependency file for `-d` option.
  90 ///
  91 /// This functionality is generally only for the benefit of the build system,
  92 /// and is modeled after the same option in TableGen.
  93 static LogicalResult
  94 createDependencyFile(StringRef outputFilename, StringRef dependencyFile,
  95                      std::set<std::string> &includedFiles) {
  96   if (outputFilename == "-") {
  97     llvm::errs() << "error: the option -d must be used together with -o\n";
  98     return failure();
  99   }
 100
 101   std::string errorMessage;
 102   std::unique_ptr<llvm::ToolOutputFile> outputFile =
 103       openOutputFile(dependencyFile, &errorMessage);
 104   if (!outputFile) {
 105     llvm::errs() << errorMessage << "\n";
 106     return failure();
 107   }
 108
 109   outputFile->os() << outputFilename << ":";
 110   for (const auto &includeFile : includedFiles)
 111     outputFile->os() << ' ' << includeFile;
 112   outputFile->os() << "\n";
 113   outputFile->keep();
 114   return success();
 115 }
 116
 117 int main(int argc, char **argv) {
 118   // FIXME: This is necessary because we link in TableGen, which defines its
 119   // options as static variables.. some of which overlap with our options.
 120   llvm::cl::ResetCommandLineParser();
 121
 122   llvm::cl::opt<std::string> inputFilename(
 123       llvm::cl::Positional, llvm::cl::desc("<input file>"), llvm::cl::init("-"),
 124       llvm::cl::value_desc("filename"));
 125
 126   llvm::cl::opt<std::string> outputFilename(
 127       "o", llvm::cl::desc("Output filename"), llvm::cl::value_desc("filename"),
 128       llvm::cl::init("-"));
 129
 130   llvm::cl::list<std::string> includeDirs(
 131       "I", llvm::cl::desc("Directory of include files"),
 132       llvm::cl::value_desc("directory"), llvm::cl::Prefix);
 133
 134   llvm::cl::opt<bool> dumpODS(
 135       "dump-ods",
 136       llvm::cl::desc(
 137           "Print out the parsed ODS information from the input file"),
 138       llvm::cl::init(false));
 139   llvm::cl::opt<bool> splitInputFile(
 140       "split-input-file",
 141       llvm::cl::desc("Split the input file into pieces and process each "
 142                      "chunk independently"),
 143       llvm::cl::init(false));
 144   llvm::cl::opt<enum OutputType> outputType(
 145       "x", llvm::cl::init(OutputType::AST),
 146       llvm::cl::desc("The type of output desired"),
 147       llvm::cl::values(clEnumValN(OutputType::AST, "ast",
 148                                   "generate the AST for the input file"),
 149                        clEnumValN(OutputType::MLIR, "mlir",
 150                                   "generate the PDL MLIR for the input file"),
 151                        clEnumValN(OutputType::CPP, "cpp",
 152                                   "generate a C++ source file containing the "
 153                                   "patterns for the input file")));
 154   llvm::cl::opt<std::string> dependencyFilename(
 155       "d", llvm::cl::desc("Dependency filename"),
 156       llvm::cl::value_desc("filename"), llvm::cl::init(""));
 157   llvm::cl::opt<bool> writeIfChanged(
 158       "write-if-changed",
 159       llvm::cl::desc("Only write to the output file if it changed"));
 160
 161   llvm::InitLLVM y(argc, argv);
 162   llvm::cl::ParseCommandLineOptions(argc, argv, "PDLL Frontend");
 163
 164   // Set up the input file.
 165   std::string errorMessage;
 166   std::unique_ptr<llvm::MemoryBuffer> inputFile =
 167       openInputFile(inputFilename, &errorMessage);
 168   if (!inputFile) {
 169     llvm::errs() << errorMessage << "\n";
 170     return 1;
 171   }
 172
 173   // If we are creating a dependency file, we'll also need to track what files
 174   // get included during processing.
 175   std::set<std::string> includedFilesStorage;
 176   std::set<std::string> *includedFiles = nullptr;
 177   if (!dependencyFilename.empty())
 178     includedFiles = &includedFilesStorage;
 179
 180   // The split-input-file mode is a very specific mode that slices the file
 181   // up into small pieces and checks each independently.
 182   std::string outputStr;
 183   llvm::raw_string_ostream outputStrOS(outputStr);
 184   auto processFn = [&](std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,
 185                        raw_ostream &os) {
 186     return processBuffer(os, std::move(chunkBuffer), outputType, includeDirs,
 187                          dumpODS, includedFiles);
 188   };
 189   if (failed(splitAndProcessBuffer(std::move(inputFile), processFn, outputStrOS,
 190                                    splitInputFile)))
 191     return 1;
 192
 193   // Write the output.
 194   bool shouldWriteOutput = true;
 195   if (writeIfChanged) {
 196     // Only update the real output file if there are any differences. This
 197     // prevents recompilation of all the files depending on it if there aren't
 198     // any.
 199     if (auto existingOrErr =
 200             llvm::MemoryBuffer::getFile(outputFilename, /*IsText=*/true))
 201       if (std::move(existingOrErr.get())->getBuffer() == outputStrOS.str())
 202         shouldWriteOutput = false;
 203   }
 204
 205   // Populate the output file if necessary.
 206   if (shouldWriteOutput) {
 207     std::unique_ptr<llvm::ToolOutputFile> outputFile =
 208         openOutputFile(outputFilename, &errorMessage);
 209     if (!outputFile) {
 210       llvm::errs() << errorMessage << "\n";
 211       return 1;
 212     }
 213     outputFile->os() << outputStrOS.str();
 214     outputFile->keep();
 215   }
 216
 217   // Always write the depfile, even if the main output hasn't changed. If it's
 218   // missing, Ninja considers the output dirty.
 219   if (!dependencyFilename.empty()) {
 220     if (failed(createDependencyFile(outputFilename, dependencyFilename,
 221                                     includedFilesStorage)))
 222       return 1;
 223   }
 224
 225   return 0;
 226 }