2137 lines · cpp
1//===- RewriterGen.cpp - MLIR pattern rewriter generator ------------------===//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// RewriterGen uses pattern rewrite definitions to generate rewriter matchers.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Support/IndentedOstream.h"14#include "mlir/TableGen/Argument.h"15#include "mlir/TableGen/Attribute.h"16#include "mlir/TableGen/CodeGenHelpers.h"17#include "mlir/TableGen/Format.h"18#include "mlir/TableGen/GenInfo.h"19#include "mlir/TableGen/Operator.h"20#include "mlir/TableGen/Pattern.h"21#include "mlir/TableGen/Predicate.h"22#include "mlir/TableGen/Property.h"23#include "mlir/TableGen/Type.h"24#include "llvm/ADT/FunctionExtras.h"25#include "llvm/ADT/SetVector.h"26#include "llvm/ADT/StringExtras.h"27#include "llvm/ADT/StringSet.h"28#include "llvm/Support/CommandLine.h"29#include "llvm/Support/Debug.h"30#include "llvm/Support/FormatAdapters.h"31#include "llvm/Support/PrettyStackTrace.h"32#include "llvm/Support/Signals.h"33#include "llvm/TableGen/Error.h"34#include "llvm/TableGen/Main.h"35#include "llvm/TableGen/Record.h"36#include "llvm/TableGen/TableGenBackend.h"37 38using namespace mlir;39using namespace mlir::tblgen;40 41using llvm::formatv;42using llvm::Record;43using llvm::RecordKeeper;44 45#define DEBUG_TYPE "mlir-tblgen-rewritergen"46 47namespace llvm {48template <>49struct format_provider<mlir::tblgen::Pattern::IdentifierLine> {50 static void format(const mlir::tblgen::Pattern::IdentifierLine &v,51 raw_ostream &os, StringRef style) {52 os << v.first << ":" << v.second;53 }54};55} // namespace llvm56 57//===----------------------------------------------------------------------===//58// PatternEmitter59//===----------------------------------------------------------------------===//60 61namespace {62 63class StaticMatcherHelper;64 65class PatternEmitter {66public:67 PatternEmitter(const Record *pat, RecordOperatorMap *mapper, raw_ostream &os,68 StaticMatcherHelper &helper);69 70 // Emits the mlir::RewritePattern struct named `rewriteName`.71 void emit(StringRef rewriteName);72 73 // Emits the static function of DAG matcher.74 void emitStaticMatcher(DagNode tree, std::string funcName);75 76private:77 // Emits the code for matching ops.78 void emitMatchLogic(DagNode tree, StringRef opName);79 80 // Emits the code for rewriting ops.81 void emitRewriteLogic();82 83 //===--------------------------------------------------------------------===//84 // Match utilities85 //===--------------------------------------------------------------------===//86 87 // Emits C++ statements for matching the DAG structure.88 void emitMatch(DagNode tree, StringRef name, int depth);89 90 // Emit C++ function call to static DAG matcher.91 void emitStaticMatchCall(DagNode tree, StringRef name);92 93 // Emit C++ function call to static type/attribute constraint function.94 void emitStaticVerifierCall(StringRef funcName, StringRef opName,95 StringRef arg, StringRef failureStr);96 97 // Emits C++ statements for matching using a native code call.98 void emitNativeCodeMatch(DagNode tree, StringRef name, int depth);99 100 // Emits C++ statements for matching the op constrained by the given DAG101 // `tree` returning the op's variable name.102 void emitOpMatch(DagNode tree, StringRef opName, int depth);103 104 // Emits C++ statements for matching the `argIndex`-th argument of the given105 // DAG `tree` as an operand. `operandName` and `operandMatcher` indicate the106 // bound name and the constraint of the operand respectively.107 void emitOperandMatch(DagNode tree, StringRef opName, StringRef operandName,108 int operandIndex, DagLeaf operandMatcher,109 StringRef argName, int argIndex,110 std::optional<int> variadicSubIndex);111 112 // Emits C++ statements for matching the operands which can be matched in113 // either order.114 void emitEitherOperandMatch(DagNode tree, DagNode eitherArgTree,115 StringRef opName, int argIndex, int &operandIndex,116 int depth);117 118 // Emits C++ statements for matching a variadic operand.119 void emitVariadicOperandMatch(DagNode tree, DagNode variadicArgTree,120 StringRef opName, int argIndex,121 int &operandIndex, int depth);122 123 // Emits C++ statements for matching the `argIndex`-th argument of the given124 // DAG `tree` as an attribute.125 void emitAttributeMatch(DagNode tree, StringRef castedName, int argIndex,126 int depth);127 128 // Emits C++ statements for matching the `argIndex`-th argument of the given129 // DAG `tree` as a property.130 void emitPropertyMatch(DagNode tree, StringRef castedName, int argIndex,131 int depth);132 133 // Emits C++ for checking a match with a corresponding match failure134 // diagnostic.135 void emitMatchCheck(StringRef opName, const FmtObjectBase &matchFmt,136 const llvm::formatv_object_base &failureFmt);137 138 // Emits C++ for checking a match with a corresponding match failure139 // diagnostics.140 void emitMatchCheck(StringRef opName, const std::string &matchStr,141 const std::string &failureStr);142 143 //===--------------------------------------------------------------------===//144 // Rewrite utilities145 //===--------------------------------------------------------------------===//146 147 // The entry point for handling a result pattern rooted at `resultTree`. This148 // method dispatches to concrete handlers according to `resultTree`'s kind and149 // returns a symbol representing the whole value pack. Callers are expected to150 // further resolve the symbol according to the specific use case.151 //152 // `depth` is the nesting level of `resultTree`; 0 means top-level result153 // pattern. For top-level result pattern, `resultIndex` indicates which result154 // of the matched root op this pattern is intended to replace, which can be155 // used to deduce the result type of the op generated from this result156 // pattern.157 std::string handleResultPattern(DagNode resultTree, int resultIndex,158 int depth);159 160 // Emits the C++ statement to replace the matched DAG with a value built via161 // calling native C++ code.162 std::string handleReplaceWithNativeCodeCall(DagNode resultTree, int depth);163 164 // Returns the symbol of the old value serving as the replacement.165 StringRef handleReplaceWithValue(DagNode tree);166 167 // Emits the C++ statement to replace the matched DAG with an array of168 // matched values.169 std::string handleVariadic(DagNode tree, int depth);170 171 // Trailing directives are used at the end of DAG node argument lists to172 // specify additional behaviour for op matchers and creators, etc.173 struct TrailingDirectives {174 // DAG node containing the `location` directive. Null if there is none.175 DagNode location;176 177 // DAG node containing the `returnType` directive. Null if there is none.178 DagNode returnType;179 180 // Number of found trailing directives.181 int numDirectives;182 };183 184 // Collect any trailing directives.185 TrailingDirectives getTrailingDirectives(DagNode tree);186 187 // Returns the location value to use.188 std::string getLocation(TrailingDirectives &tail);189 190 // Returns the location value to use.191 std::string handleLocationDirective(DagNode tree);192 193 // Emit return type argument.194 std::string handleReturnTypeArg(DagNode returnType, int i, int depth);195 196 // Emits the C++ statement to build a new op out of the given DAG `tree` and197 // returns the variable name that this op is assigned to. If the root op in198 // DAG `tree` has a specified name, the created op will be assigned to a199 // variable of the given name. Otherwise, a unique name will be used as the200 // result value name.201 std::string handleOpCreation(DagNode tree, int resultIndex, int depth);202 203 using ChildNodeIndexNameMap = DenseMap<unsigned, std::string>;204 205 // Emits a local variable for each value and attribute to be used for creating206 // an op.207 void createSeparateLocalVarsForOpArgs(DagNode node,208 ChildNodeIndexNameMap &childNodeNames);209 210 // Emits the concrete arguments used to call an op's builder.211 void supplyValuesForOpArgs(DagNode node,212 const ChildNodeIndexNameMap &childNodeNames,213 int depth);214 215 // Emits the local variables for holding all values as a whole and all named216 // attributes as a whole to be used for creating an op.217 void createAggregateLocalVarsForOpArgs(218 DagNode node, const ChildNodeIndexNameMap &childNodeNames, int depth);219 220 // Returns the C++ expression to construct a constant attribute of the given221 // `value` for the given attribute kind `attr`.222 std::string handleConstantAttr(Attribute attr, const Twine &value);223 224 // Returns the C++ expression to build an argument from the given DAG `leaf`.225 // `patArgName` is used to bound the argument to the source pattern.226 std::string handleOpArgument(DagLeaf leaf, StringRef patArgName);227 228 //===--------------------------------------------------------------------===//229 // General utilities230 //===--------------------------------------------------------------------===//231 232 // Collects all of the operations within the given dag tree.233 void collectOps(DagNode tree, llvm::SmallPtrSetImpl<const Operator *> &ops);234 235 // Returns a unique symbol for a local variable of the given `op`.236 std::string getUniqueSymbol(const Operator *op);237 238 //===--------------------------------------------------------------------===//239 // Symbol utilities240 //===--------------------------------------------------------------------===//241 242 // Returns how many static values the given DAG `node` correspond to.243 int getNodeValueCount(DagNode node);244 245private:246 // Pattern instantiation location followed by the location of multiclass247 // prototypes used. This is intended to be used as a whole to248 // PrintFatalError() on errors.249 ArrayRef<SMLoc> loc;250 251 // Op's TableGen Record to wrapper object.252 RecordOperatorMap *opMap;253 254 // Handy wrapper for pattern being emitted.255 Pattern pattern;256 257 // Map for all bound symbols' info.258 SymbolInfoMap symbolInfoMap;259 260 StaticMatcherHelper &staticMatcherHelper;261 262 // The next unused ID for newly created values.263 unsigned nextValueId = 0;264 265 raw_indented_ostream os;266 267 // Format contexts containing placeholder substitutions.268 FmtContext fmtCtx;269};270 271// Tracks DagNode's reference multiple times across patterns. Enables generating272// static matcher functions for DagNode's referenced multiple times rather than273// inlining them.274class StaticMatcherHelper {275public:276 StaticMatcherHelper(raw_ostream &os, const RecordKeeper &records,277 RecordOperatorMap &mapper);278 279 // Determine if we should inline the match logic or delegate to a static280 // function.281 bool useStaticMatcher(DagNode node) {282 // either/variadic node must be associated to the parentOp, thus we can't283 // emit a static matcher rooted at them.284 if (node.isEither() || node.isVariadic())285 return false;286 287 return refStats[node] > kStaticMatcherThreshold;288 }289 290 // Get the name of the static DAG matcher function corresponding to the node.291 std::string getMatcherName(DagNode node) {292 assert(useStaticMatcher(node));293 return matcherNames[node];294 }295 296 // Get the name of static type/attribute verification function.297 StringRef getVerifierName(DagLeaf leaf);298 299 // Collect the `Record`s, i.e., the DRR, so that we can get the information of300 // the duplicated DAGs.301 void addPattern(const Record *record);302 303 // Emit all static functions of DAG Matcher.304 void populateStaticMatchers(raw_ostream &os);305 306 // Emit all static functions for Constraints.307 void populateStaticConstraintFunctions(raw_ostream &os);308 309private:310 static constexpr unsigned kStaticMatcherThreshold = 1;311 312 // Consider two patterns as down below,313 // DagNode_Root_A DagNode_Root_B314 // \ \315 // DagNode_C DagNode_C316 // \ \317 // DagNode_D DagNode_D318 //319 // DagNode_Root_A and DagNode_Root_B share the same subtree which consists of320 // DagNode_C and DagNode_D. Both DagNode_C and DagNode_D are referenced321 // multiple times so we'll have static matchers for both of them. When we're322 // emitting the match logic for DagNode_C, we will check if DagNode_D has the323 // static matcher generated. If so, then we'll generate a call to the324 // function, inline otherwise. In this case, inlining is not what we want. As325 // a result, generate the static matcher in topological order to ensure all326 // the dependent static matchers are generated and we can avoid accidentally327 // inlining.328 //329 // The topological order of all the DagNodes among all patterns.330 SmallVector<std::pair<DagNode, const Record *>> topologicalOrder;331 332 RecordOperatorMap &opMap;333 334 // Records of the static function name of each DagNode335 DenseMap<DagNode, std::string> matcherNames;336 337 // After collecting all the DagNode in each pattern, `refStats` records the338 // number of users for each DagNode. We will generate the static matcher for a339 // DagNode while the number of users exceeds a certain threshold.340 DenseMap<DagNode, unsigned> refStats;341 342 // Number of static matcher generated. This is used to generate a unique name343 // for each DagNode.344 int staticMatcherCounter = 0;345 346 // The DagLeaf which contains type, attr, or prop constraint.347 SetVector<DagLeaf> constraints;348 349 // Static type/attribute verification function emitter.350 StaticVerifierFunctionEmitter staticVerifierEmitter;351};352 353} // namespace354 355PatternEmitter::PatternEmitter(const Record *pat, RecordOperatorMap *mapper,356 raw_ostream &os, StaticMatcherHelper &helper)357 : loc(pat->getLoc()), opMap(mapper), pattern(pat, mapper),358 symbolInfoMap(pat->getLoc()), staticMatcherHelper(helper), os(os) {359 fmtCtx.withBuilder("rewriter");360}361 362std::string PatternEmitter::handleConstantAttr(Attribute attr,363 const Twine &value) {364 if (!attr.isConstBuildable())365 PrintFatalError(loc, "Attribute " + attr.getAttrDefName() +366 " does not have the 'constBuilderCall' field");367 368 // TODO: Verify the constants here369 return std::string(tgfmt(attr.getConstBuilderTemplate(), &fmtCtx, value));370}371 372void PatternEmitter::emitStaticMatcher(DagNode tree, std::string funcName) {373 os << formatv(374 "static ::llvm::LogicalResult {0}(::mlir::PatternRewriter &rewriter, "375 "::mlir::Operation *op0, ::llvm::SmallVector<::mlir::Operation "376 "*, 4> &tblgen_ops",377 funcName);378 379 // We pass the reference of the variables that need to be captured. Hence we380 // need to collect all the symbols in the tree first.381 pattern.collectBoundSymbols(tree, symbolInfoMap, /*isSrcPattern=*/true);382 symbolInfoMap.assignUniqueAlternativeNames();383 for (const auto &info : symbolInfoMap)384 os << formatv(", {0}", info.second.getArgDecl(info.first));385 386 os << ") {\n";387 os.indent();388 os << "(void)tblgen_ops;\n";389 390 // Note that a static matcher is considered at least one step from the match391 // entry.392 emitMatch(tree, "op0", /*depth=*/1);393 394 os << "return ::mlir::success();\n";395 os.unindent();396 os << "}\n\n";397}398 399// Helper function to match patterns.400void PatternEmitter::emitMatch(DagNode tree, StringRef name, int depth) {401 if (tree.isNativeCodeCall()) {402 emitNativeCodeMatch(tree, name, depth);403 return;404 }405 406 if (tree.isOperation()) {407 emitOpMatch(tree, name, depth);408 return;409 }410 411 PrintFatalError(loc, "encountered non-op, non-NativeCodeCall match.");412}413 414void PatternEmitter::emitStaticMatchCall(DagNode tree, StringRef opName) {415 std::string funcName = staticMatcherHelper.getMatcherName(tree);416 os << formatv("if(::mlir::failed({0}(rewriter, {1}, tblgen_ops", funcName,417 opName);418 419 // TODO(chiahungduan): Add a lookupBoundSymbols() to do the subtree lookup in420 // one pass.421 422 // In general, bound symbol should have the unique name in the pattern but423 // for the operand, binding same symbol to multiple operands imply a424 // constraint at the same time. In this case, we will rename those operands425 // with different names. As a result, we need to collect all the symbolInfos426 // from the DagNode then get the updated name of the local variables from the427 // global symbolInfoMap.428 429 // Collect all the bound symbols in the Dag430 SymbolInfoMap localSymbolMap(loc);431 pattern.collectBoundSymbols(tree, localSymbolMap, /*isSrcPattern=*/true);432 433 for (const auto &info : localSymbolMap) {434 auto name = info.first;435 auto symboInfo = info.second;436 auto ret = symbolInfoMap.findBoundSymbol(name, symboInfo);437 os << formatv(", {0}", ret->second.getVarName(name));438 }439 440 os << "))) {\n";441 os.scope().os << "return ::mlir::failure();\n";442 os << "}\n";443}444 445void PatternEmitter::emitStaticVerifierCall(StringRef funcName,446 StringRef opName, StringRef arg,447 StringRef failureStr) {448 os << formatv("if(::mlir::failed({0}(rewriter, {1}, {2}, {3}))) {{\n",449 funcName, opName, arg, failureStr);450 os.scope().os << "return ::mlir::failure();\n";451 os << "}\n";452}453 454// Helper function to match patterns.455void PatternEmitter::emitNativeCodeMatch(DagNode tree, StringRef opName,456 int depth) {457 LLVM_DEBUG(llvm::dbgs() << "handle NativeCodeCall matcher pattern: ");458 LLVM_DEBUG(tree.print(llvm::dbgs()));459 LLVM_DEBUG(llvm::dbgs() << '\n');460 461 // The order of generating static matcher follows the topological order so462 // that for every dependent DagNode already have their static matcher463 // generated if needed. The reason we check if `getMatcherName(tree).empty()`464 // is when we are generating the static matcher for a DagNode itself. In this465 // case, we need to emit the function body rather than a function call.466 if (staticMatcherHelper.useStaticMatcher(tree) &&467 !staticMatcherHelper.getMatcherName(tree).empty()) {468 emitStaticMatchCall(tree, opName);469 470 // NativeCodeCall will never be at depth 0 so that we don't need to catch471 // the root operation as emitOpMatch();472 473 return;474 }475 476 // TODO(suderman): iterate through arguments, determine their types, output477 // names.478 SmallVector<std::string, 8> capture;479 480 raw_indented_ostream::DelimitedScope scope(os);481 482 for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {483 std::string argName = formatv("arg{0}_{1}", depth, i);484 if (DagNode argTree = tree.getArgAsNestedDag(i)) {485 if (argTree.isEither())486 PrintFatalError(loc, "NativeCodeCall cannot have `either` operands");487 if (argTree.isVariadic())488 PrintFatalError(loc, "NativeCodeCall cannot have `variadic` operands");489 490 os << "::mlir::Value " << argName << ";\n";491 } else {492 auto leaf = tree.getArgAsLeaf(i);493 if (leaf.isAttrMatcher() || leaf.isConstantAttr()) {494 os << "::mlir::Attribute " << argName << ";\n";495 } else if (leaf.isPropMatcher()) {496 StringRef interfaceType = leaf.getAsPropConstraint().getInterfaceType();497 if (interfaceType.empty())498 PrintFatalError(loc, "NativeCodeCall cannot have a property operand "499 "with unspecified interface type");500 os << interfaceType << " " << argName;501 if (leaf.isPropDefinition()) {502 Property propDef = leaf.getAsProperty();503 // Ensure properties that aren't zero-arg-constructable still work.504 if (propDef.hasDefaultValue())505 os << " = " << propDef.getDefaultValue();506 }507 os << ";\n";508 } else {509 os << "::mlir::Value " << argName << ";\n";510 }511 }512 513 capture.push_back(std::move(argName));514 }515 516 auto tail = getTrailingDirectives(tree);517 if (tail.returnType)518 PrintFatalError(loc, "`NativeCodeCall` cannot have return type specifier");519 auto locToUse = getLocation(tail);520 521 auto fmt = tree.getNativeCodeTemplate();522 if (fmt.count("$_self") != 1)523 PrintFatalError(loc, "NativeCodeCall must have $_self as argument for "524 "passing the defining Operation");525 526 auto nativeCodeCall = std::string(527 tgfmt(fmt, &fmtCtx.addSubst("_loc", locToUse).withSelf(opName.str()),528 static_cast<ArrayRef<std::string>>(capture)));529 530 emitMatchCheck(opName, formatv("!::mlir::failed({0})", nativeCodeCall),531 formatv("\"{0} return ::mlir::failure\"", nativeCodeCall));532 533 for (int i = 0, e = tree.getNumArgs() - tail.numDirectives; i != e; ++i) {534 auto name = tree.getArgName(i);535 if (!name.empty() && name != "_") {536 os << formatv("{0} = {1};\n", name, capture[i]);537 }538 }539 540 for (int i = 0, e = tree.getNumArgs() - tail.numDirectives; i != e; ++i) {541 std::string argName = capture[i];542 543 // Handle nested DAG construct first544 if (tree.getArgAsNestedDag(i)) {545 PrintFatalError(546 loc, formatv("Matching nested tree in NativeCodecall not support for "547 "{0} as arg {1}",548 argName, i));549 }550 551 DagLeaf leaf = tree.getArgAsLeaf(i);552 553 // The parameter for native function doesn't bind any constraints.554 if (leaf.isUnspecified())555 continue;556 557 auto constraint = leaf.getAsConstraint();558 559 std::string self;560 if (leaf.isAttrMatcher() || leaf.isConstantAttr() || leaf.isPropMatcher())561 self = argName;562 else563 self = formatv("{0}.getType()", argName);564 StringRef verifier = staticMatcherHelper.getVerifierName(leaf);565 emitStaticVerifierCall(566 verifier, opName, self,567 formatv("\"operand {0} of native code call '{1}' failed to satisfy "568 "constraint: "569 "'{2}'\"",570 i, tree.getNativeCodeTemplate(),571 escapeString(constraint.getSummary()))572 .str());573 }574 575 LLVM_DEBUG(llvm::dbgs() << "done emitting match for native code call\n");576}577 578// Helper function to match patterns.579void PatternEmitter::emitOpMatch(DagNode tree, StringRef opName, int depth) {580 Operator &op = tree.getDialectOp(opMap);581 LLVM_DEBUG(llvm::dbgs() << "start emitting match for op '"582 << op.getOperationName() << "' at depth " << depth583 << '\n');584 585 auto getCastedName = [depth]() -> std::string {586 return formatv("castedOp{0}", depth);587 };588 589 // The order of generating static matcher follows the topological order so590 // that for every dependent DagNode already have their static matcher591 // generated if needed. The reason we check if `getMatcherName(tree).empty()`592 // is when we are generating the static matcher for a DagNode itself. In this593 // case, we need to emit the function body rather than a function call.594 if (staticMatcherHelper.useStaticMatcher(tree) &&595 !staticMatcherHelper.getMatcherName(tree).empty()) {596 emitStaticMatchCall(tree, opName);597 // In the codegen of rewriter, we suppose that castedOp0 will capture the598 // root operation. Manually add it if the root DagNode is a static matcher.599 if (depth == 0)600 os << formatv("auto {2} = ::llvm::dyn_cast_or_null<{1}>({0}); "601 "(void){2};\n",602 opName, op.getQualCppClassName(), getCastedName());603 return;604 }605 606 std::string castedName = getCastedName();607 os << formatv("auto {0} = ::llvm::dyn_cast<{2}>({1}); "608 "(void){0};\n",609 castedName, opName, op.getQualCppClassName());610 611 // Skip the operand matching at depth 0 as the pattern rewriter already does.612 if (depth != 0)613 emitMatchCheck(opName, /*matchStr=*/castedName,614 formatv("\"{0} is not {1} type\"", castedName,615 op.getQualCppClassName()));616 617 // If the operand's name is set, set to that variable.618 auto name = tree.getSymbol();619 if (!name.empty())620 os << formatv("{0} = {1};\n", name, castedName);621 622 for (int i = 0, opArgIdx = 0, e = tree.getNumArgs(), nextOperand = 0; i != e;623 ++i, ++opArgIdx) {624 auto opArg = op.getArg(opArgIdx);625 std::string argName = formatv("op{0}", depth + 1);626 627 // Handle nested DAG construct first628 if (DagNode argTree = tree.getArgAsNestedDag(i)) {629 if (argTree.isEither()) {630 emitEitherOperandMatch(tree, argTree, castedName, opArgIdx, nextOperand,631 depth);632 ++opArgIdx;633 continue;634 }635 if (auto *operand =636 llvm::dyn_cast_if_present<NamedTypeConstraint *>(opArg)) {637 if (argTree.isVariadic()) {638 if (!operand->isVariadic()) {639 auto error = formatv("variadic DAG construct can't match op {0}'s "640 "non-variadic operand #{1}",641 op.getOperationName(), opArgIdx);642 PrintFatalError(loc, error);643 }644 emitVariadicOperandMatch(tree, argTree, castedName, opArgIdx,645 nextOperand, depth);646 ++nextOperand;647 continue;648 }649 if (operand->isVariableLength()) {650 auto error = formatv("use nested DAG construct to match op {0}'s "651 "variadic operand #{1} unsupported now",652 op.getOperationName(), opArgIdx);653 PrintFatalError(loc, error);654 }655 }656 657 os << "{\n";658 659 // Attributes don't count for getODSOperands.660 // TODO: Operand is a Value, check if we should remove `getDefiningOp()`.661 os.indent() << formatv(662 "auto *{0} = "663 "(*{1}.getODSOperands({2}).begin()).getDefiningOp();\n",664 argName, castedName, nextOperand);665 // Null check of operand's definingOp666 emitMatchCheck(667 castedName, /*matchStr=*/argName,668 formatv("\"There's no operation that defines operand {0} of {1}\"",669 nextOperand++, castedName));670 emitMatch(argTree, argName, depth + 1);671 os << formatv("tblgen_ops.push_back({0});\n", argName);672 os.unindent() << "}\n";673 continue;674 }675 676 // Next handle DAG leaf: operand or attribute677 if (isa<NamedTypeConstraint *>(opArg)) {678 auto operandName =679 formatv("{0}.getODSOperands({1})", castedName, nextOperand);680 emitOperandMatch(tree, castedName, operandName.str(), nextOperand,681 /*operandMatcher=*/tree.getArgAsLeaf(i),682 /*argName=*/tree.getArgName(i), opArgIdx,683 /*variadicSubIndex=*/std::nullopt);684 ++nextOperand;685 } else if (isa<NamedAttribute *>(opArg)) {686 emitAttributeMatch(tree, castedName, opArgIdx, depth);687 } else if (isa<NamedProperty *>(opArg)) {688 emitPropertyMatch(tree, castedName, opArgIdx, depth);689 } else {690 PrintFatalError(loc, "unhandled case when matching op");691 }692 }693 LLVM_DEBUG(llvm::dbgs() << "done emitting match for op '"694 << op.getOperationName() << "' at depth " << depth695 << '\n');696}697 698void PatternEmitter::emitOperandMatch(DagNode tree, StringRef opName,699 StringRef operandName, int operandIndex,700 DagLeaf operandMatcher, StringRef argName,701 int argIndex,702 std::optional<int> variadicSubIndex) {703 Operator &op = tree.getDialectOp(opMap);704 NamedTypeConstraint operand = op.getOperand(operandIndex);705 706 // If a constraint is specified, we need to generate C++ statements to707 // check the constraint.708 if (!operandMatcher.isUnspecified()) {709 if (!operandMatcher.isOperandMatcher())710 PrintFatalError(711 loc, formatv("the {1}-th argument of op '{0}' should be an operand",712 op.getOperationName(), argIndex + 1));713 714 // Only need to verify if the matcher's type is different from the one715 // of op definition.716 Constraint constraint = operandMatcher.getAsConstraint();717 if (operand.constraint != constraint) {718 if (operand.isVariableLength()) {719 auto error = formatv(720 "further constrain op {0}'s variadic operand #{1} unsupported now",721 op.getOperationName(), argIndex);722 PrintFatalError(loc, error);723 }724 auto self = formatv("(*{0}.begin()).getType()", operandName);725 StringRef verifier = staticMatcherHelper.getVerifierName(operandMatcher);726 emitStaticVerifierCall(727 verifier, opName, self.str(),728 formatv(729 "\"operand {0} of op '{1}' failed to satisfy constraint: '{2}'\"",730 operandIndex, op.getOperationName(),731 escapeString(constraint.getSummary()))732 .str());733 }734 }735 736 // Capture the value737 // `$_` is a special symbol to ignore op argument matching.738 if (!argName.empty() && argName != "_") {739 auto res = symbolInfoMap.findBoundSymbol(argName, tree, op, argIndex,740 variadicSubIndex);741 if (res == symbolInfoMap.end())742 PrintFatalError(loc, formatv("symbol not found: {0}", argName));743 744 os << formatv("{0} = {1};\n", res->second.getVarName(argName), operandName);745 }746}747 748void PatternEmitter::emitEitherOperandMatch(DagNode tree, DagNode eitherArgTree,749 StringRef opName, int argIndex,750 int &operandIndex, int depth) {751 constexpr int numEitherArgs = 2;752 if (eitherArgTree.getNumArgs() != numEitherArgs)753 PrintFatalError(loc, "`either` only supports grouping two operands");754 755 Operator &op = tree.getDialectOp(opMap);756 757 std::string codeBuffer;758 llvm::raw_string_ostream tblgenOps(codeBuffer);759 760 std::string lambda = formatv("eitherLambda{0}", depth);761 os << formatv(762 "auto {0} = [&](::mlir::OperandRange v0, ::mlir::OperandRange v1) {{\n",763 lambda);764 765 os.indent();766 767 for (int i = 0; i < numEitherArgs; ++i, ++argIndex) {768 if (DagNode argTree = eitherArgTree.getArgAsNestedDag(i)) {769 if (argTree.isEither())770 PrintFatalError(loc, "either cannot be nested");771 772 std::string argName = formatv("local_op_{0}", i).str();773 774 os << formatv("auto {0} = (*v{1}.begin()).getDefiningOp();\n", argName,775 i);776 777 // Indent emitMatchCheck and emitMatch because they declare local778 // variables.779 os << "{\n";780 os.indent();781 782 emitMatchCheck(783 opName, /*matchStr=*/argName,784 formatv("\"There's no operation that defines operand {0} of {1}\"",785 operandIndex++, opName));786 emitMatch(argTree, argName, depth + 1);787 788 os.unindent() << "}\n";789 790 // `tblgen_ops` is used to collect the matched operations. In either, we791 // need to queue the operation only if the matching success. Thus we emit792 // the code at the end.793 tblgenOps << formatv("tblgen_ops.push_back({0});\n", argName);794 } else if (isa<NamedTypeConstraint *>(op.getArg(argIndex))) {795 emitOperandMatch(tree, opName, /*operandName=*/formatv("v{0}", i).str(),796 operandIndex,797 /*operandMatcher=*/eitherArgTree.getArgAsLeaf(i),798 /*argName=*/eitherArgTree.getArgName(i), argIndex,799 /*variadicSubIndex=*/std::nullopt);800 ++operandIndex;801 } else {802 PrintFatalError(loc, "either can only be applied on operand");803 }804 }805 806 os << tblgenOps.str();807 os << "return ::mlir::success();\n";808 os.unindent() << "};\n";809 810 os << "{\n";811 os.indent();812 813 os << formatv("auto eitherOperand0 = {0}.getODSOperands({1});\n", opName,814 operandIndex - 2);815 os << formatv("auto eitherOperand1 = {0}.getODSOperands({1});\n", opName,816 operandIndex - 1);817 818 os << formatv("if(::mlir::failed({0}(eitherOperand0, eitherOperand1)) && "819 "::mlir::failed({0}(eitherOperand1, "820 "eitherOperand0)))\n",821 lambda);822 os.indent() << "return ::mlir::failure();\n";823 824 os.unindent().unindent() << "}\n";825}826 827void PatternEmitter::emitVariadicOperandMatch(DagNode tree,828 DagNode variadicArgTree,829 StringRef opName, int argIndex,830 int &operandIndex, int depth) {831 Operator &op = tree.getDialectOp(opMap);832 833 os << "{\n";834 os.indent();835 836 os << formatv("auto variadic_operand_range = {0}.getODSOperands({1});\n",837 opName, operandIndex);838 os << formatv("if (variadic_operand_range.size() != {0}) "839 "return ::mlir::failure();\n",840 variadicArgTree.getNumArgs());841 842 StringRef variadicTreeName = variadicArgTree.getSymbol();843 if (!variadicTreeName.empty()) {844 auto res =845 symbolInfoMap.findBoundSymbol(variadicTreeName, tree, op, argIndex,846 /*variadicSubIndex=*/std::nullopt);847 if (res == symbolInfoMap.end())848 PrintFatalError(loc, formatv("symbol not found: {0}", variadicTreeName));849 850 os << formatv("{0} = variadic_operand_range;\n",851 res->second.getVarName(variadicTreeName));852 }853 854 for (int i = 0; i < variadicArgTree.getNumArgs(); ++i) {855 if (DagNode argTree = variadicArgTree.getArgAsNestedDag(i)) {856 if (!argTree.isOperation())857 PrintFatalError(loc, "variadic only accepts operation sub-dags");858 859 os << "{\n";860 os.indent();861 862 std::string argName = formatv("local_op_{0}", i).str();863 os << formatv("auto *{0} = "864 "variadic_operand_range[{1}].getDefiningOp();\n",865 argName, i);866 emitMatchCheck(867 opName, /*matchStr=*/argName,868 formatv("\"There's no operation that defines variadic operand "869 "{0} (variadic sub-opearnd #{1}) of {2}\"",870 operandIndex, i, opName));871 emitMatch(argTree, argName, depth + 1);872 os << formatv("tblgen_ops.push_back({0});\n", argName);873 874 os.unindent() << "}\n";875 } else if (isa<NamedTypeConstraint *>(op.getArg(argIndex))) {876 auto operandName = formatv("variadic_operand_range.slice({0}, 1)", i);877 emitOperandMatch(tree, opName, operandName.str(), operandIndex,878 /*operandMatcher=*/variadicArgTree.getArgAsLeaf(i),879 /*argName=*/variadicArgTree.getArgName(i), argIndex, i);880 } else {881 PrintFatalError(loc, "variadic can only be applied on operand");882 }883 }884 885 os.unindent() << "}\n";886}887 888void PatternEmitter::emitAttributeMatch(DagNode tree, StringRef castedName,889 int argIndex, int depth) {890 Operator &op = tree.getDialectOp(opMap);891 auto *namedAttr = cast<NamedAttribute *>(op.getArg(argIndex));892 const auto &attr = namedAttr->attr;893 894 os << "{\n";895 if (op.getDialect().usePropertiesForAttributes()) {896 os.indent() << formatv(897 "[[maybe_unused]] auto tblgen_attr = {0}.getProperties().{1}();\n",898 castedName, op.getGetterName(namedAttr->name));899 } else {900 os.indent() << formatv("[[maybe_unused]] auto tblgen_attr = "901 "{0}->getAttrOfType<{1}>(\"{2}\");\n",902 castedName, attr.getStorageType(), namedAttr->name);903 }904 905 // TODO: This should use getter method to avoid duplication.906 if (attr.hasDefaultValue()) {907 os << "if (!tblgen_attr) tblgen_attr = "908 << std::string(tgfmt(attr.getConstBuilderTemplate(), &fmtCtx,909 tgfmt(attr.getDefaultValue(), &fmtCtx)))910 << ";\n";911 } else if (attr.isOptional()) {912 // For a missing attribute that is optional according to definition, we913 // should just capture a mlir::Attribute() to signal the missing state.914 // That is precisely what getDiscardableAttr() returns on missing915 // attributes.916 } else {917 emitMatchCheck(castedName, tgfmt("tblgen_attr", &fmtCtx),918 formatv("\"expected op '{0}' to have attribute '{1}' "919 "of type '{2}'\"",920 op.getOperationName(), namedAttr->name,921 attr.getStorageType()));922 }923 924 auto matcher = tree.getArgAsLeaf(argIndex);925 if (!matcher.isUnspecified()) {926 if (!matcher.isAttrMatcher()) {927 PrintFatalError(928 loc, formatv("the {1}-th argument of op '{0}' should be an attribute",929 op.getOperationName(), argIndex + 1));930 }931 932 // If a constraint is specified, we need to generate function call to its933 // static verifier.934 StringRef verifier = staticMatcherHelper.getVerifierName(matcher);935 if (attr.isOptional()) {936 // Avoid dereferencing null attribute. This is using a simple heuristic to937 // avoid common cases of attempting to dereference null attribute. This938 // will return where there is no check if attribute is null unless the939 // attribute's value is not used.940 // FIXME: This could be improved as some null dereferences could slip941 // through.942 if (!StringRef(matcher.getConditionTemplate()).contains("!$_self") &&943 StringRef(matcher.getConditionTemplate()).contains("$_self")) {944 os << "if (!tblgen_attr) return ::mlir::failure();\n";945 }946 }947 emitStaticVerifierCall(948 verifier, castedName, "tblgen_attr",949 formatv("\"op '{0}' attribute '{1}' failed to satisfy constraint: "950 "'{2}'\"",951 op.getOperationName(), namedAttr->name,952 escapeString(matcher.getAsConstraint().getSummary()))953 .str());954 }955 956 // Capture the value957 auto name = tree.getArgName(argIndex);958 // `$_` is a special symbol to ignore op argument matching.959 if (!name.empty() && name != "_") {960 os << formatv("{0} = tblgen_attr;\n", name);961 }962 963 os.unindent() << "}\n";964}965 966void PatternEmitter::emitPropertyMatch(DagNode tree, StringRef castedName,967 int argIndex, int depth) {968 Operator &op = tree.getDialectOp(opMap);969 auto *namedProp = cast<NamedProperty *>(op.getArg(argIndex));970 971 os << "{\n";972 os.indent() << formatv(973 "[[maybe_unused]] auto tblgen_prop = {0}.getProperties().{1}();\n",974 castedName, op.getGetterName(namedProp->name));975 976 auto matcher = tree.getArgAsLeaf(argIndex);977 if (!matcher.isUnspecified()) {978 if (!matcher.isPropMatcher()) {979 PrintFatalError(980 loc, formatv("the {1}-th argument of op '{0}' should be a property",981 op.getOperationName(), argIndex + 1));982 }983 984 // If a constraint is specified, we need to generate function call to its985 // static verifier.986 StringRef verifier = staticMatcherHelper.getVerifierName(matcher);987 emitStaticVerifierCall(988 verifier, castedName, "tblgen_prop",989 formatv("\"op '{0}' property '{1}' failed to satisfy constraint: "990 "'{2}'\"",991 op.getOperationName(), namedProp->name,992 escapeString(matcher.getAsConstraint().getSummary()))993 .str());994 }995 996 // Capture the value997 auto name = tree.getArgName(argIndex);998 // `$_` is a special symbol to ignore op argument matching.999 if (!name.empty() && name != "_") {1000 os << formatv("{0} = tblgen_prop;\n", name);1001 }1002 1003 os.unindent() << "}\n";1004}1005 1006void PatternEmitter::emitMatchCheck(1007 StringRef opName, const FmtObjectBase &matchFmt,1008 const llvm::formatv_object_base &failureFmt) {1009 emitMatchCheck(opName, matchFmt.str(), failureFmt.str());1010}1011 1012void PatternEmitter::emitMatchCheck(StringRef opName,1013 const std::string &matchStr,1014 const std::string &failureStr) {1015 1016 os << "if (!(" << matchStr << "))";1017 os.scope("{\n", "\n}\n").os << "return rewriter.notifyMatchFailure(" << opName1018 << ", [&](::mlir::Diagnostic &diag) {\n diag << "1019 << failureStr << ";\n});";1020}1021 1022void PatternEmitter::emitMatchLogic(DagNode tree, StringRef opName) {1023 LLVM_DEBUG(llvm::dbgs() << "--- start emitting match logic ---\n");1024 int depth = 0;1025 emitMatch(tree, opName, depth);1026 1027 // Some of the operands could be bound to the same symbol name, we need1028 // to enforce equality constraint on those.1029 // This has to happen before user provided constraints, which may assume the1030 // same name checks are already performed, since in the pattern source code1031 // the user provided constraints appear later.1032 // TODO: we should be able to emit equality checks early1033 // and short circuit unnecessary work if vars are not equal.1034 for (auto symbolInfoIt = symbolInfoMap.begin();1035 symbolInfoIt != symbolInfoMap.end();) {1036 auto range = symbolInfoMap.getRangeOfEqualElements(symbolInfoIt->first);1037 auto startRange = range.first;1038 auto endRange = range.second;1039 1040 auto firstOperand = symbolInfoIt->second.getVarName(symbolInfoIt->first);1041 for (++startRange; startRange != endRange; ++startRange) {1042 auto secondOperand = startRange->second.getVarName(symbolInfoIt->first);1043 emitMatchCheck(1044 opName,1045 formatv("*{0}.begin() == *{1}.begin()", firstOperand, secondOperand),1046 formatv("\"Operands '{0}' and '{1}' must be equal\"", firstOperand,1047 secondOperand));1048 }1049 1050 symbolInfoIt = endRange;1051 }1052 1053 for (auto &appliedConstraint : pattern.getConstraints()) {1054 auto &constraint = appliedConstraint.constraint;1055 auto &entities = appliedConstraint.entities;1056 1057 auto condition = constraint.getConditionTemplate();1058 if (isa<TypeConstraint>(constraint)) {1059 if (entities.size() != 1)1060 PrintFatalError(loc, "type constraint requires exactly one argument");1061 1062 auto self = formatv("({0}.getType())",1063 symbolInfoMap.getValueAndRangeUse(entities.front()));1064 emitMatchCheck(1065 opName, tgfmt(condition, &fmtCtx.withSelf(self.str())),1066 formatv("\"value entity '{0}' failed to satisfy constraint: '{1}'\"",1067 entities.front(), escapeString(constraint.getSummary())));1068 1069 } else if (isa<AttrConstraint>(constraint)) {1070 PrintFatalError(1071 loc, "cannot use AttrConstraint in Pattern multi-entity constraints");1072 } else {1073 // TODO: replace formatv arguments with the exact specified1074 // args.1075 if (entities.size() > 4) {1076 PrintFatalError(loc, "only support up to 4-entity constraints now");1077 }1078 SmallVector<std::string, 4> names;1079 int i = 0;1080 for (int e = entities.size(); i < e; ++i)1081 names.push_back(symbolInfoMap.getValueAndRangeUse(entities[i]));1082 std::string self = appliedConstraint.self;1083 if (!self.empty())1084 self = symbolInfoMap.getValueAndRangeUse(self);1085 for (; i < 4; ++i)1086 names.push_back("<unused>");1087 emitMatchCheck(opName,1088 tgfmt(condition, &fmtCtx.withSelf(self), names[0],1089 names[1], names[2], names[3]),1090 formatv("\"entities '{0}' failed to satisfy constraint: "1091 "'{1}'\"",1092 llvm::join(entities, ", "),1093 escapeString(constraint.getSummary())));1094 }1095 }1096 1097 LLVM_DEBUG(llvm::dbgs() << "--- done emitting match logic ---\n");1098}1099 1100void PatternEmitter::collectOps(DagNode tree,1101 llvm::SmallPtrSetImpl<const Operator *> &ops) {1102 // Check if this tree is an operation.1103 if (tree.isOperation()) {1104 const Operator &op = tree.getDialectOp(opMap);1105 LLVM_DEBUG(llvm::dbgs()1106 << "found operation " << op.getOperationName() << '\n');1107 ops.insert(&op);1108 }1109 1110 // Recurse the arguments of the tree.1111 for (unsigned i = 0, e = tree.getNumArgs(); i != e; ++i)1112 if (auto child = tree.getArgAsNestedDag(i))1113 collectOps(child, ops);1114}1115 1116void PatternEmitter::emit(StringRef rewriteName) {1117 // Get the DAG tree for the source pattern.1118 DagNode sourceTree = pattern.getSourcePattern();1119 1120 const Operator &rootOp = pattern.getSourceRootOp();1121 auto rootName = rootOp.getOperationName();1122 1123 // Collect the set of result operations.1124 llvm::SmallPtrSet<const Operator *, 4> resultOps;1125 LLVM_DEBUG(llvm::dbgs() << "start collecting ops used in result patterns\n");1126 for (unsigned i = 0, e = pattern.getNumResultPatterns(); i != e; ++i) {1127 collectOps(pattern.getResultPattern(i), resultOps);1128 }1129 LLVM_DEBUG(llvm::dbgs() << "done collecting ops used in result patterns\n");1130 1131 // Emit RewritePattern for Pattern.1132 auto locs = pattern.getLocation(/*forSourceOutput=*/true);1133 os << formatv("/* Generated from:\n {0:$[ instantiating\n ]}\n*/\n",1134 llvm::reverse(locs));1135 os << formatv(R"(struct {0} : public ::mlir::RewritePattern {1136 {0}(::mlir::MLIRContext *context)1137 : ::mlir::RewritePattern("{1}", {2}, context, {{)",1138 rewriteName, rootName, pattern.getBenefit());1139 // Sort result operators by name.1140 llvm::SmallVector<const Operator *, 4> sortedResultOps(resultOps.begin(),1141 resultOps.end());1142 llvm::sort(sortedResultOps, [&](const Operator *lhs, const Operator *rhs) {1143 return lhs->getOperationName() < rhs->getOperationName();1144 });1145 llvm::interleaveComma(sortedResultOps, os, [&](const Operator *op) {1146 os << '"' << op->getOperationName() << '"';1147 });1148 os << "}) {}\n";1149 1150 // Emit matchAndRewrite() function.1151 {1152 auto classScope = os.scope();1153 os.printReindented(R"(1154 ::llvm::LogicalResult matchAndRewrite(::mlir::Operation *op0,1155 ::mlir::PatternRewriter &rewriter) const override {)")1156 << '\n';1157 {1158 auto functionScope = os.scope();1159 1160 // Register all symbols bound in the source pattern.1161 pattern.collectSourcePatternBoundSymbols(symbolInfoMap);1162 1163 LLVM_DEBUG(llvm::dbgs()1164 << "start creating local variables for capturing matches\n");1165 os << "// Variables for capturing values and attributes used while "1166 "creating ops\n";1167 // Create local variables for storing the arguments and results bound1168 // to symbols.1169 for (const auto &symbolInfoPair : symbolInfoMap) {1170 const auto &symbol = symbolInfoPair.first;1171 const auto &info = symbolInfoPair.second;1172 1173 os << info.getVarDecl(symbol);1174 }1175 // TODO: capture ops with consistent numbering so that it can be1176 // reused for fused loc.1177 os << "::llvm::SmallVector<::mlir::Operation *, 4> tblgen_ops;\n\n";1178 LLVM_DEBUG(llvm::dbgs()1179 << "done creating local variables for capturing matches\n");1180 1181 os << "// Match\n";1182 os << "tblgen_ops.push_back(op0);\n";1183 emitMatchLogic(sourceTree, "op0");1184 1185 os << "\n// Rewrite\n";1186 emitRewriteLogic();1187 1188 os << "return ::mlir::success();\n";1189 }1190 os << "}\n";1191 }1192 os << "};\n\n";1193}1194 1195void PatternEmitter::emitRewriteLogic() {1196 LLVM_DEBUG(llvm::dbgs() << "--- start emitting rewrite logic ---\n");1197 const Operator &rootOp = pattern.getSourceRootOp();1198 int numExpectedResults = rootOp.getNumResults();1199 int numResultPatterns = pattern.getNumResultPatterns();1200 1201 // First register all symbols bound to ops generated in result patterns.1202 pattern.collectResultPatternBoundSymbols(symbolInfoMap);1203 1204 // Only the last N static values generated are used to replace the matched1205 // root N-result op. We need to calculate the starting index (of the results1206 // of the matched op) each result pattern is to replace.1207 SmallVector<int, 4> offsets(numResultPatterns + 1, numExpectedResults);1208 // If we don't need to replace any value at all, set the replacement starting1209 // index as the number of result patterns so we skip all of them when trying1210 // to replace the matched op's results.1211 int replStartIndex = numExpectedResults == 0 ? numResultPatterns : -1;1212 for (int i = numResultPatterns - 1; i >= 0; --i) {1213 auto numValues = getNodeValueCount(pattern.getResultPattern(i));1214 offsets[i] = offsets[i + 1] - numValues;1215 if (offsets[i] == 0) {1216 if (replStartIndex == -1)1217 replStartIndex = i;1218 } else if (offsets[i] < 0 && offsets[i + 1] > 0) {1219 auto error = formatv(1220 "cannot use the same multi-result op '{0}' to generate both "1221 "auxiliary values and values to be used for replacing the matched op",1222 pattern.getResultPattern(i).getSymbol());1223 PrintFatalError(loc, error);1224 }1225 }1226 1227 if (offsets.front() > 0) {1228 const char error[] =1229 "not enough values generated to replace the matched op";1230 PrintFatalError(loc, error);1231 }1232 1233 os << "auto odsLoc = rewriter.getFusedLoc({";1234 for (int i = 0, e = pattern.getSourcePattern().getNumOps(); i != e; ++i) {1235 os << (i ? ", " : "") << "tblgen_ops[" << i << "]->getLoc()";1236 }1237 os << "}); (void)odsLoc;\n";1238 1239 // Process auxiliary result patterns.1240 for (int i = 0; i < replStartIndex; ++i) {1241 DagNode resultTree = pattern.getResultPattern(i);1242 auto val = handleResultPattern(resultTree, offsets[i], 0);1243 // Normal op creation will be streamed to `os` by the above call; but1244 // NativeCodeCall will only be materialized to `os` if it is used. Here1245 // we are handling auxiliary patterns so we want the side effect even if1246 // NativeCodeCall is not replacing matched root op's results.1247 if (resultTree.isNativeCodeCall() &&1248 resultTree.getNumReturnsOfNativeCode() == 0)1249 os << val << ";\n";1250 }1251 1252 auto processSupplementalPatterns = [&]() {1253 int numSupplementalPatterns = pattern.getNumSupplementalPatterns();1254 for (int i = 0, offset = -numSupplementalPatterns;1255 i < numSupplementalPatterns; ++i) {1256 DagNode resultTree = pattern.getSupplementalPattern(i);1257 auto val = handleResultPattern(resultTree, offset++, 0);1258 if (resultTree.isNativeCodeCall() &&1259 resultTree.getNumReturnsOfNativeCode() == 0)1260 os << val << ";\n";1261 }1262 };1263 1264 if (numExpectedResults == 0) {1265 assert(replStartIndex >= numResultPatterns &&1266 "invalid auxiliary vs. replacement pattern division!");1267 processSupplementalPatterns();1268 // No result to replace. Just erase the op.1269 os << "rewriter.eraseOp(op0);\n";1270 } else {1271 // Process replacement result patterns.1272 os << "::llvm::SmallVector<::mlir::Value, 4> tblgen_repl_values;\n";1273 for (int i = replStartIndex; i < numResultPatterns; ++i) {1274 DagNode resultTree = pattern.getResultPattern(i);1275 auto val = handleResultPattern(resultTree, offsets[i], 0);1276 os << "\n";1277 // Resolve each symbol for all range use so that we can loop over them.1278 // We need an explicit cast to `SmallVector` to capture the cases where1279 // `{0}` resolves to an `Operation::result_range` as well as cases that1280 // are not iterable (e.g. vector that gets wrapped in additional braces by1281 // RewriterGen).1282 // TODO: Revisit the need for materializing a vector.1283 os << symbolInfoMap.getAllRangeUse(1284 val,1285 "for (auto v: ::llvm::SmallVector<::mlir::Value, 4>{ {0} }) {{\n"1286 " tblgen_repl_values.push_back(v);\n}\n",1287 "\n");1288 }1289 processSupplementalPatterns();1290 os << "\nrewriter.replaceOp(op0, tblgen_repl_values);\n";1291 }1292 1293 LLVM_DEBUG(llvm::dbgs() << "--- done emitting rewrite logic ---\n");1294}1295 1296std::string PatternEmitter::getUniqueSymbol(const Operator *op) {1297 return std::string(1298 formatv("tblgen_{0}_{1}", op->getCppClassName(), nextValueId++));1299}1300 1301std::string PatternEmitter::handleResultPattern(DagNode resultTree,1302 int resultIndex, int depth) {1303 LLVM_DEBUG(llvm::dbgs() << "handle result pattern: ");1304 LLVM_DEBUG(resultTree.print(llvm::dbgs()));1305 LLVM_DEBUG(llvm::dbgs() << '\n');1306 1307 if (resultTree.isLocationDirective()) {1308 PrintFatalError(loc,1309 "location directive can only be used with op creation");1310 }1311 1312 if (resultTree.isNativeCodeCall())1313 return handleReplaceWithNativeCodeCall(resultTree, depth);1314 1315 if (resultTree.isReplaceWithValue())1316 return handleReplaceWithValue(resultTree).str();1317 1318 if (resultTree.isVariadic())1319 return handleVariadic(resultTree, depth);1320 1321 // Normal op creation.1322 auto symbol = handleOpCreation(resultTree, resultIndex, depth);1323 if (resultTree.getSymbol().empty()) {1324 // This is an op not explicitly bound to a symbol in the rewrite rule.1325 // Register the auto-generated symbol for it.1326 symbolInfoMap.bindOpResult(symbol, pattern.getDialectOp(resultTree));1327 }1328 return symbol;1329}1330 1331std::string PatternEmitter::handleVariadic(DagNode tree, int depth) {1332 assert(tree.isVariadic());1333 1334 std::string output;1335 llvm::raw_string_ostream oss(output);1336 auto name = std::string(formatv("tblgen_variadic_values_{0}", nextValueId++));1337 symbolInfoMap.bindValue(name);1338 oss << "::llvm::SmallVector<::mlir::Value, 4> " << name << ";\n";1339 for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {1340 if (auto child = tree.getArgAsNestedDag(i)) {1341 oss << name << ".push_back(" << handleResultPattern(child, i, depth + 1)1342 << ");\n";1343 } else {1344 oss << name << ".push_back("1345 << handleOpArgument(tree.getArgAsLeaf(i), tree.getArgName(i))1346 << ");\n";1347 }1348 }1349 1350 os << oss.str();1351 return name;1352}1353 1354StringRef PatternEmitter::handleReplaceWithValue(DagNode tree) {1355 assert(tree.isReplaceWithValue());1356 1357 if (tree.getNumArgs() != 1) {1358 PrintFatalError(1359 loc, "replaceWithValue directive must take exactly one argument");1360 }1361 1362 if (!tree.getSymbol().empty()) {1363 PrintFatalError(loc, "cannot bind symbol to replaceWithValue");1364 }1365 1366 return tree.getArgName(0);1367}1368 1369std::string PatternEmitter::handleLocationDirective(DagNode tree) {1370 assert(tree.isLocationDirective());1371 auto lookUpArgLoc = [this, &tree](int idx) {1372 const auto *const lookupFmt = "{0}.getLoc()";1373 return symbolInfoMap.getValueAndRangeUse(tree.getArgName(idx), lookupFmt);1374 };1375 1376 if (tree.getNumArgs() == 0)1377 llvm::PrintFatalError(1378 "At least one argument to location directive required");1379 1380 if (!tree.getSymbol().empty())1381 PrintFatalError(loc, "cannot bind symbol to location");1382 1383 if (tree.getNumArgs() == 1) {1384 DagLeaf leaf = tree.getArgAsLeaf(0);1385 if (leaf.isStringAttr())1386 return formatv("::mlir::NameLoc::get(rewriter.getStringAttr(\"{0}\"))",1387 leaf.getStringAttr())1388 .str();1389 return lookUpArgLoc(0);1390 }1391 1392 std::string ret;1393 llvm::raw_string_ostream os(ret);1394 std::string strAttr;1395 os << "rewriter.getFusedLoc({";1396 bool first = true;1397 for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {1398 DagLeaf leaf = tree.getArgAsLeaf(i);1399 // Handle the optional string value.1400 if (leaf.isStringAttr()) {1401 if (!strAttr.empty())1402 llvm::PrintFatalError("Only one string attribute may be specified");1403 strAttr = leaf.getStringAttr();1404 continue;1405 }1406 os << (first ? "" : ", ") << lookUpArgLoc(i);1407 first = false;1408 }1409 os << "}";1410 if (!strAttr.empty()) {1411 os << ", rewriter.getStringAttr(\"" << strAttr << "\")";1412 }1413 os << ")";1414 return os.str();1415}1416 1417std::string PatternEmitter::handleReturnTypeArg(DagNode returnType, int i,1418 int depth) {1419 // Nested NativeCodeCall.1420 if (auto dagNode = returnType.getArgAsNestedDag(i)) {1421 if (!dagNode.isNativeCodeCall())1422 PrintFatalError(loc, "nested DAG in `returnType` must be a native code "1423 "call");1424 return handleReplaceWithNativeCodeCall(dagNode, depth);1425 }1426 // String literal.1427 auto dagLeaf = returnType.getArgAsLeaf(i);1428 if (dagLeaf.isStringAttr())1429 return tgfmt(dagLeaf.getStringAttr(), &fmtCtx);1430 return tgfmt(1431 "$0.getType()", &fmtCtx,1432 handleOpArgument(returnType.getArgAsLeaf(i), returnType.getArgName(i)));1433}1434 1435std::string PatternEmitter::handleOpArgument(DagLeaf leaf,1436 StringRef patArgName) {1437 if (leaf.isStringAttr())1438 PrintFatalError(loc, "raw string not supported as argument");1439 if (leaf.isConstantAttr()) {1440 auto constAttr = leaf.getAsConstantAttr();1441 return handleConstantAttr(constAttr.getAttribute(),1442 constAttr.getConstantValue());1443 }1444 if (leaf.isEnumCase()) {1445 auto enumCase = leaf.getAsEnumCase();1446 // This is an enum case backed by an IntegerAttr. We need to get its value1447 // to build the constant.1448 std::string val = std::to_string(enumCase.getValue());1449 return handleConstantAttr(Attribute(&enumCase.getDef()), val);1450 }1451 if (leaf.isConstantProp()) {1452 auto constantProp = leaf.getAsConstantProp();1453 return constantProp.getValue().str();1454 }1455 1456 LLVM_DEBUG(llvm::dbgs() << "handle argument '" << patArgName << "'\n");1457 auto argName = symbolInfoMap.getValueAndRangeUse(patArgName);1458 if (leaf.isUnspecified() || leaf.isOperandMatcher()) {1459 LLVM_DEBUG(llvm::dbgs() << "replace " << patArgName << " with '" << argName1460 << "' (via symbol ref)\n");1461 return argName;1462 }1463 if (leaf.isNativeCodeCall()) {1464 auto repl = tgfmt(leaf.getNativeCodeTemplate(), &fmtCtx.withSelf(argName));1465 LLVM_DEBUG(llvm::dbgs() << "replace " << patArgName << " with '" << repl1466 << "' (via NativeCodeCall)\n");1467 return std::string(repl);1468 }1469 PrintFatalError(loc, "unhandled case when rewriting op");1470}1471 1472std::string PatternEmitter::handleReplaceWithNativeCodeCall(DagNode tree,1473 int depth) {1474 LLVM_DEBUG(llvm::dbgs() << "handle NativeCodeCall pattern: ");1475 LLVM_DEBUG(tree.print(llvm::dbgs()));1476 LLVM_DEBUG(llvm::dbgs() << '\n');1477 1478 auto fmt = tree.getNativeCodeTemplate();1479 1480 SmallVector<std::string, 16> attrs;1481 1482 auto tail = getTrailingDirectives(tree);1483 if (tail.returnType)1484 PrintFatalError(loc, "`NativeCodeCall` cannot have return type specifier");1485 auto locToUse = getLocation(tail);1486 1487 for (int i = 0, e = tree.getNumArgs() - tail.numDirectives; i != e; ++i) {1488 if (tree.isNestedDagArg(i)) {1489 attrs.push_back(1490 handleResultPattern(tree.getArgAsNestedDag(i), i, depth + 1));1491 } else {1492 attrs.push_back(1493 handleOpArgument(tree.getArgAsLeaf(i), tree.getArgName(i)));1494 }1495 LLVM_DEBUG(llvm::dbgs() << "NativeCodeCall argument #" << i1496 << " replacement: " << attrs[i] << "\n");1497 }1498 1499 std::string symbol = tgfmt(fmt, &fmtCtx.addSubst("_loc", locToUse),1500 static_cast<ArrayRef<std::string>>(attrs));1501 1502 // In general, NativeCodeCall without naming binding don't need this. To1503 // ensure void helper function has been correctly labeled, i.e., use1504 // NativeCodeCallVoid, we cache the result to a local variable so that we will1505 // get a compilation error in the auto-generated file.1506 // Example.1507 // // In the td file1508 // Pat<(...), (NativeCodeCall<Foo> ...)>1509 //1510 // ---1511 //1512 // // In the auto-generated .cpp1513 // ...1514 // // Causes compilation error if Foo() returns void.1515 // auto nativeVar = Foo();1516 // ...1517 if (tree.getNumReturnsOfNativeCode() != 0) {1518 // Determine the local variable name for return value.1519 std::string varName =1520 SymbolInfoMap::getValuePackName(tree.getSymbol()).str();1521 if (varName.empty()) {1522 varName = formatv("nativeVar_{0}", nextValueId++);1523 // Register the local variable for later uses.1524 symbolInfoMap.bindValues(varName, tree.getNumReturnsOfNativeCode());1525 }1526 1527 // Catch the return value of helper function.1528 os << formatv("auto {0} = {1}; (void){0};\n", varName, symbol);1529 1530 if (!tree.getSymbol().empty())1531 symbol = tree.getSymbol().str();1532 else1533 symbol = varName;1534 }1535 1536 return symbol;1537}1538 1539int PatternEmitter::getNodeValueCount(DagNode node) {1540 if (node.isOperation()) {1541 // If the op is bound to a symbol in the rewrite rule, query its result1542 // count from the symbol info map.1543 auto symbol = node.getSymbol();1544 if (!symbol.empty()) {1545 return symbolInfoMap.getStaticValueCount(symbol);1546 }1547 // Otherwise this is an unbound op; we will use all its results.1548 return pattern.getDialectOp(node).getNumResults();1549 }1550 1551 if (node.isNativeCodeCall())1552 return node.getNumReturnsOfNativeCode();1553 1554 return 1;1555}1556 1557PatternEmitter::TrailingDirectives1558PatternEmitter::getTrailingDirectives(DagNode tree) {1559 TrailingDirectives tail = {DagNode(nullptr), DagNode(nullptr), 0};1560 1561 // Look backwards through the arguments.1562 auto numPatArgs = tree.getNumArgs();1563 for (int i = numPatArgs - 1; i >= 0; --i) {1564 auto dagArg = tree.getArgAsNestedDag(i);1565 // A leaf is not a directive. Stop looking.1566 if (!dagArg)1567 break;1568 1569 auto isLocation = dagArg.isLocationDirective();1570 auto isReturnType = dagArg.isReturnTypeDirective();1571 // If encountered a DAG node that isn't a trailing directive, stop looking.1572 if (!(isLocation || isReturnType))1573 break;1574 // Save the directive, but error if one of the same type was already1575 // found.1576 ++tail.numDirectives;1577 if (isLocation) {1578 if (tail.location)1579 PrintFatalError(loc, "`location` directive can only be specified "1580 "once");1581 tail.location = dagArg;1582 } else if (isReturnType) {1583 if (tail.returnType)1584 PrintFatalError(loc, "`returnType` directive can only be specified "1585 "once");1586 tail.returnType = dagArg;1587 }1588 }1589 1590 return tail;1591}1592 1593std::string1594PatternEmitter::getLocation(PatternEmitter::TrailingDirectives &tail) {1595 if (tail.location)1596 return handleLocationDirective(tail.location);1597 1598 // If no explicit location is given, use the default, all fused, location.1599 return "odsLoc";1600}1601 1602std::string PatternEmitter::handleOpCreation(DagNode tree, int resultIndex,1603 int depth) {1604 LLVM_DEBUG(llvm::dbgs() << "create op for pattern: ");1605 LLVM_DEBUG(tree.print(llvm::dbgs()));1606 LLVM_DEBUG(llvm::dbgs() << '\n');1607 1608 Operator &resultOp = tree.getDialectOp(opMap);1609 bool useProperties = resultOp.getDialect().usePropertiesForAttributes();1610 auto numOpArgs = resultOp.getNumArgs();1611 auto numPatArgs = tree.getNumArgs();1612 1613 auto tail = getTrailingDirectives(tree);1614 auto locToUse = getLocation(tail);1615 1616 auto inPattern = numPatArgs - tail.numDirectives;1617 if (numOpArgs != inPattern) {1618 PrintFatalError(loc,1619 formatv("resultant op '{0}' argument number mismatch: "1620 "{1} in pattern vs. {2} in definition",1621 resultOp.getOperationName(), inPattern, numOpArgs));1622 }1623 1624 // A map to collect all nested DAG child nodes' names, with operand index as1625 // the key. This includes both bound and unbound child nodes.1626 ChildNodeIndexNameMap childNodeNames;1627 1628 // If the argument is a type constraint, then its an operand. Check if the1629 // op's argument is variadic that the argument in the pattern is too.1630 auto checkIfMatchedVariadic = [&](int i) {1631 // FIXME: This does not yet check for variable/leaf case.1632 // FIXME: Change so that native code call can be handled.1633 const auto *operand =1634 llvm::dyn_cast_if_present<NamedTypeConstraint *>(resultOp.getArg(i));1635 if (!operand || !operand->isVariadic())1636 return;1637 1638 auto child = tree.getArgAsNestedDag(i);1639 if (!child)1640 return;1641 1642 // Skip over replaceWithValues.1643 while (child.isReplaceWithValue()) {1644 if (!(child = child.getArgAsNestedDag(0)))1645 return;1646 }1647 if (!child.isNativeCodeCall() && !child.isVariadic())1648 PrintFatalError(loc, formatv("op expects variadic operand `{0}`, while "1649 "provided is non-variadic",1650 resultOp.getArgName(i)));1651 };1652 1653 // First go through all the child nodes who are nested DAG constructs to1654 // create ops for them and remember the symbol names for them, so that we can1655 // use the results in the current node. This happens in a recursive manner.1656 for (int i = 0, e = tree.getNumArgs() - tail.numDirectives; i != e; ++i) {1657 checkIfMatchedVariadic(i);1658 if (auto child = tree.getArgAsNestedDag(i))1659 childNodeNames[i] = handleResultPattern(child, i, depth + 1);1660 }1661 1662 // The name of the local variable holding this op.1663 std::string valuePackName;1664 // The symbol for holding the result of this pattern. Note that the result of1665 // this pattern is not necessarily the same as the variable created by this1666 // pattern because we can use `__N` suffix to refer only a specific result if1667 // the generated op is a multi-result op.1668 std::string resultValue;1669 if (tree.getSymbol().empty()) {1670 // No symbol is explicitly bound to this op in the pattern. Generate a1671 // unique name.1672 valuePackName = resultValue = getUniqueSymbol(&resultOp);1673 } else {1674 resultValue = std::string(tree.getSymbol());1675 // Strip the index to get the name for the value pack and use it to name the1676 // local variable for the op.1677 valuePackName = std::string(SymbolInfoMap::getValuePackName(resultValue));1678 }1679 1680 // Create the local variable for this op.1681 os << formatv("{0} {1};\n{{\n", resultOp.getQualCppClassName(),1682 valuePackName);1683 1684 // Right now ODS don't have general type inference support. Except a few1685 // special cases listed below, DRR needs to supply types for all results1686 // when building an op.1687 bool isSameOperandsAndResultType =1688 resultOp.getTrait("::mlir::OpTrait::SameOperandsAndResultType");1689 bool useFirstAttr =1690 resultOp.getTrait("::mlir::OpTrait::FirstAttrDerivedResultType");1691 1692 if (!tail.returnType && (isSameOperandsAndResultType || useFirstAttr)) {1693 // We know how to deduce the result type for ops with these traits and we've1694 // generated builders taking aggregate parameters. Use those builders to1695 // create the ops.1696 1697 // First prepare local variables for op arguments used in builder call.1698 createAggregateLocalVarsForOpArgs(tree, childNodeNames, depth);1699 1700 // Then create the op.1701 os.scope("", "\n}\n").os1702 << formatv("{0} = {1}::create(rewriter, {2}, tblgen_values, {3});",1703 valuePackName, resultOp.getQualCppClassName(), locToUse,1704 useProperties ? "tblgen_props" : "tblgen_attrs");1705 return resultValue;1706 }1707 1708 bool usePartialResults = valuePackName != resultValue;1709 1710 if (!tail.returnType && (usePartialResults || depth > 0 || resultIndex < 0)) {1711 // For these cases (broadcastable ops, op results used both as auxiliary1712 // values and replacement values, ops in nested patterns, auxiliary ops), we1713 // still need to supply the result types when building the op. But because1714 // we don't generate a builder automatically with ODS for them, it's the1715 // developer's responsibility to make sure such a builder (with result type1716 // deduction ability) exists. We go through the separate-parameter builder1717 // here given that it's easier for developers to write compared to1718 // aggregate-parameter builders.1719 createSeparateLocalVarsForOpArgs(tree, childNodeNames);1720 1721 os.scope().os << formatv("{0} = {1}::create(rewriter, {2}", valuePackName,1722 resultOp.getQualCppClassName(), locToUse);1723 supplyValuesForOpArgs(tree, childNodeNames, depth);1724 os << "\n );\n}\n";1725 return resultValue;1726 }1727 1728 // If we are provided explicit return types, use them to build the op.1729 // However, if depth == 0 and resultIndex >= 0, it means we are replacing1730 // the values generated from the source pattern root op. Then we must use the1731 // source pattern's value types to determine the value type of the generated1732 // op here.1733 if (depth == 0 && resultIndex >= 0 && tail.returnType)1734 PrintFatalError(loc, "Cannot specify explicit return types in an op whose "1735 "return values replace the source pattern's root op");1736 1737 // First prepare local variables for op arguments used in builder call.1738 createAggregateLocalVarsForOpArgs(tree, childNodeNames, depth);1739 1740 // Then prepare the result types. We need to specify the types for all1741 // results.1742 os.indent() << formatv("::llvm::SmallVector<::mlir::Type, 4> tblgen_types; "1743 "(void)tblgen_types;\n");1744 int numResults = resultOp.getNumResults();1745 if (tail.returnType) {1746 auto numRetTys = tail.returnType.getNumArgs();1747 for (int i = 0; i < numRetTys; ++i) {1748 auto varName = handleReturnTypeArg(tail.returnType, i, depth + 1);1749 os << "tblgen_types.push_back(" << varName << ");\n";1750 }1751 } else {1752 if (numResults != 0) {1753 // Copy the result types from the source pattern.1754 for (int i = 0; i < numResults; ++i)1755 os << formatv("for (auto v: castedOp0.getODSResults({0})) {{\n"1756 " tblgen_types.push_back(v.getType());\n}\n",1757 resultIndex + i);1758 }1759 }1760 os << formatv("{0} = {1}::create(rewriter, {2}, tblgen_types, "1761 "tblgen_values, {3});\n",1762 valuePackName, resultOp.getQualCppClassName(), locToUse,1763 useProperties ? "tblgen_props" : "tblgen_attrs");1764 os.unindent() << "}\n";1765 return resultValue;1766}1767 1768void PatternEmitter::createSeparateLocalVarsForOpArgs(1769 DagNode node, ChildNodeIndexNameMap &childNodeNames) {1770 Operator &resultOp = node.getDialectOp(opMap);1771 1772 // Now prepare operands used for building this op:1773 // * If the operand is non-variadic, we create a `Value` local variable.1774 // * If the operand is variadic, we create a `SmallVector<Value>` local1775 // variable.1776 1777 int valueIndex = 0; // An index for uniquing local variable names.1778 for (int argIndex = 0, e = resultOp.getNumArgs(); argIndex < e; ++argIndex) {1779 const auto *operand = llvm::dyn_cast_if_present<NamedTypeConstraint *>(1780 resultOp.getArg(argIndex));1781 // We do not need special handling for attributes or properties.1782 if (!operand)1783 continue;1784 1785 raw_indented_ostream::DelimitedScope scope(os);1786 std::string varName;1787 if (operand->isVariadic()) {1788 varName = std::string(formatv("tblgen_values_{0}", valueIndex++));1789 os << formatv("::llvm::SmallVector<::mlir::Value, 4> {0};\n", varName);1790 std::string range;1791 if (node.isNestedDagArg(argIndex)) {1792 range = childNodeNames[argIndex];1793 } else {1794 range = std::string(node.getArgName(argIndex));1795 }1796 // Resolve the symbol for all range use so that we have a uniform way of1797 // capturing the values.1798 range = symbolInfoMap.getValueAndRangeUse(range);1799 os << formatv("for (auto v: {0}) {{\n {1}.push_back(v);\n}\n", range,1800 varName);1801 } else {1802 varName = std::string(formatv("tblgen_value_{0}", valueIndex++));1803 os << formatv("::mlir::Value {0} = ", varName);1804 if (node.isNestedDagArg(argIndex)) {1805 os << symbolInfoMap.getValueAndRangeUse(childNodeNames[argIndex]);1806 } else {1807 DagLeaf leaf = node.getArgAsLeaf(argIndex);1808 auto symbol =1809 symbolInfoMap.getValueAndRangeUse(node.getArgName(argIndex));1810 if (leaf.isNativeCodeCall()) {1811 os << std::string(1812 tgfmt(leaf.getNativeCodeTemplate(), &fmtCtx.withSelf(symbol)));1813 } else {1814 os << symbol;1815 }1816 }1817 os << ";\n";1818 }1819 1820 // Update to use the newly created local variable for building the op later.1821 childNodeNames[argIndex] = varName;1822 }1823}1824 1825void PatternEmitter::supplyValuesForOpArgs(1826 DagNode node, const ChildNodeIndexNameMap &childNodeNames, int depth) {1827 Operator &resultOp = node.getDialectOp(opMap);1828 for (int argIndex = 0, numOpArgs = resultOp.getNumArgs();1829 argIndex != numOpArgs; ++argIndex) {1830 // Start each argument on its own line.1831 os << ",\n ";1832 1833 Argument opArg = resultOp.getArg(argIndex);1834 // Handle the case of operand first.1835 if (auto *operand =1836 llvm::dyn_cast_if_present<NamedTypeConstraint *>(opArg)) {1837 if (!operand->name.empty())1838 os << "/*" << operand->name << "=*/";1839 os << childNodeNames.lookup(argIndex);1840 continue;1841 }1842 1843 // The argument in the op definition.1844 auto opArgName = resultOp.getArgName(argIndex);1845 if (auto subTree = node.getArgAsNestedDag(argIndex)) {1846 if (!subTree.isNativeCodeCall())1847 PrintFatalError(loc, "only NativeCodeCall allowed in nested dag node "1848 "for creating attributes and properties");1849 os << formatv("/*{0}=*/{1}", opArgName, childNodeNames.lookup(argIndex));1850 } else {1851 auto leaf = node.getArgAsLeaf(argIndex);1852 // The argument in the result DAG pattern.1853 auto patArgName = node.getArgName(argIndex);1854 if (leaf.isConstantAttr() || leaf.isEnumCase()) {1855 // TODO: Refactor out into map to avoid recomputing these.1856 if (!isa<NamedAttribute *>(opArg))1857 PrintFatalError(loc, Twine("expected attribute ") + Twine(argIndex));1858 if (!patArgName.empty())1859 os << "/*" << patArgName << "=*/";1860 } else if (leaf.isConstantProp()) {1861 if (!isa<NamedProperty *>(opArg))1862 PrintFatalError(loc, Twine("expected property ") + Twine(argIndex));1863 if (!patArgName.empty())1864 os << "/*" << patArgName << "=*/";1865 } else {1866 os << "/*" << opArgName << "=*/";1867 }1868 os << handleOpArgument(leaf, patArgName);1869 }1870 }1871}1872 1873void PatternEmitter::createAggregateLocalVarsForOpArgs(1874 DagNode node, const ChildNodeIndexNameMap &childNodeNames, int depth) {1875 Operator &resultOp = node.getDialectOp(opMap);1876 1877 bool useProperties = resultOp.getDialect().usePropertiesForAttributes();1878 auto scope = os.scope();1879 os << formatv("::llvm::SmallVector<::mlir::Value, 4> "1880 "tblgen_values; (void)tblgen_values;\n");1881 if (useProperties) {1882 os << formatv("{0}::Properties tblgen_props; (void)tblgen_props;\n",1883 resultOp.getQualCppClassName());1884 } else {1885 os << formatv("::llvm::SmallVector<::mlir::NamedAttribute, 4> "1886 "tblgen_attrs; (void)tblgen_attrs;\n");1887 }1888 1889 const char *setPropCmd =1890 "tblgen_props.{0} = "1891 "::llvm::dyn_cast_if_present<decltype(tblgen_props.{0})>({1});\n";1892 const char *addAttrCmd =1893 "if (auto tmpAttr = {1}) {\n"1894 " tblgen_attrs.emplace_back(rewriter.getStringAttr(\"{0}\"), "1895 "tmpAttr);\n}\n";1896 const char *setterCmd = (useProperties) ? setPropCmd : addAttrCmd;1897 const char *propSetterCmd = "tblgen_props.{0}({1});\n";1898 1899 int numVariadic = 0;1900 bool hasOperandSegmentSizes = false;1901 std::vector<std::string> sizes;1902 for (int argIndex = 0, e = resultOp.getNumArgs(); argIndex < e; ++argIndex) {1903 if (isa<NamedAttribute *>(resultOp.getArg(argIndex))) {1904 // The argument in the op definition.1905 auto opArgName = resultOp.getArgName(argIndex);1906 hasOperandSegmentSizes =1907 hasOperandSegmentSizes || opArgName == "operandSegmentSizes";1908 if (auto subTree = node.getArgAsNestedDag(argIndex)) {1909 if (!subTree.isNativeCodeCall())1910 PrintFatalError(loc, "only NativeCodeCall allowed in nested dag node "1911 "for creating attribute");1912 1913 os << formatv(setterCmd, opArgName, childNodeNames.lookup(argIndex));1914 } else {1915 auto leaf = node.getArgAsLeaf(argIndex);1916 // The argument in the result DAG pattern.1917 auto patArgName = node.getArgName(argIndex);1918 os << formatv(setterCmd, opArgName, handleOpArgument(leaf, patArgName));1919 }1920 continue;1921 }1922 1923 if (isa<NamedProperty *>(resultOp.getArg(argIndex))) {1924 // The argument in the op definition.1925 auto opArgName = resultOp.getArgName(argIndex);1926 auto setterName = resultOp.getSetterName(opArgName);1927 if (auto subTree = node.getArgAsNestedDag(argIndex)) {1928 if (!subTree.isNativeCodeCall())1929 PrintFatalError(loc, "only NativeCodeCall allowed in nested dag node "1930 "for creating property");1931 1932 os << formatv(propSetterCmd, setterName,1933 childNodeNames.lookup(argIndex));1934 } else {1935 auto leaf = node.getArgAsLeaf(argIndex);1936 // The argument in the result DAG pattern.1937 auto patArgName = node.getArgName(argIndex);1938 // The argument in the result DAG pattern.1939 os << formatv(propSetterCmd, setterName,1940 handleOpArgument(leaf, patArgName));1941 }1942 continue;1943 }1944 1945 const auto *operand =1946 cast<NamedTypeConstraint *>(resultOp.getArg(argIndex));1947 if (operand->isVariadic()) {1948 ++numVariadic;1949 std::string range;1950 if (node.isNestedDagArg(argIndex)) {1951 range = childNodeNames.lookup(argIndex);1952 } else {1953 range = std::string(node.getArgName(argIndex));1954 }1955 // Resolve the symbol for all range use so that we have a uniform way of1956 // capturing the values.1957 range = symbolInfoMap.getValueAndRangeUse(range);1958 os << formatv("for (auto v: {0}) {{\n tblgen_values.push_back(v);\n}\n",1959 range);1960 sizes.push_back(formatv("static_cast<int32_t>({0}.size())", range));1961 } else {1962 sizes.emplace_back("1");1963 os << formatv("tblgen_values.push_back(");1964 if (node.isNestedDagArg(argIndex)) {1965 os << symbolInfoMap.getValueAndRangeUse(1966 childNodeNames.lookup(argIndex));1967 } else {1968 DagLeaf leaf = node.getArgAsLeaf(argIndex);1969 if (leaf.isConstantAttr())1970 // TODO: Use better location1971 PrintFatalError(1972 loc,1973 "attribute found where value was expected, if attempting to use "1974 "constant value, construct a constant op with given attribute "1975 "instead");1976 1977 auto symbol =1978 symbolInfoMap.getValueAndRangeUse(node.getArgName(argIndex));1979 if (leaf.isNativeCodeCall()) {1980 os << std::string(1981 tgfmt(leaf.getNativeCodeTemplate(), &fmtCtx.withSelf(symbol)));1982 } else {1983 os << symbol;1984 }1985 }1986 os << ");\n";1987 }1988 }1989 1990 if (numVariadic > 1 && !hasOperandSegmentSizes) {1991 // Only set size if it can't be computed.1992 const auto *sameVariadicSize =1993 resultOp.getTrait("::mlir::OpTrait::SameVariadicOperandSize");1994 if (!sameVariadicSize) {1995 if (useProperties) {1996 const char *setSizes = R"(1997 tblgen_props.operandSegmentSizes = {{ {0} };1998 )";1999 os.printReindented(formatv(setSizes, llvm::join(sizes, ", ")).str());2000 } else {2001 const char *setSizes = R"(2002 tblgen_attrs.emplace_back(rewriter.getStringAttr("operandSegmentSizes"),2003 rewriter.getDenseI32ArrayAttr({{ {0} }));2004 )";2005 os.printReindented(formatv(setSizes, llvm::join(sizes, ", ")).str());2006 }2007 }2008 }2009}2010 2011StaticMatcherHelper::StaticMatcherHelper(raw_ostream &os,2012 const RecordKeeper &records,2013 RecordOperatorMap &mapper)2014 : opMap(mapper), staticVerifierEmitter(os, records) {}2015 2016void StaticMatcherHelper::populateStaticMatchers(raw_ostream &os) {2017 // PatternEmitter will use the static matcher if there's one generated. To2018 // ensure that all the dependent static matchers are generated before emitting2019 // the matching logic of the DagNode, we use topological order to achieve it.2020 for (auto &dagInfo : topologicalOrder) {2021 DagNode node = dagInfo.first;2022 if (!useStaticMatcher(node))2023 continue;2024 2025 std::string funcName =2026 formatv("static_dag_matcher_{0}", staticMatcherCounter++);2027 assert(!matcherNames.contains(node));2028 PatternEmitter(dagInfo.second, &opMap, os, *this)2029 .emitStaticMatcher(node, funcName);2030 matcherNames[node] = funcName;2031 }2032}2033 2034void StaticMatcherHelper::populateStaticConstraintFunctions(raw_ostream &os) {2035 staticVerifierEmitter.emitPatternConstraints(constraints.getArrayRef());2036}2037 2038void StaticMatcherHelper::addPattern(const Record *record) {2039 Pattern pat(record, &opMap);2040 2041 // While generating the function body of the DAG matcher, it may depends on2042 // other DAG matchers. To ensure the dependent matchers are ready, we compute2043 // the topological order for all the DAGs and emit the DAG matchers in this2044 // order.2045 llvm::unique_function<void(DagNode)> dfs = [&](DagNode node) {2046 ++refStats[node];2047 2048 if (refStats[node] != 1)2049 return;2050 2051 for (unsigned i = 0, e = node.getNumArgs(); i < e; ++i)2052 if (DagNode sibling = node.getArgAsNestedDag(i))2053 dfs(sibling);2054 else {2055 DagLeaf leaf = node.getArgAsLeaf(i);2056 if (!leaf.isUnspecified())2057 constraints.insert(leaf);2058 }2059 2060 topologicalOrder.push_back(std::make_pair(node, record));2061 };2062 2063 dfs(pat.getSourcePattern());2064}2065 2066StringRef StaticMatcherHelper::getVerifierName(DagLeaf leaf) {2067 if (leaf.isAttrMatcher()) {2068 std::optional<StringRef> constraint =2069 staticVerifierEmitter.getAttrConstraintFn(leaf.getAsConstraint());2070 assert(constraint && "attribute constraint was not uniqued");2071 return *constraint;2072 }2073 if (leaf.isPropMatcher()) {2074 std::optional<StringRef> constraint =2075 staticVerifierEmitter.getPropConstraintFn(leaf.getAsConstraint());2076 assert(constraint && "prop constraint was not uniqued");2077 return *constraint;2078 }2079 assert(leaf.isOperandMatcher());2080 return staticVerifierEmitter.getTypeConstraintFn(leaf.getAsConstraint());2081}2082 2083static void emitRewriters(const RecordKeeper &records, raw_ostream &os) {2084 emitSourceFileHeader("Rewriters", os, records);2085 2086 auto patterns = records.getAllDerivedDefinitions("Pattern");2087 2088 // We put the map here because it can be shared among multiple patterns.2089 RecordOperatorMap recordOpMap;2090 2091 // Exam all the patterns and generate static matcher for the duplicated2092 // DagNode.2093 StaticMatcherHelper staticMatcher(os, records, recordOpMap);2094 for (const Record *p : patterns)2095 staticMatcher.addPattern(p);2096 staticMatcher.populateStaticConstraintFunctions(os);2097 staticMatcher.populateStaticMatchers(os);2098 2099 std::vector<std::string> rewriterNames;2100 rewriterNames.reserve(patterns.size());2101 2102 std::string baseRewriterName = "GeneratedConvert";2103 int rewriterIndex = 0;2104 2105 for (const Record *p : patterns) {2106 std::string name;2107 if (p->isAnonymous()) {2108 // If no name is provided, ensure unique rewriter names simply by2109 // appending unique suffix.2110 name = baseRewriterName + llvm::utostr(rewriterIndex++);2111 } else {2112 name = std::string(p->getName());2113 }2114 LLVM_DEBUG(llvm::dbgs()2115 << "=== start generating pattern '" << name << "' ===\n");2116 PatternEmitter(p, &recordOpMap, os, staticMatcher).emit(name);2117 LLVM_DEBUG(llvm::dbgs()2118 << "=== done generating pattern '" << name << "' ===\n");2119 rewriterNames.push_back(std::move(name));2120 }2121 2122 // Emit function to add the generated matchers to the pattern list.2123 os << "[[maybe_unused]] void populateWithGenerated("2124 "::mlir::RewritePatternSet &patterns) {\n";2125 for (const auto &name : rewriterNames) {2126 os << " patterns.add<" << name << ">(patterns.getContext());\n";2127 }2128 os << "}\n";2129}2130 2131static mlir::GenRegistration2132 genRewriters("gen-rewriters", "Generate pattern rewriters",2133 [](const RecordKeeper &records, raw_ostream &os) {2134 emitRewriters(records, os);2135 return false;2136 });2137