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