summaryrefslogtreecommitdiff
path: root/mlir/lib/Transforms/Utils/DialectConversion.cpp
diff options
context:
space:
mode:
Diffstat (limited to 'mlir/lib/Transforms/Utils/DialectConversion.cpp')
-rw-r--r--mlir/lib/Transforms/Utils/DialectConversion.cpp135
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);
}