brintos

brintos / llvm-project-archived public Read only

0
0
Text · 81.8 KiB · 08d6483 Raw
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