253 lines · 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-exception6//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 "llvm/Support/VirtualFileSystem.h"23#include <set>24 25using namespace mlir;26using namespace mlir::pdll;27 28//===----------------------------------------------------------------------===//29// main30//===----------------------------------------------------------------------===//31 32/// The desired output type.33enum class OutputType {34 AST,35 MLIR,36 CPP,37};38 39static LogicalResult40processBuffer(raw_ostream &os, std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,41 OutputType outputType, std::vector<std::string> &includeDirs,42 bool dumpODS, std::set<std::string> *includedFiles) {43 llvm::SourceMgr sourceMgr;44 sourceMgr.setIncludeDirs(includeDirs);45 sourceMgr.setVirtualFileSystem(llvm::vfs::getRealFileSystem());46 sourceMgr.AddNewSourceBuffer(std::move(chunkBuffer), SMLoc());47 48 // If we are dumping ODS information, also enable documentation to ensure the49 // summary and description information is imported as well.50 bool enableDocumentation = dumpODS;51 52 ods::Context odsContext;53 ast::Context astContext(odsContext);54 FailureOr<ast::Module *> module =55 parsePDLLAST(astContext, sourceMgr, enableDocumentation);56 if (failed(module))57 return failure();58 59 // Add the files that were included to the set.60 if (includedFiles) {61 for (unsigned i = 1, e = sourceMgr.getNumBuffers(); i < e; ++i) {62 includedFiles->insert(63 sourceMgr.getMemoryBuffer(i + 1)->getBufferIdentifier().str());64 }65 }66 67 // Print out the ODS information if requested.68 if (dumpODS)69 odsContext.print(llvm::errs());70 71 // Generate the output.72 if (outputType == OutputType::AST) {73 (*module)->print(os);74 return success();75 }76 77 MLIRContext mlirContext;78 OwningOpRef<ModuleOp> pdlModule =79 codegenPDLLToMLIR(&mlirContext, astContext, sourceMgr, **module);80 if (!pdlModule)81 return failure();82 83 if (outputType == OutputType::MLIR) {84 pdlModule->print(os, OpPrintingFlags().enableDebugInfo());85 return success();86 }87 codegenPDLLToCPP(**module, *pdlModule, os);88 return success();89}90 91/// Create a dependency file for `-d` option.92///93/// This functionality is generally only for the benefit of the build system,94/// and is modeled after the same option in TableGen.95static LogicalResult96createDependencyFile(StringRef outputFilename, StringRef dependencyFile,97 std::set<std::string> &includedFiles) {98 if (outputFilename == "-") {99 llvm::errs() << "error: the option -d must be used together with -o\n";100 return failure();101 }102 103 std::string errorMessage;104 std::unique_ptr<llvm::ToolOutputFile> outputFile =105 openOutputFile(dependencyFile, &errorMessage);106 if (!outputFile) {107 llvm::errs() << errorMessage << "\n";108 return failure();109 }110 111 outputFile->os() << outputFilename << ":";112 for (const auto &includeFile : includedFiles)113 outputFile->os() << ' ' << includeFile;114 outputFile->os() << "\n";115 outputFile->keep();116 return success();117}118 119int main(int argc, char **argv) {120 // FIXME: This is necessary because we link in TableGen, which defines its121 // options as static variables.. some of which overlap with our options.122 llvm::cl::ResetCommandLineParser();123 124 llvm::cl::opt<std::string> inputFilename(125 llvm::cl::Positional, llvm::cl::desc("<input file>"), llvm::cl::init("-"),126 llvm::cl::value_desc("filename"));127 128 llvm::cl::opt<std::string> outputFilename(129 "o", llvm::cl::desc("Output filename"), llvm::cl::value_desc("filename"),130 llvm::cl::init("-"));131 132 llvm::cl::list<std::string> includeDirs(133 "I", llvm::cl::desc("Directory of include files"),134 llvm::cl::value_desc("directory"), llvm::cl::Prefix);135 136 llvm::cl::opt<bool> dumpODS(137 "dump-ods",138 llvm::cl::desc(139 "Print out the parsed ODS information from the input file"),140 llvm::cl::init(false));141 llvm::cl::opt<std::string> inputSplitMarker{142 "split-input-file", llvm::cl::ValueOptional,143 llvm::cl::callback([&](const std::string &str) {144 // Implicit value: use default marker if flag was used without value.145 if (str.empty())146 inputSplitMarker.setValue(kDefaultSplitMarker);147 }),148 llvm::cl::desc("Split the input file into chunks using the given or "149 "default marker and process each chunk independently"),150 llvm::cl::init("")};151 llvm::cl::opt<std::string> outputSplitMarker(152 "output-split-marker",153 llvm::cl::desc("Split marker to use for merging the ouput"),154 llvm::cl::init(kDefaultSplitMarker));155 llvm::cl::opt<enum OutputType> outputType(156 "x", llvm::cl::init(OutputType::AST),157 llvm::cl::desc("The type of output desired"),158 llvm::cl::values(clEnumValN(OutputType::AST, "ast",159 "generate the AST for the input file"),160 clEnumValN(OutputType::MLIR, "mlir",161 "generate the PDL MLIR for the input file"),162 clEnumValN(OutputType::CPP, "cpp",163 "generate a C++ source file containing the "164 "patterns for the input file")));165 llvm::cl::opt<std::string> dependencyFilename(166 "d", llvm::cl::desc("Dependency filename"),167 llvm::cl::value_desc("filename"), llvm::cl::init(""));168 llvm::cl::opt<bool> writeIfChanged(169 "write-if-changed",170 llvm::cl::desc("Only write to the output file if it changed"));171 172 // `ResetCommandLineParser` at the above unregistered the "D" option173 // of `llvm-tblgen`, which causes tblgen usage to fail due to174 // "Unknnown command line argument '-D...`" when a macros name is175 // present. The following is a workaround to re-register it again.176 llvm::cl::list<std::string> macroNames(177 "D",178 llvm::cl::desc("Name of the macro to be defined -- ignored by mlir-pdll"),179 llvm::cl::value_desc("macro name"), llvm::cl::Prefix);180 181 llvm::InitLLVM y(argc, argv);182 llvm::cl::ParseCommandLineOptions(argc, argv, "PDLL Frontend");183 184 // Set up the input file.185 std::string errorMessage;186 std::unique_ptr<llvm::MemoryBuffer> inputFile =187 openInputFile(inputFilename, &errorMessage);188 if (!inputFile) {189 llvm::errs() << errorMessage << "\n";190 return 1;191 }192 193 // If we are creating a dependency file, we'll also need to track what files194 // get included during processing.195 std::set<std::string> includedFilesStorage;196 std::set<std::string> *includedFiles = nullptr;197 if (!dependencyFilename.empty())198 includedFiles = &includedFilesStorage;199 200 // The split-input-file mode is a very specific mode that slices the file201 // up into small pieces and checks each independently.202 std::string outputStr;203 llvm::raw_string_ostream outputStrOS(outputStr);204 auto processFn = [&](std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,205 raw_ostream &os) {206 // Split does not guarantee null-termination. Make a copy of the buffer to207 // ensure null-termination.208 if (!chunkBuffer->getBuffer().ends_with('\0')) {209 chunkBuffer = llvm::MemoryBuffer::getMemBufferCopy(210 chunkBuffer->getBuffer(), chunkBuffer->getBufferIdentifier());211 }212 return processBuffer(os, std::move(chunkBuffer), outputType, includeDirs,213 dumpODS, includedFiles);214 };215 if (failed(splitAndProcessBuffer(std::move(inputFile), processFn, outputStrOS,216 inputSplitMarker, outputSplitMarker)))217 return 1;218 219 // Write the output.220 bool shouldWriteOutput = true;221 if (writeIfChanged) {222 // Only update the real output file if there are any differences. This223 // prevents recompilation of all the files depending on it if there aren't224 // any.225 if (auto existingOrErr =226 llvm::MemoryBuffer::getFile(outputFilename, /*IsText=*/true))227 if (std::move(existingOrErr.get())->getBuffer() == outputStr)228 shouldWriteOutput = false;229 }230 231 // Populate the output file if necessary.232 if (shouldWriteOutput) {233 std::unique_ptr<llvm::ToolOutputFile> outputFile =234 openOutputFile(outputFilename, &errorMessage);235 if (!outputFile) {236 llvm::errs() << errorMessage << "\n";237 return 1;238 }239 outputFile->os() << outputStr;240 outputFile->keep();241 }242 243 // Always write the depfile, even if the main output hasn't changed. If it's244 // missing, Ninja considers the output dirty.245 if (!dependencyFilename.empty()) {246 if (failed(createDependencyFile(outputFilename, dependencyFilename,247 includedFilesStorage)))248 return 1;249 }250 251 return 0;252}253