1 //===- PDL.cpp - Pattern Descriptor Language Dialect ----------------------===//
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 #include "mlir/Dialect/PDL/IR/PDL.h"
10 #include "mlir/Dialect/PDL/IR/PDLOps.h"
11 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
12 #include "mlir/IR/BuiltinTypes.h"
13 #include "mlir/Interfaces/InferTypeOpInterface.h"
14 #include "llvm/ADT/DenseSet.h"
15 #include "llvm/ADT/TypeSwitch.h"
16 
17 using namespace mlir;
18 using namespace mlir::pdl;
19 
20 #include "mlir/Dialect/PDL/IR/PDLOpsDialect.cpp.inc"
21 
22 //===----------------------------------------------------------------------===//
23 // PDLDialect
24 //===----------------------------------------------------------------------===//
25 
26 void PDLDialect::initialize() {
27   addOperations<
28 #define GET_OP_LIST
29 #include "mlir/Dialect/PDL/IR/PDLOps.cpp.inc"
30       >();
31   registerTypes();
32 }
33 
34 //===----------------------------------------------------------------------===//
35 // PDL Operations
36 //===----------------------------------------------------------------------===//
37 
38 /// Returns true if the given operation is used by a "binding" pdl operation.
39 static bool hasBindingUse(Operation *op) {
40   for (Operation *user : op->getUsers())
41     // A result by itself is not binding, it must also be bound.
42     if (!isa<ResultOp, ResultsOp>(user) || hasBindingUse(user))
43       return true;
44   return false;
45 }
46 
47 /// Returns success if the given operation is not in the main matcher body or
48 /// is used by a "binding" operation. On failure, emits an error.
49 static LogicalResult verifyHasBindingUse(Operation *op) {
50   // If the parent is not a pattern, there is nothing to do.
51   if (!isa<PatternOp>(op->getParentOp()))
52     return success();
53   if (hasBindingUse(op))
54     return success();
55   return op->emitOpError(
56       "expected a bindable user when defined in the matcher body of a "
57       "`pdl.pattern`");
58 }
59 
60 /// Visits all the pdl.operand(s), pdl.result(s), and pdl.operation(s)
61 /// connected to the given operation.
62 static void visit(Operation *op, DenseSet<Operation *> &visited) {
63   // If the parent is not a pattern, there is nothing to do.
64   if (!isa<PatternOp>(op->getParentOp()) || isa<RewriteOp>(op))
65     return;
66 
67   // Ignore if already visited.
68   if (visited.contains(op))
69     return;
70 
71   // Mark as visited.
72   visited.insert(op);
73 
74   // Traverse the operands / parent.
75   TypeSwitch<Operation *>(op)
76       .Case<OperationOp>([&visited](auto operation) {
77         for (Value operand : operation.operands())
78           visit(operand.getDefiningOp(), visited);
79       })
80       .Case<ResultOp, ResultsOp>([&visited](auto result) {
81         visit(result.parent().getDefiningOp(), visited);
82       });
83 
84   // Traverse the users.
85   for (Operation *user : op->getUsers())
86     visit(user, visited);
87 }
88 
89 //===----------------------------------------------------------------------===//
90 // pdl::ApplyNativeConstraintOp
91 //===----------------------------------------------------------------------===//
92 
93 static LogicalResult verify(ApplyNativeConstraintOp op) {
94   if (op.getNumOperands() == 0)
95     return op.emitOpError("expected at least one argument");
96   return success();
97 }
98 
99 //===----------------------------------------------------------------------===//
100 // pdl::ApplyNativeRewriteOp
101 //===----------------------------------------------------------------------===//
102 
103 static LogicalResult verify(ApplyNativeRewriteOp op) {
104   if (op.getNumOperands() == 0 && op.getNumResults() == 0)
105     return op.emitOpError("expected at least one argument or result");
106   return success();
107 }
108 
109 //===----------------------------------------------------------------------===//
110 // pdl::AttributeOp
111 //===----------------------------------------------------------------------===//
112 
113 static LogicalResult verify(AttributeOp op) {
114   Value attrType = op.type();
115   Optional<Attribute> attrValue = op.value();
116 
117   if (!attrValue && isa<RewriteOp>(op->getParentOp()))
118     return op.emitOpError("expected constant value when specified within a "
119                           "`pdl.rewrite`");
120   if (attrValue && attrType)
121     return op.emitOpError("expected only one of [`type`, `value`] to be set");
122   return verifyHasBindingUse(op);
123 }
124 
125 //===----------------------------------------------------------------------===//
126 // pdl::OperandOp
127 //===----------------------------------------------------------------------===//
128 
129 static LogicalResult verify(OperandOp op) { return verifyHasBindingUse(op); }
130 
131 //===----------------------------------------------------------------------===//
132 // pdl::OperandsOp
133 //===----------------------------------------------------------------------===//
134 
135 static LogicalResult verify(OperandsOp op) { return verifyHasBindingUse(op); }
136 
137 //===----------------------------------------------------------------------===//
138 // pdl::OperationOp
139 //===----------------------------------------------------------------------===//
140 
141 static ParseResult parseOperationOpAttributes(
142     OpAsmParser &p, SmallVectorImpl<OpAsmParser::OperandType> &attrOperands,
143     ArrayAttr &attrNamesAttr) {
144   Builder &builder = p.getBuilder();
145   SmallVector<Attribute, 4> attrNames;
146   if (succeeded(p.parseOptionalLBrace())) {
147     do {
148       StringAttr nameAttr;
149       OpAsmParser::OperandType operand;
150       if (p.parseAttribute(nameAttr) || p.parseEqual() ||
151           p.parseOperand(operand))
152         return failure();
153       attrNames.push_back(nameAttr);
154       attrOperands.push_back(operand);
155     } while (succeeded(p.parseOptionalComma()));
156     if (p.parseRBrace())
157       return failure();
158   }
159   attrNamesAttr = builder.getArrayAttr(attrNames);
160   return success();
161 }
162 
163 static void printOperationOpAttributes(OpAsmPrinter &p, OperationOp op,
164                                        OperandRange attrArgs,
165                                        ArrayAttr attrNames) {
166   if (attrNames.empty())
167     return;
168   p << " {";
169   interleaveComma(llvm::seq<int>(0, attrNames.size()), p,
170                   [&](int i) { p << attrNames[i] << " = " << attrArgs[i]; });
171   p << '}';
172 }
173 
174 /// Verifies that the result types of this operation, defined within a
175 /// `pdl.rewrite`, can be inferred.
176 static LogicalResult verifyResultTypesAreInferrable(OperationOp op,
177                                                     OperandRange resultTypes) {
178   // Functor that returns if the given use can be used to infer a type.
179   Block *rewriterBlock = op->getBlock();
180   auto canInferTypeFromUse = [&](OpOperand &use) {
181     // If the use is within a ReplaceOp and isn't the operation being replaced
182     // (i.e. is not the first operand of the replacement), we can infer a type.
183     ReplaceOp replOpUser = dyn_cast<ReplaceOp>(use.getOwner());
184     if (!replOpUser || use.getOperandNumber() == 0)
185       return false;
186     // Make sure the replaced operation was defined before this one.
187     Operation *replacedOp = replOpUser.operation().getDefiningOp();
188     return replacedOp->getBlock() != rewriterBlock ||
189            replacedOp->isBeforeInBlock(op);
190   };
191 
192   // Check to see if the uses of the operation itself can be used to infer
193   // types.
194   if (llvm::any_of(op.op().getUses(), canInferTypeFromUse))
195     return success();
196 
197   // Otherwise, make sure each of the types can be inferred.
198   for (auto it : llvm::enumerate(resultTypes)) {
199     Operation *resultTypeOp = it.value().getDefiningOp();
200     assert(resultTypeOp && "expected valid result type operation");
201 
202     // If the op was defined by a `apply_native_rewrite`, it is guaranteed to be
203     // usable.
204     if (isa<ApplyNativeRewriteOp>(resultTypeOp))
205       continue;
206 
207     // If the type operation was defined in the matcher and constrains the
208     // result of an input operation, it can be used.
209     auto constrainsInputOp = [rewriterBlock](Operation *user) {
210       return user->getBlock() != rewriterBlock && isa<OperationOp>(user);
211     };
212     if (TypeOp typeOp = dyn_cast<TypeOp>(resultTypeOp)) {
213       if (typeOp.type() || llvm::any_of(typeOp->getUsers(), constrainsInputOp))
214         continue;
215     } else if (TypesOp typeOp = dyn_cast<TypesOp>(resultTypeOp)) {
216       if (typeOp.types() || llvm::any_of(typeOp->getUsers(), constrainsInputOp))
217         continue;
218     }
219 
220     return op
221         .emitOpError("must have inferable or constrained result types when "
222                      "nested within `pdl.rewrite`")
223         .attachNote()
224         .append("result type #", it.index(), " was not constrained");
225   }
226   return success();
227 }
228 
229 static LogicalResult verify(OperationOp op) {
230   bool isWithinRewrite = isa<RewriteOp>(op->getParentOp());
231   if (isWithinRewrite && !op.name())
232     return op.emitOpError("must have an operation name when nested within "
233                           "a `pdl.rewrite`");
234   ArrayAttr attributeNames = op.attributeNames();
235   auto attributeValues = op.attributes();
236   if (attributeNames.size() != attributeValues.size()) {
237     return op.emitOpError()
238            << "expected the same number of attribute values and attribute "
239               "names, got "
240            << attributeNames.size() << " names and " << attributeValues.size()
241            << " values";
242   }
243 
244   // If the operation is within a rewrite body and doesn't have type inference,
245   // ensure that the result types can be resolved.
246   if (isWithinRewrite && !op.hasTypeInference()) {
247     if (failed(verifyResultTypesAreInferrable(op, op.types())))
248       return failure();
249   }
250 
251   return verifyHasBindingUse(op);
252 }
253 
254 bool OperationOp::hasTypeInference() {
255   Optional<StringRef> opName = name();
256   if (!opName)
257     return false;
258 
259   if (auto rInfo = RegisteredOperationName::lookup(*opName, getContext()))
260     return rInfo->hasInterface<InferTypeOpInterface>();
261   return false;
262 }
263 
264 //===----------------------------------------------------------------------===//
265 // pdl::PatternOp
266 //===----------------------------------------------------------------------===//
267 
268 static LogicalResult verify(PatternOp pattern) {
269   Region &body = pattern.body();
270   Operation *term = body.front().getTerminator();
271   auto rewrite_op = dyn_cast<RewriteOp>(term);
272   if (!rewrite_op) {
273     return pattern.emitOpError("expected body to terminate with `pdl.rewrite`")
274         .attachNote(term->getLoc())
275         .append("see terminator defined here");
276   }
277 
278   // Check that all values defined in the top-level pattern belong to the PDL
279   // dialect.
280   WalkResult result = body.walk([&](Operation *op) -> WalkResult {
281     if (!isa_and_nonnull<PDLDialect>(op->getDialect())) {
282       pattern
283           .emitOpError("expected only `pdl` operations within the pattern body")
284           .attachNote(op->getLoc())
285           .append("see non-`pdl` operation defined here");
286       return WalkResult::interrupt();
287     }
288     return WalkResult::advance();
289   });
290   if (result.wasInterrupted())
291     return failure();
292 
293   // Check that there is at least one operation.
294   if (body.front().getOps<OperationOp>().empty())
295     return pattern.emitOpError(
296         "the pattern must contain at least one `pdl.operation`");
297 
298   // Determine if the operations within the pdl.pattern form a connected
299   // component. This is determined by starting the search from the first
300   // operand/result/operation and visiting their users / parents / operands.
301   // We limit our attention to operations that have a user in pdl.rewrite,
302   // those that do not will be detected via other means (expected bindable
303   // user).
304   bool first = true;
305   DenseSet<Operation *> visited;
306   for (Operation &op : body.front()) {
307     // The following are the operations forming the connected component.
308     if (!isa<OperandOp, OperandsOp, ResultOp, ResultsOp, OperationOp>(op))
309       continue;
310 
311     // Determine if the operation has a user in `pdl.rewrite`.
312     bool hasUserInRewrite = false;
313     for (Operation *user : op.getUsers()) {
314       Region *region = user->getParentRegion();
315       if (isa<RewriteOp>(user) ||
316           (region && isa<RewriteOp>(region->getParentOp()))) {
317         hasUserInRewrite = true;
318         break;
319       }
320     }
321 
322     // If the operation does not have a user in `pdl.rewrite`, ignore it.
323     if (!hasUserInRewrite)
324       continue;
325 
326     if (first) {
327       // For the first operation, invoke visit.
328       visit(&op, visited);
329       first = false;
330     } else if (!visited.count(&op)) {
331       // For the subsequent operations, check if already visited.
332       return pattern
333           .emitOpError("the operations must form a connected component")
334           .attachNote(op.getLoc())
335           .append("see a disconnected value / operation here");
336     }
337   }
338 
339   return success();
340 }
341 
342 void PatternOp::build(OpBuilder &builder, OperationState &state,
343                       Optional<uint16_t> benefit, Optional<StringRef> name) {
344   build(builder, state, builder.getI16IntegerAttr(benefit ? *benefit : 0),
345         name ? builder.getStringAttr(*name) : StringAttr());
346   state.regions[0]->emplaceBlock();
347 }
348 
349 /// Returns the rewrite operation of this pattern.
350 RewriteOp PatternOp::getRewriter() {
351   return cast<RewriteOp>(body().front().getTerminator());
352 }
353 
354 //===----------------------------------------------------------------------===//
355 // pdl::ReplaceOp
356 //===----------------------------------------------------------------------===//
357 
358 static LogicalResult verify(ReplaceOp op) {
359   if (op.replOperation() && !op.replValues().empty())
360     return op.emitOpError() << "expected no replacement values to be provided"
361                                " when the replacement operation is present";
362   return success();
363 }
364 
365 //===----------------------------------------------------------------------===//
366 // pdl::ResultsOp
367 //===----------------------------------------------------------------------===//
368 
369 static ParseResult parseResultsValueType(OpAsmParser &p, IntegerAttr index,
370                                          Type &resultType) {
371   if (!index) {
372     resultType = RangeType::get(p.getBuilder().getType<ValueType>());
373     return success();
374   }
375   if (p.parseArrow() || p.parseType(resultType))
376     return failure();
377   return success();
378 }
379 
380 static void printResultsValueType(OpAsmPrinter &p, ResultsOp op,
381                                   IntegerAttr index, Type resultType) {
382   if (index)
383     p << " -> " << resultType;
384 }
385 
386 static LogicalResult verify(ResultsOp op) {
387   if (!op.index() && op.getType().isa<pdl::ValueType>()) {
388     return op.emitOpError() << "expected `pdl.range<value>` result type when "
389                                "no index is specified, but got: "
390                             << op.getType();
391   }
392   return success();
393 }
394 
395 //===----------------------------------------------------------------------===//
396 // pdl::RewriteOp
397 //===----------------------------------------------------------------------===//
398 
399 static LogicalResult verify(RewriteOp op) {
400   Region &rewriteRegion = op.body();
401 
402   // Handle the case where the rewrite is external.
403   if (op.name()) {
404     if (!rewriteRegion.empty()) {
405       return op.emitOpError()
406              << "expected rewrite region to be empty when rewrite is external";
407     }
408     return success();
409   }
410 
411   // Otherwise, check that the rewrite region only contains a single block.
412   if (rewriteRegion.empty()) {
413     return op.emitOpError() << "expected rewrite region to be non-empty if "
414                                "external name is not specified";
415   }
416 
417   // Check that no additional arguments were provided.
418   if (!op.externalArgs().empty()) {
419     return op.emitOpError() << "expected no external arguments when the "
420                                "rewrite is specified inline";
421   }
422   if (op.externalConstParams()) {
423     return op.emitOpError() << "expected no external constant parameters when "
424                                "the rewrite is specified inline";
425   }
426 
427   return success();
428 }
429 
430 //===----------------------------------------------------------------------===//
431 // pdl::TypeOp
432 //===----------------------------------------------------------------------===//
433 
434 static LogicalResult verify(TypeOp op) { return verifyHasBindingUse(op); }
435 
436 //===----------------------------------------------------------------------===//
437 // pdl::TypesOp
438 //===----------------------------------------------------------------------===//
439 
440 static LogicalResult verify(TypesOp op) { return verifyHasBindingUse(op); }
441 
442 //===----------------------------------------------------------------------===//
443 // TableGen'd op method definitions
444 //===----------------------------------------------------------------------===//
445 
446 #define GET_OP_CLASSES
447 #include "mlir/Dialect/PDL/IR/PDLOps.cpp.inc"
448