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-exception 6 // 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/Attribute.h" 15 #include "mlir/TableGen/Format.h" 16 #include "mlir/TableGen/GenInfo.h" 17 #include "mlir/TableGen/Operator.h" 18 #include "mlir/TableGen/Pattern.h" 19 #include "mlir/TableGen/Predicate.h" 20 #include "mlir/TableGen/Type.h" 21 #include "llvm/ADT/StringExtras.h" 22 #include "llvm/ADT/StringSet.h" 23 #include "llvm/Support/CommandLine.h" 24 #include "llvm/Support/Debug.h" 25 #include "llvm/Support/FormatAdapters.h" 26 #include "llvm/Support/PrettyStackTrace.h" 27 #include "llvm/Support/Signals.h" 28 #include "llvm/TableGen/Error.h" 29 #include "llvm/TableGen/Main.h" 30 #include "llvm/TableGen/Record.h" 31 #include "llvm/TableGen/TableGenBackend.h" 32 33 using namespace mlir; 34 using namespace mlir::tblgen; 35 36 using llvm::formatv; 37 using llvm::Record; 38 using llvm::RecordKeeper; 39 40 #define DEBUG_TYPE "mlir-tblgen-rewritergen" 41 42 namespace llvm { 43 template <> 44 struct format_provider<mlir::tblgen::Pattern::IdentifierLine> { 45 static void format(const mlir::tblgen::Pattern::IdentifierLine &v, 46 raw_ostream &os, StringRef style) { 47 os << v.first << ":" << v.second; 48 } 49 }; 50 } // end namespace llvm 51 52 //===----------------------------------------------------------------------===// 53 // PatternEmitter 54 //===----------------------------------------------------------------------===// 55 56 namespace { 57 class PatternEmitter { 58 public: 59 PatternEmitter(Record *pat, RecordOperatorMap *mapper, raw_ostream &os); 60 61 // Emits the mlir::RewritePattern struct named `rewriteName`. 62 void emit(StringRef rewriteName); 63 64 private: 65 // Emits the code for matching ops. 66 void emitMatchLogic(DagNode tree, StringRef opName); 67 68 // Emits the code for rewriting ops. 69 void emitRewriteLogic(); 70 71 //===--------------------------------------------------------------------===// 72 // Match utilities 73 //===--------------------------------------------------------------------===// 74 75 // Emits C++ statements for matching the DAG structure. 76 void emitMatch(DagNode tree, StringRef name, int depth); 77 78 // Emits C++ statements for matching using a native code call. 79 void emitNativeCodeMatch(DagNode tree, StringRef name, int depth); 80 81 // Emits C++ statements for matching the op constrained by the given DAG 82 // `tree` returning the op's variable name. 83 void emitOpMatch(DagNode tree, StringRef opName, int depth); 84 85 // Emits C++ statements for matching the `argIndex`-th argument of the given 86 // DAG `tree` as an operand. operandIndex is the index in the DAG excluding 87 // the preceding attributes. 88 void emitOperandMatch(DagNode tree, StringRef opName, int argIndex, 89 int operandIndex, int depth); 90 91 // Emits C++ statements for matching the `argIndex`-th argument of the given 92 // DAG `tree` as an attribute. 93 void emitAttributeMatch(DagNode tree, StringRef opName, int argIndex, 94 int depth); 95 96 // Emits C++ for checking a match with a corresponding match failure 97 // diagnostic. 98 void emitMatchCheck(StringRef opName, const FmtObjectBase &matchFmt, 99 const llvm::formatv_object_base &failureFmt); 100 101 // Emits C++ for checking a match with a corresponding match failure 102 // diagnostics. 103 void emitMatchCheck(StringRef opName, const std::string &matchStr, 104 const std::string &failureStr); 105 106 //===--------------------------------------------------------------------===// 107 // Rewrite utilities 108 //===--------------------------------------------------------------------===// 109 110 // The entry point for handling a result pattern rooted at `resultTree`. This 111 // method dispatches to concrete handlers according to `resultTree`'s kind and 112 // returns a symbol representing the whole value pack. Callers are expected to 113 // further resolve the symbol according to the specific use case. 114 // 115 // `depth` is the nesting level of `resultTree`; 0 means top-level result 116 // pattern. For top-level result pattern, `resultIndex` indicates which result 117 // of the matched root op this pattern is intended to replace, which can be 118 // used to deduce the result type of the op generated from this result 119 // pattern. 120 std::string handleResultPattern(DagNode resultTree, int resultIndex, 121 int depth); 122 123 // Emits the C++ statement to replace the matched DAG with a value built via 124 // calling native C++ code. 125 std::string handleReplaceWithNativeCodeCall(DagNode resultTree, int depth); 126 127 // Returns the symbol of the old value serving as the replacement. 128 StringRef handleReplaceWithValue(DagNode tree); 129 130 // Returns the location value to use. 131 std::pair<bool, std::string> getLocation(DagNode tree); 132 133 // Returns the location value to use. 134 std::string handleLocationDirective(DagNode tree); 135 136 // Emits the C++ statement to build a new op out of the given DAG `tree` and 137 // returns the variable name that this op is assigned to. If the root op in 138 // DAG `tree` has a specified name, the created op will be assigned to a 139 // variable of the given name. Otherwise, a unique name will be used as the 140 // result value name. 141 std::string handleOpCreation(DagNode tree, int resultIndex, int depth); 142 143 using ChildNodeIndexNameMap = DenseMap<unsigned, std::string>; 144 145 // Emits a local variable for each value and attribute to be used for creating 146 // an op. 147 void createSeparateLocalVarsForOpArgs(DagNode node, 148 ChildNodeIndexNameMap &childNodeNames); 149 150 // Emits the concrete arguments used to call an op's builder. 151 void supplyValuesForOpArgs(DagNode node, 152 const ChildNodeIndexNameMap &childNodeNames, 153 int depth); 154 155 // Emits the local variables for holding all values as a whole and all named 156 // attributes as a whole to be used for creating an op. 157 void createAggregateLocalVarsForOpArgs( 158 DagNode node, const ChildNodeIndexNameMap &childNodeNames, int depth); 159 160 // Returns the C++ expression to construct a constant attribute of the given 161 // `value` for the given attribute kind `attr`. 162 std::string handleConstantAttr(Attribute attr, StringRef value); 163 164 // Returns the C++ expression to build an argument from the given DAG `leaf`. 165 // `patArgName` is used to bound the argument to the source pattern. 166 std::string handleOpArgument(DagLeaf leaf, StringRef patArgName); 167 168 //===--------------------------------------------------------------------===// 169 // General utilities 170 //===--------------------------------------------------------------------===// 171 172 // Collects all of the operations within the given dag tree. 173 void collectOps(DagNode tree, llvm::SmallPtrSetImpl<const Operator *> &ops); 174 175 // Returns a unique symbol for a local variable of the given `op`. 176 std::string getUniqueSymbol(const Operator *op); 177 178 //===--------------------------------------------------------------------===// 179 // Symbol utilities 180 //===--------------------------------------------------------------------===// 181 182 // Returns how many static values the given DAG `node` correspond to. 183 int getNodeValueCount(DagNode node); 184 185 private: 186 // Pattern instantiation location followed by the location of multiclass 187 // prototypes used. This is intended to be used as a whole to 188 // PrintFatalError() on errors. 189 ArrayRef<llvm::SMLoc> loc; 190 191 // Op's TableGen Record to wrapper object. 192 RecordOperatorMap *opMap; 193 194 // Handy wrapper for pattern being emitted. 195 Pattern pattern; 196 197 // Map for all bound symbols' info. 198 SymbolInfoMap symbolInfoMap; 199 200 // The next unused ID for newly created values. 201 unsigned nextValueId; 202 203 raw_indented_ostream os; 204 205 // Format contexts containing placeholder substitutions. 206 FmtContext fmtCtx; 207 208 // Number of op processed. 209 int opCounter = 0; 210 }; 211 } // end anonymous namespace 212 213 PatternEmitter::PatternEmitter(Record *pat, RecordOperatorMap *mapper, 214 raw_ostream &os) 215 : loc(pat->getLoc()), opMap(mapper), pattern(pat, mapper), 216 symbolInfoMap(pat->getLoc()), nextValueId(0), os(os) { 217 fmtCtx.withBuilder("rewriter"); 218 } 219 220 std::string PatternEmitter::handleConstantAttr(Attribute attr, 221 StringRef value) { 222 if (!attr.isConstBuildable()) 223 PrintFatalError(loc, "Attribute " + attr.getAttrDefName() + 224 " does not have the 'constBuilderCall' field"); 225 226 // TODO: Verify the constants here 227 return std::string(tgfmt(attr.getConstBuilderTemplate(), &fmtCtx, value)); 228 } 229 230 // Helper function to match patterns. 231 void PatternEmitter::emitMatch(DagNode tree, StringRef name, int depth) { 232 if (tree.isNativeCodeCall()) { 233 emitNativeCodeMatch(tree, name, depth); 234 return; 235 } 236 237 if (tree.isOperation()) { 238 emitOpMatch(tree, name, depth); 239 return; 240 } 241 242 PrintFatalError(loc, "encountered non-op, non-NativeCodeCall match."); 243 } 244 245 // Helper function to match patterns. 246 void PatternEmitter::emitNativeCodeMatch(DagNode tree, StringRef opName, 247 int depth) { 248 LLVM_DEBUG(llvm::dbgs() << "handle NativeCodeCall matcher pattern: "); 249 LLVM_DEBUG(tree.print(llvm::dbgs())); 250 LLVM_DEBUG(llvm::dbgs() << '\n'); 251 252 // TODO(suderman): iterate through arguments, determine their types, output 253 // names. 254 SmallVector<std::string, 8> capture(8); 255 if (tree.getNumArgs() > 8) { 256 PrintFatalError(loc, 257 "unsupported NativeCodeCall matcher argument numbers: " + 258 Twine(tree.getNumArgs())); 259 } 260 261 raw_indented_ostream::DelimitedScope scope(os); 262 263 os << "if(!" << opName << ") return ::mlir::failure();\n"; 264 for (int i = 0, e = tree.getNumArgs(); i != e; ++i) { 265 std::string argName = formatv("arg{0}_{1}", depth, i); 266 if (DagNode argTree = tree.getArgAsNestedDag(i)) { 267 os << "Value " << argName << ";\n"; 268 } else { 269 auto leaf = tree.getArgAsLeaf(i); 270 if (leaf.isAttrMatcher() || leaf.isConstantAttr()) { 271 os << "Attribute " << argName << ";\n"; 272 } else if (leaf.isOperandMatcher()) { 273 os << "Operation " << argName << ";\n"; 274 } 275 } 276 277 capture[i] = std::move(argName); 278 } 279 280 bool hasLocationDirective; 281 std::string locToUse; 282 std::tie(hasLocationDirective, locToUse) = getLocation(tree); 283 284 auto fmt = tree.getNativeCodeTemplate(); 285 auto nativeCodeCall = std::string(tgfmt( 286 fmt, &fmtCtx.addSubst("_loc", locToUse), opName, capture[0], capture[1], 287 capture[2], capture[3], capture[4], capture[5], capture[6], capture[7])); 288 289 os << "if (failed(" << nativeCodeCall << ")) return ::mlir::failure();\n"; 290 291 for (int i = 0, e = tree.getNumArgs(); i != e; ++i) { 292 auto name = tree.getArgName(i); 293 if (!name.empty() && name != "_") { 294 os << formatv("{0} = {1};\n", name, capture[i]); 295 } 296 } 297 298 for (int i = 0, e = tree.getNumArgs(); i != e; ++i) { 299 std::string argName = capture[i]; 300 301 // Handle nested DAG construct first 302 if (DagNode argTree = tree.getArgAsNestedDag(i)) { 303 PrintFatalError( 304 loc, formatv("Matching nested tree in NativeCodecall not support for " 305 "{0} as arg {1}", 306 argName, i)); 307 } 308 309 DagLeaf leaf = tree.getArgAsLeaf(i); 310 auto constraint = leaf.getAsConstraint(); 311 312 auto self = formatv("{0}", argName); 313 emitMatchCheck( 314 opName, 315 tgfmt(constraint.getConditionTemplate(), &fmtCtx.withSelf(self)), 316 formatv("\"operand {0} of native code call '{1}' failed to satisfy " 317 "constraint: " 318 "'{2}'\"", 319 i, tree.getNativeCodeTemplate(), constraint.getSummary())); 320 } 321 322 LLVM_DEBUG(llvm::dbgs() << "done emitting match for native code call\n"); 323 } 324 325 // Helper function to match patterns. 326 void PatternEmitter::emitOpMatch(DagNode tree, StringRef opName, int depth) { 327 Operator &op = tree.getDialectOp(opMap); 328 LLVM_DEBUG(llvm::dbgs() << "start emitting match for op '" 329 << op.getOperationName() << "' at depth " << depth 330 << '\n'); 331 332 std::string castedName = formatv("castedOp{0}", depth); 333 os << formatv("auto {0} = ::llvm::dyn_cast_or_null<{2}>({1}); " 334 "(void){0};\n", 335 castedName, opName, op.getQualCppClassName()); 336 // Skip the operand matching at depth 0 as the pattern rewriter already does. 337 if (depth != 0) { 338 // Skip if there is no defining operation (e.g., arguments to function). 339 os << formatv("if (!{0}) return ::mlir::failure();\n", castedName); 340 } 341 if (tree.getNumArgs() != op.getNumArgs()) { 342 PrintFatalError(loc, formatv("op '{0}' argument number mismatch: {1} in " 343 "pattern vs. {2} in definition", 344 op.getOperationName(), tree.getNumArgs(), 345 op.getNumArgs())); 346 } 347 348 // If the operand's name is set, set to that variable. 349 auto name = tree.getSymbol(); 350 if (!name.empty()) 351 os << formatv("{0} = {1};\n", name, castedName); 352 353 for (int i = 0, e = tree.getNumArgs(), nextOperand = 0; i != e; ++i) { 354 auto opArg = op.getArg(i); 355 std::string argName = formatv("op{0}", depth + 1); 356 357 // Handle nested DAG construct first 358 if (DagNode argTree = tree.getArgAsNestedDag(i)) { 359 if (auto *operand = opArg.dyn_cast<NamedTypeConstraint *>()) { 360 if (operand->isVariableLength()) { 361 auto error = formatv("use nested DAG construct to match op {0}'s " 362 "variadic operand #{1} unsupported now", 363 op.getOperationName(), i); 364 PrintFatalError(loc, error); 365 } 366 } 367 os << "{\n"; 368 369 // Attributes don't count for getODSOperands. 370 os.indent() << formatv( 371 "auto *{0} = " 372 "(*{1}.getODSOperands({2}).begin()).getDefiningOp();\n", 373 argName, castedName, nextOperand++); 374 emitMatch(argTree, argName, depth + 1); 375 os << formatv("tblgen_ops[{0}] = {1};\n", ++opCounter, argName); 376 os.unindent() << "}\n"; 377 continue; 378 } 379 380 // Next handle DAG leaf: operand or attribute 381 if (opArg.is<NamedTypeConstraint *>()) { 382 // emitOperandMatch's argument indexing counts attributes. 383 emitOperandMatch(tree, castedName, i, nextOperand, depth); 384 ++nextOperand; 385 } else if (opArg.is<NamedAttribute *>()) { 386 emitAttributeMatch(tree, opName, i, depth); 387 } else { 388 PrintFatalError(loc, "unhandled case when matching op"); 389 } 390 } 391 LLVM_DEBUG(llvm::dbgs() << "done emitting match for op '" 392 << op.getOperationName() << "' at depth " << depth 393 << '\n'); 394 } 395 396 void PatternEmitter::emitOperandMatch(DagNode tree, StringRef opName, 397 int argIndex, int operandIndex, 398 int depth) { 399 Operator &op = tree.getDialectOp(opMap); 400 auto *operand = op.getArg(argIndex).get<NamedTypeConstraint *>(); 401 auto matcher = tree.getArgAsLeaf(argIndex); 402 403 // If a constraint is specified, we need to generate C++ statements to 404 // check the constraint. 405 if (!matcher.isUnspecified()) { 406 if (!matcher.isOperandMatcher()) { 407 PrintFatalError( 408 loc, formatv("the {1}-th argument of op '{0}' should be an operand", 409 op.getOperationName(), argIndex + 1)); 410 } 411 412 // Only need to verify if the matcher's type is different from the one 413 // of op definition. 414 Constraint constraint = matcher.getAsConstraint(); 415 if (operand->constraint != constraint) { 416 if (operand->isVariableLength()) { 417 auto error = formatv( 418 "further constrain op {0}'s variadic operand #{1} unsupported now", 419 op.getOperationName(), argIndex); 420 PrintFatalError(loc, error); 421 } 422 auto self = formatv("(*{0}.getODSOperands({1}).begin()).getType()", 423 opName, operandIndex); 424 emitMatchCheck( 425 opName, 426 tgfmt(constraint.getConditionTemplate(), &fmtCtx.withSelf(self)), 427 formatv("\"operand {0} of op '{1}' failed to satisfy constraint: " 428 "'{2}'\"", 429 operand - op.operand_begin(), op.getOperationName(), 430 constraint.getSummary())); 431 } 432 } 433 434 // Capture the value 435 auto name = tree.getArgName(argIndex); 436 // `$_` is a special symbol to ignore op argument matching. 437 if (!name.empty() && name != "_") { 438 // We need to subtract the number of attributes before this operand to get 439 // the index in the operand list. 440 auto numPrevAttrs = std::count_if( 441 op.arg_begin(), op.arg_begin() + argIndex, 442 [](const Argument &arg) { return arg.is<NamedAttribute *>(); }); 443 444 auto res = symbolInfoMap.findBoundSymbol(name, op, argIndex); 445 os << formatv("{0} = {1}.getODSOperands({2});\n", 446 res->second.getVarName(name), opName, 447 argIndex - numPrevAttrs); 448 } 449 } 450 451 void PatternEmitter::emitAttributeMatch(DagNode tree, StringRef opName, 452 int argIndex, int depth) { 453 Operator &op = tree.getDialectOp(opMap); 454 auto *namedAttr = op.getArg(argIndex).get<NamedAttribute *>(); 455 const auto &attr = namedAttr->attr; 456 457 os << "{\n"; 458 os.indent() << formatv("auto tblgen_attr = {0}->getAttrOfType<{1}>(\"{2}\");" 459 "(void)tblgen_attr;\n", 460 opName, attr.getStorageType(), namedAttr->name); 461 462 // TODO: This should use getter method to avoid duplication. 463 if (attr.hasDefaultValue()) { 464 os << "if (!tblgen_attr) tblgen_attr = " 465 << std::string(tgfmt(attr.getConstBuilderTemplate(), &fmtCtx, 466 attr.getDefaultValue())) 467 << ";\n"; 468 } else if (attr.isOptional()) { 469 // For a missing attribute that is optional according to definition, we 470 // should just capture a mlir::Attribute() to signal the missing state. 471 // That is precisely what getAttr() returns on missing attributes. 472 } else { 473 emitMatchCheck(opName, tgfmt("tblgen_attr", &fmtCtx), 474 formatv("\"expected op '{0}' to have attribute '{1}' " 475 "of type '{2}'\"", 476 op.getOperationName(), namedAttr->name, 477 attr.getStorageType())); 478 } 479 480 auto matcher = tree.getArgAsLeaf(argIndex); 481 if (!matcher.isUnspecified()) { 482 if (!matcher.isAttrMatcher()) { 483 PrintFatalError( 484 loc, formatv("the {1}-th argument of op '{0}' should be an attribute", 485 op.getOperationName(), argIndex + 1)); 486 } 487 488 // If a constraint is specified, we need to generate C++ statements to 489 // check the constraint. 490 emitMatchCheck( 491 opName, 492 tgfmt(matcher.getConditionTemplate(), &fmtCtx.withSelf("tblgen_attr")), 493 formatv("\"op '{0}' attribute '{1}' failed to satisfy constraint: " 494 "{2}\"", 495 op.getOperationName(), namedAttr->name, 496 matcher.getAsConstraint().getSummary())); 497 } 498 499 // Capture the value 500 auto name = tree.getArgName(argIndex); 501 // `$_` is a special symbol to ignore op argument matching. 502 if (!name.empty() && name != "_") { 503 os << formatv("{0} = tblgen_attr;\n", name); 504 } 505 506 os.unindent() << "}\n"; 507 } 508 509 void PatternEmitter::emitMatchCheck( 510 StringRef opName, const FmtObjectBase &matchFmt, 511 const llvm::formatv_object_base &failureFmt) { 512 emitMatchCheck(opName, matchFmt.str(), failureFmt.str()); 513 } 514 515 void PatternEmitter::emitMatchCheck(StringRef opName, 516 const std::string &matchStr, 517 const std::string &failureStr) { 518 519 os << "if (!(" << matchStr << "))"; 520 os.scope("{\n", "\n}\n").os << "return rewriter.notifyMatchFailure(" << opName 521 << ", [&](::mlir::Diagnostic &diag) {\n diag << " 522 << failureStr << ";\n});"; 523 } 524 525 void PatternEmitter::emitMatchLogic(DagNode tree, StringRef opName) { 526 LLVM_DEBUG(llvm::dbgs() << "--- start emitting match logic ---\n"); 527 int depth = 0; 528 emitMatch(tree, opName, depth); 529 530 for (auto &appliedConstraint : pattern.getConstraints()) { 531 auto &constraint = appliedConstraint.constraint; 532 auto &entities = appliedConstraint.entities; 533 534 auto condition = constraint.getConditionTemplate(); 535 if (isa<TypeConstraint>(constraint)) { 536 auto self = formatv("({0}.getType())", 537 symbolInfoMap.getValueAndRangeUse(entities.front())); 538 emitMatchCheck( 539 opName, tgfmt(condition, &fmtCtx.withSelf(self.str())), 540 formatv("\"value entity '{0}' failed to satisfy constraint: {1}\"", 541 entities.front(), constraint.getSummary())); 542 543 } else if (isa<AttrConstraint>(constraint)) { 544 PrintFatalError( 545 loc, "cannot use AttrConstraint in Pattern multi-entity constraints"); 546 } else { 547 // TODO: replace formatv arguments with the exact specified 548 // args. 549 if (entities.size() > 4) { 550 PrintFatalError(loc, "only support up to 4-entity constraints now"); 551 } 552 SmallVector<std::string, 4> names; 553 int i = 0; 554 for (int e = entities.size(); i < e; ++i) 555 names.push_back(symbolInfoMap.getValueAndRangeUse(entities[i])); 556 std::string self = appliedConstraint.self; 557 if (!self.empty()) 558 self = symbolInfoMap.getValueAndRangeUse(self); 559 for (; i < 4; ++i) 560 names.push_back("<unused>"); 561 emitMatchCheck(opName, 562 tgfmt(condition, &fmtCtx.withSelf(self), names[0], 563 names[1], names[2], names[3]), 564 formatv("\"entities '{0}' failed to satisfy constraint: " 565 "{1}\"", 566 llvm::join(entities, ", "), 567 constraint.getSummary())); 568 } 569 } 570 571 // Some of the operands could be bound to the same symbol name, we need 572 // to enforce equality constraint on those. 573 // TODO: we should be able to emit equality checks early 574 // and short circuit unnecessary work if vars are not equal. 575 for (auto symbolInfoIt = symbolInfoMap.begin(); 576 symbolInfoIt != symbolInfoMap.end();) { 577 auto range = symbolInfoMap.getRangeOfEqualElements(symbolInfoIt->first); 578 auto startRange = range.first; 579 auto endRange = range.second; 580 581 auto firstOperand = symbolInfoIt->second.getVarName(symbolInfoIt->first); 582 for (++startRange; startRange != endRange; ++startRange) { 583 auto secondOperand = startRange->second.getVarName(symbolInfoIt->first); 584 emitMatchCheck( 585 opName, 586 formatv("*{0}.begin() == *{1}.begin()", firstOperand, secondOperand), 587 formatv("\"Operands '{0}' and '{1}' must be equal\"", firstOperand, 588 secondOperand)); 589 } 590 591 symbolInfoIt = endRange; 592 } 593 594 LLVM_DEBUG(llvm::dbgs() << "--- done emitting match logic ---\n"); 595 } 596 597 void PatternEmitter::collectOps(DagNode tree, 598 llvm::SmallPtrSetImpl<const Operator *> &ops) { 599 // Check if this tree is an operation. 600 if (tree.isOperation()) { 601 const Operator &op = tree.getDialectOp(opMap); 602 LLVM_DEBUG(llvm::dbgs() 603 << "found operation " << op.getOperationName() << '\n'); 604 ops.insert(&op); 605 } 606 607 // Recurse the arguments of the tree. 608 for (unsigned i = 0, e = tree.getNumArgs(); i != e; ++i) 609 if (auto child = tree.getArgAsNestedDag(i)) 610 collectOps(child, ops); 611 } 612 613 void PatternEmitter::emit(StringRef rewriteName) { 614 // Get the DAG tree for the source pattern. 615 DagNode sourceTree = pattern.getSourcePattern(); 616 617 const Operator &rootOp = pattern.getSourceRootOp(); 618 auto rootName = rootOp.getOperationName(); 619 620 // Collect the set of result operations. 621 llvm::SmallPtrSet<const Operator *, 4> resultOps; 622 LLVM_DEBUG(llvm::dbgs() << "start collecting ops used in result patterns\n"); 623 for (unsigned i = 0, e = pattern.getNumResultPatterns(); i != e; ++i) { 624 collectOps(pattern.getResultPattern(i), resultOps); 625 } 626 LLVM_DEBUG(llvm::dbgs() << "done collecting ops used in result patterns\n"); 627 628 // Emit RewritePattern for Pattern. 629 auto locs = pattern.getLocation(); 630 os << formatv("/* Generated from:\n {0:$[ instantiating\n ]}\n*/\n", 631 make_range(locs.rbegin(), locs.rend())); 632 os << formatv(R"(struct {0} : public ::mlir::RewritePattern { 633 {0}(::mlir::MLIRContext *context) 634 : ::mlir::RewritePattern("{1}", {{)", 635 rewriteName, rootName); 636 // Sort result operators by name. 637 llvm::SmallVector<const Operator *, 4> sortedResultOps(resultOps.begin(), 638 resultOps.end()); 639 llvm::sort(sortedResultOps, [&](const Operator *lhs, const Operator *rhs) { 640 return lhs->getOperationName() < rhs->getOperationName(); 641 }); 642 llvm::interleaveComma(sortedResultOps, os, [&](const Operator *op) { 643 os << '"' << op->getOperationName() << '"'; 644 }); 645 os << formatv(R"(}, {0}, context) {{})", pattern.getBenefit()) << "\n"; 646 647 // Emit matchAndRewrite() function. 648 { 649 auto classScope = os.scope(); 650 os.reindent(R"( 651 ::mlir::LogicalResult matchAndRewrite(::mlir::Operation *op0, 652 ::mlir::PatternRewriter &rewriter) const override {)") 653 << '\n'; 654 { 655 auto functionScope = os.scope(); 656 657 // Register all symbols bound in the source pattern. 658 pattern.collectSourcePatternBoundSymbols(symbolInfoMap); 659 660 LLVM_DEBUG(llvm::dbgs() 661 << "start creating local variables for capturing matches\n"); 662 os << "// Variables for capturing values and attributes used while " 663 "creating ops\n"; 664 // Create local variables for storing the arguments and results bound 665 // to symbols. 666 for (const auto &symbolInfoPair : symbolInfoMap) { 667 const auto &symbol = symbolInfoPair.first; 668 const auto &info = symbolInfoPair.second; 669 670 os << info.getVarDecl(symbol); 671 } 672 // TODO: capture ops with consistent numbering so that it can be 673 // reused for fused loc. 674 os << formatv("::mlir::Operation *tblgen_ops[{0}];\n\n", 675 pattern.getSourcePattern().getNumOps()); 676 LLVM_DEBUG(llvm::dbgs() 677 << "done creating local variables for capturing matches\n"); 678 679 os << "// Match\n"; 680 os << "tblgen_ops[0] = op0;\n"; 681 emitMatchLogic(sourceTree, "op0"); 682 683 os << "\n// Rewrite\n"; 684 emitRewriteLogic(); 685 686 os << "return ::mlir::success();\n"; 687 } 688 os << "};\n"; 689 } 690 os << "};\n\n"; 691 } 692 693 void PatternEmitter::emitRewriteLogic() { 694 LLVM_DEBUG(llvm::dbgs() << "--- start emitting rewrite logic ---\n"); 695 const Operator &rootOp = pattern.getSourceRootOp(); 696 int numExpectedResults = rootOp.getNumResults(); 697 int numResultPatterns = pattern.getNumResultPatterns(); 698 699 // First register all symbols bound to ops generated in result patterns. 700 pattern.collectResultPatternBoundSymbols(symbolInfoMap); 701 702 // Only the last N static values generated are used to replace the matched 703 // root N-result op. We need to calculate the starting index (of the results 704 // of the matched op) each result pattern is to replace. 705 SmallVector<int, 4> offsets(numResultPatterns + 1, numExpectedResults); 706 // If we don't need to replace any value at all, set the replacement starting 707 // index as the number of result patterns so we skip all of them when trying 708 // to replace the matched op's results. 709 int replStartIndex = numExpectedResults == 0 ? numResultPatterns : -1; 710 for (int i = numResultPatterns - 1; i >= 0; --i) { 711 auto numValues = getNodeValueCount(pattern.getResultPattern(i)); 712 offsets[i] = offsets[i + 1] - numValues; 713 if (offsets[i] == 0) { 714 if (replStartIndex == -1) 715 replStartIndex = i; 716 } else if (offsets[i] < 0 && offsets[i + 1] > 0) { 717 auto error = formatv( 718 "cannot use the same multi-result op '{0}' to generate both " 719 "auxiliary values and values to be used for replacing the matched op", 720 pattern.getResultPattern(i).getSymbol()); 721 PrintFatalError(loc, error); 722 } 723 } 724 725 if (offsets.front() > 0) { 726 const char error[] = "no enough values generated to replace the matched op"; 727 PrintFatalError(loc, error); 728 } 729 730 os << "auto odsLoc = rewriter.getFusedLoc({"; 731 for (int i = 0, e = pattern.getSourcePattern().getNumOps(); i != e; ++i) { 732 os << (i ? ", " : "") << "tblgen_ops[" << i << "]->getLoc()"; 733 } 734 os << "}); (void)odsLoc;\n"; 735 736 // Process auxiliary result patterns. 737 for (int i = 0; i < replStartIndex; ++i) { 738 DagNode resultTree = pattern.getResultPattern(i); 739 auto val = handleResultPattern(resultTree, offsets[i], 0); 740 // Normal op creation will be streamed to `os` by the above call; but 741 // NativeCodeCall will only be materialized to `os` if it is used. Here 742 // we are handling auxiliary patterns so we want the side effect even if 743 // NativeCodeCall is not replacing matched root op's results. 744 if (resultTree.isNativeCodeCall()) 745 os << val << ";\n"; 746 } 747 748 if (numExpectedResults == 0) { 749 assert(replStartIndex >= numResultPatterns && 750 "invalid auxiliary vs. replacement pattern division!"); 751 // No result to replace. Just erase the op. 752 os << "rewriter.eraseOp(op0);\n"; 753 } else { 754 // Process replacement result patterns. 755 os << "::llvm::SmallVector<::mlir::Value, 4> tblgen_repl_values;\n"; 756 for (int i = replStartIndex; i < numResultPatterns; ++i) { 757 DagNode resultTree = pattern.getResultPattern(i); 758 auto val = handleResultPattern(resultTree, offsets[i], 0); 759 os << "\n"; 760 // Resolve each symbol for all range use so that we can loop over them. 761 // We need an explicit cast to `SmallVector` to capture the cases where 762 // `{0}` resolves to an `Operation::result_range` as well as cases that 763 // are not iterable (e.g. vector that gets wrapped in additional braces by 764 // RewriterGen). 765 // TODO: Revisit the need for materializing a vector. 766 os << symbolInfoMap.getAllRangeUse( 767 val, 768 "for (auto v: ::llvm::SmallVector<::mlir::Value, 4>{ {0} }) {{\n" 769 " tblgen_repl_values.push_back(v);\n}\n", 770 "\n"); 771 } 772 os << "\nrewriter.replaceOp(op0, tblgen_repl_values);\n"; 773 } 774 775 LLVM_DEBUG(llvm::dbgs() << "--- done emitting rewrite logic ---\n"); 776 } 777 778 std::string PatternEmitter::getUniqueSymbol(const Operator *op) { 779 return std::string( 780 formatv("tblgen_{0}_{1}", op->getCppClassName(), nextValueId++)); 781 } 782 783 std::string PatternEmitter::handleResultPattern(DagNode resultTree, 784 int resultIndex, int depth) { 785 LLVM_DEBUG(llvm::dbgs() << "handle result pattern: "); 786 LLVM_DEBUG(resultTree.print(llvm::dbgs())); 787 LLVM_DEBUG(llvm::dbgs() << '\n'); 788 789 if (resultTree.isLocationDirective()) { 790 PrintFatalError(loc, 791 "location directive can only be used with op creation"); 792 } 793 794 if (resultTree.isNativeCodeCall()) { 795 auto symbol = handleReplaceWithNativeCodeCall(resultTree, depth); 796 symbolInfoMap.bindValue(symbol); 797 return symbol; 798 } 799 800 if (resultTree.isReplaceWithValue()) 801 return handleReplaceWithValue(resultTree).str(); 802 803 // Normal op creation. 804 auto symbol = handleOpCreation(resultTree, resultIndex, depth); 805 if (resultTree.getSymbol().empty()) { 806 // This is an op not explicitly bound to a symbol in the rewrite rule. 807 // Register the auto-generated symbol for it. 808 symbolInfoMap.bindOpResult(symbol, pattern.getDialectOp(resultTree)); 809 } 810 return symbol; 811 } 812 813 StringRef PatternEmitter::handleReplaceWithValue(DagNode tree) { 814 assert(tree.isReplaceWithValue()); 815 816 if (tree.getNumArgs() != 1) { 817 PrintFatalError( 818 loc, "replaceWithValue directive must take exactly one argument"); 819 } 820 821 if (!tree.getSymbol().empty()) { 822 PrintFatalError(loc, "cannot bind symbol to replaceWithValue"); 823 } 824 825 return tree.getArgName(0); 826 } 827 828 std::string PatternEmitter::handleLocationDirective(DagNode tree) { 829 assert(tree.isLocationDirective()); 830 auto lookUpArgLoc = [this, &tree](int idx) { 831 const auto *const lookupFmt = "(*{0}.begin()).getLoc()"; 832 return symbolInfoMap.getAllRangeUse(tree.getArgName(idx), lookupFmt); 833 }; 834 835 if (tree.getNumArgs() == 0) 836 llvm::PrintFatalError( 837 "At least one argument to location directive required"); 838 839 if (!tree.getSymbol().empty()) 840 PrintFatalError(loc, "cannot bind symbol to location"); 841 842 if (tree.getNumArgs() == 1) { 843 DagLeaf leaf = tree.getArgAsLeaf(0); 844 if (leaf.isStringAttr()) 845 return formatv("::mlir::NameLoc::get(rewriter.getIdentifier(\"{0}\"))", 846 leaf.getStringAttr()) 847 .str(); 848 return lookUpArgLoc(0); 849 } 850 851 std::string ret; 852 llvm::raw_string_ostream os(ret); 853 std::string strAttr; 854 os << "rewriter.getFusedLoc({"; 855 bool first = true; 856 for (int i = 0, e = tree.getNumArgs(); i != e; ++i) { 857 DagLeaf leaf = tree.getArgAsLeaf(i); 858 // Handle the optional string value. 859 if (leaf.isStringAttr()) { 860 if (!strAttr.empty()) 861 llvm::PrintFatalError("Only one string attribute may be specified"); 862 strAttr = leaf.getStringAttr(); 863 continue; 864 } 865 os << (first ? "" : ", ") << lookUpArgLoc(i); 866 first = false; 867 } 868 os << "}"; 869 if (!strAttr.empty()) { 870 os << ", rewriter.getStringAttr(\"" << strAttr << "\")"; 871 } 872 os << ")"; 873 return os.str(); 874 } 875 876 std::string PatternEmitter::handleOpArgument(DagLeaf leaf, 877 StringRef patArgName) { 878 if (leaf.isStringAttr()) 879 PrintFatalError(loc, "raw string not supported as argument"); 880 if (leaf.isConstantAttr()) { 881 auto constAttr = leaf.getAsConstantAttr(); 882 return handleConstantAttr(constAttr.getAttribute(), 883 constAttr.getConstantValue()); 884 } 885 if (leaf.isEnumAttrCase()) { 886 auto enumCase = leaf.getAsEnumAttrCase(); 887 if (enumCase.isStrCase()) 888 return handleConstantAttr(enumCase, enumCase.getSymbol()); 889 // This is an enum case backed by an IntegerAttr. We need to get its value 890 // to build the constant. 891 std::string val = std::to_string(enumCase.getValue()); 892 return handleConstantAttr(enumCase, val); 893 } 894 895 LLVM_DEBUG(llvm::dbgs() << "handle argument '" << patArgName << "'\n"); 896 auto argName = symbolInfoMap.getValueAndRangeUse(patArgName); 897 if (leaf.isUnspecified() || leaf.isOperandMatcher()) { 898 LLVM_DEBUG(llvm::dbgs() << "replace " << patArgName << " with '" << argName 899 << "' (via symbol ref)\n"); 900 return argName; 901 } 902 if (leaf.isNativeCodeCall()) { 903 auto repl = tgfmt(leaf.getNativeCodeTemplate(), &fmtCtx.withSelf(argName)); 904 LLVM_DEBUG(llvm::dbgs() << "replace " << patArgName << " with '" << repl 905 << "' (via NativeCodeCall)\n"); 906 return std::string(repl); 907 } 908 PrintFatalError(loc, "unhandled case when rewriting op"); 909 } 910 911 std::string PatternEmitter::handleReplaceWithNativeCodeCall(DagNode tree, 912 int depth) { 913 LLVM_DEBUG(llvm::dbgs() << "handle NativeCodeCall pattern: "); 914 LLVM_DEBUG(tree.print(llvm::dbgs())); 915 LLVM_DEBUG(llvm::dbgs() << '\n'); 916 917 auto fmt = tree.getNativeCodeTemplate(); 918 // TODO: replace formatv arguments with the exact specified args. 919 SmallVector<std::string, 8> attrs(8); 920 if (tree.getNumArgs() > 8) { 921 PrintFatalError(loc, 922 "unsupported NativeCodeCall replace argument numbers: " + 923 Twine(tree.getNumArgs())); 924 } 925 bool hasLocationDirective; 926 std::string locToUse; 927 std::tie(hasLocationDirective, locToUse) = getLocation(tree); 928 929 for (int i = 0, e = tree.getNumArgs() - hasLocationDirective; i != e; ++i) { 930 if (tree.isNestedDagArg(i)) { 931 attrs[i] = handleResultPattern(tree.getArgAsNestedDag(i), i, depth + 1); 932 } else { 933 attrs[i] = handleOpArgument(tree.getArgAsLeaf(i), tree.getArgName(i)); 934 } 935 LLVM_DEBUG(llvm::dbgs() << "NativeCodeCall argument #" << i 936 << " replacement: " << attrs[i] << "\n"); 937 } 938 return std::string(tgfmt(fmt, &fmtCtx.addSubst("_loc", locToUse), attrs[0], 939 attrs[1], attrs[2], attrs[3], attrs[4], attrs[5], 940 attrs[6], attrs[7])); 941 } 942 943 int PatternEmitter::getNodeValueCount(DagNode node) { 944 if (node.isOperation()) { 945 // If the op is bound to a symbol in the rewrite rule, query its result 946 // count from the symbol info map. 947 auto symbol = node.getSymbol(); 948 if (!symbol.empty()) { 949 return symbolInfoMap.getStaticValueCount(symbol); 950 } 951 // Otherwise this is an unbound op; we will use all its results. 952 return pattern.getDialectOp(node).getNumResults(); 953 } 954 // TODO: This considers all NativeCodeCall as returning one 955 // value. Enhance if multi-value ones are needed. 956 return 1; 957 } 958 959 std::pair<bool, std::string> PatternEmitter::getLocation(DagNode tree) { 960 auto numPatArgs = tree.getNumArgs(); 961 962 if (numPatArgs != 0) { 963 if (auto lastArg = tree.getArgAsNestedDag(numPatArgs - 1)) 964 if (lastArg.isLocationDirective()) { 965 return std::make_pair(true, handleLocationDirective(lastArg)); 966 } 967 } 968 969 // If no explicit location is given, use the default, all fused, location. 970 return std::make_pair(false, "odsLoc"); 971 } 972 973 std::string PatternEmitter::handleOpCreation(DagNode tree, int resultIndex, 974 int depth) { 975 LLVM_DEBUG(llvm::dbgs() << "create op for pattern: "); 976 LLVM_DEBUG(tree.print(llvm::dbgs())); 977 LLVM_DEBUG(llvm::dbgs() << '\n'); 978 979 Operator &resultOp = tree.getDialectOp(opMap); 980 auto numOpArgs = resultOp.getNumArgs(); 981 auto numPatArgs = tree.getNumArgs(); 982 983 bool hasLocationDirective; 984 std::string locToUse; 985 std::tie(hasLocationDirective, locToUse) = getLocation(tree); 986 987 auto inPattern = numPatArgs - hasLocationDirective; 988 if (numOpArgs != inPattern) { 989 PrintFatalError(loc, 990 formatv("resultant op '{0}' argument number mismatch: " 991 "{1} in pattern vs. {2} in definition", 992 resultOp.getOperationName(), inPattern, numOpArgs)); 993 } 994 995 // A map to collect all nested DAG child nodes' names, with operand index as 996 // the key. This includes both bound and unbound child nodes. 997 ChildNodeIndexNameMap childNodeNames; 998 999 // First go through all the child nodes who are nested DAG constructs to 1000 // create ops for them and remember the symbol names for them, so that we can 1001 // use the results in the current node. This happens in a recursive manner. 1002 for (int i = 0, e = tree.getNumArgs() - hasLocationDirective; i != e; ++i) { 1003 if (auto child = tree.getArgAsNestedDag(i)) 1004 childNodeNames[i] = handleResultPattern(child, i, depth + 1); 1005 } 1006 1007 // The name of the local variable holding this op. 1008 std::string valuePackName; 1009 // The symbol for holding the result of this pattern. Note that the result of 1010 // this pattern is not necessarily the same as the variable created by this 1011 // pattern because we can use `__N` suffix to refer only a specific result if 1012 // the generated op is a multi-result op. 1013 std::string resultValue; 1014 if (tree.getSymbol().empty()) { 1015 // No symbol is explicitly bound to this op in the pattern. Generate a 1016 // unique name. 1017 valuePackName = resultValue = getUniqueSymbol(&resultOp); 1018 } else { 1019 resultValue = std::string(tree.getSymbol()); 1020 // Strip the index to get the name for the value pack and use it to name the 1021 // local variable for the op. 1022 valuePackName = std::string(SymbolInfoMap::getValuePackName(resultValue)); 1023 } 1024 1025 // Create the local variable for this op. 1026 os << formatv("{0} {1};\n{{\n", resultOp.getQualCppClassName(), 1027 valuePackName); 1028 1029 // Right now ODS don't have general type inference support. Except a few 1030 // special cases listed below, DRR needs to supply types for all results 1031 // when building an op. 1032 bool isSameOperandsAndResultType = 1033 resultOp.getTrait("::mlir::OpTrait::SameOperandsAndResultType"); 1034 bool useFirstAttr = 1035 resultOp.getTrait("::mlir::OpTrait::FirstAttrDerivedResultType"); 1036 1037 if (isSameOperandsAndResultType || useFirstAttr) { 1038 // We know how to deduce the result type for ops with these traits and we've 1039 // generated builders taking aggregate parameters. Use those builders to 1040 // create the ops. 1041 1042 // First prepare local variables for op arguments used in builder call. 1043 createAggregateLocalVarsForOpArgs(tree, childNodeNames, depth); 1044 1045 // Then create the op. 1046 os.scope("", "\n}\n").os << formatv( 1047 "{0} = rewriter.create<{1}>({2}, tblgen_values, tblgen_attrs);", 1048 valuePackName, resultOp.getQualCppClassName(), locToUse); 1049 return resultValue; 1050 } 1051 1052 bool usePartialResults = valuePackName != resultValue; 1053 1054 if (usePartialResults || depth > 0 || resultIndex < 0) { 1055 // For these cases (broadcastable ops, op results used both as auxiliary 1056 // values and replacement values, ops in nested patterns, auxiliary ops), we 1057 // still need to supply the result types when building the op. But because 1058 // we don't generate a builder automatically with ODS for them, it's the 1059 // developer's responsibility to make sure such a builder (with result type 1060 // deduction ability) exists. We go through the separate-parameter builder 1061 // here given that it's easier for developers to write compared to 1062 // aggregate-parameter builders. 1063 createSeparateLocalVarsForOpArgs(tree, childNodeNames); 1064 1065 os.scope().os << formatv("{0} = rewriter.create<{1}>({2}", valuePackName, 1066 resultOp.getQualCppClassName(), locToUse); 1067 supplyValuesForOpArgs(tree, childNodeNames, depth); 1068 os << "\n );\n}\n"; 1069 return resultValue; 1070 } 1071 1072 // If depth == 0 and resultIndex >= 0, it means we are replacing the values 1073 // generated from the source pattern root op. Then we can use the source 1074 // pattern's value types to determine the value type of the generated op 1075 // here. 1076 1077 // First prepare local variables for op arguments used in builder call. 1078 createAggregateLocalVarsForOpArgs(tree, childNodeNames, depth); 1079 1080 // Then prepare the result types. We need to specify the types for all 1081 // results. 1082 os.indent() << formatv("::mlir::SmallVector<::mlir::Type, 4> tblgen_types; " 1083 "(void)tblgen_types;\n"); 1084 int numResults = resultOp.getNumResults(); 1085 if (numResults != 0) { 1086 for (int i = 0; i < numResults; ++i) 1087 os << formatv("for (auto v: castedOp0.getODSResults({0})) {{\n" 1088 " tblgen_types.push_back(v.getType());\n}\n", 1089 resultIndex + i); 1090 } 1091 os << formatv("{0} = rewriter.create<{1}>({2}, tblgen_types, " 1092 "tblgen_values, tblgen_attrs);\n", 1093 valuePackName, resultOp.getQualCppClassName(), locToUse); 1094 os.unindent() << "}\n"; 1095 return resultValue; 1096 } 1097 1098 void PatternEmitter::createSeparateLocalVarsForOpArgs( 1099 DagNode node, ChildNodeIndexNameMap &childNodeNames) { 1100 Operator &resultOp = node.getDialectOp(opMap); 1101 1102 // Now prepare operands used for building this op: 1103 // * If the operand is non-variadic, we create a `Value` local variable. 1104 // * If the operand is variadic, we create a `SmallVector<Value>` local 1105 // variable. 1106 1107 int valueIndex = 0; // An index for uniquing local variable names. 1108 for (int argIndex = 0, e = resultOp.getNumArgs(); argIndex < e; ++argIndex) { 1109 const auto *operand = 1110 resultOp.getArg(argIndex).dyn_cast<NamedTypeConstraint *>(); 1111 // We do not need special handling for attributes. 1112 if (!operand) 1113 continue; 1114 1115 raw_indented_ostream::DelimitedScope scope(os); 1116 std::string varName; 1117 if (operand->isVariadic()) { 1118 varName = std::string(formatv("tblgen_values_{0}", valueIndex++)); 1119 os << formatv("::mlir::SmallVector<::mlir::Value, 4> {0};\n", varName); 1120 std::string range; 1121 if (node.isNestedDagArg(argIndex)) { 1122 range = childNodeNames[argIndex]; 1123 } else { 1124 range = std::string(node.getArgName(argIndex)); 1125 } 1126 // Resolve the symbol for all range use so that we have a uniform way of 1127 // capturing the values. 1128 range = symbolInfoMap.getValueAndRangeUse(range); 1129 os << formatv("for (auto v: {0}) {{\n {1}.push_back(v);\n}\n", range, 1130 varName); 1131 } else { 1132 varName = std::string(formatv("tblgen_value_{0}", valueIndex++)); 1133 os << formatv("::mlir::Value {0} = ", varName); 1134 if (node.isNestedDagArg(argIndex)) { 1135 os << symbolInfoMap.getValueAndRangeUse(childNodeNames[argIndex]); 1136 } else { 1137 DagLeaf leaf = node.getArgAsLeaf(argIndex); 1138 auto symbol = 1139 symbolInfoMap.getValueAndRangeUse(node.getArgName(argIndex)); 1140 if (leaf.isNativeCodeCall()) { 1141 os << std::string( 1142 tgfmt(leaf.getNativeCodeTemplate(), &fmtCtx.withSelf(symbol))); 1143 } else { 1144 os << symbol; 1145 } 1146 } 1147 os << ";\n"; 1148 } 1149 1150 // Update to use the newly created local variable for building the op later. 1151 childNodeNames[argIndex] = varName; 1152 } 1153 } 1154 1155 void PatternEmitter::supplyValuesForOpArgs( 1156 DagNode node, const ChildNodeIndexNameMap &childNodeNames, int depth) { 1157 Operator &resultOp = node.getDialectOp(opMap); 1158 for (int argIndex = 0, numOpArgs = resultOp.getNumArgs(); 1159 argIndex != numOpArgs; ++argIndex) { 1160 // Start each argument on its own line. 1161 os << ",\n "; 1162 1163 Argument opArg = resultOp.getArg(argIndex); 1164 // Handle the case of operand first. 1165 if (auto *operand = opArg.dyn_cast<NamedTypeConstraint *>()) { 1166 if (!operand->name.empty()) 1167 os << "/*" << operand->name << "=*/"; 1168 os << childNodeNames.lookup(argIndex); 1169 continue; 1170 } 1171 1172 // The argument in the op definition. 1173 auto opArgName = resultOp.getArgName(argIndex); 1174 if (auto subTree = node.getArgAsNestedDag(argIndex)) { 1175 if (!subTree.isNativeCodeCall()) 1176 PrintFatalError(loc, "only NativeCodeCall allowed in nested dag node " 1177 "for creating attribute"); 1178 os << formatv("/*{0}=*/{1}", opArgName, 1179 handleReplaceWithNativeCodeCall(subTree, depth)); 1180 } else { 1181 auto leaf = node.getArgAsLeaf(argIndex); 1182 // The argument in the result DAG pattern. 1183 auto patArgName = node.getArgName(argIndex); 1184 if (leaf.isConstantAttr() || leaf.isEnumAttrCase()) { 1185 // TODO: Refactor out into map to avoid recomputing these. 1186 if (!opArg.is<NamedAttribute *>()) 1187 PrintFatalError(loc, Twine("expected attribute ") + Twine(argIndex)); 1188 if (!patArgName.empty()) 1189 os << "/*" << patArgName << "=*/"; 1190 } else { 1191 os << "/*" << opArgName << "=*/"; 1192 } 1193 os << handleOpArgument(leaf, patArgName); 1194 } 1195 } 1196 } 1197 1198 void PatternEmitter::createAggregateLocalVarsForOpArgs( 1199 DagNode node, const ChildNodeIndexNameMap &childNodeNames, int depth) { 1200 Operator &resultOp = node.getDialectOp(opMap); 1201 1202 auto scope = os.scope(); 1203 os << formatv("::mlir::SmallVector<::mlir::Value, 4> " 1204 "tblgen_values; (void)tblgen_values;\n"); 1205 os << formatv("::mlir::SmallVector<::mlir::NamedAttribute, 4> " 1206 "tblgen_attrs; (void)tblgen_attrs;\n"); 1207 1208 const char *addAttrCmd = 1209 "if (auto tmpAttr = {1}) {\n" 1210 " tblgen_attrs.emplace_back(rewriter.getIdentifier(\"{0}\"), " 1211 "tmpAttr);\n}\n"; 1212 for (int argIndex = 0, e = resultOp.getNumArgs(); argIndex < e; ++argIndex) { 1213 if (resultOp.getArg(argIndex).is<NamedAttribute *>()) { 1214 // The argument in the op definition. 1215 auto opArgName = resultOp.getArgName(argIndex); 1216 if (auto subTree = node.getArgAsNestedDag(argIndex)) { 1217 if (!subTree.isNativeCodeCall()) 1218 PrintFatalError(loc, "only NativeCodeCall allowed in nested dag node " 1219 "for creating attribute"); 1220 os << formatv(addAttrCmd, opArgName, 1221 handleReplaceWithNativeCodeCall(subTree, depth + 1)); 1222 } else { 1223 auto leaf = node.getArgAsLeaf(argIndex); 1224 // The argument in the result DAG pattern. 1225 auto patArgName = node.getArgName(argIndex); 1226 os << formatv(addAttrCmd, opArgName, 1227 handleOpArgument(leaf, patArgName)); 1228 } 1229 continue; 1230 } 1231 1232 const auto *operand = 1233 resultOp.getArg(argIndex).get<NamedTypeConstraint *>(); 1234 std::string varName; 1235 if (operand->isVariadic()) { 1236 std::string range; 1237 if (node.isNestedDagArg(argIndex)) { 1238 range = childNodeNames.lookup(argIndex); 1239 } else { 1240 range = std::string(node.getArgName(argIndex)); 1241 } 1242 // Resolve the symbol for all range use so that we have a uniform way of 1243 // capturing the values. 1244 range = symbolInfoMap.getValueAndRangeUse(range); 1245 os << formatv("for (auto v: {0}) {{\n tblgen_values.push_back(v);\n}\n", 1246 range); 1247 } else { 1248 os << formatv("tblgen_values.push_back("); 1249 if (node.isNestedDagArg(argIndex)) { 1250 os << symbolInfoMap.getValueAndRangeUse( 1251 childNodeNames.lookup(argIndex)); 1252 } else { 1253 DagLeaf leaf = node.getArgAsLeaf(argIndex); 1254 auto symbol = 1255 symbolInfoMap.getValueAndRangeUse(node.getArgName(argIndex)); 1256 if (leaf.isNativeCodeCall()) { 1257 os << std::string( 1258 tgfmt(leaf.getNativeCodeTemplate(), &fmtCtx.withSelf(symbol))); 1259 } else { 1260 os << symbol; 1261 } 1262 } 1263 os << ");\n"; 1264 } 1265 } 1266 } 1267 1268 static void emitRewriters(const RecordKeeper &recordKeeper, raw_ostream &os) { 1269 emitSourceFileHeader("Rewriters", os); 1270 1271 const auto &patterns = recordKeeper.getAllDerivedDefinitions("Pattern"); 1272 auto numPatterns = patterns.size(); 1273 1274 // We put the map here because it can be shared among multiple patterns. 1275 RecordOperatorMap recordOpMap; 1276 1277 std::vector<std::string> rewriterNames; 1278 rewriterNames.reserve(numPatterns); 1279 1280 std::string baseRewriterName = "GeneratedConvert"; 1281 int rewriterIndex = 0; 1282 1283 for (Record *p : patterns) { 1284 std::string name; 1285 if (p->isAnonymous()) { 1286 // If no name is provided, ensure unique rewriter names simply by 1287 // appending unique suffix. 1288 name = baseRewriterName + llvm::utostr(rewriterIndex++); 1289 } else { 1290 name = std::string(p->getName()); 1291 } 1292 LLVM_DEBUG(llvm::dbgs() 1293 << "=== start generating pattern '" << name << "' ===\n"); 1294 PatternEmitter(p, &recordOpMap, os).emit(name); 1295 LLVM_DEBUG(llvm::dbgs() 1296 << "=== done generating pattern '" << name << "' ===\n"); 1297 rewriterNames.push_back(std::move(name)); 1298 } 1299 1300 // Emit function to add the generated matchers to the pattern list. 1301 os << "void LLVM_ATTRIBUTE_UNUSED populateWithGenerated(::mlir::MLIRContext " 1302 "*context, ::mlir::OwningRewritePatternList &patterns) {\n"; 1303 for (const auto &name : rewriterNames) { 1304 os << " patterns.insert<" << name << ">(context);\n"; 1305 } 1306 os << "}\n"; 1307 } 1308 1309 static mlir::GenRegistration 1310 genRewriters("gen-rewriters", "Generate pattern rewriters", 1311 [](const RecordKeeper &records, raw_ostream &os) { 1312 emitRewriters(records, os); 1313 return false; 1314 }); 1315