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(F: std::move(chunkBuffer), IncludeLoc: 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(ctx&: astContext, sourceMgr, enableDocumentation); |
54 | if (failed(result: 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 | x: sourceMgr.getMemoryBuffer(i: i + 1)->getBufferIdentifier().str()); |
62 | } |
63 | } |
64 | |
65 | // Print out the ODS information if requested. |
66 | if (dumpODS) |
67 | odsContext.print(os&: 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(outputFilename: dependencyFile, errorMessage: &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(Val: "-" ), |
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(Val: "-" )); |
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(Val: false)); |
139 | llvm::cl::opt<std::string> inputSplitMarker{ |
140 | "split-input-file" , llvm::cl::ValueOptional, |
141 | llvm::cl::callback(CB: [&](const std::string &str) { |
142 | // Implicit value: use default marker if flag was used without value. |
143 | if (str.empty()) |
144 | inputSplitMarker.setValue(V: kDefaultSplitMarker); |
145 | }), |
146 | llvm::cl::desc("Split the input file into chunks using the given or " |
147 | "default marker and process each chunk independently" ), |
148 | llvm::cl::init(Val: "" )}; |
149 | llvm::cl::opt<std::string> outputSplitMarker( |
150 | "output-split-marker" , |
151 | llvm::cl::desc("Split marker to use for merging the ouput" ), |
152 | llvm::cl::init(Val: kDefaultSplitMarker)); |
153 | llvm::cl::opt<enum OutputType> outputType( |
154 | "x" , llvm::cl::init(Val: OutputType::AST), |
155 | llvm::cl::desc("The type of output desired" ), |
156 | llvm::cl::values(clEnumValN(OutputType::AST, "ast" , |
157 | "generate the AST for the input file" ), |
158 | clEnumValN(OutputType::MLIR, "mlir" , |
159 | "generate the PDL MLIR for the input file" ), |
160 | clEnumValN(OutputType::CPP, "cpp" , |
161 | "generate a C++ source file containing the " |
162 | "patterns for the input file" ))); |
163 | llvm::cl::opt<std::string> dependencyFilename( |
164 | "d" , llvm::cl::desc("Dependency filename" ), |
165 | llvm::cl::value_desc("filename" ), llvm::cl::init(Val: "" )); |
166 | llvm::cl::opt<bool> writeIfChanged( |
167 | "write-if-changed" , |
168 | llvm::cl::desc("Only write to the output file if it changed" )); |
169 | |
170 | llvm::InitLLVM y(argc, argv); |
171 | llvm::cl::ParseCommandLineOptions(argc, argv, Overview: "PDLL Frontend" ); |
172 | |
173 | // Set up the input file. |
174 | std::string errorMessage; |
175 | std::unique_ptr<llvm::MemoryBuffer> inputFile = |
176 | openInputFile(inputFilename, errorMessage: &errorMessage); |
177 | if (!inputFile) { |
178 | llvm::errs() << errorMessage << "\n" ; |
179 | return 1; |
180 | } |
181 | |
182 | // If we are creating a dependency file, we'll also need to track what files |
183 | // get included during processing. |
184 | std::set<std::string> includedFilesStorage; |
185 | std::set<std::string> *includedFiles = nullptr; |
186 | if (!dependencyFilename.empty()) |
187 | includedFiles = &includedFilesStorage; |
188 | |
189 | // The split-input-file mode is a very specific mode that slices the file |
190 | // up into small pieces and checks each independently. |
191 | std::string outputStr; |
192 | llvm::raw_string_ostream outputStrOS(outputStr); |
193 | auto processFn = [&](std::unique_ptr<llvm::MemoryBuffer> chunkBuffer, |
194 | raw_ostream &os) { |
195 | return processBuffer(os, chunkBuffer: std::move(chunkBuffer), outputType, includeDirs, |
196 | dumpODS, includedFiles); |
197 | }; |
198 | if (failed(result: splitAndProcessBuffer(originalBuffer: std::move(inputFile), processChunkBuffer: processFn, os&: outputStrOS, |
199 | inputSplitMarker, outputSplitMarker))) |
200 | return 1; |
201 | |
202 | // Write the output. |
203 | bool shouldWriteOutput = true; |
204 | if (writeIfChanged) { |
205 | // Only update the real output file if there are any differences. This |
206 | // prevents recompilation of all the files depending on it if there aren't |
207 | // any. |
208 | if (auto existingOrErr = |
209 | llvm::MemoryBuffer::getFile(Filename: outputFilename, /*IsText=*/true)) |
210 | if (std::move(existingOrErr.get())->getBuffer() == outputStrOS.str()) |
211 | shouldWriteOutput = false; |
212 | } |
213 | |
214 | // Populate the output file if necessary. |
215 | if (shouldWriteOutput) { |
216 | std::unique_ptr<llvm::ToolOutputFile> outputFile = |
217 | openOutputFile(outputFilename, errorMessage: &errorMessage); |
218 | if (!outputFile) { |
219 | llvm::errs() << errorMessage << "\n" ; |
220 | return 1; |
221 | } |
222 | outputFile->os() << outputStrOS.str(); |
223 | outputFile->keep(); |
224 | } |
225 | |
226 | // Always write the depfile, even if the main output hasn't changed. If it's |
227 | // missing, Ninja considers the output dirty. |
228 | if (!dependencyFilename.empty()) { |
229 | if (failed(result: createDependencyFile(outputFilename, dependencyFile: dependencyFilename, |
230 | includedFiles&: includedFilesStorage))) |
231 | return 1; |
232 | } |
233 | |
234 | return 0; |
235 | } |
236 | |