//===- DialectConversion.cpp - MLIR dialect conversion generic pass -------===// // // Copyright 2019 The MLIR Authors. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= #include "mlir/Transforms/DialectConversion.h" #include "mlir/IR/Block.h" #include "mlir/IR/BlockAndValueMapping.h" #include "mlir/IR/Builders.h" #include "mlir/IR/Function.h" #include "mlir/IR/Module.h" #include "mlir/Transforms/Utils.h" #include "llvm/ADT/SetVector.h" #include "llvm/ADT/SmallPtrSet.h" #include "llvm/Support/Debug.h" #include "llvm/Support/raw_ostream.h" using namespace mlir; using namespace mlir::detail; #define DEBUG_TYPE "dialect-conversion" //===----------------------------------------------------------------------===// // Multi-Level Value Mapper //===----------------------------------------------------------------------===// namespace { /// This class wraps a BlockAndValueMapping to provide recursive lookup /// functionality, i.e. we will traverse if the mapped value also has a mapping. struct ConversionValueMapping { /// Lookup a mapped value within the map. If a mapping for the provided value /// does not exist then return the provided value. Value *lookupOrDefault(Value *from) const; /// Map a value to the one provided. void map(Value *oldVal, Value *newVal) { mapping.map(oldVal, newVal); } /// Drop the last mapping for the given value. void erase(Value *value) { mapping.erase(value); } private: /// Current value mappings. BlockAndValueMapping mapping; }; } // end anonymous namespace /// Lookup a mapped value within the map. If a mapping for the provided value /// does not exist then return the provided value. Value *ConversionValueMapping::lookupOrDefault(Value *from) const { // If this value had a valid mapping, unmap that value as well in the case // that it was also replaced. while (auto *mappedValue = mapping.lookupOrNull(from)) from = mappedValue; return from; } //===----------------------------------------------------------------------===// // ArgConverter //===----------------------------------------------------------------------===// namespace { /// This class provides a simple interface for converting the types of block /// arguments. This is done by inserting fake cast operations that map from the /// illegal type to the original type to allow for undoing pending rewrites in /// the case of failure. struct ArgConverter { ArgConverter(TypeConverter *typeConverter, PatternRewriter &rewriter) : castOpName(kCastName, rewriter.getContext()), loc(rewriter.getUnknownLoc()), typeConverter(typeConverter), rewriter(rewriter) {} /// Erase any rewrites registered for arguments to blocks within the given /// region. This function is called when the given region is to be destroyed. void cancelPendingRewrites(Block *block); /// Cleanup and undo any generated conversions for the arguments of block. /// This method differs from 'cancelPendingRewrites' in that it returns the /// block signature to its original state. void discardPendingRewrites(Block *block); /// Replace usages of the cast operations with the argument directly. void applyRewrites(); /// Return if the signature of the given block has already been converted. bool hasBeenConverted(Block *block) const { return argMapping.count(block); } /// Attempt to convert the signature of the given block. LogicalResult convertSignature(Block *block, ConversionValueMapping &mapping); /// Apply the given signature conversion on the given block. void applySignatureConversion( Block *block, TypeConverter::SignatureConversion &signatureConversion, ConversionValueMapping &mapping); /// Convert the given block argument given the provided set of new argument /// values that are to replace it. This function returns the operation used /// to perform the conversion. Operation *convertArgument(BlockArgument *origArg, ArrayRef newValues, ConversionValueMapping &mapping); /// A utility function used to create a conversion cast operation with the /// given input and result types. Operation *createCast(ArrayRef inputs, Type outputType); /// This is an operation name for a fake operation that is inserted during the /// conversion process. Operations of this type are guaranteed to never escape /// the converter. static constexpr StringLiteral kCastName = "__mlir_conversion.cast"; OperationName castOpName; /// This is a collection of cast operations that were generated during the /// conversion process when converting the types of block arguments. llvm::MapVector> argMapping; /// An instance of the unknown location that is used when generating /// producers. Location loc; /// The type converter to use when changing types. TypeConverter *typeConverter; /// The pattern rewriter to use when materializing conversions. PatternRewriter &rewriter; }; } // end anonymous namespace constexpr StringLiteral ArgConverter::kCastName; /// Erase any rewrites registered for arguments to the given block. void ArgConverter::cancelPendingRewrites(Block *block) { auto it = argMapping.find(block); if (it == argMapping.end()) return; for (auto *op : it->second) { op->dropAllDefinedValueUses(); op->erase(); } argMapping.erase(it); } /// Cleanup and undo any generated conversions for the arguments of block. /// This method differs from 'cancelPendingRewrites' in that it returns the /// block signature to its original state. void ArgConverter::discardPendingRewrites(Block *block) { auto it = argMapping.find(block); if (it == argMapping.end()) return; // Erase all of the new arguments. for (int i = block->getNumArguments() - 1; i >= 0; --i) { block->getArgument(i)->dropAllUses(); block->eraseArgument(i, /*updatePredTerms=*/false); } // Re-instate the old arguments. auto &mapping = it->second; for (unsigned i = 0, e = mapping.size(); i != e; ++i) { auto *op = mapping[i]; auto *arg = block->addArgument(op->getResult(0)->getType()); op->getResult(0)->replaceAllUsesWith(arg); // If this operation is within a block, it will be cleaned up automatically. if (!op->getBlock()) op->erase(); } argMapping.erase(it); } /// Replace usages of the cast operations with the argument directly. void ArgConverter::applyRewrites() { Block *block; ArrayRef argOps; for (auto &mapping : argMapping) { std::tie(block, argOps) = mapping; // Process the remapping for each of the original arguments. for (unsigned i = 0, e = argOps.size(); i != e; ++i) { auto *op = argOps[i]; // Handle the case of a 1->N value mapping. if (op->getNumOperands() > 1) { // If all of the uses were removed, we can drop this op. Otherwise, // keep the operation alive and let the user handle any remaining // usages. if (op->use_empty()) op->erase(); continue; } // If mapping is 1-1, replace the remaining uses and drop the cast // operation. // FIXME(riverriddle) This should check that the result type and operand // type are the same, otherwise it should force a conversion to be // materialized. This works around a current limitation with regards to // region entry argument type conversion. if (op->getNumOperands() == 1) { op->getResult(0)->replaceAllUsesWith(op->getOperand(0)); op->destroy(); continue; } // Otherwise, if there are any dangling uses then replace the fake // conversion operation with one generated by the type converter. This // is necessary as the cast must persist in the IR after conversion. auto *opResult = op->getResult(0); if (!opResult->use_empty()) { rewriter.setInsertionPointToStart(block); SmallVector operands(op->getOperands()); auto *newOp = typeConverter->materializeConversion( rewriter, opResult->getType(), operands, op->getLoc()); opResult->replaceAllUsesWith(newOp->getResult(0)); } op->destroy(); } } } /// Converts the signature of the given entry block. LogicalResult ArgConverter::convertSignature(Block *block, ConversionValueMapping &mapping) { if (auto conversion = typeConverter->convertBlockSignature(block)) return applySignatureConversion(block, *conversion, mapping), success(); return failure(); } /// Apply the given signature conversion on the given block. void ArgConverter::applySignatureConversion( Block *block, TypeConverter::SignatureConversion &signatureConversion, ConversionValueMapping &mapping) { unsigned origArgCount = block->getNumArguments(); auto convertedTypes = signatureConversion.getConvertedTypes(); if (origArgCount == 0 && convertedTypes.empty()) return; SmallVector newArgRange(block->addArguments(convertedTypes)); ArrayRef newArgRef(newArgRange); // Remap each of the original arguments as determined by the signature // conversion. auto &newArgMapping = argMapping[block]; rewriter.setInsertionPointToStart(block); for (unsigned i = 0; i != origArgCount; ++i) { ArrayRef remappedValues; if (auto inputMap = signatureConversion.getInputMapping(i)) remappedValues = newArgRef.slice(inputMap->inputNo, inputMap->size); BlockArgument *arg = block->getArgument(i); newArgMapping.push_back(convertArgument(arg, remappedValues, mapping)); } // Erase all of the original arguments. for (unsigned i = 0; i != origArgCount; ++i) block->eraseArgument(0, /*updatePredTerms=*/false); } /// Convert the given block argument given the provided set of new argument /// values that are to replace it. This function returns the operation used /// to perform the conversion. Operation *ArgConverter::convertArgument(BlockArgument *origArg, ArrayRef newValues, ConversionValueMapping &mapping) { // Handle the cases of 1->0 or 1->1 mappings. if (newValues.size() < 2) { // Create a temporary producer for the argument during the conversion // process. auto *cast = createCast(newValues, origArg->getType()); origArg->replaceAllUsesWith(cast->getResult(0)); // Insert a mapping between this argument and the one that is replacing // it. if (!newValues.empty()) mapping.map(cast->getResult(0), newValues[0]); return cast; } // Otherwise, this is a 1->N mapping. Call into the provided type converter // to pack the new values. auto *cast = typeConverter->materializeConversion( rewriter, origArg->getType(), newValues, loc); assert(cast->getNumResults() == 1 && cast->getNumOperands() == newValues.size()); origArg->replaceAllUsesWith(cast->getResult(0)); return cast; } /// A utility function used to create a conversion cast operation with the /// given input and result types. Operation *ArgConverter::createCast(ArrayRef inputs, Type outputType) { return Operation::create(loc, castOpName, inputs, outputType, llvm::None, llvm::None, 0, false); } //===----------------------------------------------------------------------===// // ConversionPatternRewriterImpl //===----------------------------------------------------------------------===// namespace { /// This class contains a snapshot of the current conversion rewriter state. /// This is useful when saving and undoing a set of rewrites. struct RewriterState { RewriterState(unsigned numCreatedOperations, unsigned numReplacements, unsigned numBlockActions) : numCreatedOperations(numCreatedOperations), numReplacements(numReplacements), numBlockActions(numBlockActions) {} /// The current number of created operations. unsigned numCreatedOperations; /// The current number of replacements queued. unsigned numReplacements; /// The current number of block actions performed. unsigned numBlockActions; }; } // end anonymous namespace namespace mlir { namespace detail { struct ConversionPatternRewriterImpl { /// This class represents one requested operation replacement via 'replaceOp'. struct OpReplacement { OpReplacement() = default; OpReplacement(Operation *op, ArrayRef newValues) : op(op), newValues(newValues.begin(), newValues.end()) {} Operation *op; SmallVector newValues; }; /// The kind of the block action performed during the rewrite. Actions can be /// undone if the conversion fails. enum class BlockActionKind { Split, Move, TypeConversion }; /// Original position of the given block in its parent region. We cannot use /// a region iterator because it could have been invalidated by other region /// operations since the position was stored. struct BlockPosition { Region *region; Region::iterator::difference_type position; }; /// The storage class for an undoable block action (one of BlockActionKind), /// contains the information necessary to undo this action. struct BlockAction { static BlockAction getSplit(Block *block, Block *originalBlock) { BlockAction action{BlockActionKind::Split, block, {}}; action.originalBlock = originalBlock; return action; } static BlockAction getMove(Block *block, BlockPosition originalPos) { return {BlockActionKind::Move, block, {originalPos}}; } static BlockAction getTypeConversion(Block *block) { return BlockAction{BlockActionKind::TypeConversion, block, {}}; } // The action kind. BlockActionKind kind; // A pointer to the block that was created by the action. Block *block; union { // In use if kind == BlockActionKind::Move and contains a pointer to the // region that originally contained the block as well as the position of // the block in that region. BlockPosition originalPosition; // In use if kind == BlockActionKind::Split and contains a pointer to the // block that was split into two parts. Block *originalBlock; }; }; ConversionPatternRewriterImpl(PatternRewriter &rewriter, TypeConverter *converter) : argConverter(converter, rewriter) {} /// Return the current state of the rewriter. RewriterState getCurrentState(); /// Reset the state of the rewriter to a previously saved point. void resetState(RewriterState state); /// Undo the block actions (motions, splits) one by one in reverse order until /// "numActionsToKeep" actions remains. void undoBlockActions(unsigned numActionsToKeep = 0); /// Cleanup and destroy any generated rewrite operations. This method is /// invoked when the conversion process fails. void discardRewrites(); /// Apply all requested operation rewrites. This method is invoked when the /// conversion process succeeds. void applyRewrites(); /// Convert the signature of the given block. LogicalResult convertBlockSignature(Block *block); /// Apply a signature conversion on the given region. void applySignatureConversion(Region *region, TypeConverter::SignatureConversion &conversion); /// PatternRewriter hook for replacing the results of an operation. void replaceOp(Operation *op, ArrayRef newValues, ArrayRef valuesToRemoveIfDead); /// Notifies that a block was split. void notifySplitBlock(Block *block, Block *continuation); /// Notifies that the blocks of a region are about to be moved. void notifyRegionIsBeingInlinedBefore(Region ®ion, Region &parent, Region::iterator before); /// Remap the given operands to those with potentially different types. void remapValues(Operation::operand_range operands, SmallVectorImpl &remapped); // Mapping between replaced values that differ in type. This happens when // replacing a value with one of a different type. ConversionValueMapping mapping; /// Utility used to convert block arguments. ArgConverter argConverter; /// Ordered vector of all of the newly created operations during conversion. SmallVector createdOps; /// Ordered vector of any requested operation replacements. SmallVector replacements; /// Ordered list of block operations (creations, splits, motions). SmallVector blockActions; }; } // end namespace detail } // end namespace mlir RewriterState ConversionPatternRewriterImpl::getCurrentState() { return RewriterState(createdOps.size(), replacements.size(), blockActions.size()); } void ConversionPatternRewriterImpl::resetState(RewriterState state) { // Undo any block actions. undoBlockActions(state.numBlockActions); // Reset any replaced operations and undo any saved mappings. for (auto &repl : llvm::drop_begin(replacements, state.numReplacements)) for (auto *result : repl.op->getResults()) mapping.erase(result); replacements.resize(state.numReplacements); // Pop all of the newly created operations. while (createdOps.size() != state.numCreatedOperations) createdOps.pop_back_val()->erase(); } void ConversionPatternRewriterImpl::undoBlockActions( unsigned numActionsToKeep) { for (auto &action : llvm::reverse(llvm::drop_begin(blockActions, numActionsToKeep))) { switch (action.kind) { // Merge back the block that was split out. case BlockActionKind::Split: { action.originalBlock->getOperations().splice( action.originalBlock->end(), action.block->getOperations()); action.block->dropAllUses(); action.block->erase(); break; } // Move the block back to its original position. case BlockActionKind::Move: { Region *originalRegion = action.originalPosition.region; originalRegion->getBlocks().splice( std::next(originalRegion->begin(), action.originalPosition.position), action.block->getParent()->getBlocks(), action.block); break; } // Undo the type conversion. case BlockActionKind::TypeConversion: { argConverter.discardPendingRewrites(action.block); break; } } } blockActions.resize(numActionsToKeep); } void ConversionPatternRewriterImpl::discardRewrites() { undoBlockActions(); // Remove any newly created ops. for (auto *op : llvm::reverse(createdOps)) op->erase(); } void ConversionPatternRewriterImpl::applyRewrites() { // Apply all of the rewrites replacements requested during conversion. for (auto &repl : replacements) { for (unsigned i = 0, e = repl.newValues.size(); i != e; ++i) repl.op->getResult(i)->replaceAllUsesWith( mapping.lookupOrDefault(repl.newValues[i])); // If this operation defines any regions, drop any pending argument // rewrites. if (argConverter.typeConverter && repl.op->getNumRegions()) { for (auto ®ion : repl.op->getRegions()) for (auto &block : region) argConverter.cancelPendingRewrites(&block); } } // In a second pass, erase all of the replaced operations in reverse. This // allows processing nested operations before their parent region is // destroyed. for (auto &repl : llvm::reverse(replacements)) repl.op->erase(); argConverter.applyRewrites(); } LogicalResult ConversionPatternRewriterImpl::convertBlockSignature(Block *block) { // Check to see if this block should not be converted: // * There is no type converter. // * The block has already been converted. // * This is an entry block, these are converted explicitly via patterns. if (!argConverter.typeConverter || argConverter.hasBeenConverted(block) || block->isEntryBlock()) return success(); // Otherwise, try to convert the block signature. if (failed(argConverter.convertSignature(block, mapping))) return failure(); blockActions.push_back(BlockAction::getTypeConversion(block)); return success(); } void ConversionPatternRewriterImpl::applySignatureConversion( Region *region, TypeConverter::SignatureConversion &conversion) { if (!region->empty()) { argConverter.applySignatureConversion(®ion->front(), conversion, mapping); blockActions.push_back(BlockAction::getTypeConversion(®ion->front())); } } void ConversionPatternRewriterImpl::replaceOp( Operation *op, ArrayRef newValues, ArrayRef valuesToRemoveIfDead) { assert(newValues.size() == op->getNumResults()); // Create mappings for each of the new result values. for (unsigned i = 0, e = newValues.size(); i < e; ++i) { assert((newValues[i] || op->getResult(i)->use_empty()) && "result value has remaining uses that must be replaced"); if (newValues[i]) mapping.map(op->getResult(i), newValues[i]); } // Record the requested operation replacement. replacements.emplace_back(op, newValues); } void ConversionPatternRewriterImpl::notifySplitBlock(Block *block, Block *continuation) { blockActions.push_back(BlockAction::getSplit(continuation, block)); } void ConversionPatternRewriterImpl::notifyRegionIsBeingInlinedBefore( Region ®ion, Region &parent, Region::iterator before) { for (auto &pair : llvm::enumerate(region)) { Block &block = pair.value(); unsigned position = pair.index(); blockActions.push_back(BlockAction::getMove(&block, {®ion, position})); } } void ConversionPatternRewriterImpl::remapValues( Operation::operand_range operands, SmallVectorImpl &remapped) { remapped.reserve(llvm::size(operands)); for (Value *operand : operands) remapped.push_back(mapping.lookupOrDefault(operand)); } //===----------------------------------------------------------------------===// // ConversionPatternRewriter //===----------------------------------------------------------------------===// ConversionPatternRewriter::ConversionPatternRewriter(MLIRContext *ctx, TypeConverter *converter) : PatternRewriter(ctx), impl(new detail::ConversionPatternRewriterImpl(*this, converter)) {} ConversionPatternRewriter::~ConversionPatternRewriter() {} /// PatternRewriter hook for replacing the results of an operation. void ConversionPatternRewriter::replaceOp( Operation *op, ArrayRef newValues, ArrayRef valuesToRemoveIfDead) { LLVM_DEBUG(llvm::dbgs() << "** Replacing operation : " << op->getName() << "\n"); impl->replaceOp(op, newValues, valuesToRemoveIfDead); } /// Apply a signature conversion to the entry block of the given region. void ConversionPatternRewriter::applySignatureConversion( Region *region, TypeConverter::SignatureConversion &conversion) { impl->applySignatureConversion(region, conversion); } /// Clone the given operation without cloning its regions. Operation *ConversionPatternRewriter::cloneWithoutRegions(Operation *op) { Operation *newOp = OpBuilder::cloneWithoutRegions(*op); impl->createdOps.push_back(newOp); return newOp; } /// PatternRewriter hook for splitting a block into two parts. Block *ConversionPatternRewriter::splitBlock(Block *block, Block::iterator before) { auto *continuation = PatternRewriter::splitBlock(block, before); impl->notifySplitBlock(block, continuation); return continuation; } /// PatternRewriter hook for moving blocks out of a region. void ConversionPatternRewriter::inlineRegionBefore(Region ®ion, Region &parent, Region::iterator before) { impl->notifyRegionIsBeingInlinedBefore(region, parent, before); PatternRewriter::inlineRegionBefore(region, parent, before); } /// PatternRewriter hook for creating a new operation. Operation * ConversionPatternRewriter::createOperation(const OperationState &state) { LLVM_DEBUG(llvm::dbgs() << "** Creating operation : " << state.name << "\n"); auto *result = OpBuilder::createOperation(state); impl->createdOps.push_back(result); return result; } /// PatternRewriter hook for updating the root operation in-place. void ConversionPatternRewriter::notifyRootUpdated(Operation *op) { // The rewriter caches changes to the IR to allow for operating in-place and // backtracking. The rewriter is currently not capable of backtracking // in-place modifications. llvm_unreachable("in-place operation updates are not supported"); } /// Return a reference to the internal implementation. detail::ConversionPatternRewriterImpl &ConversionPatternRewriter::getImpl() { return *impl; } //===----------------------------------------------------------------------===// // Conversion Patterns //===----------------------------------------------------------------------===// /// Attempt to match and rewrite the IR root at the specified operation. PatternMatchResult ConversionPattern::matchAndRewrite(Operation *op, PatternRewriter &rewriter) const { SmallVector operands; auto &dialectRewriter = static_cast(rewriter); dialectRewriter.getImpl().remapValues(op->getOperands(), operands); // If this operation has no successors, invoke the rewrite directly. if (op->getNumSuccessors() == 0) return matchAndRewrite(op, operands, dialectRewriter); // Otherwise, we need to remap the successors. SmallVector destinations; destinations.reserve(op->getNumSuccessors()); SmallVector, 2> operandsPerDestination; unsigned firstSuccessorOperand = op->getSuccessorOperandIndex(0); for (unsigned i = 0, seen = 0, e = op->getNumSuccessors(); i < e; ++i) { destinations.push_back(op->getSuccessor(i)); // Lookup the successors operands. unsigned n = op->getNumSuccessorOperands(i); operandsPerDestination.push_back( llvm::makeArrayRef(operands.data() + firstSuccessorOperand + seen, n)); seen += n; } // Rewrite the operation. return matchAndRewrite( op, llvm::makeArrayRef(operands.data(), operands.data() + firstSuccessorOperand), destinations, operandsPerDestination, dialectRewriter); } //===----------------------------------------------------------------------===// // OperationLegalizer //===----------------------------------------------------------------------===// namespace { /// A set of rewrite patterns that can be used to legalize a given operation. using LegalizationPatterns = SmallVector; /// This class defines a recursive operation legalizer. class OperationLegalizer { public: using LegalizationAction = ConversionTarget::LegalizationAction; OperationLegalizer(ConversionTarget &targetInfo, const OwningRewritePatternList &patterns) : target(targetInfo) { buildLegalizationGraph(patterns); computeLegalizationGraphBenefit(); } /// Returns if the given operation is known to be illegal on the target. bool isIllegal(Operation *op) const; /// Attempt to legalize the given operation. Returns success if the operation /// was legalized, failure otherwise. LogicalResult legalize(Operation *op, ConversionPatternRewriter &rewriter); private: /// Attempt to legalize the given operation by applying the provided pattern. /// Returns success if the operation was legalized, failure otherwise. LogicalResult legalizePattern(Operation *op, RewritePattern *pattern, ConversionPatternRewriter &rewriter); /// Build an optimistic legalization graph given the provided patterns. This /// function populates 'legalizerPatterns' with the operations that are not /// directly legal, but may be transitively legal for the current target given /// the provided patterns. void buildLegalizationGraph(const OwningRewritePatternList &patterns); /// Compute the benefit of each node within the computed legalization graph. /// This orders the patterns within 'legalizerPatterns' based upon two /// criteria: /// 1) Prefer patterns that have the lowest legalization depth, i.e. /// represent the more direct mapping to the target. /// 2) When comparing patterns with the same legalization depth, prefer the /// pattern with the highest PatternBenefit. This allows for users to /// prefer specific legalizations over others. void computeLegalizationGraphBenefit(); /// The current set of patterns that have been applied. llvm::SmallPtrSet appliedPatterns; /// The set of legality information for operations transitively supported by /// the target. DenseMap legalizerPatterns; /// The legalization information provided by the target. ConversionTarget ⌖ }; } // namespace bool OperationLegalizer::isIllegal(Operation *op) const { // Check if the target explicitly marked this operation as illegal. if (auto action = target.getOpAction(op->getName())) return action == LegalizationAction::Illegal; return false; } LogicalResult OperationLegalizer::legalize(Operation *op, ConversionPatternRewriter &rewriter) { LLVM_DEBUG(llvm::dbgs() << "Legalizing operation : " << op->getName() << "\n"); // Check if this operation is legal on the target. if (target.isLegal(op)) { LLVM_DEBUG(llvm::dbgs() << "-- Success : Operation marked legal by the target\n"); return success(); } // Otherwise, we need to apply a legalization pattern to this operation. auto it = legalizerPatterns.find(op->getName()); if (it == legalizerPatterns.end()) { LLVM_DEBUG(llvm::dbgs() << "-- FAIL : no known legalization path.\n"); return failure(); } // The patterns are sorted by expected benefit, so try to apply each in-order. for (auto *pattern : it->second) if (succeeded(legalizePattern(op, pattern, rewriter))) return success(); LLVM_DEBUG(llvm::dbgs() << "-- FAIL : no matched legalization pattern.\n"); return failure(); } LogicalResult OperationLegalizer::legalizePattern(Operation *op, RewritePattern *pattern, ConversionPatternRewriter &rewriter) { LLVM_DEBUG({ llvm::dbgs() << "-* Applying rewrite pattern '" << op->getName() << " -> ("; interleaveComma(pattern->getGeneratedOps(), llvm::dbgs()); llvm::dbgs() << ")'.\n"; }); // Ensure that we don't cycle by not allowing the same pattern to be // applied twice in the same recursion stack. // TODO(riverriddle) We could eventually converge, but that requires more // complicated analysis. if (!appliedPatterns.insert(pattern).second) { LLVM_DEBUG(llvm::dbgs() << "-- FAIL: Pattern was already applied.\n"); return failure(); } auto &rewriterImpl = rewriter.getImpl(); RewriterState curState = rewriterImpl.getCurrentState(); auto cleanupFailure = [&] { // Reset the rewriter state and pop this pattern. rewriterImpl.resetState(curState); appliedPatterns.erase(pattern); return failure(); }; // Try to rewrite with the given pattern. rewriter.setInsertionPoint(op); if (!pattern->matchAndRewrite(op, rewriter)) { LLVM_DEBUG(llvm::dbgs() << "-- FAIL: Pattern failed to match.\n"); return cleanupFailure(); } // If the pattern moved any blocks, try to legalize their types. This ensures // that the types of the block arguments are legal for the region they were // moved into. for (unsigned i = curState.numBlockActions, e = rewriterImpl.blockActions.size(); i != e; ++i) { auto &action = rewriterImpl.blockActions[i]; if (action.kind != ConversionPatternRewriterImpl::BlockActionKind::Move) continue; // Convert the block signature. if (failed(rewriterImpl.convertBlockSignature(action.block))) { LLVM_DEBUG(llvm::dbgs() << "-- FAIL: failed to convert types of moved block.\n"); return cleanupFailure(); } } // Recursively legalize each of the new operations. for (unsigned i = curState.numCreatedOperations, e = rewriterImpl.createdOps.size(); i != e; ++i) { Operation *op = rewriterImpl.createdOps[i]; if (failed(legalize(op, rewriter))) { LLVM_DEBUG(llvm::dbgs() << "-- FAIL: Generated operation '" << op->getName() << "' was illegal.\n"); return cleanupFailure(); } } appliedPatterns.erase(pattern); return success(); } void OperationLegalizer::buildLegalizationGraph( const OwningRewritePatternList &patterns) { // A mapping between an operation and a set of operations that can be used to // generate it. DenseMap> parentOps; // A mapping between an operation and any currently invalid patterns it has. DenseMap> invalidPatterns; // A worklist of patterns to consider for legality. llvm::SetVector patternWorklist; // Build the mapping from operations to the parent ops that may generate them. for (auto &pattern : patterns) { auto root = pattern->getRootKind(); // Skip operations that are always known to be legal. if (target.getOpAction(root) == LegalizationAction::Legal) continue; // Add this pattern to the invalid set for the root op and record this root // as a parent for any generated operations. invalidPatterns[root].insert(pattern.get()); for (auto op : pattern->getGeneratedOps()) parentOps[op].insert(root); // Add this pattern to the worklist. patternWorklist.insert(pattern.get()); } while (!patternWorklist.empty()) { auto *pattern = patternWorklist.pop_back_val(); // Check to see if any of the generated operations are invalid. if (llvm::any_of(pattern->getGeneratedOps(), [&](OperationName op) { auto action = target.getOpAction(op); return !legalizerPatterns.count(op) && (!action || action == LegalizationAction::Illegal); })) continue; // Otherwise, if all of the generated operation are valid, this op is now // legal so add all of the child patterns to the worklist. legalizerPatterns[pattern->getRootKind()].push_back(pattern); invalidPatterns[pattern->getRootKind()].erase(pattern); // Add any invalid patterns of the parent operations to see if they have now // become legal. for (auto op : parentOps[pattern->getRootKind()]) patternWorklist.set_union(invalidPatterns[op]); } } void OperationLegalizer::computeLegalizationGraphBenefit() { // The smallest pattern depth, when legalizing an operation. DenseMap minPatternDepth; // Compute the minimum legalization depth for a given operation. std::function computeDepth = [&](OperationName op) { // Check for existing depth. auto depthIt = minPatternDepth.find(op); if (depthIt != minPatternDepth.end()) return depthIt->second; // If a mapping for this operation does not exist, then this operation // is always legal. Return 0 as the depth for a directly legal operation. auto opPatternsIt = legalizerPatterns.find(op); if (opPatternsIt == legalizerPatterns.end() || opPatternsIt->second.empty()) return 0u; // Initialize the depth to the maximum value. unsigned minDepth = std::numeric_limits::max(); // Record this initial depth in case we encounter this op again when // recursively computing the depth. minPatternDepth.try_emplace(op, minDepth); // Compute the depth for each pattern used to legalize this operation. SmallVector, 4> patternsByDepth; patternsByDepth.reserve(opPatternsIt->second.size()); for (RewritePattern *pattern : opPatternsIt->second) { unsigned depth = 0; for (auto generatedOp : pattern->getGeneratedOps()) depth = std::max(depth, computeDepth(generatedOp) + 1); patternsByDepth.emplace_back(pattern, depth); // Update the min depth for this operation. minDepth = std::min(minDepth, depth); } // Update the pattern depth. minPatternDepth[op] = minDepth; // If the operation only has one legalization pattern, there is no need to // sort them. if (patternsByDepth.size() == 1) return minDepth; // Sort the patterns by those likely to be the most beneficial. llvm::array_pod_sort( patternsByDepth.begin(), patternsByDepth.end(), [](const std::pair *lhs, const std::pair *rhs) { // First sort by the smaller pattern legalization depth. if (lhs->second != rhs->second) return llvm::array_pod_sort_comparator(&lhs->second, &rhs->second); // Then sort by the larger pattern benefit. auto lhsBenefit = lhs->first->getBenefit(); auto rhsBenefit = rhs->first->getBenefit(); return llvm::array_pod_sort_comparator(&rhsBenefit, &lhsBenefit); }); // Update the legalization pattern to use the new sorted list. opPatternsIt->second.clear(); for (auto &patternIt : patternsByDepth) opPatternsIt->second.push_back(patternIt.first); return minDepth; }; // For each operation that is transitively legal, compute a cost for it. for (auto &opIt : legalizerPatterns) if (!minPatternDepth.count(opIt.first)) computeDepth(opIt.first); } //===----------------------------------------------------------------------===// // OperationConverter //===----------------------------------------------------------------------===// namespace { enum OpConversionMode { // In this mode, the conversion will ignore failed conversions to allow // illegal operations to co-exist in the IR. Partial, // In this mode, all operations must be legal for the given target for the // conversion to succeeed. Full, // In this mode, operations are analyzed for legality. No actual rewrites are // applied to the operations on success. Analysis, }; // This class converts operations to a given conversion target via a set of // rewrite patterns. The conversion behaves differently depending on the // conversion mode. struct OperationConverter { explicit OperationConverter(ConversionTarget &target, const OwningRewritePatternList &patterns, OpConversionMode mode, DenseSet *legalizableOps = nullptr) : opLegalizer(target, patterns), mode(mode), legalizableOps(legalizableOps) {} /// Converts the given operations to the conversion target. LogicalResult convertOperations(ArrayRef ops, TypeConverter *typeConverter); private: /// Converts an operation with the given rewriter. LogicalResult convert(ConversionPatternRewriter &rewriter, Operation *op); /// Recursively collect all of the operations to convert from within 'region'. LogicalResult computeConversionSet(Region ®ion, std::vector &toConvert); /// Converts the type signatures of the blocks nested within 'op'. LogicalResult convertBlockSignatures(ConversionPatternRewriter &rewriter, Operation *op); /// The legalizer to use when converting operations. OperationLegalizer opLegalizer; /// The conversion mode to use when legalizing operations. OpConversionMode mode; /// A set of pre-existing operations that were found to be legalizable to the /// target. This field is only used when mode == OpConversionMode::Analysis. DenseSet *legalizableOps; }; } // end anonymous namespace LogicalResult OperationConverter::convertBlockSignatures(ConversionPatternRewriter &rewriter, Operation *op) { // Check to see if type signatures need to be converted. if (!rewriter.getImpl().argConverter.typeConverter) return success(); for (auto ®ion : op->getRegions()) { for (auto &block : region) if (failed(rewriter.getImpl().convertBlockSignature(&block))) return failure(); } return success(); } LogicalResult OperationConverter::computeConversionSet(Region ®ion, std::vector &toConvert) { if (region.empty()) return success(); // Traverse starting from the entry block. SmallVector worklist(1, ®ion.front()); DenseSet visitedBlocks; visitedBlocks.insert(®ion.front()); while (!worklist.empty()) { auto *block = worklist.pop_back_val(); // Compute the conversion set of each of the nested operations. for (auto &op : *block) { toConvert.emplace_back(&op); for (auto ®ion : op.getRegions()) computeConversionSet(region, toConvert); } // Recurse to children that haven't been visited. for (Block *succ : block->getSuccessors()) if (visitedBlocks.insert(succ).second) worklist.push_back(succ); } // Check that all blocks in the region were visited. if (llvm::any_of(llvm::drop_begin(region.getBlocks(), 1), [&](Block &block) { return !visitedBlocks.count(&block); })) return emitError(region.getLoc(), "unreachable blocks were not converted"); return success(); } LogicalResult OperationConverter::convert(ConversionPatternRewriter &rewriter, Operation *op) { // Legalize the given operation. if (failed(opLegalizer.legalize(op, rewriter))) { // Handle the case of a failed conversion for each of the different modes. /// Full conversions expect all operations to be converted. if (mode == OpConversionMode::Full) return op->emitError() << "failed to legalize operation '" << op->getName() << "'"; /// Partial conversions allow conversions to fail iff the operation was not /// explicitly marked as illegal. if (mode == OpConversionMode::Partial && opLegalizer.isIllegal(op)) return op->emitError() << "failed to legalize operation '" << op->getName() << "' that was explicitly marked illegal"; } else { /// Analysis conversions don't fail if any operations fail to legalize, /// they are only interested in the operations that were successfully /// legalized. if (mode == OpConversionMode::Analysis) legalizableOps->insert(op); // If legalization succeeded, convert the types any of the blocks within // this operation. if (failed(convertBlockSignatures(rewriter, op))) return failure(); } return success(); } LogicalResult OperationConverter::convertOperations(ArrayRef ops, TypeConverter *typeConverter) { if (ops.empty()) return success(); /// Compute the set of operations and blocks to convert. std::vector toConvert; for (auto *op : ops) { toConvert.emplace_back(op); for (auto ®ion : op->getRegions()) if (failed(computeConversionSet(region, toConvert))) return failure(); } // Convert each operation and discard rewrites on failure. ConversionPatternRewriter rewriter(ops.front()->getContext(), typeConverter); for (auto *op : toConvert) if (failed(convert(rewriter, op))) return rewriter.getImpl().discardRewrites(), failure(); // Otherwise, the body conversion succeeded. Apply rewrites if this is not an // analysis conversion. if (mode == OpConversionMode::Analysis) rewriter.getImpl().discardRewrites(); else rewriter.getImpl().applyRewrites(); return success(); } //===----------------------------------------------------------------------===// // Type Conversion //===----------------------------------------------------------------------===// /// Remap an input of the original signature with a new set of types. The /// new types are appended to the new signature conversion. void TypeConverter::SignatureConversion::addInputs(unsigned origInputNo, ArrayRef types) { assert(!types.empty() && "expected valid types"); remapInput(origInputNo, /*newInputNo=*/argTypes.size(), types.size()); addInputs(types); } /// Append new input types to the signature conversion, this should only be /// used if the new types are not intended to remap an existing input. void TypeConverter::SignatureConversion::addInputs(ArrayRef types) { assert(!types.empty() && "1->0 type remappings don't need to be added explicitly"); argTypes.append(types.begin(), types.end()); } /// Remap an input of the original signature with a range of types in the /// new signature. void TypeConverter::SignatureConversion::remapInput(unsigned origInputNo, unsigned newInputNo, unsigned newInputCount) { assert(!remappedInputs[origInputNo] && "input has already been remapped"); assert(newInputCount != 0 && "expected valid input count"); remappedInputs[origInputNo] = InputMapping{newInputNo, newInputCount}; } /// This hooks allows for converting a type. LogicalResult TypeConverter::convertType(Type t, SmallVectorImpl &results) { if (auto newT = convertType(t)) { results.push_back(newT); return success(); } return failure(); } /// Convert the given set of types, filling 'results' as necessary. This /// returns failure if the conversion of any of the types fails, success /// otherwise. LogicalResult TypeConverter::convertTypes(ArrayRef types, SmallVectorImpl &results) { for (auto type : types) if (failed(convertType(type, results))) return failure(); return success(); } /// Return true if the given type is legal for this type converter, i.e. the /// type converts to itself. bool TypeConverter::isLegal(Type type) { SmallVector results; return succeeded(convertType(type, results)) && results.size() == 1 && results.front() == type; } /// Return true if the inputs and outputs of the given function type are /// legal. bool TypeConverter::isSignatureLegal(FunctionType funcType) { return llvm::all_of( llvm::concat(funcType.getInputs(), funcType.getResults()), [this](Type type) { return isLegal(type); }); } /// This hook allows for converting a specific argument of a signature. LogicalResult TypeConverter::convertSignatureArg(unsigned inputNo, Type type, SignatureConversion &result) { // Try to convert the given input type. SmallVector convertedTypes; if (failed(convertType(type, convertedTypes))) return failure(); // If this argument is being dropped, there is nothing left to do. if (convertedTypes.empty()) return success(); // Otherwise, add the new inputs. result.addInputs(inputNo, convertedTypes); return success(); } /// Create a default conversion pattern that rewrites the type signature of a /// FuncOp. namespace { struct FuncOpSignatureConversion : public ConversionPattern { FuncOpSignatureConversion(MLIRContext *ctx, TypeConverter &converter) : ConversionPattern(FuncOp::getOperationName(), 1, ctx), converter(converter) {} /// Hook for derived classes to implement combined matching and rewriting. PatternMatchResult matchAndRewrite(Operation *op, ArrayRef operands, ConversionPatternRewriter &rewriter) const override { auto funcOp = cast(op); FunctionType type = funcOp.getType(); // Convert the original function arguments. TypeConverter::SignatureConversion result(type.getNumInputs()); for (unsigned i = 0, e = type.getNumInputs(); i != e; ++i) if (failed(converter.convertSignatureArg(i, type.getInput(i), result))) return matchFailure(); // Convert the original function results. SmallVector convertedResults; if (failed(converter.convertTypes(type.getResults(), convertedResults))) return matchFailure(); // Create a new function with an updated signature. auto newFuncOp = rewriter.cloneWithoutRegions(funcOp); rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(), newFuncOp.end()); newFuncOp.setType(FunctionType::get(result.getConvertedTypes(), convertedResults, funcOp.getContext())); // Tell the rewriter to convert the region signature. rewriter.applySignatureConversion(&newFuncOp.getBody(), result); rewriter.replaceOp(op, llvm::None); return matchSuccess(); } /// The type converter to use when rewriting the signature. TypeConverter &converter; }; } // end anonymous namespace void mlir::populateFuncOpTypeConversionPattern( OwningRewritePatternList &patterns, MLIRContext *ctx, TypeConverter &converter) { patterns.insert(ctx, converter); } /// This function converts the type signature of the given block, by invoking /// 'convertSignatureArg' for each argument. This function should return a valid /// conversion for the signature on success, None otherwise. auto TypeConverter::convertBlockSignature(Block *block) -> llvm::Optional { SignatureConversion conversion(block->getNumArguments()); for (unsigned i = 0, e = block->getNumArguments(); i != e; ++i) if (failed(convertSignatureArg(i, block->getArgument(i)->getType(), conversion))) return llvm::None; return conversion; } //===----------------------------------------------------------------------===// // ConversionTarget //===----------------------------------------------------------------------===// /// Register a legality action for the given operation. void ConversionTarget::setOpAction(OperationName op, LegalizationAction action) { legalOperations[op] = action; } /// Register a legality action for the given dialects. void ConversionTarget::setDialectAction(ArrayRef dialectNames, LegalizationAction action) { for (StringRef dialect : dialectNames) legalDialects[dialect] = action; } /// Get the legality action for the given operation. auto ConversionTarget::getOpAction(OperationName op) const -> llvm::Optional { // Check for an action for this specific operation. auto it = legalOperations.find(op); if (it != legalOperations.end()) return it->second; // Otherwise, default to checking for an action on the parent dialect. auto dialectIt = legalDialects.find(op.getDialect()); if (dialectIt != legalDialects.end()) return dialectIt->second; return llvm::None; } /// Return if the given operation instance is legal on this target. bool ConversionTarget::isLegal(Operation *op) const { auto action = getOpAction(op->getName()); // Handle dynamic legality. if (action == LegalizationAction::Dynamic) { // Check for callbacks on the operation or dialect. auto opFn = opLegalityFns.find(op->getName()); if (opFn != opLegalityFns.end()) return opFn->second(op); auto dialectFn = dialectLegalityFns.find(op->getName().getDialect()); if (dialectFn != dialectLegalityFns.end()) return dialectFn->second(op); // Otherwise, invoke the hook on the derived instance. return isDynamicallyLegal(op); } // Otherwise, the operation is only legal if it was marked 'Legal'. return action == LegalizationAction::Legal; } /// Set the dynamic legality callback for the given operation. void ConversionTarget::setLegalityCallback( OperationName name, const DynamicLegalityCallbackFn &callback) { assert(callback && "expected valid legality callback"); opLegalityFns[name] = callback; } /// Set the dynamic legality callback for the given dialects. void ConversionTarget::setLegalityCallback( ArrayRef dialects, const DynamicLegalityCallbackFn &callback) { assert(callback && "expected valid legality callback"); for (StringRef dialect : dialects) dialectLegalityFns[dialect] = callback; } //===----------------------------------------------------------------------===// // Op Conversion Entry Points //===----------------------------------------------------------------------===// /// Apply a partial conversion on the given operations, and all nested /// operations. This method converts as many operations to the target as /// possible, ignoring operations that failed to legalize. LogicalResult mlir::applyPartialConversion( ArrayRef ops, ConversionTarget &target, const OwningRewritePatternList &patterns, TypeConverter *converter) { OperationConverter opConverter(target, patterns, OpConversionMode::Partial); return opConverter.convertOperations(ops, converter); } LogicalResult mlir::applyPartialConversion(Operation *op, ConversionTarget &target, const OwningRewritePatternList &patterns, TypeConverter *converter) { return applyPartialConversion(llvm::makeArrayRef(op), target, patterns, converter); } /// Apply a complete conversion on the given operations, and all nested /// operations. This method will return failure if the conversion of any /// operation fails. LogicalResult mlir::applyFullConversion(ArrayRef ops, ConversionTarget &target, const OwningRewritePatternList &patterns, TypeConverter *converter) { OperationConverter opConverter(target, patterns, OpConversionMode::Full); return opConverter.convertOperations(ops, converter); } LogicalResult mlir::applyFullConversion(Operation *op, ConversionTarget &target, const OwningRewritePatternList &patterns, TypeConverter *converter) { return applyFullConversion(llvm::makeArrayRef(op), target, patterns, converter); } /// Apply an analysis conversion on the given operations, and all nested /// operations. This method analyzes which operations would be successfully /// converted to the target if a conversion was applied. All operations that /// were found to be legalizable to the given 'target' are placed within the /// provided 'convertedOps' set; note that no actual rewrites are applied to the /// operations on success and only pre-existing operations are added to the set. LogicalResult mlir::applyAnalysisConversion( ArrayRef ops, ConversionTarget &target, const OwningRewritePatternList &patterns, DenseSet &convertedOps, TypeConverter *converter) { OperationConverter opConverter(target, patterns, OpConversionMode::Analysis, &convertedOps); return opConverter.convertOperations(ops, converter); } LogicalResult mlir::applyAnalysisConversion(Operation *op, ConversionTarget &target, const OwningRewritePatternList &patterns, DenseSet &convertedOps, TypeConverter *converter) { return applyAnalysisConversion(llvm::makeArrayRef(op), target, patterns, convertedOps, converter); }