1 //===- TransformDialect.cpp - Transform dialect operations ----------------===//
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/Transform/IR/TransformOps.h"
10 #include "mlir/Dialect/PDL/IR/PDLOps.h"
11 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
12 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
13 #include "mlir/IR/Builders.h"
14 #include "mlir/IR/OpImplementation.h"
15 #include "mlir/IR/PatternMatch.h"
16 #include "mlir/Interfaces/ControlFlowInterfaces.h"
17 #include "mlir/Rewrite/FrozenRewritePatternSet.h"
18 #include "mlir/Rewrite/PatternApplicator.h"
19 #include "llvm/ADT/ScopeExit.h"
20 
21 using namespace mlir;
22 
23 #define GET_OP_CLASSES
24 #include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
25 
26 //===----------------------------------------------------------------------===//
27 // PatternApplicatorExtension
28 //===----------------------------------------------------------------------===//
29 
30 namespace {
31 /// A simple pattern rewriter that can be constructed from a context. This is
32 /// necessary to apply patterns to a specific op locally.
33 class TrivialPatternRewriter : public PatternRewriter {
34 public:
35   explicit TrivialPatternRewriter(MLIRContext *context)
36       : PatternRewriter(context) {}
37 };
38 
39 /// A TransformState extension that keeps track of compiled PDL pattern sets.
40 /// This is intended to be used along the WithPDLPatterns op. The extension
41 /// can be constructed given an operation that has a SymbolTable trait and
42 /// contains pdl::PatternOp instances. The patterns are compiled lazily and one
43 /// by one when requested; this behavior is subject to change.
44 class PatternApplicatorExtension : public transform::TransformState::Extension {
45 public:
46   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PatternApplicatorExtension)
47 
48   /// Creates the extension for patterns contained in `patternContainer`.
49   explicit PatternApplicatorExtension(transform::TransformState &state,
50                                       Operation *patternContainer)
51       : Extension(state), patterns(patternContainer) {}
52 
53   /// Appends to `results` the operations contained in `root` that matched the
54   /// PDL pattern with the given name. Note that `root` may or may not be the
55   /// operation that contains PDL patterns. Reports an error if the pattern
56   /// cannot be found. Note that when no operations are matched, this still
57   /// succeeds as long as the pattern exists.
58   LogicalResult findAllMatches(StringRef patternName, Operation *root,
59                                SmallVectorImpl<Operation *> &results);
60 
61 private:
62   /// Map from the pattern name to a singleton set of rewrite patterns that only
63   /// contains the pattern with this name. Populated when the pattern is first
64   /// requested.
65   // TODO: reconsider the efficiency of this storage when more usage data is
66   // available. Storing individual patterns in a set and triggering compilation
67   // for each of them has overhead. So does compiling a large set of patterns
68   // only to apply a handlful of them.
69   llvm::StringMap<FrozenRewritePatternSet> compiledPatterns;
70 
71   /// A symbol table operation containing the relevant PDL patterns.
72   SymbolTable patterns;
73 };
74 
75 LogicalResult PatternApplicatorExtension::findAllMatches(
76     StringRef patternName, Operation *root,
77     SmallVectorImpl<Operation *> &results) {
78   auto it = compiledPatterns.find(patternName);
79   if (it == compiledPatterns.end()) {
80     auto patternOp = patterns.lookup<pdl::PatternOp>(patternName);
81     if (!patternOp)
82       return failure();
83 
84     OwningOpRef<ModuleOp> pdlModuleOp = ModuleOp::create(patternOp.getLoc());
85     patternOp->moveBefore(pdlModuleOp->getBody(),
86                           pdlModuleOp->getBody()->end());
87     PDLPatternModule patternModule(std::move(pdlModuleOp));
88 
89     // Merge in the hooks owned by the dialect. Make a copy as they may be
90     // also used by the following operations.
91     auto *dialect =
92         root->getContext()->getLoadedDialect<transform::TransformDialect>();
93     for (const auto &pair : dialect->getPDLConstraintHooks())
94       patternModule.registerConstraintFunction(pair.first(), pair.second);
95 
96     // Register a noop rewriter because PDL requires patterns to end with some
97     // rewrite call.
98     patternModule.registerRewriteFunction(
99         "transform.dialect", [](PatternRewriter &, Operation *) {});
100 
101     it = compiledPatterns
102              .try_emplace(patternOp.getName(), std::move(patternModule))
103              .first;
104   }
105 
106   PatternApplicator applicator(it->second);
107   TrivialPatternRewriter rewriter(root->getContext());
108   applicator.applyDefaultCostModel();
109   root->walk([&](Operation *op) {
110     if (succeeded(applicator.matchAndRewrite(op, rewriter)))
111       results.push_back(op);
112   });
113 
114   return success();
115 }
116 } // namespace
117 
118 //===----------------------------------------------------------------------===//
119 // PDLMatchOp
120 //===----------------------------------------------------------------------===//
121 
122 LogicalResult transform::PDLMatchOp::apply(transform::TransformResults &results,
123                                            transform::TransformState &state) {
124   auto *extension = state.getExtension<PatternApplicatorExtension>();
125   assert(extension &&
126          "expected PatternApplicatorExtension to be attached by the parent op");
127   SmallVector<Operation *> targets;
128   for (Operation *root : state.getPayloadOps(getRoot())) {
129     if (failed(extension->findAllMatches(
130             getPatternName().getLeafReference().getValue(), root, targets))) {
131       return emitOpError() << "could not find pattern '" << getPatternName()
132                            << "'";
133     }
134   }
135   results.set(getResult().cast<OpResult>(), targets);
136   return success();
137 }
138 
139 //===----------------------------------------------------------------------===//
140 // SequenceOp
141 //===----------------------------------------------------------------------===//
142 
143 LogicalResult transform::SequenceOp::apply(transform::TransformResults &results,
144                                            transform::TransformState &state) {
145   // Map the entry block argument to the list of operations.
146   auto scope = state.make_region_scope(*getBodyBlock()->getParent());
147   if (failed(mapBlockArguments(state)))
148     return failure();
149 
150   // Apply the sequenced ops one by one.
151   for (Operation &transform : getBodyBlock()->without_terminator())
152     if (failed(state.applyTransform(cast<TransformOpInterface>(transform))))
153       return failure();
154 
155   // Forward the operation mapping for values yielded from the sequence to the
156   // values produced by the sequence op.
157   for (const auto &pair :
158        llvm::zip(getBodyBlock()->getTerminator()->getOperands(),
159                  getOperation()->getOpResults())) {
160     Value terminatorOperand = std::get<0>(pair);
161     OpResult result = std::get<1>(pair);
162     results.set(result, state.getPayloadOps(terminatorOperand));
163   }
164 
165   return success();
166 }
167 
168 /// Returns `true` if the given op operand may be consuming the handle value in
169 /// the Transform IR. That is, if it may have a Free effect on it.
170 static bool isValueUsePotentialConsumer(OpOperand &use) {
171   // Conservatively assume the effect being present in absence of the interface.
172   auto memEffectInterface = dyn_cast<MemoryEffectOpInterface>(use.getOwner());
173   if (!memEffectInterface)
174     return true;
175 
176   SmallVector<MemoryEffects::EffectInstance, 2> effects;
177   memEffectInterface.getEffectsOnValue(use.get(), effects);
178   return llvm::any_of(effects, [](const MemoryEffects::EffectInstance &effect) {
179     return isa<transform::TransformMappingResource>(effect.getResource()) &&
180            isa<MemoryEffects::Free>(effect.getEffect());
181   });
182 }
183 
184 LogicalResult
185 checkDoubleConsume(Value value,
186                    function_ref<InFlightDiagnostic()> reportError) {
187   OpOperand *potentialConsumer = nullptr;
188   for (OpOperand &use : value.getUses()) {
189     if (!isValueUsePotentialConsumer(use))
190       continue;
191 
192     if (!potentialConsumer) {
193       potentialConsumer = &use;
194       continue;
195     }
196 
197     InFlightDiagnostic diag = reportError()
198                               << " has more than one potential consumer";
199     diag.attachNote(potentialConsumer->getOwner()->getLoc())
200         << "used here as operand #" << potentialConsumer->getOperandNumber();
201     diag.attachNote(use.getOwner()->getLoc())
202         << "used here as operand #" << use.getOperandNumber();
203     return diag;
204   }
205 
206   return success();
207 }
208 
209 LogicalResult transform::SequenceOp::verify() {
210   // Check if the block argument has more than one consuming use.
211   for (BlockArgument argument : getBodyBlock()->getArguments()) {
212     auto report = [&]() {
213       return (emitOpError() << "block argument #" << argument.getArgNumber());
214     };
215     if (failed(checkDoubleConsume(argument, report)))
216       return failure();
217   }
218 
219   // Check properties of the nested operations they cannot check themselves.
220   for (Operation &child : *getBodyBlock()) {
221     if (!isa<TransformOpInterface>(child) &&
222         &child != &getBodyBlock()->back()) {
223       InFlightDiagnostic diag =
224           emitOpError()
225           << "expected children ops to implement TransformOpInterface";
226       diag.attachNote(child.getLoc()) << "op without interface";
227       return diag;
228     }
229 
230     for (OpResult result : child.getResults()) {
231       auto report = [&]() {
232         return (child.emitError() << "result #" << result.getResultNumber());
233       };
234       if (failed(checkDoubleConsume(result, report)))
235         return failure();
236     }
237   }
238 
239   if (getBodyBlock()->getTerminator()->getOperandTypes() !=
240       getOperation()->getResultTypes()) {
241     InFlightDiagnostic diag = emitOpError()
242                               << "expects the types of the terminator operands "
243                                  "to match the types of the result";
244     diag.attachNote(getBodyBlock()->getTerminator()->getLoc()) << "terminator";
245     return diag;
246   }
247   return success();
248 }
249 
250 void transform::SequenceOp::getEffects(
251     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
252   auto *mappingResource = TransformMappingResource::get();
253   effects.emplace_back(MemoryEffects::Read::get(), getRoot(), mappingResource);
254 
255   for (Value result : getResults()) {
256     effects.emplace_back(MemoryEffects::Allocate::get(), result,
257                          mappingResource);
258     effects.emplace_back(MemoryEffects::Write::get(), result, mappingResource);
259   }
260 
261   if (!getRoot()) {
262     for (Operation &op : *getBodyBlock()) {
263       auto iface = dyn_cast<MemoryEffectOpInterface>(&op);
264       if (!iface) {
265         // TODO: fill all possible effects; or require ops to actually implement
266         // the memory effect interface always
267         assert(false);
268       }
269 
270       SmallVector<MemoryEffects::EffectInstance, 2> nestedEffects;
271       iface.getEffects(effects);
272     }
273     return;
274   }
275 
276   // Carry over all effects on the argument of the entry block as those on the
277   // operand, this is the same value just remapped.
278   for (Operation &op : *getBodyBlock()) {
279     auto iface = dyn_cast<MemoryEffectOpInterface>(&op);
280     if (!iface) {
281       // TODO: fill all possible effects; or require ops to actually implement
282       // the memory effect interface always
283       assert(false);
284     }
285 
286     SmallVector<MemoryEffects::EffectInstance, 2> nestedEffects;
287     iface.getEffectsOnValue(getBodyBlock()->getArgument(0), nestedEffects);
288     for (const auto &effect : nestedEffects)
289       effects.emplace_back(effect.getEffect(), getRoot(), effect.getResource());
290   }
291 }
292 
293 OperandRange transform::SequenceOp::getSuccessorEntryOperands(unsigned index) {
294   assert(index == 0 && "unexpected region index");
295   if (getOperation()->getNumOperands() == 1)
296     return getOperation()->getOperands();
297   return OperandRange(getOperation()->operand_end(),
298                       getOperation()->operand_end());
299 }
300 
301 void transform::SequenceOp::getSuccessorRegions(
302     Optional<unsigned> index, ArrayRef<Attribute> operands,
303     SmallVectorImpl<RegionSuccessor> &regions) {
304   if (!index.hasValue()) {
305     Region *bodyRegion = &getBody();
306     regions.emplace_back(bodyRegion, !operands.empty()
307                                          ? bodyRegion->getArguments()
308                                          : Block::BlockArgListType());
309     return;
310   }
311 
312   assert(*index == 0 && "unexpected region index");
313   regions.emplace_back(getOperation()->getResults());
314 }
315 
316 void transform::SequenceOp::getRegionInvocationBounds(
317     ArrayRef<Attribute> operands, SmallVectorImpl<InvocationBounds> &bounds) {
318   (void)operands;
319   bounds.emplace_back(1, 1);
320 }
321 
322 //===----------------------------------------------------------------------===//
323 // WithPDLPatternsOp
324 //===----------------------------------------------------------------------===//
325 
326 LogicalResult
327 transform::WithPDLPatternsOp::apply(transform::TransformResults &results,
328                                     transform::TransformState &state) {
329   OwningOpRef<ModuleOp> pdlModuleOp =
330       ModuleOp::create(getOperation()->getLoc());
331   TransformOpInterface transformOp = nullptr;
332   for (Operation &nested : getBody().front()) {
333     if (!isa<pdl::PatternOp>(nested)) {
334       transformOp = cast<TransformOpInterface>(nested);
335       break;
336     }
337   }
338 
339   state.addExtension<PatternApplicatorExtension>(getOperation());
340   auto guard = llvm::make_scope_exit(
341       [&]() { state.removeExtension<PatternApplicatorExtension>(); });
342 
343   auto scope = state.make_region_scope(getBody());
344   if (failed(mapBlockArguments(state)))
345     return failure();
346   return state.applyTransform(transformOp);
347 }
348 
349 LogicalResult transform::WithPDLPatternsOp::verify() {
350   Block *body = getBodyBlock();
351   Operation *topLevelOp = nullptr;
352   for (Operation &op : body->getOperations()) {
353     if (isa<pdl::PatternOp>(op))
354       continue;
355 
356     if (op.hasTrait<::mlir::transform::PossibleTopLevelTransformOpTrait>()) {
357       if (topLevelOp) {
358         InFlightDiagnostic diag =
359             emitOpError() << "expects only one non-pattern op in its body";
360         diag.attachNote(topLevelOp->getLoc()) << "first non-pattern op";
361         diag.attachNote(op.getLoc()) << "second non-pattern op";
362         return diag;
363       }
364       topLevelOp = &op;
365       continue;
366     }
367 
368     InFlightDiagnostic diag =
369         emitOpError()
370         << "expects only pattern and top-level transform ops in its body";
371     diag.attachNote(op.getLoc()) << "offending op";
372     return diag;
373   }
374 
375   if (auto parent = getOperation()->getParentOfType<WithPDLPatternsOp>()) {
376     InFlightDiagnostic diag = emitOpError() << "cannot be nested";
377     diag.attachNote(parent.getLoc()) << "parent operation";
378     return diag;
379   }
380 
381   return success();
382 }
383