394 lines · cpp
1//===- mlir-rewrite.cpp - MLIR Rewrite Driver -----------------------------===//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// Main entry function for mlir-rewrite.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/AsmParser/AsmParser.h"14#include "mlir/AsmParser/AsmParserState.h"15#include "mlir/IR/AsmState.h"16#include "mlir/IR/Dialect.h"17#include "mlir/IR/MLIRContext.h"18#include "mlir/InitAllDialects.h"19#include "mlir/Pass/Pass.h"20#include "mlir/Pass/PassManager.h"21#include "mlir/Support/FileUtilities.h"22#include "mlir/Tools/ParseUtilities.h"23#include "llvm/ADT/RewriteBuffer.h"24#include "llvm/Support/CommandLine.h"25#include "llvm/Support/InitLLVM.h"26#include "llvm/Support/LineIterator.h"27#include "llvm/Support/ManagedStatic.h"28#include "llvm/Support/Regex.h"29#include "llvm/Support/SourceMgr.h"30#include "llvm/Support/ToolOutputFile.h"31 32using namespace mlir;33 34namespace mlir {35using OperationDefinition = AsmParserState::OperationDefinition;36 37/// Return the source code associated with the OperationDefinition.38static SMRange getOpRange(const OperationDefinition &op) {39 const char *startOp = op.scopeLoc.Start.getPointer();40 const char *endOp = op.scopeLoc.End.getPointer();41 42 for (const auto &res : op.resultGroups) {43 SMRange range = res.definition.loc;44 startOp = std::min(startOp, range.Start.getPointer());45 }46 return {SMLoc::getFromPointer(startOp), SMLoc::getFromPointer(endOp)};47}48 49/// Helper to simplify rewriting the source file.50class RewritePad {51public:52 static std::unique_ptr<RewritePad> init(StringRef inputFilename,53 StringRef outputFilename);54 55 /// Return the context the file was parsed into.56 MLIRContext *getContext() { return &context; }57 58 /// Return the OperationDefinition's of the operation's parsed.59 iterator_range<AsmParserState::OperationDefIterator> getOpDefs() {60 return asmState.getOpDefs();61 }62 63 /// Insert the specified string at the specified location in the original64 /// buffer.65 void insertText(SMLoc pos, StringRef str, bool insertAfter = true) {66 rewriteBuffer.InsertText(pos.getPointer() - start, str, insertAfter);67 }68 69 /// Replace the range of the source text with the corresponding string in the70 /// output.71 void replaceRange(SMRange range, StringRef str) {72 rewriteBuffer.ReplaceText(range.Start.getPointer() - start,73 range.End.getPointer() - range.Start.getPointer(),74 str);75 }76 77 /// Replace the range of the operation in the source text with the78 /// corresponding string in the output.79 void replaceDef(const OperationDefinition &opDef, StringRef newDef) {80 replaceRange(getOpRange(opDef), newDef);81 }82 83 /// Return the source string corresponding to the source range.84 StringRef getSourceString(SMRange range) {85 return StringRef(range.Start.getPointer(),86 range.End.getPointer() - range.Start.getPointer());87 }88 89 /// Return the source string corresponding to operation definition.90 StringRef getSourceString(const OperationDefinition &opDef) {91 auto range = getOpRange(opDef);92 return getSourceString(range);93 }94 95 /// Write to stream the result of applying all changes to the96 /// original buffer.97 /// Note that it isn't safe to use this function to overwrite memory mapped98 /// files in-place (PR17960).99 ///100 /// The original buffer is not actually changed.101 raw_ostream &write(raw_ostream &stream) const {102 return rewriteBuffer.write(stream);103 }104 105 /// Return lines that are purely comments.106 SmallVector<SMRange> getSingleLineComments() {107 unsigned curBuf = sourceMgr.getMainFileID();108 const llvm::MemoryBuffer *curMB = sourceMgr.getMemoryBuffer(curBuf);109 llvm::line_iterator lineIterator(*curMB);110 SmallVector<SMRange> ret;111 for (; !lineIterator.is_at_end(); ++lineIterator) {112 StringRef trimmed = lineIterator->ltrim();113 if (trimmed.starts_with("//")) {114 ret.emplace_back(115 SMLoc::getFromPointer(trimmed.data()),116 SMLoc::getFromPointer(trimmed.data() + trimmed.size()));117 }118 }119 return ret;120 }121 122 /// Return the IR from parsed file.123 Block *getParsed() { return &parsedIR; }124 125 /// Return the definition for the given operation, or nullptr if the given126 /// operation does not have a definition.127 const OperationDefinition &getOpDef(Operation *op) const {128 return *asmState.getOpDef(op);129 }130 131private:132 // The context and state required to parse.133 MLIRContext context;134 llvm::SourceMgr sourceMgr;135 DialectRegistry registry;136 FallbackAsmResourceMap fallbackResourceMap;137 138 // Storage of textual parsing results.139 AsmParserState asmState;140 141 // Parsed IR.142 Block parsedIR;143 144 // The RewriteBuffer is doing most of the real work.145 llvm::RewriteBuffer rewriteBuffer;146 147 // Start of the original input, used to compute offset.148 const char *start;149};150 151std::unique_ptr<RewritePad> RewritePad::init(StringRef inputFilename,152 StringRef outputFilename) {153 std::unique_ptr<RewritePad> r = std::make_unique<RewritePad>();154 155 // Register all the dialects needed.156 registerAllDialects(r->registry);157 158 // Set up the input file.159 std::string errorMessage;160 std::unique_ptr<llvm::MemoryBuffer> file =161 openInputFile(inputFilename, &errorMessage);162 if (!file) {163 llvm::errs() << errorMessage << "\n";164 return nullptr;165 }166 r->sourceMgr.AddNewSourceBuffer(std::move(file), SMLoc());167 168 // Set up the MLIR context and error handling.169 r->context.appendDialectRegistry(r->registry);170 171 // Record the start of the buffer to compute offsets with.172 unsigned curBuf = r->sourceMgr.getMainFileID();173 const llvm::MemoryBuffer *curMB = r->sourceMgr.getMemoryBuffer(curBuf);174 r->start = curMB->getBufferStart();175 r->rewriteBuffer.Initialize(curMB->getBuffer());176 177 // Parse and populate the AsmParserState.178 ParserConfig parseConfig(&r->context, /*verifyAfterParse=*/true,179 &r->fallbackResourceMap);180 // Always allow unregistered.181 r->context.allowUnregisteredDialects(true);182 if (failed(parseAsmSourceFile(r->sourceMgr, &r->parsedIR, parseConfig,183 &r->asmState)))184 return nullptr;185 186 return r;187}188 189/// Return the source code associated with the operation name.190static SMRange getOpNameRange(const OperationDefinition &op) { return op.loc; }191 192/// Return whether the operation was printed using generic syntax in original193/// buffer.194static bool isGeneric(const OperationDefinition &op) {195 return op.loc.Start.getPointer()[0] == '"';196}197 198static inline int asMainReturnCode(LogicalResult r) {199 return r.succeeded() ? EXIT_SUCCESS : EXIT_FAILURE;200}201 202/// Reriter function to invoke.203using RewriterFunction = std::function<mlir::LogicalResult(204 mlir::RewritePad &rewriteState, llvm::raw_ostream &os)>;205 206/// Structure to group information about a rewriter (argument to invoke via207/// mlir-tblgen, description, and rewriter function).208class RewriterInfo {209public:210 /// RewriterInfo constructor should not be invoked directly, instead use211 /// RewriterRegistration or registerRewriter.212 RewriterInfo(StringRef arg, StringRef description, RewriterFunction rewriter)213 : arg(arg), description(description), rewriter(std::move(rewriter)) {}214 215 /// Invokes the rewriter and returns whether the rewriter failed.216 LogicalResult invoke(mlir::RewritePad &rewriteState, raw_ostream &os) const {217 assert(rewriter && "Cannot call rewriter with null rewriter");218 return rewriter(rewriteState, os);219 }220 221 /// Returns the command line option that may be passed to 'mlir-rewrite' to222 /// invoke this rewriter.223 StringRef getRewriterArgument() const { return arg; }224 225 /// Returns a description for the rewriter.226 StringRef getRewriterDescription() const { return description; }227 228private:229 // The argument with which to invoke the rewriter via mlir-tblgen.230 StringRef arg;231 232 // Description of the rewriter.233 StringRef description;234 235 // Rewritererator function.236 RewriterFunction rewriter;237};238 239static llvm::ManagedStatic<std::vector<RewriterInfo>> rewriterRegistry;240 241/// Adds command line option for each registered rewriter.242struct RewriterNameParser : public llvm::cl::parser<const RewriterInfo *> {243 RewriterNameParser(llvm::cl::Option &opt);244 245 void printOptionInfo(const llvm::cl::Option &o,246 size_t globalWidth) const override;247};248 249/// RewriterRegistration provides a global initializer that registers a rewriter250/// function.251struct RewriterRegistration {252 RewriterRegistration(StringRef arg, StringRef description,253 const RewriterFunction &function);254};255 256RewriterRegistration::RewriterRegistration(StringRef arg, StringRef description,257 const RewriterFunction &function) {258 rewriterRegistry->emplace_back(arg, description, function);259}260 261RewriterNameParser::RewriterNameParser(llvm::cl::Option &opt)262 : llvm::cl::parser<const RewriterInfo *>(opt) {263 for (const auto &kv : *rewriterRegistry) {264 addLiteralOption(kv.getRewriterArgument(), &kv,265 kv.getRewriterDescription());266 }267}268 269void RewriterNameParser::printOptionInfo(const llvm::cl::Option &o,270 size_t globalWidth) const {271 RewriterNameParser *tp = const_cast<RewriterNameParser *>(this);272 llvm::array_pod_sort(tp->Values.begin(), tp->Values.end(),273 [](const RewriterNameParser::OptionInfo *vT1,274 const RewriterNameParser::OptionInfo *vT2) {275 return vT1->Name.compare(vT2->Name);276 });277 using llvm::cl::parser;278 parser<const RewriterInfo *>::printOptionInfo(o, globalWidth);279}280 281} // namespace mlir282 283// TODO: Make these injectable too in non-global way.284static llvm::cl::OptionCategory clSimpleRenameCategory{"simple-rename options"};285static llvm::cl::opt<std::string> simpleRenameOpName{286 "simple-rename-op-name", llvm::cl::desc("Name of op to match on"),287 llvm::cl::cat(clSimpleRenameCategory)};288static llvm::cl::opt<std::string> simpleRenameMatch{289 "simple-rename-match", llvm::cl::desc("Match string for rename"),290 llvm::cl::cat(clSimpleRenameCategory)};291static llvm::cl::opt<std::string> simpleRenameReplace{292 "simple-rename-replace", llvm::cl::desc("Replace string for rename"),293 llvm::cl::cat(clSimpleRenameCategory)};294 295// Rewriter that does simple renames.296static LogicalResult simpleRename(RewritePad &rewriteState, raw_ostream &os) {297 StringRef opName = simpleRenameOpName;298 StringRef match = simpleRenameMatch;299 StringRef replace = simpleRenameReplace;300 llvm::Regex regex(match);301 302 rewriteState.getParsed()->walk([&](Operation *op) {303 if (op->getName().getStringRef() != opName)304 return;305 306 const OperationDefinition &opDef = rewriteState.getOpDef(op);307 SMRange range = getOpRange(opDef);308 // This is a little bit overkill for simple.309 std::string str = regex.sub(replace, rewriteState.getSourceString(range));310 rewriteState.replaceRange(range, str);311 });312 return success();313}314 315static mlir::RewriterRegistration rewriteSimpleRename("simple-rename",316 "Perform a simple rename",317 simpleRename);318 319// Rewriter that insert range markers.320static LogicalResult markRanges(RewritePad &rewriteState, raw_ostream &os) {321 for (const auto &it : rewriteState.getOpDefs()) {322 auto [startOp, endOp] = getOpRange(it);323 324 rewriteState.insertText(startOp, "<");325 rewriteState.insertText(endOp, ">");326 327 auto nameRange = getOpNameRange(it);328 329 if (isGeneric(it)) {330 rewriteState.insertText(nameRange.Start, "[");331 rewriteState.insertText(nameRange.End, "]");332 } else {333 rewriteState.insertText(nameRange.Start, "![");334 rewriteState.insertText(nameRange.End, "]!");335 }336 }337 338 // Highlight all comment lines.339 // TODO: Could be replaced if this is kept in memory.340 for (auto commentLine : rewriteState.getSingleLineComments()) {341 rewriteState.insertText(commentLine.Start, "{");342 rewriteState.insertText(commentLine.End, "}");343 }344 345 return success();346}347 348static mlir::RewriterRegistration349 rewriteMarkRanges("mark-ranges", "Indicate ranges parsed", markRanges);350 351int main(int argc, char **argv) {352 llvm::cl::opt<std::string> inputFilename(llvm::cl::Positional,353 llvm::cl::desc("<input file>"),354 llvm::cl::init("-"));355 356 llvm::cl::opt<std::string> outputFilename(357 "o", llvm::cl::desc("Output filename"), llvm::cl::value_desc("filename"),358 llvm::cl::init("-"));359 360 llvm::cl::opt<const mlir::RewriterInfo *, false, mlir::RewriterNameParser>361 rewriter("", llvm::cl::desc("Rewriter to run"));362 363 std::string helpHeader = "mlir-rewrite";364 365 llvm::cl::ParseCommandLineOptions(argc, argv, helpHeader);366 367 // If no rewriter has been selected, exit with error code. Could also just368 // return but its unlikely this was intentionally being used as `cp`.369 if (!rewriter) {370 llvm::errs() << "No rewriter selected!\n";371 return mlir::asMainReturnCode(mlir::failure());372 }373 374 // Set up rewrite buffer.375 auto rewriterOr = RewritePad::init(inputFilename, outputFilename);376 if (!rewriterOr)377 return mlir::asMainReturnCode(mlir::failure());378 379 // Set up the output file.380 std::string errorMessage;381 auto output = openOutputFile(outputFilename, &errorMessage);382 if (!output) {383 llvm::errs() << errorMessage << "\n";384 return mlir::asMainReturnCode(mlir::failure());385 }386 387 LogicalResult result = rewriter->invoke(*rewriterOr, output->os());388 if (succeeded(result)) {389 rewriterOr->write(output->os());390 output->keep();391 }392 return mlir::asMainReturnCode(result);393}394