[Mlir-commits] [mlir] c0cba25 - [mlir][Transforms] Dialect conversion: Hardening `replaceOp` (#109540)

[llvmlistbot at llvm.org]

Author: Matthias Springer
Date: 2024-10-29T21:13:54+09:00
New Revision: c0cba25cdd06d700bdc15e9ae48c1fcadd0963bd

URL: https://github.com/llvm/llvm-project/commit/c0cba25cdd06d700bdc15e9ae48c1fcadd0963bd
DIFF: https://github.com/llvm/llvm-project/commit/c0cba25cdd06d700bdc15e9ae48c1fcadd0963bd.diff

LOG: [mlir][Transforms] Dialect conversion: Hardening `replaceOp` (#109540)

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 #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.

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/DialectConversion.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 9f8a482d6e2d22..44cf8331d55a73 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1382,16 +1382,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".
@@ -1400,10 +1405,20 @@ void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op,
           result.getLoc(), /*inputs=*/ValueRange(),
           /*outputType=*/result.getType(), /*originalType=*/Type(),
           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);