1 //===- DialectConversion.cpp - MLIR dialect conversion generic pass -------===//
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/Transforms/DialectConversion.h"
10 #include "mlir/IR/Block.h"
11 #include "mlir/IR/BlockAndValueMapping.h"
12 #include "mlir/IR/Builders.h"
13 #include "mlir/IR/BuiltinOps.h"
14 #include "mlir/Rewrite/PatternApplicator.h"
15 #include "mlir/Transforms/Utils.h"
16 #include "llvm/ADT/SetVector.h"
17 #include "llvm/ADT/SmallPtrSet.h"
18 #include "llvm/Support/Debug.h"
19 #include "llvm/Support/FormatVariadic.h"
20 #include "llvm/Support/SaveAndRestore.h"
21 #include "llvm/Support/ScopedPrinter.h"
22 
23 using namespace mlir;
24 using namespace mlir::detail;
25 
26 #define DEBUG_TYPE "dialect-conversion"
27 
28 /// Recursively collect all of the operations to convert from within 'region'.
29 /// If 'target' is nonnull, operations that are recursively legal have their
30 /// regions pre-filtered to avoid considering them for legalization.
31 static LogicalResult
32 computeConversionSet(iterator_range<Region::iterator> region,
33                      Location regionLoc, std::vector<Operation *> &toConvert,
34                      ConversionTarget *target = nullptr) {
35   if (llvm::empty(region))
36     return success();
37 
38   // Traverse starting from the entry block.
39   SmallVector<Block *, 16> worklist(1, &*region.begin());
40   DenseSet<Block *> visitedBlocks;
41   visitedBlocks.insert(worklist.front());
42   while (!worklist.empty()) {
43     Block *block = worklist.pop_back_val();
44 
45     // Compute the conversion set of each of the nested operations.
46     for (Operation &op : *block) {
47       toConvert.emplace_back(&op);
48 
49       // Don't check this operation's children for conversion if the operation
50       // is recursively legal.
51       auto legalityInfo = target ? target->isLegal(&op)
52                                  : Optional<ConversionTarget::LegalOpDetails>();
53       if (legalityInfo && legalityInfo->isRecursivelyLegal)
54         continue;
55       for (auto &region : op.getRegions()) {
56         if (failed(computeConversionSet(region.getBlocks(), region.getLoc(),
57                                         toConvert, target)))
58           return failure();
59       }
60     }
61 
62     // Recurse to children that haven't been visited.
63     for (Block *succ : block->getSuccessors())
64       if (visitedBlocks.insert(succ).second)
65         worklist.push_back(succ);
66   }
67 
68   // Check that all blocks in the region were visited.
69   if (llvm::any_of(llvm::drop_begin(region, 1),
70                    [&](Block &block) { return !visitedBlocks.count(&block); }))
71     return emitError(regionLoc, "unreachable blocks were not converted");
72   return success();
73 }
74 
75 /// A utility function to log a successful result for the given reason.
76 template <typename... Args>
77 static void logSuccess(llvm::ScopedPrinter &os, StringRef fmt, Args &&...args) {
78   LLVM_DEBUG({
79     os.unindent();
80     os.startLine() << "} -> SUCCESS";
81     if (!fmt.empty())
82       os.getOStream() << " : "
83                       << llvm::formatv(fmt.data(), std::forward<Args>(args)...);
84     os.getOStream() << "\n";
85   });
86 }
87 
88 /// A utility function to log a failure result for the given reason.
89 template <typename... Args>
90 static void logFailure(llvm::ScopedPrinter &os, StringRef fmt, Args &&...args) {
91   LLVM_DEBUG({
92     os.unindent();
93     os.startLine() << "} -> FAILURE : "
94                    << llvm::formatv(fmt.data(), std::forward<Args>(args)...)
95                    << "\n";
96   });
97 }
98 
99 //===----------------------------------------------------------------------===//
100 // ConversionValueMapping
101 //===----------------------------------------------------------------------===//
102 
103 namespace {
104 /// This class wraps a BlockAndValueMapping to provide recursive lookup
105 /// functionality, i.e. we will traverse if the mapped value also has a mapping.
106 struct ConversionValueMapping {
107   /// Lookup a mapped value within the map. If a mapping for the provided value
108   /// does not exist then return the provided value. If `desiredType` is
109   /// non-null, returns the most recently mapped value with that type. If an
110   /// operand of that type does not exist, defaults to normal behavior.
111   Value lookupOrDefault(Value from, Type desiredType = nullptr) const;
112 
113   /// Lookup a mapped value within the map, or return null if a mapping does not
114   /// exist. If a mapping exists, this follows the same behavior of
115   /// `lookupOrDefault`.
116   Value lookupOrNull(Value from) const;
117 
118   /// Map a value to the one provided.
119   void map(Value oldVal, Value newVal) { mapping.map(oldVal, newVal); }
120 
121   /// Drop the last mapping for the given value.
122   void erase(Value value) { mapping.erase(value); }
123 
124 private:
125   /// Current value mappings.
126   BlockAndValueMapping mapping;
127 };
128 } // end anonymous namespace
129 
130 Value ConversionValueMapping::lookupOrDefault(Value from,
131                                               Type desiredType) const {
132   // If there was no desired type, simply find the leaf value.
133   if (!desiredType) {
134     // If this value had a valid mapping, unmap that value as well in the case
135     // that it was also replaced.
136     while (auto mappedValue = mapping.lookupOrNull(from))
137       from = mappedValue;
138     return from;
139   }
140 
141   // Otherwise, try to find the deepest value that has the desired type.
142   Value desiredValue;
143   do {
144     if (from.getType() == desiredType)
145       desiredValue = from;
146 
147     Value mappedValue = mapping.lookupOrNull(from);
148     if (!mappedValue)
149       break;
150     from = mappedValue;
151   } while (true);
152 
153   // If the desired value was found use it, otherwise default to the leaf value.
154   return desiredValue ? desiredValue : from;
155 }
156 
157 Value ConversionValueMapping::lookupOrNull(Value from) const {
158   Value result = lookupOrDefault(from);
159   return result == from ? nullptr : result;
160 }
161 
162 //===----------------------------------------------------------------------===//
163 // ArgConverter
164 //===----------------------------------------------------------------------===//
165 namespace {
166 /// This class provides a simple interface for converting the types of block
167 /// arguments. This is done by creating a new block that contains the new legal
168 /// types and extracting the block that contains the old illegal types to allow
169 /// for undoing pending rewrites in the case of failure.
170 struct ArgConverter {
171   ArgConverter(PatternRewriter &rewriter) : rewriter(rewriter) {}
172 
173   /// This structure contains the information pertaining to an argument that has
174   /// been converted.
175   struct ConvertedArgInfo {
176     ConvertedArgInfo(unsigned newArgIdx, unsigned newArgSize,
177                      Value castValue = nullptr)
178         : newArgIdx(newArgIdx), newArgSize(newArgSize), castValue(castValue) {}
179 
180     /// The start index of in the new argument list that contains arguments that
181     /// replace the original.
182     unsigned newArgIdx;
183 
184     /// The number of arguments that replaced the original argument.
185     unsigned newArgSize;
186 
187     /// The cast value that was created to cast from the new arguments to the
188     /// old. This only used if 'newArgSize' > 1.
189     Value castValue;
190   };
191 
192   /// This structure contains information pertaining to a block that has had its
193   /// signature converted.
194   struct ConvertedBlockInfo {
195     ConvertedBlockInfo(Block *origBlock, TypeConverter &converter)
196         : origBlock(origBlock), converter(&converter) {}
197 
198     /// The original block that was requested to have its signature converted.
199     Block *origBlock;
200 
201     /// The conversion information for each of the arguments. The information is
202     /// None if the argument was dropped during conversion.
203     SmallVector<Optional<ConvertedArgInfo>, 1> argInfo;
204 
205     /// The type converter used to convert the arguments.
206     TypeConverter *converter;
207   };
208 
209   /// Return if the signature of the given block has already been converted.
210   bool hasBeenConverted(Block *block) const {
211     return conversionInfo.count(block) || convertedBlocks.count(block);
212   }
213 
214   /// Set the type converter to use for the given region.
215   void setConverter(Region *region, TypeConverter *typeConverter) {
216     assert(typeConverter && "expected valid type converter");
217     regionToConverter[region] = typeConverter;
218   }
219 
220   /// Return the type converter to use for the given region, or null if there
221   /// isn't one.
222   TypeConverter *getConverter(Region *region) {
223     return regionToConverter.lookup(region);
224   }
225 
226   //===--------------------------------------------------------------------===//
227   // Rewrite Application
228   //===--------------------------------------------------------------------===//
229 
230   /// Erase any rewrites registered for the blocks within the given operation
231   /// which is about to be removed. This merely drops the rewrites without
232   /// undoing them.
233   void notifyOpRemoved(Operation *op);
234 
235   /// Cleanup and undo any generated conversions for the arguments of block.
236   /// This method replaces the new block with the original, reverting the IR to
237   /// its original state.
238   void discardRewrites(Block *block);
239 
240   /// Fully replace uses of the old arguments with the new.
241   void applyRewrites(ConversionValueMapping &mapping);
242 
243   /// Materialize any necessary conversions for converted arguments that have
244   /// live users, using the provided `findLiveUser` to search for a user that
245   /// survives the conversion process.
246   LogicalResult
247   materializeLiveConversions(ConversionValueMapping &mapping,
248                              OpBuilder &builder,
249                              function_ref<Operation *(Value)> findLiveUser);
250 
251   //===--------------------------------------------------------------------===//
252   // Conversion
253   //===--------------------------------------------------------------------===//
254 
255   /// Attempt to convert the signature of the given block, if successful a new
256   /// block is returned containing the new arguments. Returns `block` if it did
257   /// not require conversion.
258   FailureOr<Block *> convertSignature(Block *block, TypeConverter &converter,
259                                       ConversionValueMapping &mapping);
260 
261   /// Apply the given signature conversion on the given block. The new block
262   /// containing the updated signature is returned. If no conversions were
263   /// necessary, e.g. if the block has no arguments, `block` is returned.
264   /// `converter` is used to generate any necessary cast operations that
265   /// translate between the origin argument types and those specified in the
266   /// signature conversion.
267   Block *applySignatureConversion(
268       Block *block, TypeConverter &converter,
269       TypeConverter::SignatureConversion &signatureConversion,
270       ConversionValueMapping &mapping);
271 
272   /// Insert a new conversion into the cache.
273   void insertConversion(Block *newBlock, ConvertedBlockInfo &&info);
274 
275   /// A collection of blocks that have had their arguments converted. This is a
276   /// map from the new replacement block, back to the original block.
277   llvm::MapVector<Block *, ConvertedBlockInfo> conversionInfo;
278 
279   /// The set of original blocks that were converted.
280   DenseSet<Block *> convertedBlocks;
281 
282   /// A mapping from valid regions, to those containing the original blocks of a
283   /// conversion.
284   DenseMap<Region *, std::unique_ptr<Region>> regionMapping;
285 
286   /// A mapping of regions to type converters that should be used when
287   /// converting the arguments of blocks within that region.
288   DenseMap<Region *, TypeConverter *> regionToConverter;
289 
290   /// The pattern rewriter to use when materializing conversions.
291   PatternRewriter &rewriter;
292 };
293 } // end anonymous namespace
294 
295 //===----------------------------------------------------------------------===//
296 // Rewrite Application
297 
298 void ArgConverter::notifyOpRemoved(Operation *op) {
299   if (conversionInfo.empty())
300     return;
301 
302   for (Region &region : op->getRegions()) {
303     for (Block &block : region) {
304       // Drop any rewrites from within.
305       for (Operation &nestedOp : block)
306         if (nestedOp.getNumRegions())
307           notifyOpRemoved(&nestedOp);
308 
309       // Check if this block was converted.
310       auto it = conversionInfo.find(&block);
311       if (it == conversionInfo.end())
312         continue;
313 
314       // Drop all uses of the original arguments and delete the original block.
315       Block *origBlock = it->second.origBlock;
316       for (BlockArgument arg : origBlock->getArguments())
317         arg.dropAllUses();
318       conversionInfo.erase(it);
319     }
320   }
321 }
322 
323 void ArgConverter::discardRewrites(Block *block) {
324   auto it = conversionInfo.find(block);
325   if (it == conversionInfo.end())
326     return;
327   Block *origBlock = it->second.origBlock;
328 
329   // Drop all uses of the new block arguments and replace uses of the new block.
330   for (int i = block->getNumArguments() - 1; i >= 0; --i)
331     block->getArgument(i).dropAllUses();
332   block->replaceAllUsesWith(origBlock);
333 
334   // Move the operations back the original block and the delete the new block.
335   origBlock->getOperations().splice(origBlock->end(), block->getOperations());
336   origBlock->moveBefore(block);
337   block->erase();
338 
339   convertedBlocks.erase(origBlock);
340   conversionInfo.erase(it);
341 }
342 
343 void ArgConverter::applyRewrites(ConversionValueMapping &mapping) {
344   for (auto &info : conversionInfo) {
345     ConvertedBlockInfo &blockInfo = info.second;
346     Block *origBlock = blockInfo.origBlock;
347 
348     // Process the remapping for each of the original arguments.
349     for (unsigned i = 0, e = origBlock->getNumArguments(); i != e; ++i) {
350       Optional<ConvertedArgInfo> &argInfo = blockInfo.argInfo[i];
351       BlockArgument origArg = origBlock->getArgument(i);
352 
353       // Handle the case of a 1->0 value mapping.
354       if (!argInfo) {
355         if (Value newArg = mapping.lookupOrNull(origArg))
356           origArg.replaceAllUsesWith(newArg);
357         continue;
358       }
359 
360       // Otherwise this is a 1->1+ value mapping.
361       Value castValue = argInfo->castValue;
362       assert(argInfo->newArgSize >= 1 && castValue && "expected 1->1+ mapping");
363 
364       // If the argument is still used, replace it with the generated cast.
365       if (!origArg.use_empty())
366         origArg.replaceAllUsesWith(mapping.lookupOrDefault(castValue));
367     }
368   }
369 }
370 
371 LogicalResult ArgConverter::materializeLiveConversions(
372     ConversionValueMapping &mapping, OpBuilder &builder,
373     function_ref<Operation *(Value)> findLiveUser) {
374   for (auto &info : conversionInfo) {
375     Block *newBlock = info.first;
376     ConvertedBlockInfo &blockInfo = info.second;
377     Block *origBlock = blockInfo.origBlock;
378 
379     // Process the remapping for each of the original arguments.
380     for (unsigned i = 0, e = origBlock->getNumArguments(); i != e; ++i) {
381       // FIXME: We should run the below checks even if the type conversion was
382       // 1->N, but a lot of existing lowering rely on the block argument being
383       // blindly replaced. Those usages should be updated, and this if should be
384       // removed.
385       if (blockInfo.argInfo[i])
386         continue;
387 
388       // If the type of this argument changed and the argument is still live, we
389       // need to materialize a conversion.
390       BlockArgument origArg = origBlock->getArgument(i);
391       auto argReplacementValue = mapping.lookupOrDefault(origArg);
392       bool isDroppedArg = argReplacementValue == origArg;
393       if (argReplacementValue.getType() == origArg.getType() && !isDroppedArg)
394         continue;
395       Operation *liveUser = findLiveUser(origArg);
396       if (!liveUser)
397         continue;
398 
399       if (OpResult result = argReplacementValue.dyn_cast<OpResult>())
400         rewriter.setInsertionPointAfter(result.getOwner());
401       else
402         rewriter.setInsertionPointToStart(newBlock);
403       Value newArg = blockInfo.converter->materializeSourceConversion(
404           rewriter, origArg.getLoc(), origArg.getType(),
405           isDroppedArg ? ValueRange() : ValueRange(argReplacementValue));
406       if (!newArg) {
407         InFlightDiagnostic diag =
408             emitError(origArg.getLoc())
409             << "failed to materialize conversion for block argument #" << i
410             << " that remained live after conversion, type was "
411             << origArg.getType();
412         if (!isDroppedArg)
413           diag << ", with target type " << argReplacementValue.getType();
414         diag.attachNote(liveUser->getLoc())
415             << "see existing live user here: " << *liveUser;
416         return failure();
417       }
418       mapping.map(origArg, newArg);
419     }
420   }
421   return success();
422 }
423 
424 //===----------------------------------------------------------------------===//
425 // Conversion
426 
427 FailureOr<Block *>
428 ArgConverter::convertSignature(Block *block, TypeConverter &converter,
429                                ConversionValueMapping &mapping) {
430   // Check if the block was already converted. If the block is detached,
431   // conservatively assume it is going to be deleted.
432   if (hasBeenConverted(block) || !block->getParent())
433     return block;
434 
435   // Try to convert the signature for the block with the provided converter.
436   if (auto conversion = converter.convertBlockSignature(block))
437     return applySignatureConversion(block, converter, *conversion, mapping);
438   return failure();
439 }
440 
441 Block *ArgConverter::applySignatureConversion(
442     Block *block, TypeConverter &converter,
443     TypeConverter::SignatureConversion &signatureConversion,
444     ConversionValueMapping &mapping) {
445   // If no arguments are being changed or added, there is nothing to do.
446   unsigned origArgCount = block->getNumArguments();
447   auto convertedTypes = signatureConversion.getConvertedTypes();
448   if (origArgCount == 0 && convertedTypes.empty())
449     return block;
450 
451   // Split the block at the beginning to get a new block to use for the updated
452   // signature.
453   Block *newBlock = block->splitBlock(block->begin());
454   block->replaceAllUsesWith(newBlock);
455 
456   SmallVector<Value, 4> newArgRange(newBlock->addArguments(convertedTypes));
457   ArrayRef<Value> newArgs(newArgRange);
458 
459   // Remap each of the original arguments as determined by the signature
460   // conversion.
461   ConvertedBlockInfo info(block, converter);
462   info.argInfo.resize(origArgCount);
463 
464   OpBuilder::InsertionGuard guard(rewriter);
465   rewriter.setInsertionPointToStart(newBlock);
466   for (unsigned i = 0; i != origArgCount; ++i) {
467     auto inputMap = signatureConversion.getInputMapping(i);
468     if (!inputMap)
469       continue;
470     BlockArgument origArg = block->getArgument(i);
471 
472     // If inputMap->replacementValue is not nullptr, then the argument is
473     // dropped and a replacement value is provided to be the remappedValue.
474     if (inputMap->replacementValue) {
475       assert(inputMap->size == 0 &&
476              "invalid to provide a replacement value when the argument isn't "
477              "dropped");
478       mapping.map(origArg, inputMap->replacementValue);
479       continue;
480     }
481 
482     // Otherwise, this is a 1->1+ mapping. Call into the provided type converter
483     // to pack the new values. For 1->1 mappings, if there is no materialization
484     // provided, use the argument directly instead.
485     auto replArgs = newArgs.slice(inputMap->inputNo, inputMap->size);
486     Value newArg = converter.materializeArgumentConversion(
487         rewriter, origArg.getLoc(), origArg.getType(), replArgs);
488     if (!newArg) {
489       assert(replArgs.size() == 1 &&
490              "couldn't materialize the result of 1->N conversion");
491       newArg = replArgs.front();
492     }
493     mapping.map(origArg, newArg);
494     info.argInfo[i] =
495         ConvertedArgInfo(inputMap->inputNo, inputMap->size, newArg);
496   }
497 
498   // Remove the original block from the region and return the new one.
499   insertConversion(newBlock, std::move(info));
500   return newBlock;
501 }
502 
503 void ArgConverter::insertConversion(Block *newBlock,
504                                     ConvertedBlockInfo &&info) {
505   // Get a region to insert the old block.
506   Region *region = newBlock->getParent();
507   std::unique_ptr<Region> &mappedRegion = regionMapping[region];
508   if (!mappedRegion)
509     mappedRegion = std::make_unique<Region>(region->getParentOp());
510 
511   // Move the original block to the mapped region and emplace the conversion.
512   mappedRegion->getBlocks().splice(mappedRegion->end(), region->getBlocks(),
513                                    info.origBlock->getIterator());
514   convertedBlocks.insert(info.origBlock);
515   conversionInfo.insert({newBlock, std::move(info)});
516 }
517 
518 //===----------------------------------------------------------------------===//
519 // Rewriter and Translation State
520 //===----------------------------------------------------------------------===//
521 namespace {
522 /// This class contains a snapshot of the current conversion rewriter state.
523 /// This is useful when saving and undoing a set of rewrites.
524 struct RewriterState {
525   RewriterState(unsigned numCreatedOps, unsigned numReplacements,
526                 unsigned numArgReplacements, unsigned numBlockActions,
527                 unsigned numIgnoredOperations, unsigned numRootUpdates)
528       : numCreatedOps(numCreatedOps), numReplacements(numReplacements),
529         numArgReplacements(numArgReplacements),
530         numBlockActions(numBlockActions),
531         numIgnoredOperations(numIgnoredOperations),
532         numRootUpdates(numRootUpdates) {}
533 
534   /// The current number of created operations.
535   unsigned numCreatedOps;
536 
537   /// The current number of replacements queued.
538   unsigned numReplacements;
539 
540   /// The current number of argument replacements queued.
541   unsigned numArgReplacements;
542 
543   /// The current number of block actions performed.
544   unsigned numBlockActions;
545 
546   /// The current number of ignored operations.
547   unsigned numIgnoredOperations;
548 
549   /// The current number of operations that were updated in place.
550   unsigned numRootUpdates;
551 };
552 
553 /// The state of an operation that was updated by a pattern in-place. This
554 /// contains all of the necessary information to reconstruct an operation that
555 /// was updated in place.
556 class OperationTransactionState {
557 public:
558   OperationTransactionState() = default;
559   OperationTransactionState(Operation *op)
560       : op(op), loc(op->getLoc()), attrs(op->getAttrDictionary()),
561         operands(op->operand_begin(), op->operand_end()),
562         successors(op->successor_begin(), op->successor_end()) {}
563 
564   /// Discard the transaction state and reset the state of the original
565   /// operation.
566   void resetOperation() const {
567     op->setLoc(loc);
568     op->setAttrs(attrs);
569     op->setOperands(operands);
570     for (auto it : llvm::enumerate(successors))
571       op->setSuccessor(it.value(), it.index());
572   }
573 
574   /// Return the original operation of this state.
575   Operation *getOperation() const { return op; }
576 
577 private:
578   Operation *op;
579   LocationAttr loc;
580   DictionaryAttr attrs;
581   SmallVector<Value, 8> operands;
582   SmallVector<Block *, 2> successors;
583 };
584 
585 /// This class represents one requested operation replacement via 'replaceOp' or
586 /// 'eraseOp`.
587 struct OpReplacement {
588   OpReplacement() = default;
589   OpReplacement(TypeConverter *converter) : converter(converter) {}
590 
591   /// An optional type converter that can be used to materialize conversions
592   /// between the new and old values if necessary.
593   TypeConverter *converter = nullptr;
594 };
595 
596 /// The kind of the block action performed during the rewrite.  Actions can be
597 /// undone if the conversion fails.
598 enum class BlockActionKind {
599   Create,
600   Erase,
601   Merge,
602   Move,
603   Split,
604   TypeConversion
605 };
606 
607 /// Original position of the given block in its parent region. During undo
608 /// actions, the block needs to be placed after `insertAfterBlock`.
609 struct BlockPosition {
610   Region *region;
611   Block *insertAfterBlock;
612 };
613 
614 /// Information needed to undo the merge actions.
615 /// - the source block, and
616 /// - the Operation that was the last operation in the dest block before the
617 ///   merge (could be null if the dest block was empty).
618 struct MergeInfo {
619   Block *sourceBlock;
620   Operation *destBlockLastInst;
621 };
622 
623 /// The storage class for an undoable block action (one of BlockActionKind),
624 /// contains the information necessary to undo this action.
625 struct BlockAction {
626   static BlockAction getCreate(Block *block) {
627     return {BlockActionKind::Create, block, {}};
628   }
629   static BlockAction getErase(Block *block, BlockPosition originalPosition) {
630     return {BlockActionKind::Erase, block, {originalPosition}};
631   }
632   static BlockAction getMerge(Block *block, Block *sourceBlock) {
633     BlockAction action{BlockActionKind::Merge, block, {}};
634     action.mergeInfo = {sourceBlock, block->empty() ? nullptr : &block->back()};
635     return action;
636   }
637   static BlockAction getMove(Block *block, BlockPosition originalPosition) {
638     return {BlockActionKind::Move, block, {originalPosition}};
639   }
640   static BlockAction getSplit(Block *block, Block *originalBlock) {
641     BlockAction action{BlockActionKind::Split, block, {}};
642     action.originalBlock = originalBlock;
643     return action;
644   }
645   static BlockAction getTypeConversion(Block *block) {
646     return BlockAction{BlockActionKind::TypeConversion, block, {}};
647   }
648 
649   // The action kind.
650   BlockActionKind kind;
651 
652   // A pointer to the block that was created by the action.
653   Block *block;
654 
655   union {
656     // In use if kind == BlockActionKind::Move or BlockActionKind::Erase, and
657     // contains a pointer to the region that originally contained the block as
658     // well as the position of the block in that region.
659     BlockPosition originalPosition;
660     // In use if kind == BlockActionKind::Split and contains a pointer to the
661     // block that was split into two parts.
662     Block *originalBlock;
663     // In use if kind == BlockActionKind::Merge, and contains the information
664     // needed to undo the merge.
665     MergeInfo mergeInfo;
666   };
667 };
668 } // end anonymous namespace
669 
670 //===----------------------------------------------------------------------===//
671 // ConversionPatternRewriterImpl
672 //===----------------------------------------------------------------------===//
673 namespace mlir {
674 namespace detail {
675 struct ConversionPatternRewriterImpl {
676   ConversionPatternRewriterImpl(PatternRewriter &rewriter)
677       : argConverter(rewriter) {}
678 
679   /// Cleanup and destroy any generated rewrite operations. This method is
680   /// invoked when the conversion process fails.
681   void discardRewrites();
682 
683   /// Apply all requested operation rewrites. This method is invoked when the
684   /// conversion process succeeds.
685   void applyRewrites();
686 
687   //===--------------------------------------------------------------------===//
688   // State Management
689   //===--------------------------------------------------------------------===//
690 
691   /// Return the current state of the rewriter.
692   RewriterState getCurrentState();
693 
694   /// Reset the state of the rewriter to a previously saved point.
695   void resetState(RewriterState state);
696 
697   /// Erase any blocks that were unlinked from their regions and stored in block
698   /// actions.
699   void eraseDanglingBlocks();
700 
701   /// Undo the block actions (motions, splits) one by one in reverse order until
702   /// "numActionsToKeep" actions remains.
703   void undoBlockActions(unsigned numActionsToKeep = 0);
704 
705   /// Remap the given operands to those with potentially different types. The
706   /// provided type converter is used to ensure that the remapped types are
707   /// legal. Returns success if the operands could be remapped, failure
708   /// otherwise.
709   LogicalResult remapValues(Location loc, PatternRewriter &rewriter,
710                             TypeConverter *converter,
711                             Operation::operand_range operands,
712                             SmallVectorImpl<Value> &remapped);
713 
714   /// Returns true if the given operation is ignored, and does not need to be
715   /// converted.
716   bool isOpIgnored(Operation *op) const;
717 
718   /// Recursively marks the nested operations under 'op' as ignored. This
719   /// removes them from being considered for legalization.
720   void markNestedOpsIgnored(Operation *op);
721 
722   //===--------------------------------------------------------------------===//
723   // Type Conversion
724   //===--------------------------------------------------------------------===//
725 
726   /// Convert the signature of the given block.
727   FailureOr<Block *> convertBlockSignature(
728       Block *block, TypeConverter &converter,
729       TypeConverter::SignatureConversion *conversion = nullptr);
730 
731   /// Apply a signature conversion on the given region.
732   Block *
733   applySignatureConversion(Region *region,
734                            TypeConverter::SignatureConversion &conversion);
735 
736   /// Convert the types of block arguments within the given region.
737   FailureOr<Block *>
738   convertRegionTypes(Region *region, TypeConverter &converter,
739                      TypeConverter::SignatureConversion *entryConversion);
740 
741   //===--------------------------------------------------------------------===//
742   // Rewriter Notification Hooks
743   //===--------------------------------------------------------------------===//
744 
745   /// PatternRewriter hook for replacing the results of an operation.
746   void notifyOpReplaced(Operation *op, ValueRange newValues);
747 
748   /// Notifies that a block is about to be erased.
749   void notifyBlockIsBeingErased(Block *block);
750 
751   /// Notifies that a block was created.
752   void notifyCreatedBlock(Block *block);
753 
754   /// Notifies that a block was split.
755   void notifySplitBlock(Block *block, Block *continuation);
756 
757   /// Notifies that `block` is being merged with `srcBlock`.
758   void notifyBlocksBeingMerged(Block *block, Block *srcBlock);
759 
760   /// Notifies that the blocks of a region are about to be moved.
761   void notifyRegionIsBeingInlinedBefore(Region &region, Region &parent,
762                                         Region::iterator before);
763 
764   /// Notifies that the blocks of a region were cloned into another.
765   void notifyRegionWasClonedBefore(iterator_range<Region::iterator> &blocks,
766                                    Location origRegionLoc);
767 
768   /// Notifies that a pattern match failed for the given reason.
769   LogicalResult
770   notifyMatchFailure(Location loc,
771                      function_ref<void(Diagnostic &)> reasonCallback);
772 
773   //===--------------------------------------------------------------------===//
774   // State
775   //===--------------------------------------------------------------------===//
776 
777   // Mapping between replaced values that differ in type. This happens when
778   // replacing a value with one of a different type.
779   ConversionValueMapping mapping;
780 
781   /// Utility used to convert block arguments.
782   ArgConverter argConverter;
783 
784   /// Ordered vector of all of the newly created operations during conversion.
785   std::vector<Operation *> createdOps;
786 
787   /// Ordered map of requested operation replacements.
788   llvm::MapVector<Operation *, OpReplacement> replacements;
789 
790   /// Ordered vector of any requested block argument replacements.
791   SmallVector<BlockArgument, 4> argReplacements;
792 
793   /// Ordered list of block operations (creations, splits, motions).
794   SmallVector<BlockAction, 4> blockActions;
795 
796   /// A set of operations that should no longer be considered for legalization,
797   /// but were not directly replace/erased/etc. by a pattern. These are
798   /// generally child operations of other operations who were
799   /// replaced/erased/etc. This is not meant to be an exhaustive list of all
800   /// operations, but the minimal set that can be used to detect if a given
801   /// operation should be `ignored`. For example, we may add the operations that
802   /// define non-empty regions to the set, but not any of the others. This
803   /// simplifies the amount of memory needed as we can query if the parent
804   /// operation was ignored.
805   llvm::SetVector<Operation *> ignoredOps;
806 
807   /// A transaction state for each of operations that were updated in-place.
808   SmallVector<OperationTransactionState, 4> rootUpdates;
809 
810   /// A vector of indices into `replacements` of operations that were replaced
811   /// with values with different result types than the original operation, e.g.
812   /// 1->N conversion of some kind.
813   SmallVector<unsigned, 4> operationsWithChangedResults;
814 
815   /// A default type converter, used when block conversions do not have one
816   /// explicitly provided.
817   TypeConverter defaultTypeConverter;
818 
819   /// The current conversion pattern that is being rewritten, or nullptr if
820   /// called from outside of a conversion pattern rewrite.
821   const ConversionPattern *currentConversionPattern = nullptr;
822 
823 #ifndef NDEBUG
824   /// A set of operations that have pending updates. This tracking isn't
825   /// strictly necessary, and is thus only active during debug builds for extra
826   /// verification.
827   SmallPtrSet<Operation *, 1> pendingRootUpdates;
828 
829   /// A logger used to emit diagnostics during the conversion process.
830   llvm::ScopedPrinter logger{llvm::dbgs()};
831 #endif
832 };
833 } // end namespace detail
834 } // end namespace mlir
835 
836 /// Detach any operations nested in the given operation from their parent
837 /// blocks, and erase the given operation. This can be used when the nested
838 /// operations are scheduled for erasure themselves, so deleting the regions of
839 /// the given operation together with their content would result in double-free.
840 /// This happens, for example, when rolling back op creation in the reverse
841 /// order and if the nested ops were created before the parent op. This function
842 /// does not need to collect nested ops recursively because it is expected to
843 /// also be called for each nested op when it is about to be deleted.
844 static void detachNestedAndErase(Operation *op) {
845   for (Region &region : op->getRegions()) {
846     for (Block &block : region.getBlocks()) {
847       while (!block.getOperations().empty())
848         block.getOperations().remove(block.getOperations().begin());
849       block.dropAllDefinedValueUses();
850     }
851   }
852   op->dropAllUses();
853   op->erase();
854 }
855 
856 void ConversionPatternRewriterImpl::discardRewrites() {
857   // Reset any operations that were updated in place.
858   for (auto &state : rootUpdates)
859     state.resetOperation();
860 
861   undoBlockActions();
862 
863   // Remove any newly created ops.
864   for (auto *op : llvm::reverse(createdOps))
865     detachNestedAndErase(op);
866 }
867 
868 void ConversionPatternRewriterImpl::applyRewrites() {
869   // Apply all of the rewrites replacements requested during conversion.
870   for (auto &repl : replacements) {
871     for (OpResult result : repl.first->getResults())
872       if (Value newValue = mapping.lookupOrNull(result))
873         result.replaceAllUsesWith(newValue);
874 
875     // If this operation defines any regions, drop any pending argument
876     // rewrites.
877     if (repl.first->getNumRegions())
878       argConverter.notifyOpRemoved(repl.first);
879   }
880 
881   // Apply all of the requested argument replacements.
882   for (BlockArgument arg : argReplacements) {
883     Value repl = mapping.lookupOrDefault(arg);
884     if (repl.isa<BlockArgument>()) {
885       arg.replaceAllUsesWith(repl);
886       continue;
887     }
888 
889     // If the replacement value is an operation, we check to make sure that we
890     // don't replace uses that are within the parent operation of the
891     // replacement value.
892     Operation *replOp = repl.cast<OpResult>().getOwner();
893     Block *replBlock = replOp->getBlock();
894     arg.replaceUsesWithIf(repl, [&](OpOperand &operand) {
895       Operation *user = operand.getOwner();
896       return user->getBlock() != replBlock || replOp->isBeforeInBlock(user);
897     });
898   }
899 
900   // In a second pass, erase all of the replaced operations in reverse. This
901   // allows processing nested operations before their parent region is
902   // destroyed.
903   for (auto &repl : llvm::reverse(replacements))
904     repl.first->erase();
905 
906   argConverter.applyRewrites(mapping);
907 
908   // Now that the ops have been erased, also erase dangling blocks.
909   eraseDanglingBlocks();
910 }
911 
912 //===----------------------------------------------------------------------===//
913 // State Management
914 
915 RewriterState ConversionPatternRewriterImpl::getCurrentState() {
916   return RewriterState(createdOps.size(), replacements.size(),
917                        argReplacements.size(), blockActions.size(),
918                        ignoredOps.size(), rootUpdates.size());
919 }
920 
921 void ConversionPatternRewriterImpl::resetState(RewriterState state) {
922   // Reset any operations that were updated in place.
923   for (unsigned i = state.numRootUpdates, e = rootUpdates.size(); i != e; ++i)
924     rootUpdates[i].resetOperation();
925   rootUpdates.resize(state.numRootUpdates);
926 
927   // Reset any replaced arguments.
928   for (BlockArgument replacedArg :
929        llvm::drop_begin(argReplacements, state.numArgReplacements))
930     mapping.erase(replacedArg);
931   argReplacements.resize(state.numArgReplacements);
932 
933   // Undo any block actions.
934   undoBlockActions(state.numBlockActions);
935 
936   // Reset any replaced operations and undo any saved mappings.
937   for (auto &repl : llvm::drop_begin(replacements, state.numReplacements))
938     for (auto result : repl.first->getResults())
939       mapping.erase(result);
940   while (replacements.size() != state.numReplacements)
941     replacements.pop_back();
942 
943   // Pop all of the newly created operations.
944   while (createdOps.size() != state.numCreatedOps) {
945     detachNestedAndErase(createdOps.back());
946     createdOps.pop_back();
947   }
948 
949   // Pop all of the recorded ignored operations that are no longer valid.
950   while (ignoredOps.size() != state.numIgnoredOperations)
951     ignoredOps.pop_back();
952 
953   // Reset operations with changed results.
954   while (!operationsWithChangedResults.empty() &&
955          operationsWithChangedResults.back() >= state.numReplacements)
956     operationsWithChangedResults.pop_back();
957 }
958 
959 void ConversionPatternRewriterImpl::eraseDanglingBlocks() {
960   for (auto &action : blockActions)
961     if (action.kind == BlockActionKind::Erase)
962       delete action.block;
963 }
964 
965 void ConversionPatternRewriterImpl::undoBlockActions(
966     unsigned numActionsToKeep) {
967   for (auto &action :
968        llvm::reverse(llvm::drop_begin(blockActions, numActionsToKeep))) {
969     switch (action.kind) {
970     // Delete the created block.
971     case BlockActionKind::Create: {
972       // Unlink all of the operations within this block, they will be deleted
973       // separately.
974       auto &blockOps = action.block->getOperations();
975       while (!blockOps.empty())
976         blockOps.remove(blockOps.begin());
977       action.block->dropAllDefinedValueUses();
978       action.block->erase();
979       break;
980     }
981     // Put the block (owned by action) back into its original position.
982     case BlockActionKind::Erase: {
983       auto &blockList = action.originalPosition.region->getBlocks();
984       Block *insertAfterBlock = action.originalPosition.insertAfterBlock;
985       blockList.insert((insertAfterBlock
986                             ? std::next(Region::iterator(insertAfterBlock))
987                             : blockList.begin()),
988                        action.block);
989       break;
990     }
991     // Split the block at the position which was originally the end of the
992     // destination block (owned by action), and put the instructions back into
993     // the block used before the merge.
994     case BlockActionKind::Merge: {
995       Block *sourceBlock = action.mergeInfo.sourceBlock;
996       Block::iterator splitPoint =
997           (action.mergeInfo.destBlockLastInst
998                ? ++Block::iterator(action.mergeInfo.destBlockLastInst)
999                : action.block->begin());
1000       sourceBlock->getOperations().splice(sourceBlock->begin(),
1001                                           action.block->getOperations(),
1002                                           splitPoint, action.block->end());
1003       break;
1004     }
1005     // Move the block back to its original position.
1006     case BlockActionKind::Move: {
1007       Region *originalRegion = action.originalPosition.region;
1008       Block *insertAfterBlock = action.originalPosition.insertAfterBlock;
1009       originalRegion->getBlocks().splice(
1010           (insertAfterBlock ? std::next(Region::iterator(insertAfterBlock))
1011                             : originalRegion->end()),
1012           action.block->getParent()->getBlocks(), action.block);
1013       break;
1014     }
1015     // Merge back the block that was split out.
1016     case BlockActionKind::Split: {
1017       action.originalBlock->getOperations().splice(
1018           action.originalBlock->end(), action.block->getOperations());
1019       action.block->dropAllDefinedValueUses();
1020       action.block->erase();
1021       break;
1022     }
1023     // Undo the type conversion.
1024     case BlockActionKind::TypeConversion: {
1025       argConverter.discardRewrites(action.block);
1026       break;
1027     }
1028     }
1029   }
1030   blockActions.resize(numActionsToKeep);
1031 }
1032 
1033 LogicalResult ConversionPatternRewriterImpl::remapValues(
1034     Location loc, PatternRewriter &rewriter, TypeConverter *converter,
1035     Operation::operand_range operands, SmallVectorImpl<Value> &remapped) {
1036   remapped.reserve(llvm::size(operands));
1037 
1038   SmallVector<Type, 1> legalTypes;
1039   for (auto it : llvm::enumerate(operands)) {
1040     Value operand = it.value();
1041     Type origType = operand.getType();
1042 
1043     // If a converter was provided, get the desired legal types for this
1044     // operand.
1045     Type desiredType;
1046     if (converter) {
1047       // If there is no legal conversion, fail to match this pattern.
1048       legalTypes.clear();
1049       if (failed(converter->convertType(origType, legalTypes))) {
1050         return notifyMatchFailure(loc, [=](Diagnostic &diag) {
1051           diag << "unable to convert type for operand #" << it.index()
1052                << ", type was " << origType;
1053         });
1054       }
1055       // TODO: There currently isn't any mechanism to do 1->N type conversion
1056       // via the PatternRewriter replacement API, so for now we just ignore it.
1057       if (legalTypes.size() == 1)
1058         desiredType = legalTypes.front();
1059     } else {
1060       // TODO: What we should do here is just set `desiredType` to `origType`
1061       // and then handle the necessary type conversions after the conversion
1062       // process has finished. Unfortunately a lot of patterns currently rely on
1063       // receiving the new operands even if the types change, so we keep the
1064       // original behavior here for now until all of the patterns relying on
1065       // this get updated.
1066     }
1067     Value newOperand = mapping.lookupOrDefault(operand, desiredType);
1068 
1069     // Handle the case where the conversion was 1->1 and the new operand type
1070     // isn't legal.
1071     Type newOperandType = newOperand.getType();
1072     if (converter && desiredType && newOperandType != desiredType) {
1073       // Attempt to materialize a conversion for this new value.
1074       newOperand = converter->materializeTargetConversion(
1075           rewriter, loc, desiredType, newOperand);
1076       if (!newOperand) {
1077         return notifyMatchFailure(loc, [=](Diagnostic &diag) {
1078           diag << "unable to materialize a conversion for "
1079                   "operand #"
1080                << it.index() << ", from " << newOperandType << " to "
1081                << desiredType;
1082         });
1083       }
1084     }
1085     remapped.push_back(newOperand);
1086   }
1087   return success();
1088 }
1089 
1090 bool ConversionPatternRewriterImpl::isOpIgnored(Operation *op) const {
1091   // Check to see if this operation was replaced or its parent ignored.
1092   return replacements.count(op) || ignoredOps.count(op->getParentOp());
1093 }
1094 
1095 void ConversionPatternRewriterImpl::markNestedOpsIgnored(Operation *op) {
1096   // Walk this operation and collect nested operations that define non-empty
1097   // regions. We mark such operations as 'ignored' so that we know we don't have
1098   // to convert them, or their nested ops.
1099   if (op->getNumRegions() == 0)
1100     return;
1101   op->walk([&](Operation *op) {
1102     if (llvm::any_of(op->getRegions(),
1103                      [](Region &region) { return !region.empty(); }))
1104       ignoredOps.insert(op);
1105   });
1106 }
1107 
1108 //===----------------------------------------------------------------------===//
1109 // Type Conversion
1110 
1111 FailureOr<Block *> ConversionPatternRewriterImpl::convertBlockSignature(
1112     Block *block, TypeConverter &converter,
1113     TypeConverter::SignatureConversion *conversion) {
1114   FailureOr<Block *> result =
1115       conversion ? argConverter.applySignatureConversion(block, converter,
1116                                                          *conversion, mapping)
1117                  : argConverter.convertSignature(block, converter, mapping);
1118   if (Block *newBlock = result.getValue()) {
1119     if (newBlock != block)
1120       blockActions.push_back(BlockAction::getTypeConversion(newBlock));
1121   }
1122   return result;
1123 }
1124 
1125 Block *ConversionPatternRewriterImpl::applySignatureConversion(
1126     Region *region, TypeConverter::SignatureConversion &conversion) {
1127   if (!region->empty()) {
1128     return *convertBlockSignature(&region->front(), defaultTypeConverter,
1129                                   &conversion);
1130   }
1131   return nullptr;
1132 }
1133 
1134 FailureOr<Block *> ConversionPatternRewriterImpl::convertRegionTypes(
1135     Region *region, TypeConverter &converter,
1136     TypeConverter::SignatureConversion *entryConversion) {
1137   argConverter.setConverter(region, &converter);
1138   if (region->empty())
1139     return nullptr;
1140 
1141   // Convert the arguments of each block within the region.
1142   FailureOr<Block *> newEntry =
1143       convertBlockSignature(&region->front(), converter, entryConversion);
1144   for (Block &block : llvm::make_early_inc_range(llvm::drop_begin(*region, 1)))
1145     if (failed(convertBlockSignature(&block, converter)))
1146       return failure();
1147   return newEntry;
1148 }
1149 
1150 //===----------------------------------------------------------------------===//
1151 // Rewriter Notification Hooks
1152 
1153 void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op,
1154                                                      ValueRange newValues) {
1155   assert(newValues.size() == op->getNumResults());
1156   assert(!replacements.count(op) && "operation was already replaced");
1157 
1158   // Track if any of the results changed, e.g. erased and replaced with null.
1159   bool resultChanged = false;
1160 
1161   // Create mappings for each of the new result values.
1162   Value newValue, result;
1163   for (auto it : llvm::zip(newValues, op->getResults())) {
1164     std::tie(newValue, result) = it;
1165     if (!newValue) {
1166       resultChanged = true;
1167       continue;
1168     }
1169     // Remap, and check for any result type changes.
1170     mapping.map(result, newValue);
1171     resultChanged |= (newValue.getType() != result.getType());
1172   }
1173   if (resultChanged)
1174     operationsWithChangedResults.push_back(replacements.size());
1175 
1176   // Record the requested operation replacement.
1177   TypeConverter *converter = nullptr;
1178   if (currentConversionPattern)
1179     converter = currentConversionPattern->getTypeConverter();
1180   replacements.insert(std::make_pair(op, OpReplacement(converter)));
1181 
1182   // Mark this operation as recursively ignored so that we don't need to
1183   // convert any nested operations.
1184   markNestedOpsIgnored(op);
1185 }
1186 
1187 void ConversionPatternRewriterImpl::notifyBlockIsBeingErased(Block *block) {
1188   Region *region = block->getParent();
1189   Block *origPrevBlock = block->getPrevNode();
1190   blockActions.push_back(BlockAction::getErase(block, {region, origPrevBlock}));
1191 }
1192 
1193 void ConversionPatternRewriterImpl::notifyCreatedBlock(Block *block) {
1194   blockActions.push_back(BlockAction::getCreate(block));
1195 }
1196 
1197 void ConversionPatternRewriterImpl::notifySplitBlock(Block *block,
1198                                                      Block *continuation) {
1199   blockActions.push_back(BlockAction::getSplit(continuation, block));
1200 }
1201 
1202 void ConversionPatternRewriterImpl::notifyBlocksBeingMerged(Block *block,
1203                                                             Block *srcBlock) {
1204   blockActions.push_back(BlockAction::getMerge(block, srcBlock));
1205 }
1206 
1207 void ConversionPatternRewriterImpl::notifyRegionIsBeingInlinedBefore(
1208     Region &region, Region &parent, Region::iterator before) {
1209   if (region.empty())
1210     return;
1211   Block *laterBlock = &region.back();
1212   for (auto &earlierBlock : llvm::drop_begin(llvm::reverse(region), 1)) {
1213     blockActions.push_back(
1214         BlockAction::getMove(laterBlock, {&region, &earlierBlock}));
1215     laterBlock = &earlierBlock;
1216   }
1217   blockActions.push_back(BlockAction::getMove(laterBlock, {&region, nullptr}));
1218 }
1219 
1220 void ConversionPatternRewriterImpl::notifyRegionWasClonedBefore(
1221     iterator_range<Region::iterator> &blocks, Location origRegionLoc) {
1222   for (Block &block : blocks)
1223     blockActions.push_back(BlockAction::getCreate(&block));
1224 
1225   // Compute the conversion set for the inlined region.
1226   auto result = computeConversionSet(blocks, origRegionLoc, createdOps);
1227 
1228   // This original region has already had its conversion set computed, so there
1229   // shouldn't be any new failures.
1230   (void)result;
1231   assert(succeeded(result) && "expected region to have no unreachable blocks");
1232 }
1233 
1234 LogicalResult ConversionPatternRewriterImpl::notifyMatchFailure(
1235     Location loc, function_ref<void(Diagnostic &)> reasonCallback) {
1236   LLVM_DEBUG({
1237     Diagnostic diag(loc, DiagnosticSeverity::Remark);
1238     reasonCallback(diag);
1239     logger.startLine() << "** Failure : " << diag.str() << "\n";
1240   });
1241   return failure();
1242 }
1243 
1244 //===----------------------------------------------------------------------===//
1245 // ConversionPatternRewriter
1246 //===----------------------------------------------------------------------===//
1247 
1248 ConversionPatternRewriter::ConversionPatternRewriter(MLIRContext *ctx)
1249     : PatternRewriter(ctx),
1250       impl(new detail::ConversionPatternRewriterImpl(*this)) {}
1251 ConversionPatternRewriter::~ConversionPatternRewriter() {}
1252 
1253 /// PatternRewriter hook for replacing the results of an operation when the
1254 /// given functor returns true.
1255 void ConversionPatternRewriter::replaceOpWithIf(
1256     Operation *op, ValueRange newValues, bool *allUsesReplaced,
1257     llvm::unique_function<bool(OpOperand &) const> functor) {
1258   // TODO: To support this we will need to rework a bit of how replacements are
1259   // tracked, given that this isn't guranteed to replace all of the uses of an
1260   // operation. The main change is that now an operation can be replaced
1261   // multiple times, in parts. The current "set" based tracking is mainly useful
1262   // for tracking if a replaced operation should be ignored, i.e. if all of the
1263   // uses will be replaced.
1264   llvm_unreachable(
1265       "replaceOpWithIf is currently not supported by DialectConversion");
1266 }
1267 
1268 /// PatternRewriter hook for replacing the results of an operation.
1269 void ConversionPatternRewriter::replaceOp(Operation *op, ValueRange newValues) {
1270   LLVM_DEBUG({
1271     impl->logger.startLine()
1272         << "** Replace : '" << op->getName() << "'(" << op << ")\n";
1273   });
1274   impl->notifyOpReplaced(op, newValues);
1275 }
1276 
1277 /// PatternRewriter hook for erasing a dead operation. The uses of this
1278 /// operation *must* be made dead by the end of the conversion process,
1279 /// otherwise an assert will be issued.
1280 void ConversionPatternRewriter::eraseOp(Operation *op) {
1281   LLVM_DEBUG({
1282     impl->logger.startLine()
1283         << "** Erase   : '" << op->getName() << "'(" << op << ")\n";
1284   });
1285   SmallVector<Value, 1> nullRepls(op->getNumResults(), nullptr);
1286   impl->notifyOpReplaced(op, nullRepls);
1287 }
1288 
1289 void ConversionPatternRewriter::eraseBlock(Block *block) {
1290   impl->notifyBlockIsBeingErased(block);
1291 
1292   // Mark all ops for erasure.
1293   for (Operation &op : *block)
1294     eraseOp(&op);
1295 
1296   // Unlink the block from its parent region. The block is kept in the block
1297   // action and will be actually destroyed when rewrites are applied. This
1298   // allows us to keep the operations in the block live and undo the removal by
1299   // re-inserting the block.
1300   block->getParent()->getBlocks().remove(block);
1301 }
1302 
1303 Block *ConversionPatternRewriter::applySignatureConversion(
1304     Region *region, TypeConverter::SignatureConversion &conversion) {
1305   return impl->applySignatureConversion(region, conversion);
1306 }
1307 
1308 FailureOr<Block *> ConversionPatternRewriter::convertRegionTypes(
1309     Region *region, TypeConverter &converter,
1310     TypeConverter::SignatureConversion *entryConversion) {
1311   return impl->convertRegionTypes(region, converter, entryConversion);
1312 }
1313 
1314 void ConversionPatternRewriter::replaceUsesOfBlockArgument(BlockArgument from,
1315                                                            Value to) {
1316   LLVM_DEBUG({
1317     Operation *parentOp = from.getOwner()->getParentOp();
1318     impl->logger.startLine() << "** Replace Argument : '" << from
1319                              << "'(in region of '" << parentOp->getName()
1320                              << "'(" << from.getOwner()->getParentOp() << ")\n";
1321   });
1322   impl->argReplacements.push_back(from);
1323   impl->mapping.map(impl->mapping.lookupOrDefault(from), to);
1324 }
1325 
1326 /// Return the converted value that replaces 'key'. Return 'key' if there is
1327 /// no such a converted value.
1328 Value ConversionPatternRewriter::getRemappedValue(Value key) {
1329   return impl->mapping.lookupOrDefault(key);
1330 }
1331 
1332 /// PatternRewriter hook for creating a new block with the given arguments.
1333 void ConversionPatternRewriter::notifyBlockCreated(Block *block) {
1334   impl->notifyCreatedBlock(block);
1335 }
1336 
1337 /// PatternRewriter hook for splitting a block into two parts.
1338 Block *ConversionPatternRewriter::splitBlock(Block *block,
1339                                              Block::iterator before) {
1340   auto *continuation = PatternRewriter::splitBlock(block, before);
1341   impl->notifySplitBlock(block, continuation);
1342   return continuation;
1343 }
1344 
1345 /// PatternRewriter hook for merging a block into another.
1346 void ConversionPatternRewriter::mergeBlocks(Block *source, Block *dest,
1347                                             ValueRange argValues) {
1348   impl->notifyBlocksBeingMerged(dest, source);
1349   assert(llvm::all_of(source->getPredecessors(),
1350                       [dest](Block *succ) { return succ == dest; }) &&
1351          "expected 'source' to have no predecessors or only 'dest'");
1352   assert(argValues.size() == source->getNumArguments() &&
1353          "incorrect # of argument replacement values");
1354   for (auto it : llvm::zip(source->getArguments(), argValues))
1355     replaceUsesOfBlockArgument(std::get<0>(it), std::get<1>(it));
1356   dest->getOperations().splice(dest->end(), source->getOperations());
1357   eraseBlock(source);
1358 }
1359 
1360 /// PatternRewriter hook for moving blocks out of a region.
1361 void ConversionPatternRewriter::inlineRegionBefore(Region &region,
1362                                                    Region &parent,
1363                                                    Region::iterator before) {
1364   impl->notifyRegionIsBeingInlinedBefore(region, parent, before);
1365   PatternRewriter::inlineRegionBefore(region, parent, before);
1366 }
1367 
1368 /// PatternRewriter hook for cloning blocks of one region into another.
1369 void ConversionPatternRewriter::cloneRegionBefore(
1370     Region &region, Region &parent, Region::iterator before,
1371     BlockAndValueMapping &mapping) {
1372   if (region.empty())
1373     return;
1374   PatternRewriter::cloneRegionBefore(region, parent, before, mapping);
1375 
1376   // Collect the range of the cloned blocks.
1377   auto clonedBeginIt = mapping.lookup(&region.front())->getIterator();
1378   auto clonedBlocks = llvm::make_range(clonedBeginIt, before);
1379   impl->notifyRegionWasClonedBefore(clonedBlocks, region.getLoc());
1380 }
1381 
1382 /// PatternRewriter hook for creating a new operation.
1383 void ConversionPatternRewriter::notifyOperationInserted(Operation *op) {
1384   LLVM_DEBUG({
1385     impl->logger.startLine()
1386         << "** Insert  : '" << op->getName() << "'(" << op << ")\n";
1387   });
1388   impl->createdOps.push_back(op);
1389 }
1390 
1391 /// PatternRewriter hook for updating the root operation in-place.
1392 void ConversionPatternRewriter::startRootUpdate(Operation *op) {
1393 #ifndef NDEBUG
1394   impl->pendingRootUpdates.insert(op);
1395 #endif
1396   impl->rootUpdates.emplace_back(op);
1397 }
1398 
1399 /// PatternRewriter hook for updating the root operation in-place.
1400 void ConversionPatternRewriter::finalizeRootUpdate(Operation *op) {
1401   // There is nothing to do here, we only need to track the operation at the
1402   // start of the update.
1403 #ifndef NDEBUG
1404   assert(impl->pendingRootUpdates.erase(op) &&
1405          "operation did not have a pending in-place update");
1406 #endif
1407 }
1408 
1409 /// PatternRewriter hook for updating the root operation in-place.
1410 void ConversionPatternRewriter::cancelRootUpdate(Operation *op) {
1411 #ifndef NDEBUG
1412   assert(impl->pendingRootUpdates.erase(op) &&
1413          "operation did not have a pending in-place update");
1414 #endif
1415   // Erase the last update for this operation.
1416   auto stateHasOp = [op](const auto &it) { return it.getOperation() == op; };
1417   auto &rootUpdates = impl->rootUpdates;
1418   auto it = llvm::find_if(llvm::reverse(rootUpdates), stateHasOp);
1419   rootUpdates.erase(rootUpdates.begin() + (rootUpdates.rend() - it));
1420 }
1421 
1422 /// PatternRewriter hook for notifying match failure reasons.
1423 LogicalResult ConversionPatternRewriter::notifyMatchFailure(
1424     Operation *op, function_ref<void(Diagnostic &)> reasonCallback) {
1425   return impl->notifyMatchFailure(op->getLoc(), reasonCallback);
1426 }
1427 
1428 /// Return a reference to the internal implementation.
1429 detail::ConversionPatternRewriterImpl &ConversionPatternRewriter::getImpl() {
1430   return *impl;
1431 }
1432 
1433 //===----------------------------------------------------------------------===//
1434 // ConversionPattern
1435 //===----------------------------------------------------------------------===//
1436 
1437 /// Attempt to match and rewrite the IR root at the specified operation.
1438 LogicalResult
1439 ConversionPattern::matchAndRewrite(Operation *op,
1440                                    PatternRewriter &rewriter) const {
1441   auto &dialectRewriter = static_cast<ConversionPatternRewriter &>(rewriter);
1442   auto &rewriterImpl = dialectRewriter.getImpl();
1443 
1444   // Track the current conversion pattern in the rewriter.
1445   assert(!rewriterImpl.currentConversionPattern &&
1446          "already inside of a pattern rewrite");
1447   llvm::SaveAndRestore<const ConversionPattern *> currentPatternGuard(
1448       rewriterImpl.currentConversionPattern, this);
1449 
1450   // Remap the operands of the operation.
1451   SmallVector<Value, 4> operands;
1452   if (failed(rewriterImpl.remapValues(op->getLoc(), rewriter,
1453                                       getTypeConverter(), op->getOperands(),
1454                                       operands))) {
1455     return failure();
1456   }
1457   return matchAndRewrite(op, operands, dialectRewriter);
1458 }
1459 
1460 //===----------------------------------------------------------------------===//
1461 // OperationLegalizer
1462 //===----------------------------------------------------------------------===//
1463 
1464 namespace {
1465 /// A set of rewrite patterns that can be used to legalize a given operation.
1466 using LegalizationPatterns = SmallVector<const Pattern *, 1>;
1467 
1468 /// This class defines a recursive operation legalizer.
1469 class OperationLegalizer {
1470 public:
1471   using LegalizationAction = ConversionTarget::LegalizationAction;
1472 
1473   OperationLegalizer(ConversionTarget &targetInfo,
1474                      const FrozenRewritePatternList &patterns);
1475 
1476   /// Returns true if the given operation is known to be illegal on the target.
1477   bool isIllegal(Operation *op) const;
1478 
1479   /// Attempt to legalize the given operation. Returns success if the operation
1480   /// was legalized, failure otherwise.
1481   LogicalResult legalize(Operation *op, ConversionPatternRewriter &rewriter);
1482 
1483   /// Returns the conversion target in use by the legalizer.
1484   ConversionTarget &getTarget() { return target; }
1485 
1486 private:
1487   /// Attempt to legalize the given operation by folding it.
1488   LogicalResult legalizeWithFold(Operation *op,
1489                                  ConversionPatternRewriter &rewriter);
1490 
1491   /// Attempt to legalize the given operation by applying a pattern. Returns
1492   /// success if the operation was legalized, failure otherwise.
1493   LogicalResult legalizeWithPattern(Operation *op,
1494                                     ConversionPatternRewriter &rewriter);
1495 
1496   /// Return true if the given pattern may be applied to the given operation,
1497   /// false otherwise.
1498   bool canApplyPattern(Operation *op, const Pattern &pattern,
1499                        ConversionPatternRewriter &rewriter);
1500 
1501   /// Legalize the resultant IR after successfully applying the given pattern.
1502   LogicalResult legalizePatternResult(Operation *op, const Pattern &pattern,
1503                                       ConversionPatternRewriter &rewriter,
1504                                       RewriterState &curState);
1505 
1506   /// Legalizes the actions registered during the execution of a pattern.
1507   LogicalResult legalizePatternBlockActions(Operation *op,
1508                                             ConversionPatternRewriter &rewriter,
1509                                             ConversionPatternRewriterImpl &impl,
1510                                             RewriterState &state,
1511                                             RewriterState &newState);
1512   LogicalResult legalizePatternCreatedOperations(
1513       ConversionPatternRewriter &rewriter, ConversionPatternRewriterImpl &impl,
1514       RewriterState &state, RewriterState &newState);
1515   LogicalResult legalizePatternRootUpdates(ConversionPatternRewriter &rewriter,
1516                                            ConversionPatternRewriterImpl &impl,
1517                                            RewriterState &state,
1518                                            RewriterState &newState);
1519 
1520   //===--------------------------------------------------------------------===//
1521   // Cost Model
1522   //===--------------------------------------------------------------------===//
1523 
1524   /// Build an optimistic legalization graph given the provided patterns. This
1525   /// function populates 'anyOpLegalizerPatterns' and 'legalizerPatterns' with
1526   /// patterns for operations that are not directly legal, but may be
1527   /// transitively legal for the current target given the provided patterns.
1528   void buildLegalizationGraph(
1529       LegalizationPatterns &anyOpLegalizerPatterns,
1530       DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns);
1531 
1532   /// Compute the benefit of each node within the computed legalization graph.
1533   /// This orders the patterns within 'legalizerPatterns' based upon two
1534   /// criteria:
1535   ///  1) Prefer patterns that have the lowest legalization depth, i.e.
1536   ///     represent the more direct mapping to the target.
1537   ///  2) When comparing patterns with the same legalization depth, prefer the
1538   ///     pattern with the highest PatternBenefit. This allows for users to
1539   ///     prefer specific legalizations over others.
1540   void computeLegalizationGraphBenefit(
1541       LegalizationPatterns &anyOpLegalizerPatterns,
1542       DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns);
1543 
1544   /// Compute the legalization depth when legalizing an operation of the given
1545   /// type.
1546   unsigned computeOpLegalizationDepth(
1547       OperationName op, DenseMap<OperationName, unsigned> &minOpPatternDepth,
1548       DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns);
1549 
1550   /// Apply the conversion cost model to the given set of patterns, and return
1551   /// the smallest legalization depth of any of the patterns. See
1552   /// `computeLegalizationGraphBenefit` for the breakdown of the cost model.
1553   unsigned applyCostModelToPatterns(
1554       LegalizationPatterns &patterns,
1555       DenseMap<OperationName, unsigned> &minOpPatternDepth,
1556       DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns);
1557 
1558   /// The current set of patterns that have been applied.
1559   SmallPtrSet<const Pattern *, 8> appliedPatterns;
1560 
1561   /// The legalization information provided by the target.
1562   ConversionTarget &target;
1563 
1564   /// The pattern applicator to use for conversions.
1565   PatternApplicator applicator;
1566 };
1567 } // namespace
1568 
1569 OperationLegalizer::OperationLegalizer(ConversionTarget &targetInfo,
1570                                        const FrozenRewritePatternList &patterns)
1571     : target(targetInfo), applicator(patterns) {
1572   // The set of patterns that can be applied to illegal operations to transform
1573   // them into legal ones.
1574   DenseMap<OperationName, LegalizationPatterns> legalizerPatterns;
1575   LegalizationPatterns anyOpLegalizerPatterns;
1576 
1577   buildLegalizationGraph(anyOpLegalizerPatterns, legalizerPatterns);
1578   computeLegalizationGraphBenefit(anyOpLegalizerPatterns, legalizerPatterns);
1579 }
1580 
1581 bool OperationLegalizer::isIllegal(Operation *op) const {
1582   // Check if the target explicitly marked this operation as illegal.
1583   return target.getOpAction(op->getName()) == LegalizationAction::Illegal;
1584 }
1585 
1586 LogicalResult
1587 OperationLegalizer::legalize(Operation *op,
1588                              ConversionPatternRewriter &rewriter) {
1589 #ifndef NDEBUG
1590   const char *logLineComment =
1591       "//===-------------------------------------------===//\n";
1592 
1593   auto &rewriterImpl = rewriter.getImpl();
1594 #endif
1595   LLVM_DEBUG({
1596     auto &os = rewriterImpl.logger;
1597     os.getOStream() << "\n";
1598     os.startLine() << logLineComment;
1599     os.startLine() << "Legalizing operation : '" << op->getName() << "'(" << op
1600                    << ") {\n";
1601     os.indent();
1602 
1603     // If the operation has no regions, just print it here.
1604     if (op->getNumRegions() == 0) {
1605       op->print(os.startLine(), OpPrintingFlags().printGenericOpForm());
1606       os.getOStream() << "\n\n";
1607     }
1608   });
1609 
1610   // Check if this operation is legal on the target.
1611   if (auto legalityInfo = target.isLegal(op)) {
1612     LLVM_DEBUG({
1613       logSuccess(
1614           rewriterImpl.logger, "operation marked legal by the target{0}",
1615           legalityInfo->isRecursivelyLegal
1616               ? "; NOTE: operation is recursively legal; skipping internals"
1617               : "");
1618       rewriterImpl.logger.startLine() << logLineComment;
1619     });
1620 
1621     // If this operation is recursively legal, mark its children as ignored so
1622     // that we don't consider them for legalization.
1623     if (legalityInfo->isRecursivelyLegal)
1624       rewriter.getImpl().markNestedOpsIgnored(op);
1625     return success();
1626   }
1627 
1628   // Check to see if the operation is ignored and doesn't need to be converted.
1629   if (rewriter.getImpl().isOpIgnored(op)) {
1630     LLVM_DEBUG({
1631       logSuccess(rewriterImpl.logger,
1632                  "operation marked 'ignored' during conversion");
1633       rewriterImpl.logger.startLine() << logLineComment;
1634     });
1635     return success();
1636   }
1637 
1638   // If the operation isn't legal, try to fold it in-place.
1639   // TODO: Should we always try to do this, even if the op is
1640   // already legal?
1641   if (succeeded(legalizeWithFold(op, rewriter))) {
1642     LLVM_DEBUG({
1643       logSuccess(rewriterImpl.logger, "operation was folded");
1644       rewriterImpl.logger.startLine() << logLineComment;
1645     });
1646     return success();
1647   }
1648 
1649   // Otherwise, we need to apply a legalization pattern to this operation.
1650   if (succeeded(legalizeWithPattern(op, rewriter))) {
1651     LLVM_DEBUG({
1652       logSuccess(rewriterImpl.logger, "");
1653       rewriterImpl.logger.startLine() << logLineComment;
1654     });
1655     return success();
1656   }
1657 
1658   LLVM_DEBUG({
1659     logFailure(rewriterImpl.logger, "no matched legalization pattern");
1660     rewriterImpl.logger.startLine() << logLineComment;
1661   });
1662   return failure();
1663 }
1664 
1665 LogicalResult
1666 OperationLegalizer::legalizeWithFold(Operation *op,
1667                                      ConversionPatternRewriter &rewriter) {
1668   auto &rewriterImpl = rewriter.getImpl();
1669   RewriterState curState = rewriterImpl.getCurrentState();
1670 
1671   LLVM_DEBUG({
1672     rewriterImpl.logger.startLine() << "* Fold {\n";
1673     rewriterImpl.logger.indent();
1674   });
1675 
1676   // Try to fold the operation.
1677   SmallVector<Value, 2> replacementValues;
1678   rewriter.setInsertionPoint(op);
1679   if (failed(rewriter.tryFold(op, replacementValues))) {
1680     LLVM_DEBUG(logFailure(rewriterImpl.logger, "unable to fold"));
1681     return failure();
1682   }
1683 
1684   // Insert a replacement for 'op' with the folded replacement values.
1685   rewriter.replaceOp(op, replacementValues);
1686 
1687   // Recursively legalize any new constant operations.
1688   for (unsigned i = curState.numCreatedOps, e = rewriterImpl.createdOps.size();
1689        i != e; ++i) {
1690     Operation *cstOp = rewriterImpl.createdOps[i];
1691     if (failed(legalize(cstOp, rewriter))) {
1692       LLVM_DEBUG(logFailure(rewriterImpl.logger,
1693                             "generated constant '{0}' was illegal",
1694                             cstOp->getName()));
1695       rewriterImpl.resetState(curState);
1696       return failure();
1697     }
1698   }
1699 
1700   LLVM_DEBUG(logSuccess(rewriterImpl.logger, ""));
1701   return success();
1702 }
1703 
1704 LogicalResult
1705 OperationLegalizer::legalizeWithPattern(Operation *op,
1706                                         ConversionPatternRewriter &rewriter) {
1707   auto &rewriterImpl = rewriter.getImpl();
1708 
1709   // Functor that returns if the given pattern may be applied.
1710   auto canApply = [&](const Pattern &pattern) {
1711     return canApplyPattern(op, pattern, rewriter);
1712   };
1713 
1714   // Functor that cleans up the rewriter state after a pattern failed to match.
1715   RewriterState curState = rewriterImpl.getCurrentState();
1716   auto onFailure = [&](const Pattern &pattern) {
1717     LLVM_DEBUG(logFailure(rewriterImpl.logger, "pattern failed to match"));
1718     rewriterImpl.resetState(curState);
1719     appliedPatterns.erase(&pattern);
1720   };
1721 
1722   // Functor that performs additional legalization when a pattern is
1723   // successfully applied.
1724   auto onSuccess = [&](const Pattern &pattern) {
1725     auto result = legalizePatternResult(op, pattern, rewriter, curState);
1726     appliedPatterns.erase(&pattern);
1727     if (failed(result))
1728       rewriterImpl.resetState(curState);
1729     return result;
1730   };
1731 
1732   // Try to match and rewrite a pattern on this operation.
1733   return applicator.matchAndRewrite(op, rewriter, canApply, onFailure,
1734                                     onSuccess);
1735 }
1736 
1737 bool OperationLegalizer::canApplyPattern(Operation *op, const Pattern &pattern,
1738                                          ConversionPatternRewriter &rewriter) {
1739   LLVM_DEBUG({
1740     auto &os = rewriter.getImpl().logger;
1741     os.getOStream() << "\n";
1742     os.startLine() << "* Pattern : '" << op->getName() << " -> (";
1743     llvm::interleaveComma(pattern.getGeneratedOps(), llvm::dbgs());
1744     os.getOStream() << ")' {\n";
1745     os.indent();
1746   });
1747 
1748   // Ensure that we don't cycle by not allowing the same pattern to be
1749   // applied twice in the same recursion stack if it is not known to be safe.
1750   if (!pattern.hasBoundedRewriteRecursion() &&
1751       !appliedPatterns.insert(&pattern).second) {
1752     LLVM_DEBUG(
1753         logFailure(rewriter.getImpl().logger, "pattern was already applied"));
1754     return false;
1755   }
1756   return true;
1757 }
1758 
1759 LogicalResult
1760 OperationLegalizer::legalizePatternResult(Operation *op, const Pattern &pattern,
1761                                           ConversionPatternRewriter &rewriter,
1762                                           RewriterState &curState) {
1763   auto &impl = rewriter.getImpl();
1764 
1765 #ifndef NDEBUG
1766   assert(impl.pendingRootUpdates.empty() && "dangling root updates");
1767 #endif
1768 
1769   // Check that the root was either replaced or updated in place.
1770   auto replacedRoot = [&] {
1771     return llvm::any_of(
1772         llvm::drop_begin(impl.replacements, curState.numReplacements),
1773         [op](auto &it) { return it.first == op; });
1774   };
1775   auto updatedRootInPlace = [&] {
1776     return llvm::any_of(
1777         llvm::drop_begin(impl.rootUpdates, curState.numRootUpdates),
1778         [op](auto &state) { return state.getOperation() == op; });
1779   };
1780   (void)replacedRoot;
1781   (void)updatedRootInPlace;
1782   assert((replacedRoot() || updatedRootInPlace()) &&
1783          "expected pattern to replace the root operation");
1784 
1785   // Legalize each of the actions registered during application.
1786   RewriterState newState = impl.getCurrentState();
1787   if (failed(legalizePatternBlockActions(op, rewriter, impl, curState,
1788                                          newState)) ||
1789       failed(legalizePatternRootUpdates(rewriter, impl, curState, newState)) ||
1790       failed(legalizePatternCreatedOperations(rewriter, impl, curState,
1791                                               newState))) {
1792     return failure();
1793   }
1794 
1795   LLVM_DEBUG(logSuccess(impl.logger, "pattern applied successfully"));
1796   return success();
1797 }
1798 
1799 LogicalResult OperationLegalizer::legalizePatternBlockActions(
1800     Operation *op, ConversionPatternRewriter &rewriter,
1801     ConversionPatternRewriterImpl &impl, RewriterState &state,
1802     RewriterState &newState) {
1803   SmallPtrSet<Operation *, 16> operationsToIgnore;
1804 
1805   // If the pattern moved or created any blocks, make sure the types of block
1806   // arguments get legalized.
1807   for (int i = state.numBlockActions, e = newState.numBlockActions; i != e;
1808        ++i) {
1809     auto &action = impl.blockActions[i];
1810     if (action.kind == BlockActionKind::TypeConversion ||
1811         action.kind == BlockActionKind::Erase)
1812       continue;
1813     // Only check blocks outside of the current operation.
1814     Operation *parentOp = action.block->getParentOp();
1815     if (!parentOp || parentOp == op || action.block->getNumArguments() == 0)
1816       continue;
1817 
1818     // If the region of the block has a type converter, try to convert the block
1819     // directly.
1820     if (auto *converter =
1821             impl.argConverter.getConverter(action.block->getParent())) {
1822       if (failed(impl.convertBlockSignature(action.block, *converter))) {
1823         LLVM_DEBUG(logFailure(impl.logger, "failed to convert types of moved "
1824                                            "block"));
1825         return failure();
1826       }
1827       continue;
1828     }
1829 
1830     // Otherwise, check that this operation isn't one generated by this pattern.
1831     // This is because we will attempt to legalize the parent operation, and
1832     // blocks in regions created by this pattern will already be legalized later
1833     // on. If we haven't built the set yet, build it now.
1834     if (operationsToIgnore.empty()) {
1835       auto createdOps = ArrayRef<Operation *>(impl.createdOps)
1836                             .drop_front(state.numCreatedOps);
1837       operationsToIgnore.insert(createdOps.begin(), createdOps.end());
1838     }
1839 
1840     // If this operation should be considered for re-legalization, try it.
1841     if (operationsToIgnore.insert(parentOp).second &&
1842         failed(legalize(parentOp, rewriter))) {
1843       LLVM_DEBUG(logFailure(
1844           impl.logger, "operation '{0}'({1}) became illegal after block action",
1845           parentOp->getName(), parentOp));
1846       return failure();
1847     }
1848   }
1849   return success();
1850 }
1851 LogicalResult OperationLegalizer::legalizePatternCreatedOperations(
1852     ConversionPatternRewriter &rewriter, ConversionPatternRewriterImpl &impl,
1853     RewriterState &state, RewriterState &newState) {
1854   for (int i = state.numCreatedOps, e = newState.numCreatedOps; i != e; ++i) {
1855     Operation *op = impl.createdOps[i];
1856     if (failed(legalize(op, rewriter))) {
1857       LLVM_DEBUG(logFailure(impl.logger,
1858                             "generated operation '{0}'({1}) was illegal",
1859                             op->getName(), op));
1860       return failure();
1861     }
1862   }
1863   return success();
1864 }
1865 LogicalResult OperationLegalizer::legalizePatternRootUpdates(
1866     ConversionPatternRewriter &rewriter, ConversionPatternRewriterImpl &impl,
1867     RewriterState &state, RewriterState &newState) {
1868   for (int i = state.numRootUpdates, e = newState.numRootUpdates; i != e; ++i) {
1869     Operation *op = impl.rootUpdates[i].getOperation();
1870     if (failed(legalize(op, rewriter))) {
1871       LLVM_DEBUG(logFailure(impl.logger,
1872                             "operation updated in-place '{0}' was illegal",
1873                             op->getName()));
1874       return failure();
1875     }
1876   }
1877   return success();
1878 }
1879 
1880 //===----------------------------------------------------------------------===//
1881 // Cost Model
1882 
1883 void OperationLegalizer::buildLegalizationGraph(
1884     LegalizationPatterns &anyOpLegalizerPatterns,
1885     DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns) {
1886   // A mapping between an operation and a set of operations that can be used to
1887   // generate it.
1888   DenseMap<OperationName, SmallPtrSet<OperationName, 2>> parentOps;
1889   // A mapping between an operation and any currently invalid patterns it has.
1890   DenseMap<OperationName, SmallPtrSet<const Pattern *, 2>> invalidPatterns;
1891   // A worklist of patterns to consider for legality.
1892   llvm::SetVector<const Pattern *> patternWorklist;
1893 
1894   // Build the mapping from operations to the parent ops that may generate them.
1895   applicator.walkAllPatterns([&](const Pattern &pattern) {
1896     Optional<OperationName> root = pattern.getRootKind();
1897 
1898     // If the pattern has no specific root, we can't analyze the relationship
1899     // between the root op and generated operations. Given that, add all such
1900     // patterns to the legalization set.
1901     if (!root) {
1902       anyOpLegalizerPatterns.push_back(&pattern);
1903       return;
1904     }
1905 
1906     // Skip operations that are always known to be legal.
1907     if (target.getOpAction(*root) == LegalizationAction::Legal)
1908       return;
1909 
1910     // Add this pattern to the invalid set for the root op and record this root
1911     // as a parent for any generated operations.
1912     invalidPatterns[*root].insert(&pattern);
1913     for (auto op : pattern.getGeneratedOps())
1914       parentOps[op].insert(*root);
1915 
1916     // Add this pattern to the worklist.
1917     patternWorklist.insert(&pattern);
1918   });
1919 
1920   // If there are any patterns that don't have a specific root kind, we can't
1921   // make direct assumptions about what operations will never be legalized.
1922   // Note: Technically we could, but it would require an analysis that may
1923   // recurse into itself. It would be better to perform this kind of filtering
1924   // at a higher level than here anyways.
1925   if (!anyOpLegalizerPatterns.empty()) {
1926     for (const Pattern *pattern : patternWorklist)
1927       legalizerPatterns[*pattern->getRootKind()].push_back(pattern);
1928     return;
1929   }
1930 
1931   while (!patternWorklist.empty()) {
1932     auto *pattern = patternWorklist.pop_back_val();
1933 
1934     // Check to see if any of the generated operations are invalid.
1935     if (llvm::any_of(pattern->getGeneratedOps(), [&](OperationName op) {
1936           Optional<LegalizationAction> action = target.getOpAction(op);
1937           return !legalizerPatterns.count(op) &&
1938                  (!action || action == LegalizationAction::Illegal);
1939         }))
1940       continue;
1941 
1942     // Otherwise, if all of the generated operation are valid, this op is now
1943     // legal so add all of the child patterns to the worklist.
1944     legalizerPatterns[*pattern->getRootKind()].push_back(pattern);
1945     invalidPatterns[*pattern->getRootKind()].erase(pattern);
1946 
1947     // Add any invalid patterns of the parent operations to see if they have now
1948     // become legal.
1949     for (auto op : parentOps[*pattern->getRootKind()])
1950       patternWorklist.set_union(invalidPatterns[op]);
1951   }
1952 }
1953 
1954 void OperationLegalizer::computeLegalizationGraphBenefit(
1955     LegalizationPatterns &anyOpLegalizerPatterns,
1956     DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns) {
1957   // The smallest pattern depth, when legalizing an operation.
1958   DenseMap<OperationName, unsigned> minOpPatternDepth;
1959 
1960   // For each operation that is transitively legal, compute a cost for it.
1961   for (auto &opIt : legalizerPatterns)
1962     if (!minOpPatternDepth.count(opIt.first))
1963       computeOpLegalizationDepth(opIt.first, minOpPatternDepth,
1964                                  legalizerPatterns);
1965 
1966   // Apply the cost model to the patterns that can match any operation. Those
1967   // with a specific operation type are already resolved when computing the op
1968   // legalization depth.
1969   if (!anyOpLegalizerPatterns.empty())
1970     applyCostModelToPatterns(anyOpLegalizerPatterns, minOpPatternDepth,
1971                              legalizerPatterns);
1972 
1973   // Apply a cost model to the pattern applicator. We order patterns first by
1974   // depth then benefit. `legalizerPatterns` contains per-op patterns by
1975   // decreasing benefit.
1976   applicator.applyCostModel([&](const Pattern &pattern) {
1977     ArrayRef<const Pattern *> orderedPatternList;
1978     if (Optional<OperationName> rootName = pattern.getRootKind())
1979       orderedPatternList = legalizerPatterns[*rootName];
1980     else
1981       orderedPatternList = anyOpLegalizerPatterns;
1982 
1983     // If the pattern is not found, then it was removed and cannot be matched.
1984     auto it = llvm::find(orderedPatternList, &pattern);
1985     if (it == orderedPatternList.end())
1986       return PatternBenefit::impossibleToMatch();
1987 
1988     // Patterns found earlier in the list have higher benefit.
1989     return PatternBenefit(std::distance(it, orderedPatternList.end()));
1990   });
1991 }
1992 
1993 unsigned OperationLegalizer::computeOpLegalizationDepth(
1994     OperationName op, DenseMap<OperationName, unsigned> &minOpPatternDepth,
1995     DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns) {
1996   // Check for existing depth.
1997   auto depthIt = minOpPatternDepth.find(op);
1998   if (depthIt != minOpPatternDepth.end())
1999     return depthIt->second;
2000 
2001   // If a mapping for this operation does not exist, then this operation
2002   // is always legal. Return 0 as the depth for a directly legal operation.
2003   auto opPatternsIt = legalizerPatterns.find(op);
2004   if (opPatternsIt == legalizerPatterns.end() || opPatternsIt->second.empty())
2005     return 0u;
2006 
2007   // Record this initial depth in case we encounter this op again when
2008   // recursively computing the depth.
2009   minOpPatternDepth.try_emplace(op, std::numeric_limits<unsigned>::max());
2010 
2011   // Apply the cost model to the operation patterns, and update the minimum
2012   // depth.
2013   unsigned minDepth = applyCostModelToPatterns(
2014       opPatternsIt->second, minOpPatternDepth, legalizerPatterns);
2015   minOpPatternDepth[op] = minDepth;
2016   return minDepth;
2017 }
2018 
2019 unsigned OperationLegalizer::applyCostModelToPatterns(
2020     LegalizationPatterns &patterns,
2021     DenseMap<OperationName, unsigned> &minOpPatternDepth,
2022     DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns) {
2023   unsigned minDepth = std::numeric_limits<unsigned>::max();
2024 
2025   // Compute the depth for each pattern within the set.
2026   SmallVector<std::pair<const Pattern *, unsigned>, 4> patternsByDepth;
2027   patternsByDepth.reserve(patterns.size());
2028   for (const Pattern *pattern : patterns) {
2029     unsigned depth = 0;
2030     for (auto generatedOp : pattern->getGeneratedOps()) {
2031       unsigned generatedOpDepth = computeOpLegalizationDepth(
2032           generatedOp, minOpPatternDepth, legalizerPatterns);
2033       depth = std::max(depth, generatedOpDepth + 1);
2034     }
2035     patternsByDepth.emplace_back(pattern, depth);
2036 
2037     // Update the minimum depth of the pattern list.
2038     minDepth = std::min(minDepth, depth);
2039   }
2040 
2041   // If the operation only has one legalization pattern, there is no need to
2042   // sort them.
2043   if (patternsByDepth.size() == 1)
2044     return minDepth;
2045 
2046   // Sort the patterns by those likely to be the most beneficial.
2047   llvm::array_pod_sort(patternsByDepth.begin(), patternsByDepth.end(),
2048                        [](const std::pair<const Pattern *, unsigned> *lhs,
2049                           const std::pair<const Pattern *, unsigned> *rhs) {
2050                          // First sort by the smaller pattern legalization
2051                          // depth.
2052                          if (lhs->second != rhs->second)
2053                            return llvm::array_pod_sort_comparator<unsigned>(
2054                                &lhs->second, &rhs->second);
2055 
2056                          // Then sort by the larger pattern benefit.
2057                          auto lhsBenefit = lhs->first->getBenefit();
2058                          auto rhsBenefit = rhs->first->getBenefit();
2059                          return llvm::array_pod_sort_comparator<PatternBenefit>(
2060                              &rhsBenefit, &lhsBenefit);
2061                        });
2062 
2063   // Update the legalization pattern to use the new sorted list.
2064   patterns.clear();
2065   for (auto &patternIt : patternsByDepth)
2066     patterns.push_back(patternIt.first);
2067   return minDepth;
2068 }
2069 
2070 //===----------------------------------------------------------------------===//
2071 // OperationConverter
2072 //===----------------------------------------------------------------------===//
2073 namespace {
2074 enum OpConversionMode {
2075   // In this mode, the conversion will ignore failed conversions to allow
2076   // illegal operations to co-exist in the IR.
2077   Partial,
2078 
2079   // In this mode, all operations must be legal for the given target for the
2080   // conversion to succeed.
2081   Full,
2082 
2083   // In this mode, operations are analyzed for legality. No actual rewrites are
2084   // applied to the operations on success.
2085   Analysis,
2086 };
2087 
2088 // This class converts operations to a given conversion target via a set of
2089 // rewrite patterns. The conversion behaves differently depending on the
2090 // conversion mode.
2091 struct OperationConverter {
2092   explicit OperationConverter(ConversionTarget &target,
2093                               const FrozenRewritePatternList &patterns,
2094                               OpConversionMode mode,
2095                               DenseSet<Operation *> *trackedOps = nullptr)
2096       : opLegalizer(target, patterns), mode(mode), trackedOps(trackedOps) {}
2097 
2098   /// Converts the given operations to the conversion target.
2099   LogicalResult convertOperations(ArrayRef<Operation *> ops);
2100 
2101 private:
2102   /// Converts an operation with the given rewriter.
2103   LogicalResult convert(ConversionPatternRewriter &rewriter, Operation *op);
2104 
2105   /// This method is called after the conversion process to legalize any
2106   /// remaining artifacts and complete the conversion.
2107   LogicalResult finalize(ConversionPatternRewriter &rewriter);
2108 
2109   /// Legalize the types of converted block arguments.
2110   LogicalResult
2111   legalizeConvertedArgumentTypes(ConversionPatternRewriter &rewriter,
2112                                  ConversionPatternRewriterImpl &rewriterImpl);
2113 
2114   /// Legalize an operation result that was marked as "erased".
2115   LogicalResult
2116   legalizeErasedResult(Operation *op, OpResult result,
2117                        ConversionPatternRewriterImpl &rewriterImpl);
2118 
2119   /// Legalize an operation result that was replaced with a value of a different
2120   /// type.
2121   LogicalResult
2122   legalizeChangedResultType(Operation *op, OpResult result, Value newValue,
2123                             TypeConverter *replConverter,
2124                             ConversionPatternRewriter &rewriter,
2125                             ConversionPatternRewriterImpl &rewriterImpl);
2126 
2127   /// The legalizer to use when converting operations.
2128   OperationLegalizer opLegalizer;
2129 
2130   /// The conversion mode to use when legalizing operations.
2131   OpConversionMode mode;
2132 
2133   /// A set of pre-existing operations. When mode == OpConversionMode::Analysis,
2134   /// this is populated with ops found to be legalizable to the target.
2135   /// When mode == OpConversionMode::Partial, this is populated with ops found
2136   /// *not* to be legalizable to the target.
2137   DenseSet<Operation *> *trackedOps;
2138 };
2139 } // end anonymous namespace
2140 
2141 LogicalResult OperationConverter::convert(ConversionPatternRewriter &rewriter,
2142                                           Operation *op) {
2143   // Legalize the given operation.
2144   if (failed(opLegalizer.legalize(op, rewriter))) {
2145     // Handle the case of a failed conversion for each of the different modes.
2146     // Full conversions expect all operations to be converted.
2147     if (mode == OpConversionMode::Full)
2148       return op->emitError()
2149              << "failed to legalize operation '" << op->getName() << "'";
2150     // Partial conversions allow conversions to fail iff the operation was not
2151     // explicitly marked as illegal. If the user provided a nonlegalizableOps
2152     // set, non-legalizable ops are included.
2153     if (mode == OpConversionMode::Partial) {
2154       if (opLegalizer.isIllegal(op))
2155         return op->emitError()
2156                << "failed to legalize operation '" << op->getName()
2157                << "' that was explicitly marked illegal";
2158       if (trackedOps)
2159         trackedOps->insert(op);
2160     }
2161   } else if (mode == OpConversionMode::Analysis) {
2162     // Analysis conversions don't fail if any operations fail to legalize,
2163     // they are only interested in the operations that were successfully
2164     // legalized.
2165     trackedOps->insert(op);
2166   }
2167   return success();
2168 }
2169 
2170 LogicalResult OperationConverter::convertOperations(ArrayRef<Operation *> ops) {
2171   if (ops.empty())
2172     return success();
2173   ConversionTarget &target = opLegalizer.getTarget();
2174 
2175   // Compute the set of operations and blocks to convert.
2176   std::vector<Operation *> toConvert;
2177   for (auto *op : ops) {
2178     toConvert.emplace_back(op);
2179     for (auto &region : op->getRegions())
2180       if (failed(computeConversionSet(region.getBlocks(), region.getLoc(),
2181                                       toConvert, &target)))
2182         return failure();
2183   }
2184 
2185   // Convert each operation and discard rewrites on failure.
2186   ConversionPatternRewriter rewriter(ops.front()->getContext());
2187   ConversionPatternRewriterImpl &rewriterImpl = rewriter.getImpl();
2188   for (auto *op : toConvert)
2189     if (failed(convert(rewriter, op)))
2190       return rewriterImpl.discardRewrites(), failure();
2191 
2192   // Now that all of the operations have been converted, finalize the conversion
2193   // process to ensure any lingering conversion artifacts are cleaned up and
2194   // legalized.
2195   if (failed(finalize(rewriter)))
2196     return rewriterImpl.discardRewrites(), failure();
2197 
2198   // After a successful conversion, apply rewrites if this is not an analysis
2199   // conversion.
2200   if (mode == OpConversionMode::Analysis)
2201     rewriterImpl.discardRewrites();
2202   else
2203     rewriterImpl.applyRewrites();
2204   return success();
2205 }
2206 
2207 LogicalResult
2208 OperationConverter::finalize(ConversionPatternRewriter &rewriter) {
2209   ConversionPatternRewriterImpl &rewriterImpl = rewriter.getImpl();
2210 
2211   // Legalize converted block arguments.
2212   if (failed(legalizeConvertedArgumentTypes(rewriter, rewriterImpl)))
2213     return failure();
2214 
2215   // Process requested operation replacements.
2216   for (unsigned i = 0, e = rewriterImpl.operationsWithChangedResults.size();
2217        i != e; ++i) {
2218     unsigned replIdx = rewriterImpl.operationsWithChangedResults[i];
2219     auto &repl = *(rewriterImpl.replacements.begin() + replIdx);
2220     for (OpResult result : repl.first->getResults()) {
2221       Value newValue = rewriterImpl.mapping.lookupOrNull(result);
2222 
2223       // If the operation result was replaced with null, all of the uses of this
2224       // value should be replaced.
2225       if (!newValue) {
2226         if (failed(legalizeErasedResult(repl.first, result, rewriterImpl)))
2227           return failure();
2228         continue;
2229       }
2230 
2231       // Otherwise, check to see if the type of the result changed.
2232       if (result.getType() == newValue.getType())
2233         continue;
2234 
2235       // Legalize this result.
2236       rewriter.setInsertionPoint(repl.first);
2237       if (failed(legalizeChangedResultType(repl.first, result, newValue,
2238                                            repl.second.converter, rewriter,
2239                                            rewriterImpl)))
2240         return failure();
2241 
2242       // Update the end iterator for this loop in the case it was updated
2243       // when legalizing generated conversion operations.
2244       e = rewriterImpl.operationsWithChangedResults.size();
2245     }
2246   }
2247   return success();
2248 }
2249 
2250 LogicalResult OperationConverter::legalizeConvertedArgumentTypes(
2251     ConversionPatternRewriter &rewriter,
2252     ConversionPatternRewriterImpl &rewriterImpl) {
2253   // Functor used to check if all users of a value will be dead after
2254   // conversion.
2255   auto findLiveUser = [&](Value val) {
2256     auto liveUserIt = llvm::find_if_not(val.getUsers(), [&](Operation *user) {
2257       return rewriterImpl.isOpIgnored(user);
2258     });
2259     return liveUserIt == val.user_end() ? nullptr : *liveUserIt;
2260   };
2261 
2262   // Materialize any necessary conversions for converted block arguments that
2263   // are still live.
2264   size_t numCreatedOps = rewriterImpl.createdOps.size();
2265   if (failed(rewriterImpl.argConverter.materializeLiveConversions(
2266           rewriterImpl.mapping, rewriter, findLiveUser)))
2267     return failure();
2268 
2269   // Legalize any newly created operations during argument materialization.
2270   for (int i : llvm::seq<int>(numCreatedOps, rewriterImpl.createdOps.size())) {
2271     if (failed(opLegalizer.legalize(rewriterImpl.createdOps[i], rewriter))) {
2272       return rewriterImpl.createdOps[i]->emitError()
2273              << "failed to legalize conversion operation generated for block "
2274                 "argument that remained live after conversion";
2275     }
2276   }
2277   return success();
2278 }
2279 
2280 LogicalResult OperationConverter::legalizeErasedResult(
2281     Operation *op, OpResult result,
2282     ConversionPatternRewriterImpl &rewriterImpl) {
2283   // If the operation result was replaced with null, all of the uses of this
2284   // value should be replaced.
2285   auto liveUserIt = llvm::find_if_not(result.getUsers(), [&](Operation *user) {
2286     return rewriterImpl.isOpIgnored(user);
2287   });
2288   if (liveUserIt != result.user_end()) {
2289     InFlightDiagnostic diag = op->emitError("failed to legalize operation '")
2290                               << op->getName() << "' marked as erased";
2291     diag.attachNote(liveUserIt->getLoc())
2292         << "found live user of result #" << result.getResultNumber() << ": "
2293         << *liveUserIt;
2294     return failure();
2295   }
2296   return success();
2297 }
2298 
2299 LogicalResult OperationConverter::legalizeChangedResultType(
2300     Operation *op, OpResult result, Value newValue,
2301     TypeConverter *replConverter, ConversionPatternRewriter &rewriter,
2302     ConversionPatternRewriterImpl &rewriterImpl) {
2303   // Walk the users of this value to see if there are any live users that
2304   // weren't replaced during conversion.
2305   auto liveUserIt = llvm::find_if_not(result.getUsers(), [&](Operation *user) {
2306     return rewriterImpl.isOpIgnored(user);
2307   });
2308   if (liveUserIt == result.user_end())
2309     return success();
2310 
2311   // If the replacement has a type converter, attempt to materialize a
2312   // conversion back to the original type.
2313   if (!replConverter) {
2314     // TODO: We should emit an error here, similarly to the case where the
2315     // result is replaced with null. Unfortunately a lot of existing
2316     // patterns rely on this behavior, so until those patterns are updated
2317     // we keep the legacy behavior here of just forwarding the new value.
2318     return success();
2319   }
2320 
2321   // Track the number of created operations so that new ones can be legalized.
2322   size_t numCreatedOps = rewriterImpl.createdOps.size();
2323 
2324   // Materialize a conversion for this live result value.
2325   Type resultType = result.getType();
2326   Value convertedValue = replConverter->materializeSourceConversion(
2327       rewriter, op->getLoc(), resultType, newValue);
2328   if (!convertedValue) {
2329     InFlightDiagnostic diag = op->emitError()
2330                               << "failed to materialize conversion for result #"
2331                               << result.getResultNumber() << " of operation '"
2332                               << op->getName()
2333                               << "' that remained live after conversion";
2334     diag.attachNote(liveUserIt->getLoc())
2335         << "see existing live user here: " << *liveUserIt;
2336     return failure();
2337   }
2338 
2339   // Legalize all of the newly created conversion operations.
2340   for (int i : llvm::seq<int>(numCreatedOps, rewriterImpl.createdOps.size())) {
2341     if (failed(opLegalizer.legalize(rewriterImpl.createdOps[i], rewriter))) {
2342       return op->emitError("failed to legalize conversion operation generated ")
2343              << "for result #" << result.getResultNumber() << " of operation '"
2344              << op->getName() << "' that remained live after conversion";
2345     }
2346   }
2347 
2348   rewriterImpl.mapping.map(result, convertedValue);
2349   return success();
2350 }
2351 
2352 //===----------------------------------------------------------------------===//
2353 // Type Conversion
2354 //===----------------------------------------------------------------------===//
2355 
2356 /// Remap an input of the original signature with a new set of types. The
2357 /// new types are appended to the new signature conversion.
2358 void TypeConverter::SignatureConversion::addInputs(unsigned origInputNo,
2359                                                    ArrayRef<Type> types) {
2360   assert(!types.empty() && "expected valid types");
2361   remapInput(origInputNo, /*newInputNo=*/argTypes.size(), types.size());
2362   addInputs(types);
2363 }
2364 
2365 /// Append new input types to the signature conversion, this should only be
2366 /// used if the new types are not intended to remap an existing input.
2367 void TypeConverter::SignatureConversion::addInputs(ArrayRef<Type> types) {
2368   assert(!types.empty() &&
2369          "1->0 type remappings don't need to be added explicitly");
2370   argTypes.append(types.begin(), types.end());
2371 }
2372 
2373 /// Remap an input of the original signature with a range of types in the
2374 /// new signature.
2375 void TypeConverter::SignatureConversion::remapInput(unsigned origInputNo,
2376                                                     unsigned newInputNo,
2377                                                     unsigned newInputCount) {
2378   assert(!remappedInputs[origInputNo] && "input has already been remapped");
2379   assert(newInputCount != 0 && "expected valid input count");
2380   remappedInputs[origInputNo] =
2381       InputMapping{newInputNo, newInputCount, /*replacementValue=*/nullptr};
2382 }
2383 
2384 /// Remap an input of the original signature to another `replacementValue`
2385 /// value. This would make the signature converter drop this argument.
2386 void TypeConverter::SignatureConversion::remapInput(unsigned origInputNo,
2387                                                     Value replacementValue) {
2388   assert(!remappedInputs[origInputNo] && "input has already been remapped");
2389   remappedInputs[origInputNo] =
2390       InputMapping{origInputNo, /*size=*/0, replacementValue};
2391 }
2392 
2393 /// This hooks allows for converting a type.
2394 LogicalResult TypeConverter::convertType(Type t,
2395                                          SmallVectorImpl<Type> &results) {
2396   auto existingIt = cachedDirectConversions.find(t);
2397   if (existingIt != cachedDirectConversions.end()) {
2398     if (existingIt->second)
2399       results.push_back(existingIt->second);
2400     return success(existingIt->second != nullptr);
2401   }
2402   auto multiIt = cachedMultiConversions.find(t);
2403   if (multiIt != cachedMultiConversions.end()) {
2404     results.append(multiIt->second.begin(), multiIt->second.end());
2405     return success();
2406   }
2407 
2408   // Walk the added converters in reverse order to apply the most recently
2409   // registered first.
2410   size_t currentCount = results.size();
2411   for (ConversionCallbackFn &converter : llvm::reverse(conversions)) {
2412     if (Optional<LogicalResult> result = converter(t, results)) {
2413       if (!succeeded(*result)) {
2414         cachedDirectConversions.try_emplace(t, nullptr);
2415         return failure();
2416       }
2417       auto newTypes = ArrayRef<Type>(results).drop_front(currentCount);
2418       if (newTypes.size() == 1)
2419         cachedDirectConversions.try_emplace(t, newTypes.front());
2420       else
2421         cachedMultiConversions.try_emplace(t, llvm::to_vector<2>(newTypes));
2422       return success();
2423     }
2424   }
2425   return failure();
2426 }
2427 
2428 /// This hook simplifies defining 1-1 type conversions. This function returns
2429 /// the type to convert to on success, and a null type on failure.
2430 Type TypeConverter::convertType(Type t) {
2431   // Use the multi-type result version to convert the type.
2432   SmallVector<Type, 1> results;
2433   if (failed(convertType(t, results)))
2434     return nullptr;
2435 
2436   // Check to ensure that only one type was produced.
2437   return results.size() == 1 ? results.front() : nullptr;
2438 }
2439 
2440 /// Convert the given set of types, filling 'results' as necessary. This
2441 /// returns failure if the conversion of any of the types fails, success
2442 /// otherwise.
2443 LogicalResult TypeConverter::convertTypes(ArrayRef<Type> types,
2444                                           SmallVectorImpl<Type> &results) {
2445   for (auto type : types)
2446     if (failed(convertType(type, results)))
2447       return failure();
2448   return success();
2449 }
2450 
2451 /// Return true if the given type is legal for this type converter, i.e. the
2452 /// type converts to itself.
2453 bool TypeConverter::isLegal(Type type) { return convertType(type) == type; }
2454 /// Return true if the given operation has legal operand and result types.
2455 bool TypeConverter::isLegal(Operation *op) {
2456   return isLegal(op->getOperandTypes()) && isLegal(op->getResultTypes());
2457 }
2458 
2459 /// Return true if the types of block arguments within the region are legal.
2460 bool TypeConverter::isLegal(Region *region) {
2461   return llvm::all_of(*region, [this](Block &block) {
2462     return isLegal(block.getArgumentTypes());
2463   });
2464 }
2465 
2466 /// Return true if the inputs and outputs of the given function type are
2467 /// legal.
2468 bool TypeConverter::isSignatureLegal(FunctionType ty) {
2469   return isLegal(llvm::concat<const Type>(ty.getInputs(), ty.getResults()));
2470 }
2471 
2472 /// This hook allows for converting a specific argument of a signature.
2473 LogicalResult TypeConverter::convertSignatureArg(unsigned inputNo, Type type,
2474                                                  SignatureConversion &result) {
2475   // Try to convert the given input type.
2476   SmallVector<Type, 1> convertedTypes;
2477   if (failed(convertType(type, convertedTypes)))
2478     return failure();
2479 
2480   // If this argument is being dropped, there is nothing left to do.
2481   if (convertedTypes.empty())
2482     return success();
2483 
2484   // Otherwise, add the new inputs.
2485   result.addInputs(inputNo, convertedTypes);
2486   return success();
2487 }
2488 LogicalResult TypeConverter::convertSignatureArgs(TypeRange types,
2489                                                   SignatureConversion &result,
2490                                                   unsigned origInputOffset) {
2491   for (unsigned i = 0, e = types.size(); i != e; ++i)
2492     if (failed(convertSignatureArg(origInputOffset + i, types[i], result)))
2493       return failure();
2494   return success();
2495 }
2496 
2497 Value TypeConverter::materializeConversion(
2498     MutableArrayRef<MaterializationCallbackFn> materializations,
2499     OpBuilder &builder, Location loc, Type resultType, ValueRange inputs) {
2500   for (MaterializationCallbackFn &fn : llvm::reverse(materializations))
2501     if (Optional<Value> result = fn(builder, resultType, inputs, loc))
2502       return result.getValue();
2503   return nullptr;
2504 }
2505 
2506 /// This function converts the type signature of the given block, by invoking
2507 /// 'convertSignatureArg' for each argument. This function should return a valid
2508 /// conversion for the signature on success, None otherwise.
2509 auto TypeConverter::convertBlockSignature(Block *block)
2510     -> Optional<SignatureConversion> {
2511   SignatureConversion conversion(block->getNumArguments());
2512   if (failed(convertSignatureArgs(block->getArgumentTypes(), conversion)))
2513     return llvm::None;
2514   return conversion;
2515 }
2516 
2517 /// Create a default conversion pattern that rewrites the type signature of a
2518 /// FuncOp.
2519 namespace {
2520 struct FuncOpSignatureConversion : public OpConversionPattern<FuncOp> {
2521   FuncOpSignatureConversion(MLIRContext *ctx, TypeConverter &converter)
2522       : OpConversionPattern(converter, ctx) {}
2523 
2524   /// Hook for derived classes to implement combined matching and rewriting.
2525   LogicalResult
2526   matchAndRewrite(FuncOp funcOp, ArrayRef<Value> operands,
2527                   ConversionPatternRewriter &rewriter) const override {
2528     FunctionType type = funcOp.getType();
2529 
2530     // Convert the original function types.
2531     TypeConverter::SignatureConversion result(type.getNumInputs());
2532     SmallVector<Type, 1> newResults;
2533     if (failed(typeConverter->convertSignatureArgs(type.getInputs(), result)) ||
2534         failed(typeConverter->convertTypes(type.getResults(), newResults)) ||
2535         failed(rewriter.convertRegionTypes(&funcOp.getBody(), *typeConverter,
2536                                            &result)))
2537       return failure();
2538 
2539     // Update the function signature in-place.
2540     rewriter.updateRootInPlace(funcOp, [&] {
2541       funcOp.setType(FunctionType::get(funcOp.getContext(),
2542                                        result.getConvertedTypes(), newResults));
2543     });
2544     return success();
2545   }
2546 };
2547 } // end anonymous namespace
2548 
2549 void mlir::populateFuncOpTypeConversionPattern(
2550     OwningRewritePatternList &patterns, MLIRContext *ctx,
2551     TypeConverter &converter) {
2552   patterns.insert<FuncOpSignatureConversion>(ctx, converter);
2553 }
2554 
2555 //===----------------------------------------------------------------------===//
2556 // ConversionTarget
2557 //===----------------------------------------------------------------------===//
2558 
2559 /// Register a legality action for the given operation.
2560 void ConversionTarget::setOpAction(OperationName op,
2561                                    LegalizationAction action) {
2562   legalOperations[op] = {action, /*isRecursivelyLegal=*/false, llvm::None};
2563 }
2564 
2565 /// Register a legality action for the given dialects.
2566 void ConversionTarget::setDialectAction(ArrayRef<StringRef> dialectNames,
2567                                         LegalizationAction action) {
2568   for (StringRef dialect : dialectNames)
2569     legalDialects[dialect] = action;
2570 }
2571 
2572 /// Get the legality action for the given operation.
2573 auto ConversionTarget::getOpAction(OperationName op) const
2574     -> Optional<LegalizationAction> {
2575   Optional<LegalizationInfo> info = getOpInfo(op);
2576   return info ? info->action : Optional<LegalizationAction>();
2577 }
2578 
2579 /// If the given operation instance is legal on this target, a structure
2580 /// containing legality information is returned. If the operation is not legal,
2581 /// None is returned.
2582 auto ConversionTarget::isLegal(Operation *op) const
2583     -> Optional<LegalOpDetails> {
2584   Optional<LegalizationInfo> info = getOpInfo(op->getName());
2585   if (!info)
2586     return llvm::None;
2587 
2588   // Returns true if this operation instance is known to be legal.
2589   auto isOpLegal = [&] {
2590     // Handle dynamic legality either with the provided legality function, or
2591     // the default hook on the derived instance.
2592     if (info->action == LegalizationAction::Dynamic)
2593       return info->legalityFn ? (*info->legalityFn)(op)
2594                               : isDynamicallyLegal(op);
2595 
2596     // Otherwise, the operation is only legal if it was marked 'Legal'.
2597     return info->action == LegalizationAction::Legal;
2598   };
2599   if (!isOpLegal())
2600     return llvm::None;
2601 
2602   // This operation is legal, compute any additional legality information.
2603   LegalOpDetails legalityDetails;
2604   if (info->isRecursivelyLegal) {
2605     auto legalityFnIt = opRecursiveLegalityFns.find(op->getName());
2606     if (legalityFnIt != opRecursiveLegalityFns.end())
2607       legalityDetails.isRecursivelyLegal = legalityFnIt->second(op);
2608     else
2609       legalityDetails.isRecursivelyLegal = true;
2610   }
2611   return legalityDetails;
2612 }
2613 
2614 /// Set the dynamic legality callback for the given operation.
2615 void ConversionTarget::setLegalityCallback(
2616     OperationName name, const DynamicLegalityCallbackFn &callback) {
2617   assert(callback && "expected valid legality callback");
2618   auto infoIt = legalOperations.find(name);
2619   assert(infoIt != legalOperations.end() &&
2620          infoIt->second.action == LegalizationAction::Dynamic &&
2621          "expected operation to already be marked as dynamically legal");
2622   infoIt->second.legalityFn = callback;
2623 }
2624 
2625 /// Set the recursive legality callback for the given operation and mark the
2626 /// operation as recursively legal.
2627 void ConversionTarget::markOpRecursivelyLegal(
2628     OperationName name, const DynamicLegalityCallbackFn &callback) {
2629   auto infoIt = legalOperations.find(name);
2630   assert(infoIt != legalOperations.end() &&
2631          infoIt->second.action != LegalizationAction::Illegal &&
2632          "expected operation to already be marked as legal");
2633   infoIt->second.isRecursivelyLegal = true;
2634   if (callback)
2635     opRecursiveLegalityFns[name] = callback;
2636   else
2637     opRecursiveLegalityFns.erase(name);
2638 }
2639 
2640 /// Set the dynamic legality callback for the given dialects.
2641 void ConversionTarget::setLegalityCallback(
2642     ArrayRef<StringRef> dialects, const DynamicLegalityCallbackFn &callback) {
2643   assert(callback && "expected valid legality callback");
2644   for (StringRef dialect : dialects)
2645     dialectLegalityFns[dialect] = callback;
2646 }
2647 
2648 /// Get the legalization information for the given operation.
2649 auto ConversionTarget::getOpInfo(OperationName op) const
2650     -> Optional<LegalizationInfo> {
2651   // Check for info for this specific operation.
2652   auto it = legalOperations.find(op);
2653   if (it != legalOperations.end())
2654     return it->second;
2655   // Check for info for the parent dialect.
2656   auto dialectIt = legalDialects.find(op.getDialect());
2657   if (dialectIt != legalDialects.end()) {
2658     Optional<DynamicLegalityCallbackFn> callback;
2659     auto dialectFn = dialectLegalityFns.find(op.getDialect());
2660     if (dialectFn != dialectLegalityFns.end())
2661       callback = dialectFn->second;
2662     return LegalizationInfo{dialectIt->second, /*isRecursivelyLegal=*/false,
2663                             callback};
2664   }
2665   // Otherwise, check if we mark unknown operations as dynamic.
2666   if (unknownOpsDynamicallyLegal)
2667     return LegalizationInfo{LegalizationAction::Dynamic,
2668                             /*isRecursivelyLegal=*/false, unknownLegalityFn};
2669   return llvm::None;
2670 }
2671 
2672 //===----------------------------------------------------------------------===//
2673 // Op Conversion Entry Points
2674 //===----------------------------------------------------------------------===//
2675 
2676 /// Apply a partial conversion on the given operations and all nested
2677 /// operations. This method converts as many operations to the target as
2678 /// possible, ignoring operations that failed to legalize. This method only
2679 /// returns failure if there ops explicitly marked as illegal.
2680 /// If an `unconvertedOps` set is provided, all operations that are found not
2681 /// to be legalizable to the given `target` are placed within that set. (Note
2682 /// that if there is an op explicitly marked as illegal, the conversion
2683 /// terminates and the `unconvertedOps` set will not necessarily be complete.)
2684 LogicalResult
2685 mlir::applyPartialConversion(ArrayRef<Operation *> ops,
2686                              ConversionTarget &target,
2687                              const FrozenRewritePatternList &patterns,
2688                              DenseSet<Operation *> *unconvertedOps) {
2689   OperationConverter opConverter(target, patterns, OpConversionMode::Partial,
2690                                  unconvertedOps);
2691   return opConverter.convertOperations(ops);
2692 }
2693 LogicalResult
2694 mlir::applyPartialConversion(Operation *op, ConversionTarget &target,
2695                              const FrozenRewritePatternList &patterns,
2696                              DenseSet<Operation *> *unconvertedOps) {
2697   return applyPartialConversion(llvm::makeArrayRef(op), target, patterns,
2698                                 unconvertedOps);
2699 }
2700 
2701 /// Apply a complete conversion on the given operations, and all nested
2702 /// operations. This method will return failure if the conversion of any
2703 /// operation fails.
2704 LogicalResult
2705 mlir::applyFullConversion(ArrayRef<Operation *> ops, ConversionTarget &target,
2706                           const FrozenRewritePatternList &patterns) {
2707   OperationConverter opConverter(target, patterns, OpConversionMode::Full);
2708   return opConverter.convertOperations(ops);
2709 }
2710 LogicalResult
2711 mlir::applyFullConversion(Operation *op, ConversionTarget &target,
2712                           const FrozenRewritePatternList &patterns) {
2713   return applyFullConversion(llvm::makeArrayRef(op), target, patterns);
2714 }
2715 
2716 /// Apply an analysis conversion on the given operations, and all nested
2717 /// operations. This method analyzes which operations would be successfully
2718 /// converted to the target if a conversion was applied. All operations that
2719 /// were found to be legalizable to the given 'target' are placed within the
2720 /// provided 'convertedOps' set; note that no actual rewrites are applied to the
2721 /// operations on success and only pre-existing operations are added to the set.
2722 LogicalResult
2723 mlir::applyAnalysisConversion(ArrayRef<Operation *> ops,
2724                               ConversionTarget &target,
2725                               const FrozenRewritePatternList &patterns,
2726                               DenseSet<Operation *> &convertedOps) {
2727   OperationConverter opConverter(target, patterns, OpConversionMode::Analysis,
2728                                  &convertedOps);
2729   return opConverter.convertOperations(ops);
2730 }
2731 LogicalResult
2732 mlir::applyAnalysisConversion(Operation *op, ConversionTarget &target,
2733                               const FrozenRewritePatternList &patterns,
2734                               DenseSet<Operation *> &convertedOps) {
2735   return applyAnalysisConversion(llvm::makeArrayRef(op), target, patterns,
2736                                  convertedOps);
2737 }
2738