[mlir][Transforms] Dialect conversion: Hardening replaceOp

[matthias-springer]

This commit adds extra checks/assertions to the ConversionPatternRewriterImpl::notifyOpReplaced to improve its robustness.

1. Replacing an unrealized_conversion_cast op that was created by the driver is now forbidden and caught early during replaceOp. It may work in some cases, but it is generally dangerous because the conversion driver keeps track of these ops and performs some extra legalization steps during the "finalize" phase. (Erasing is them is fine.)
2. null replacement values are no longer registered in the ConversionValueMapping. This was an oversight in [mlir][Transforms] Dialect conversion: Align handling of dropped values #106760. There is no benefit in having null values in the ConversionValueMapping. (It may even cause problems.)

This change is in preparation of merging the 1:1 and 1:N dialect conversion drivers.

[llvmbot]

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

This commit adds extra checks/assertions to the ConversionPatternRewriterImpl::notifyOpReplaced to improve its robustness.

1. Replacing an unrealized_conversion_cast op that was created by the driver is now forbidden and caught early during replaceOp. It may work in some cases, but it is generally dangerous because the conversion driver keeps track of these ops and performs some extra legalization steps during the "finalize" phase. (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).

2. null replacement values are no longer registered in the ConversionValueMapping. This was an oversight in #106760. null values in the mapping could cause crashes when using the ConversionValueMapping API.

---

Full diff: https://github.com/llvm/llvm-project/pull/109540.diff

1 Files Affected:

• (modified) mlir/lib/Transforms/Utils/DialectConversion.cpp (+21-8)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 69036e947ebdb0..4b24d7639e435e 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1361,16 +1361,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".
@@ -1378,10 +1383,18 @@ 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 is fine and can happen
+      // when the user is erasing an entire block (including its body).
+      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);

[zero9178]

Trying to play devils advocate: How difficult would it be to support replacing unrealized casts as well? What about in the one-shot dialect conversion? Trying to think of this in terms of pattern rewrite API where in an ideal world the dialect conversion would not need any asterisks stating "well yes, it derives fromRewritePattern and RewriterBase but its not an actual subtype as these specific things are not allowed".

[matthias-springer]

Trying to play devils advocate: How difficult would it be to support replacing unrealized casts as well? What about in the one-shot dialect conversion? Trying to think of this in terms of pattern rewrite API where in an ideal world the dialect conversion would not need any asterisks stating "well yes, it derives fromRewritePattern and RewriterBase but its not an actual subtype as these specific things are not allowed".

The One-Shot Dialect Conversion will support it. (Because much less bookkeeping is necessary and most internal state will entirely disappear. The One-Shot driver will just look at the actual IR in most cases. And then it doesn't matter how we ended up with that IR.) It could even be supported in the current dialect conversion driver, but some internal state would have to be updated. I'm trying to keep the driver as simple as possible before the switch to One-Shot. (And all code that I would write to support this case, will be deleted anyway.)

This change is mainly to help with debugging if something goes wrong in downstream projects during the migration. I plan to send a few more PRs that just assertions for unsupported API usage.

[ingomueller-net]

LGTM!