[Mlir-commits] [mlir] [mlir][Transforms] Dialect conversion: Hardening `replaceOp` (PR #109540)
Matthias Springer
llvmlistbot at llvm.org
Sun Sep 29 02:34:34 PDT 2024
https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/109540
>From c5d6e99ec4c762c0a3f3ca1b2ffe5f6690d5e080 Mon Sep 17 00:00:00 2001
From: Matthias Springer <mspringer at nvidia.com>
Date: Sun, 29 Sep 2024 10:34:05 +0200
Subject: [PATCH 1/2] [mlir][Transforms][NFC] Dialect Conversion: Update docs
for `remapValues`
---
.../Transforms/Utils/DialectConversion.cpp | 70 ++++++++++---------
1 file changed, 38 insertions(+), 32 deletions(-)
diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 4693edadfb5eec..b5aab2416c3eb9 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1092,44 +1092,50 @@ LogicalResult ConversionPatternRewriterImpl::remapValues(
SmallVectorImpl<Value> &remapped) {
remapped.reserve(llvm::size(values));
- SmallVector<Type, 1> legalTypes;
for (const auto &it : llvm::enumerate(values)) {
Value operand = it.value();
Type origType = operand.getType();
+ Location operandLoc = inputLoc ? *inputLoc : operand.getLoc();
- // If a converter was provided, get the desired legal types for this
- // operand.
- Type desiredType;
- if (currentTypeConverter) {
- // If there is no legal conversion, fail to match this pattern.
- legalTypes.clear();
- if (failed(currentTypeConverter->convertType(origType, legalTypes))) {
- Location operandLoc = inputLoc ? *inputLoc : operand.getLoc();
- notifyMatchFailure(operandLoc, [=](Diagnostic &diag) {
- diag << "unable to convert type for " << valueDiagTag << " #"
- << it.index() << ", type was " << origType;
- });
- return failure();
- }
- // TODO: There currently isn't any mechanism to do 1->N type conversion
- // via the PatternRewriter replacement API, so for now we just ignore it.
- if (legalTypes.size() == 1)
- desiredType = legalTypes.front();
- } else {
- // TODO: What we should do here is just set `desiredType` to `origType`
- // and then handle the necessary type conversions after the conversion
- // process has finished. Unfortunately a lot of patterns currently rely on
- // receiving the new operands even if the types change, so we keep the
- // original behavior here for now until all of the patterns relying on
- // this get updated.
+ if (!currentTypeConverter) {
+ // The current pattern does not have a type converter. I.e., it does not
+ // distinguish between legal and illegal types. For each operand, simply
+ // pass through the most recently mapped value.
+ remapped.push_back(mapping.lookupOrDefault(operand));
+ continue;
+ }
+
+ // If there is no legal conversion, fail to match this pattern.
+ SmallVector<Type, 1> legalTypes;
+ if (failed(currentTypeConverter->convertType(origType, legalTypes))) {
+ notifyMatchFailure(operandLoc, [=](Diagnostic &diag) {
+ diag << "unable to convert type for " << valueDiagTag << " #"
+ << it.index() << ", type was " << origType;
+ });
+ return failure();
}
- Value newOperand = mapping.lookupOrDefault(operand, desiredType);
- // Handle the case where the conversion was 1->1 and the new operand type
- // isn't legal.
- Type newOperandType = newOperand.getType();
- if (currentTypeConverter && desiredType && newOperandType != desiredType) {
- Location operandLoc = inputLoc ? *inputLoc : operand.getLoc();
+ if (legalTypes.size() != 1) {
+ // TODO: Parts of the dialect conversion infrastructure do not support
+ // 1->N type conversions yet. Therefore, if a type is converted to 0 or
+ // multiple types, the only thing that we can do for now is passing
+ // through the most recently mapped value. Fixing this requires
+ // improvements to the `ConversionValueMapping` (to be able to store 1:N
+ // mappings) and to the `ConversionPattern` adaptor handling (to be able
+ // to pass multiple remapped values for a single operand to the adaptor).
+ remapped.push_back(mapping.lookupOrDefault(operand));
+ continue;
+ }
+
+ // Handle 1->1 type conversions.
+ Type desiredType = legalTypes.front();
+ // Try to find a mapped value with the desired type. (Or the operand itself
+ // if the value is not mapped at all.)
+ Value newOperand = mapping.lookupOrDefault(operand, desiredType);
+ if (newOperand.getType() != desiredType) {
+ // If the looked up value's type does not have the desired type, it means
+ // that the value was replaced with a value of different type and no
+ // source materialization was created yet.
Value castValue = buildUnresolvedMaterialization(
MaterializationKind::Target, computeInsertPoint(newOperand),
operandLoc, /*inputs=*/newOperand, /*outputType=*/desiredType,
>From 3399448ea348a5f1aa9fdff2e8b83833667e1b32 Mon Sep 17 00:00:00 2001
From: Matthias Springer <mspringer at nvidia.com>
Date: Sat, 21 Sep 2024 20:03:17 +0200
Subject: [PATCH 2/2] [mlir][Transforms] Dialect conversion: Extra checks
during `replaceOp`
This commit adds extra checks/assertions to the `ConversionPatternRewriterImpl::notifyOpReplaced` to improve its robustness.
Replacing an `unrealized_conversion_cast` op that was created by the driver is forbidden and is now caught early during `replaceOp`. It may work in some cases, but is generally dangerous because the conversion driver keeps track of these ops. (Erasing is them is fine.) This change is also in preparation of a subsequent commit that splits the `ConversionValueMapping` into replacements and materializations (with the goal of simplifying block signature conversions).
`null` replacement values are no longer registered in the `ConversionValueMapping`. This was an oversight in #106760. `null` values in the mapping could result in crashes when using the `ConversionValueMapping` API.
---
.../Transforms/Utils/DialectConversion.cpp | 31 ++++++++++++++-----
1 file changed, 23 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index b5aab2416c3eb9..9535dcafcd3b3b 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1365,16 +1365,21 @@ void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op,
assert(newValues.size() == op->getNumResults());
assert(!ignoredOps.contains(op) && "operation was already replaced");
+ // Check if replaced op is an unresolved materialization, i.e., an
+ // unrealized_conversion_cast op that was created by the conversion driver.
+ bool isUnresolvedMaterialization = false;
+ if (auto castOp = dyn_cast<UnrealizedConversionCastOp>(op))
+ if (unresolvedMaterializations.contains(castOp))
+ isUnresolvedMaterialization = true;
+
// Create mappings for each of the new result values.
for (auto [newValue, result] : llvm::zip(newValues, op->getResults())) {
if (!newValue) {
// This result was dropped and no replacement value was provided.
- if (auto castOp = dyn_cast<UnrealizedConversionCastOp>(op)) {
- if (unresolvedMaterializations.contains(castOp)) {
- // Do not create another materializations if we are erasing a
- // materialization.
- continue;
- }
+ if (isUnresolvedMaterialization) {
+ // Do not create another materializations if we are erasing a
+ // materialization.
+ continue;
}
// Materialize a replacement value "out of thin air".
@@ -1382,10 +1387,20 @@ void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op,
MaterializationKind::Source, computeInsertPoint(result),
result.getLoc(), /*inputs=*/ValueRange(),
/*outputType=*/result.getType(), currentTypeConverter);
+ } else {
+ // Make sure that the user does not mess with unresolved materializations
+ // that were inserted by the conversion driver. We keep track of these
+ // ops in internal data structures. Erasing them must be allowed because
+ // this can happen when the user is erasing an entire block (including
+ // its body). But replacing them with another value should be forbidden
+ // to avoid problems with the `mapping`.
+ assert(!isUnresolvedMaterialization &&
+ "attempting to replace an unresolved materialization");
}
- // Remap, and check for any result type changes.
- mapping.map(result, newValue);
+ // Remap result to replacement value.
+ if (newValue)
+ mapping.map(result, newValue);
}
appendRewrite<ReplaceOperationRewrite>(op, currentTypeConverter);
More information about the Mlir-commits
mailing list