1 //===- Pattern.cpp - Pattern wrapper class --------------------------------===//
2 //
3 // Copyright 2019 The MLIR Authors.
4 //
5 // Licensed under the Apache License, Version 2.0 (the "License");
6 // you may not use this file except in compliance with the License.
7 // You may obtain a copy of the License at
8 //
9 //   http://www.apache.org/licenses/LICENSE-2.0
10 //
11 // Unless required by applicable law or agreed to in writing, software
12 // distributed under the License is distributed on an "AS IS" BASIS,
13 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 // See the License for the specific language governing permissions and
15 // limitations under the License.
16 // =============================================================================
17 //
18 // Pattern wrapper class to simplify using TableGen Record defining a MLIR
19 // Pattern.
20 //
21 //===----------------------------------------------------------------------===//
22 
23 #include "mlir/TableGen/Pattern.h"
24 #include "llvm/ADT/Twine.h"
25 #include "llvm/Support/FormatVariadic.h"
26 #include "llvm/TableGen/Error.h"
27 #include "llvm/TableGen/Record.h"
28 
29 using namespace mlir;
30 
31 using llvm::formatv;
32 using mlir::tblgen::Operator;
33 
34 //===----------------------------------------------------------------------===//
35 // DagLeaf
36 //===----------------------------------------------------------------------===//
37 
38 bool tblgen::DagLeaf::isUnspecified() const {
39   return dyn_cast_or_null<llvm::UnsetInit>(def);
40 }
41 
42 bool tblgen::DagLeaf::isOperandMatcher() const {
43   // Operand matchers specify a type constraint.
44   return isSubClassOf("TypeConstraint");
45 }
46 
47 bool tblgen::DagLeaf::isAttrMatcher() const {
48   // Attribute matchers specify an attribute constraint.
49   return isSubClassOf("AttrConstraint");
50 }
51 
52 bool tblgen::DagLeaf::isNativeCodeCall() const {
53   return isSubClassOf("NativeCodeCall");
54 }
55 
56 bool tblgen::DagLeaf::isConstantAttr() const {
57   return isSubClassOf("ConstantAttr");
58 }
59 
60 bool tblgen::DagLeaf::isEnumAttrCase() const {
61   return isSubClassOf("EnumAttrCaseInfo");
62 }
63 
64 tblgen::Constraint tblgen::DagLeaf::getAsConstraint() const {
65   assert((isOperandMatcher() || isAttrMatcher()) &&
66          "the DAG leaf must be operand or attribute");
67   return Constraint(cast<llvm::DefInit>(def)->getDef());
68 }
69 
70 tblgen::ConstantAttr tblgen::DagLeaf::getAsConstantAttr() const {
71   assert(isConstantAttr() && "the DAG leaf must be constant attribute");
72   return ConstantAttr(cast<llvm::DefInit>(def));
73 }
74 
75 tblgen::EnumAttrCase tblgen::DagLeaf::getAsEnumAttrCase() const {
76   assert(isEnumAttrCase() && "the DAG leaf must be an enum attribute case");
77   return EnumAttrCase(cast<llvm::DefInit>(def));
78 }
79 
80 std::string tblgen::DagLeaf::getConditionTemplate() const {
81   return getAsConstraint().getConditionTemplate();
82 }
83 
84 llvm::StringRef tblgen::DagLeaf::getNativeCodeTemplate() const {
85   assert(isNativeCodeCall() && "the DAG leaf must be NativeCodeCall");
86   return cast<llvm::DefInit>(def)->getDef()->getValueAsString("expression");
87 }
88 
89 bool tblgen::DagLeaf::isSubClassOf(StringRef superclass) const {
90   if (auto *defInit = dyn_cast_or_null<llvm::DefInit>(def))
91     return defInit->getDef()->isSubClassOf(superclass);
92   return false;
93 }
94 
95 //===----------------------------------------------------------------------===//
96 // DagNode
97 //===----------------------------------------------------------------------===//
98 
99 bool tblgen::DagNode::isNativeCodeCall() const {
100   if (auto *defInit = dyn_cast_or_null<llvm::DefInit>(node->getOperator()))
101     return defInit->getDef()->isSubClassOf("NativeCodeCall");
102   return false;
103 }
104 
105 bool tblgen::DagNode::isOperation() const {
106   return !(isNativeCodeCall() || isReplaceWithValue());
107 }
108 
109 llvm::StringRef tblgen::DagNode::getNativeCodeTemplate() const {
110   assert(isNativeCodeCall() && "the DAG leaf must be NativeCodeCall");
111   return cast<llvm::DefInit>(node->getOperator())
112       ->getDef()
113       ->getValueAsString("expression");
114 }
115 
116 llvm::StringRef tblgen::DagNode::getSymbol() const {
117   return node->getNameStr();
118 }
119 
120 Operator &tblgen::DagNode::getDialectOp(RecordOperatorMap *mapper) const {
121   llvm::Record *opDef = cast<llvm::DefInit>(node->getOperator())->getDef();
122   auto it = mapper->find(opDef);
123   if (it != mapper->end())
124     return *it->second;
125   return *mapper->try_emplace(opDef, std::make_unique<Operator>(opDef))
126               .first->second;
127 }
128 
129 int tblgen::DagNode::getNumOps() const {
130   int count = isReplaceWithValue() ? 0 : 1;
131   for (int i = 0, e = getNumArgs(); i != e; ++i) {
132     if (auto child = getArgAsNestedDag(i))
133       count += child.getNumOps();
134   }
135   return count;
136 }
137 
138 int tblgen::DagNode::getNumArgs() const { return node->getNumArgs(); }
139 
140 bool tblgen::DagNode::isNestedDagArg(unsigned index) const {
141   return isa<llvm::DagInit>(node->getArg(index));
142 }
143 
144 tblgen::DagNode tblgen::DagNode::getArgAsNestedDag(unsigned index) const {
145   return DagNode(dyn_cast_or_null<llvm::DagInit>(node->getArg(index)));
146 }
147 
148 tblgen::DagLeaf tblgen::DagNode::getArgAsLeaf(unsigned index) const {
149   assert(!isNestedDagArg(index));
150   return DagLeaf(node->getArg(index));
151 }
152 
153 StringRef tblgen::DagNode::getArgName(unsigned index) const {
154   return node->getArgNameStr(index);
155 }
156 
157 bool tblgen::DagNode::isReplaceWithValue() const {
158   auto *dagOpDef = cast<llvm::DefInit>(node->getOperator())->getDef();
159   return dagOpDef->getName() == "replaceWithValue";
160 }
161 
162 //===----------------------------------------------------------------------===//
163 // SymbolInfoMap
164 //===----------------------------------------------------------------------===//
165 
166 StringRef tblgen::SymbolInfoMap::getValuePackName(StringRef symbol,
167                                                   int *index) {
168   StringRef name, indexStr;
169   int idx = -1;
170   std::tie(name, indexStr) = symbol.rsplit("__");
171 
172   if (indexStr.consumeInteger(10, idx)) {
173     // The second part is not an index; we return the whole symbol as-is.
174     return symbol;
175   }
176   if (index) {
177     *index = idx;
178   }
179   return name;
180 }
181 
182 tblgen::SymbolInfoMap::SymbolInfo::SymbolInfo(const Operator *op,
183                                               SymbolInfo::Kind kind,
184                                               Optional<int> index)
185     : op(op), kind(kind), argIndex(index) {}
186 
187 int tblgen::SymbolInfoMap::SymbolInfo::getStaticValueCount() const {
188   switch (kind) {
189   case Kind::Attr:
190   case Kind::Operand:
191   case Kind::Value:
192     return 1;
193   case Kind::Result:
194     return op->getNumResults();
195   }
196   llvm_unreachable("unknown kind");
197 }
198 
199 std::string
200 tblgen::SymbolInfoMap::SymbolInfo::getVarDecl(StringRef name) const {
201   switch (kind) {
202   case Kind::Attr: {
203     auto type =
204         op->getArg(*argIndex).get<NamedAttribute *>()->attr.getStorageType();
205     return formatv("{0} {1};\n", type, name);
206   }
207   case Kind::Operand:
208   case Kind::Value: {
209     return formatv("Value *{0};\n", name);
210   }
211   case Kind::Result: {
212     // Use the op itself for the results.
213     return formatv("{0} {1};\n", op->getQualCppClassName(), name);
214   }
215   }
216   llvm_unreachable("unknown kind");
217 }
218 
219 std::string
220 tblgen::SymbolInfoMap::SymbolInfo::getValueAndRangeUse(StringRef name,
221                                                        int index) const {
222   switch (kind) {
223   case Kind::Attr:
224   case Kind::Operand: {
225     assert(index < 0 && "only allowed for symbol bound to result");
226     return name;
227   }
228   case Kind::Result: {
229     // TODO(b/133341698): The following is incorrect for variadic results. We
230     // should use getODSResults().
231     if (index >= 0) {
232       return formatv("{0}.getOperation()->getResult({1})", name, index);
233     }
234 
235     // If referencing multiple results, compose a comma-separated list.
236     SmallVector<std::string, 4> values;
237     for (int i = 0, e = op->getNumResults(); i < e; ++i) {
238       values.push_back(formatv("{0}.getOperation()->getResult({1})", name, i));
239     }
240     return llvm::join(values, ", ");
241   }
242   case Kind::Value: {
243     assert(index < 0 && "only allowed for symbol bound to result");
244     assert(op == nullptr);
245     return name;
246   }
247   }
248   llvm_unreachable("unknown kind");
249 }
250 
251 bool tblgen::SymbolInfoMap::bindOpArgument(StringRef symbol, const Operator &op,
252                                            int argIndex) {
253   StringRef name = getValuePackName(symbol);
254   if (name != symbol) {
255     auto error = formatv(
256         "symbol '{0}' with trailing index cannot bind to op argument", symbol);
257     PrintFatalError(loc, error);
258   }
259 
260   auto symInfo = op.getArg(argIndex).is<NamedAttribute *>()
261                      ? SymbolInfo::getAttr(&op, argIndex)
262                      : SymbolInfo::getOperand(&op, argIndex);
263 
264   return symbolInfoMap.insert({symbol, symInfo}).second;
265 }
266 
267 bool tblgen::SymbolInfoMap::bindOpResult(StringRef symbol, const Operator &op) {
268   StringRef name = getValuePackName(symbol);
269   return symbolInfoMap.insert({name, SymbolInfo::getResult(&op)}).second;
270 }
271 
272 bool tblgen::SymbolInfoMap::bindValue(StringRef symbol) {
273   return symbolInfoMap.insert({symbol, SymbolInfo::getValue()}).second;
274 }
275 
276 bool tblgen::SymbolInfoMap::contains(StringRef symbol) const {
277   return find(symbol) != symbolInfoMap.end();
278 }
279 
280 tblgen::SymbolInfoMap::const_iterator
281 tblgen::SymbolInfoMap::find(StringRef key) const {
282   StringRef name = getValuePackName(key);
283   return symbolInfoMap.find(name);
284 }
285 
286 int tblgen::SymbolInfoMap::getStaticValueCount(StringRef symbol) const {
287   StringRef name = getValuePackName(symbol);
288   if (name != symbol) {
289     // If there is a trailing index inside symbol, it references just one
290     // static value.
291     return 1;
292   }
293   // Otherwise, find how many it represents by querying the symbol's info.
294   return find(name)->getValue().getStaticValueCount();
295 }
296 
297 std::string tblgen::SymbolInfoMap::getValueAndRangeUse(StringRef symbol) const {
298   int index = -1;
299   StringRef name = getValuePackName(symbol, &index);
300 
301   auto it = symbolInfoMap.find(name);
302   if (it == symbolInfoMap.end()) {
303     auto error = formatv("referencing unbound symbol '{0}'", symbol);
304     PrintFatalError(loc, error);
305   }
306 
307   return it->getValue().getValueAndRangeUse(name, index);
308 }
309 
310 //===----------------------------------------------------------------------===//
311 // Pattern
312 //==----------------------------------------------------------------------===//
313 
314 tblgen::Pattern::Pattern(const llvm::Record *def, RecordOperatorMap *mapper)
315     : def(*def), recordOpMap(mapper) {}
316 
317 tblgen::DagNode tblgen::Pattern::getSourcePattern() const {
318   return tblgen::DagNode(def.getValueAsDag("sourcePattern"));
319 }
320 
321 int tblgen::Pattern::getNumResultPatterns() const {
322   auto *results = def.getValueAsListInit("resultPatterns");
323   return results->size();
324 }
325 
326 tblgen::DagNode tblgen::Pattern::getResultPattern(unsigned index) const {
327   auto *results = def.getValueAsListInit("resultPatterns");
328   return tblgen::DagNode(cast<llvm::DagInit>(results->getElement(index)));
329 }
330 
331 void tblgen::Pattern::collectSourcePatternBoundSymbols(
332     tblgen::SymbolInfoMap &infoMap) {
333   collectBoundSymbols(getSourcePattern(), infoMap, /*isSrcPattern=*/true);
334 }
335 
336 void tblgen::Pattern::collectResultPatternBoundSymbols(
337     tblgen::SymbolInfoMap &infoMap) {
338   for (int i = 0, e = getNumResultPatterns(); i < e; ++i) {
339     auto pattern = getResultPattern(i);
340     collectBoundSymbols(pattern, infoMap, /*isSrcPattern=*/false);
341   }
342 }
343 
344 const tblgen::Operator &tblgen::Pattern::getSourceRootOp() {
345   return getSourcePattern().getDialectOp(recordOpMap);
346 }
347 
348 tblgen::Operator &tblgen::Pattern::getDialectOp(DagNode node) {
349   return node.getDialectOp(recordOpMap);
350 }
351 
352 std::vector<tblgen::AppliedConstraint> tblgen::Pattern::getConstraints() const {
353   auto *listInit = def.getValueAsListInit("constraints");
354   std::vector<tblgen::AppliedConstraint> ret;
355   ret.reserve(listInit->size());
356 
357   for (auto it : *listInit) {
358     auto *dagInit = dyn_cast<llvm::DagInit>(it);
359     if (!dagInit)
360       PrintFatalError(def.getLoc(), "all elemements in Pattern multi-entity "
361                                     "constraints should be DAG nodes");
362 
363     std::vector<std::string> entities;
364     entities.reserve(dagInit->arg_size());
365     for (auto *argName : dagInit->getArgNames())
366       entities.push_back(argName->getValue());
367 
368     ret.emplace_back(cast<llvm::DefInit>(dagInit->getOperator())->getDef(),
369                      dagInit->getNameStr(), std::move(entities));
370   }
371   return ret;
372 }
373 
374 int tblgen::Pattern::getBenefit() const {
375   // The initial benefit value is a heuristic with number of ops in the source
376   // pattern.
377   int initBenefit = getSourcePattern().getNumOps();
378   llvm::DagInit *delta = def.getValueAsDag("benefitDelta");
379   if (delta->getNumArgs() != 1 || !isa<llvm::IntInit>(delta->getArg(0))) {
380     PrintFatalError(def.getLoc(),
381                     "The 'addBenefit' takes and only takes one integer value");
382   }
383   return initBenefit + dyn_cast<llvm::IntInit>(delta->getArg(0))->getValue();
384 }
385 
386 std::vector<tblgen::Pattern::IdentifierLine>
387 tblgen::Pattern::getLocation() const {
388   std::vector<std::pair<StringRef, unsigned>> result;
389   result.reserve(def.getLoc().size());
390   for (auto loc : def.getLoc()) {
391     unsigned buf = llvm::SrcMgr.FindBufferContainingLoc(loc);
392     assert(buf && "invalid source location");
393     result.emplace_back(
394         llvm::SrcMgr.getBufferInfo(buf).Buffer->getBufferIdentifier(),
395         llvm::SrcMgr.getLineAndColumn(loc, buf).first);
396   }
397   return result;
398 }
399 
400 void tblgen::Pattern::collectBoundSymbols(DagNode tree, SymbolInfoMap &infoMap,
401                                           bool isSrcPattern) {
402   auto treeName = tree.getSymbol();
403   if (!tree.isOperation()) {
404     if (!treeName.empty()) {
405       PrintFatalError(
406           def.getLoc(),
407           formatv("binding symbol '{0}' to non-operation unsupported right now",
408                   treeName));
409     }
410     return;
411   }
412 
413   auto &op = getDialectOp(tree);
414   auto numOpArgs = op.getNumArgs();
415   auto numTreeArgs = tree.getNumArgs();
416 
417   if (numOpArgs != numTreeArgs) {
418     auto err = formatv("op '{0}' argument number mismatch: "
419                        "{1} in pattern vs. {2} in definition",
420                        op.getOperationName(), numTreeArgs, numOpArgs);
421     PrintFatalError(def.getLoc(), err);
422   }
423 
424   // The name attached to the DAG node's operator is for representing the
425   // results generated from this op. It should be remembered as bound results.
426   if (!treeName.empty()) {
427     if (!infoMap.bindOpResult(treeName, op))
428       PrintFatalError(def.getLoc(),
429                       formatv("symbol '{0}' bound more than once", treeName));
430   }
431 
432   for (int i = 0; i != numTreeArgs; ++i) {
433     if (auto treeArg = tree.getArgAsNestedDag(i)) {
434       // This DAG node argument is a DAG node itself. Go inside recursively.
435       collectBoundSymbols(treeArg, infoMap, isSrcPattern);
436     } else if (isSrcPattern) {
437       // We can only bind symbols to op arguments in source pattern. Those
438       // symbols are referenced in result patterns.
439       auto treeArgName = tree.getArgName(i);
440       if (!treeArgName.empty()) {
441         if (!infoMap.bindOpArgument(treeArgName, op, i)) {
442           auto err = formatv("symbol '{0}' bound more than once", treeArgName);
443           PrintFatalError(def.getLoc(), err);
444         }
445       }
446     }
447   }
448 }
449