diff options
Diffstat (limited to 'mlir/lib/Transforms/Utils/DialectConversion.cpp')
| -rw-r--r-- | mlir/lib/Transforms/Utils/DialectConversion.cpp | 135 |
1 files changed, 97 insertions, 38 deletions
diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp index 0a62628b9ad2..5b2cfd370900 100644 --- a/mlir/lib/Transforms/Utils/DialectConversion.cpp +++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp @@ -67,6 +67,10 @@ static OpBuilder::InsertPoint computeInsertPoint(Value value) { // ConversionValueMapping //===----------------------------------------------------------------------===// +/// A list of replacement SSA values. Optimized for the common case of a single +/// SSA value. +using ReplacementValues = SmallVector<Value, 1>; + namespace { /// This class wraps a IRMapping to provide recursive lookup /// functionality, i.e. we will traverse if the mapped value also has a mapping. @@ -818,6 +822,22 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener { Type originalType, const TypeConverter *converter); + /// Build an N:1 materialization for the given original value that was + /// replaced with the given replacement values. + /// + /// This is a workaround around incomplete 1:N support in the dialect + /// conversion driver. The conversion mapping can store only 1:1 replacements + /// and the conversion patterns only support single Value replacements in the + /// adaptor, so N values must be converted back to a single value. This + /// function will be deleted when full 1:N support has been added. + /// + /// This function inserts an argument materialization back to the original + /// type, followed by a target materialization to the legalized type (if + /// applicable). + void insertNTo1Materialization(OpBuilder::InsertPoint ip, Location loc, + ValueRange replacements, Value originalValue, + const TypeConverter *converter); + //===--------------------------------------------------------------------===// // Rewriter Notification Hooks //===--------------------------------------------------------------------===// @@ -827,7 +847,7 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener { OpBuilder::InsertPoint previous) override; /// Notifies that an op is about to be replaced with the given values. - void notifyOpReplaced(Operation *op, ValueRange newValues); + void notifyOpReplaced(Operation *op, ArrayRef<ReplacementValues> newValues); /// Notifies that a block is about to be erased. void notifyBlockIsBeingErased(Block *block); @@ -1148,7 +1168,8 @@ LogicalResult ConversionPatternRewriterImpl::remapValues( // source materialization was created yet. Value castValue = buildUnresolvedMaterialization( MaterializationKind::Target, computeInsertPoint(newOperand), - operandLoc, /*inputs=*/newOperand, /*outputType=*/desiredType, + operandLoc, + /*inputs=*/newOperand, /*outputType=*/desiredType, /*originalType=*/origType, currentTypeConverter); mapping.map(newOperand, castValue); newOperand = castValue; @@ -1287,33 +1308,9 @@ Block *ConversionPatternRewriterImpl::applySignatureConversion( // used as a replacement. auto replArgs = newBlock->getArguments().slice(inputMap->inputNo, inputMap->size); - Value argMat = buildUnresolvedMaterialization( - MaterializationKind::Argument, + insertNTo1Materialization( OpBuilder::InsertPoint(newBlock, newBlock->begin()), origArg.getLoc(), - /*inputs=*/replArgs, /*outputType=*/origArgType, - /*originalType=*/Type(), converter); - mapping.map(origArg, argMat); - - Type legalOutputType; - if (converter) { - legalOutputType = converter->convertType(origArgType); - } else if (replArgs.size() == 1) { - // When there is no type converter, assume that the new block argument - // types are legal. This is reasonable to assume because they were - // specified by the user. - // FIXME: This won't work for 1->N conversions because multiple output - // types are not supported in parts of the dialect conversion. In such a - // case, we currently use the original block argument type (produced by - // the argument materialization). - legalOutputType = replArgs[0].getType(); - } - if (legalOutputType && legalOutputType != origArgType) { - Value targetMat = buildUnresolvedMaterialization( - MaterializationKind::Target, computeInsertPoint(argMat), - origArg.getLoc(), /*inputs=*/argMat, /*outputType=*/legalOutputType, - /*originalType=*/origArgType, converter); - mapping.map(argMat, targetMat); - } + /*replacements=*/replArgs, /*outputValue=*/origArg, converter); appendRewrite<ReplaceBlockArgRewrite>(block, origArg); } @@ -1354,6 +1351,39 @@ Value ConversionPatternRewriterImpl::buildUnresolvedMaterialization( return convertOp.getResult(0); } +void ConversionPatternRewriterImpl::insertNTo1Materialization( + OpBuilder::InsertPoint ip, Location loc, ValueRange replacements, + Value originalValue, const TypeConverter *converter) { + // Insert argument materialization back to the original type. + Type originalType = originalValue.getType(); + Value argMat = + buildUnresolvedMaterialization(MaterializationKind::Argument, ip, loc, + /*inputs=*/replacements, originalType, + /*originalType=*/Type(), converter); + mapping.map(originalValue, argMat); + + // Insert target materialization to the legalized type. + Type legalOutputType; + if (converter) { + legalOutputType = converter->convertType(originalType); + } else if (replacements.size() == 1) { + // When there is no type converter, assume that the replacement value + // types are legal. This is reasonable to assume because they were + // specified by the user. + // FIXME: This won't work for 1->N conversions because multiple output + // types are not supported in parts of the dialect conversion. In such a + // case, we currently use the original value type. + legalOutputType = replacements[0].getType(); + } + if (legalOutputType && legalOutputType != originalType) { + Value targetMat = buildUnresolvedMaterialization( + MaterializationKind::Target, computeInsertPoint(argMat), loc, + /*inputs=*/argMat, /*outputType=*/legalOutputType, + /*originalType=*/originalType, converter); + mapping.map(argMat, targetMat); + } +} + //===----------------------------------------------------------------------===// // Rewriter Notification Hooks @@ -1377,8 +1407,8 @@ void ConversionPatternRewriterImpl::notifyOperationInserted( appendRewrite<MoveOperationRewrite>(op, previous.getBlock(), prevOp); } -void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op, - ValueRange newValues) { +void ConversionPatternRewriterImpl::notifyOpReplaced( + Operation *op, ArrayRef<ReplacementValues> newValues) { assert(newValues.size() == op->getNumResults()); assert(!ignoredOps.contains(op) && "operation was already replaced"); @@ -1390,8 +1420,9 @@ void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op, isUnresolvedMaterialization = true; // Create mappings for each of the new result values. - for (auto [newValue, result] : llvm::zip(newValues, op->getResults())) { - if (!newValue) { + for (auto [n, result] : llvm::zip_equal(newValues, op->getResults())) { + ReplacementValues repl = n; + if (repl.empty()) { // This result was dropped and no replacement value was provided. if (isUnresolvedMaterialization) { // Do not create another materializations if we are erasing a @@ -1400,11 +1431,12 @@ void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op, } // Materialize a replacement value "out of thin air". - newValue = buildUnresolvedMaterialization( + Value sourceMat = buildUnresolvedMaterialization( MaterializationKind::Source, computeInsertPoint(result), result.getLoc(), /*inputs=*/ValueRange(), /*outputType=*/result.getType(), /*originalType=*/Type(), currentTypeConverter); + repl.push_back(sourceMat); } else { // Make sure that the user does not mess with unresolved materializations // that were inserted by the conversion driver. We keep track of these @@ -1417,12 +1449,21 @@ void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op, } // Remap result to replacement value. - if (newValue) - mapping.map(result, newValue); + if (repl.empty()) + continue; + + if (repl.size() == 1) { + // Single replacement value: replace directly. + mapping.map(result, repl.front()); + } else { + // Multiple replacement values: insert N:1 materialization. + insertNTo1Materialization(computeInsertPoint(result), result.getLoc(), + /*replacements=*/repl, /*outputValue=*/result, + currentTypeConverter); + } } appendRewrite<ReplaceOperationRewrite>(op, currentTypeConverter); - // Mark this operation and all nested ops as replaced. op->walk([&](Operation *op) { replacedOps.insert(op); }); } @@ -1497,7 +1538,25 @@ void ConversionPatternRewriter::replaceOp(Operation *op, ValueRange newValues) { impl->logger.startLine() << "** Replace : '" << op->getName() << "'(" << op << ")\n"; }); - impl->notifyOpReplaced(op, newValues); + SmallVector<ReplacementValues> newVals(newValues.size()); + for (auto [index, val] : llvm::enumerate(newValues)) + if (val) + newVals[index].push_back(val); + impl->notifyOpReplaced(op, newVals); +} + +void ConversionPatternRewriter::replaceOpWithMultiple( + Operation *op, ArrayRef<ValueRange> newValues) { + assert(op->getNumResults() == newValues.size() && + "incorrect # of replacement values"); + LLVM_DEBUG({ + impl->logger.startLine() + << "** Replace : '" << op->getName() << "'(" << op << ")\n"; + }); + SmallVector<ReplacementValues> newVals(newValues.size(), {}); + for (auto [index, val] : llvm::enumerate(newValues)) + llvm::append_range(newVals[index], val); + impl->notifyOpReplaced(op, newVals); } void ConversionPatternRewriter::eraseOp(Operation *op) { @@ -1505,7 +1564,7 @@ void ConversionPatternRewriter::eraseOp(Operation *op) { impl->logger.startLine() << "** Erase : '" << op->getName() << "'(" << op << ")\n"; }); - SmallVector<Value, 1> nullRepls(op->getNumResults(), nullptr); + SmallVector<ReplacementValues> nullRepls(op->getNumResults(), {}); impl->notifyOpReplaced(op, nullRepls); } |
