[mlir][Transforms] Dialect conversion: Align handling of dropped values

[matthias-springer]

Handle dropped block arguments and dropped op results in the same way: build a source materialization (that may fold away if unused). This simplifies the code base a bit and makes it possible to merge legalizeConvertedArgumentTypes and legalizeConvertedOpResultTypes in a future commit. These two functions are almost doing the same thing now.

As a side effect, this commit also changes the dialect conversion such that temporary circular cast ops are no longer generated. (There was a workaround in #107109 that can now be removed again.) Example:

%0 = "builtin.unrealized_conversion_cast"(%1) : (!a) -> !b
%1 = "builtin.unrealized_conversion_cast"(%0) : (!b) -> !a
// No further uses of %0, %1.

This happened when:

1. An op was erased. (No replacement values provided.)
2. A conversion pattern for another op builds a replacement value for the erased op's results (first cast op) during remapValues, but that SSA value is not used during the pattern application.
3. During the finalization phase, legalizeConvertedOpResultTypes thinks that the erased op is alive because of the cast op that was built in Step 2. It builds a cast from that replacement value to the original type.
4. During the commit phase, all uses of the original op are replaced with the casted value produced in Step 3. We have generated circular IR.

This problem can be avoided by making sure that source materializations are generated for all dropped results. This ensures that we always have some replacement SSA value in the mapping. Previously, we sometimes had a value mapped and sometimes not. (No more special casing is needed anymore to distinguish between "value dropped" or "value replaced with SSA value".)

[llvmbot]

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

Handle dropped block arguments and dropped op results in the same way: build a source materialization (that may fold away if unused). This simplifies the code base a bit and makes it possible to merge legalizeConvertedArgumentTypes and legalizeConvertedOpResultTypes in a future commit. These two functions are almost doing the same thing now.

This commit also fixes a bug where circular materializations were built, e.g.:

%0 = "builtin.unrealized_conversion_cast"(%1) : (!a) -> !b
%1 = "builtin.unrealized_conversion_cast"(%0) : (!b) -> !a
// No further uses of %0, %1.

This happened when:

1. An op was erased. (No replacement values provided.)
2. A conversion pattern for another op builds a replacement value (first cast op) during remapValues, but that SSA value is not used during the pattern application.
3. During the finalization phase, legalizeConvertedOpResultTypes thinks that the erased op is alive because of the cast op that was built in Step 2. It builds a cast from that replacement value to the original type.
4. During the commit phase, all uses of the original op are repalced with the casted value produced in Step 3. We have generated circular IR.

---

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

2 Files Affected:

• (modified) mlir/lib/Transforms/Utils/DialectConversion.cpp (+8-34)
• (modified) mlir/test/Transforms/test-legalize-erased-op-with-uses.mlir (+2-2)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index cc9c9495e5155c..ebcc9f948df35c 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1385,9 +1385,14 @@ void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op,
   // Create mappings for each of the new result values.
   for (auto [newValue, result] : llvm::zip(newValues, op->getResults())) {
     if (!newValue) {
-      resultChanged = true;
-      continue;
+      // This result was dropped and no replacement value was provided.
+      // Materialize a replacement value "out of thin air".
+      newValue = buildUnresolvedMaterialization(
+          MaterializationKind::Source, computeInsertPoint(result),
+          result.getLoc(), /*inputs=*/ValueRange(),
+          /*outputType=*/result.getType(), currentTypeConverter);
     }
+
     // Remap, and check for any result type changes.
     mapping.map(result, newValue);
     resultChanged |= (newValue.getType() != result.getType());
@@ -2359,11 +2364,6 @@ struct OperationConverter {
       ConversionPatternRewriterImpl &rewriterImpl,
       DenseMap<Value, SmallVector<Value>> &inverseMapping);
 
-  /// Legalize an operation result that was marked as "erased".
-  LogicalResult
-  legalizeErasedResult(Operation *op, OpResult result,
-                       ConversionPatternRewriterImpl &rewriterImpl);
-
   /// Dialect conversion configuration.
   ConversionConfig config;
 
@@ -2579,14 +2579,7 @@ LogicalResult OperationConverter::legalizeConvertedOpResultTypes(
     Operation *op = opReplacement->getOperation();
     for (OpResult result : op->getResults()) {
       Value newValue = rewriterImpl.mapping.lookupOrNull(result);
-
-      // If the operation result was replaced with null, all of the uses of this
-      // value should be replaced.
-      if (!newValue) {
-        if (failed(legalizeErasedResult(op, result, rewriterImpl)))
-          return failure();
-        continue;
-      }
+      assert(newValue && "replacement value not found");
 
       // Otherwise, check to see if the type of the result changed.
       if (result.getType() == newValue.getType())
@@ -2655,25 +2648,6 @@ LogicalResult OperationConverter::legalizeConvertedArgumentTypes(
   return success();
 }
 
-LogicalResult OperationConverter::legalizeErasedResult(
-    Operation *op, OpResult result,
-    ConversionPatternRewriterImpl &rewriterImpl) {
-  // If the operation result was replaced with null, all of the uses of this
-  // value should be replaced.
-  auto liveUserIt = llvm::find_if_not(result.getUsers(), [&](Operation *user) {
-    return rewriterImpl.isOpIgnored(user);
-  });
-  if (liveUserIt != result.user_end()) {
-    InFlightDiagnostic diag = op->emitError("failed to legalize operation '")
-                              << op->getName() << "' marked as erased";
-    diag.attachNote(liveUserIt->getLoc())
-        << "found live user of result #" << result.getResultNumber() << ": "
-        << *liveUserIt;
-    return failure();
-  }
-  return success();
-}
-
 //===----------------------------------------------------------------------===//
 // Reconcile Unrealized Casts
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Transforms/test-legalize-erased-op-with-uses.mlir b/mlir/test/Transforms/test-legalize-erased-op-with-uses.mlir
index 49275e8008e749..b536b80ce58309 100644
--- a/mlir/test/Transforms/test-legalize-erased-op-with-uses.mlir
+++ b/mlir/test/Transforms/test-legalize-erased-op-with-uses.mlir
@@ -3,8 +3,8 @@
 // Test that an error is emitted when an operation is marked as "erased", but
 // has users that live across the conversion.
 func.func @remove_all_ops(%arg0: i32) -> i32 {
-  // expected-error@below {{failed to legalize operation 'test.illegal_op_a' marked as erased}}
+  // expected-error@below {{failed to legalize unresolved materialization from () to 'i32' that remained live after conversion}}
   %0 = "test.illegal_op_a"() : () -> i32
-  // expected-note@below {{found live user of result #0: func.return %0 : i32}}
+  // expected-note@below {{see existing live user here: func.return %0 : i32}}
   return %0 : i32
 }

[zero9178]

Is my intuition correct that this is rather difficult to test until the other PR relands but could then be tested by checking for a () -> i32 unrealized_conversion_cast remaining in the IR (in the scenario showin in @remove_all_ops)?

Or could this maybe be tested right now be checking scenarios where previously two materializations would be created but now only one is?

[matthias-springer]

@zero9178 I added a test case for the circular casts to #107109. I added an extra piece of code there to handle the circular casts. This piece of code can probably be deleted again with this PR.

I'd like to land #106760 first because I already ran it by Google and IREE folks and it should be fine (they have much more code that stress-tests the dialect conversion). #106760 is also needed more urgently by CIRCT people in their LLVM integration (see comments here: #106778 (comment)). I'd like to keep a few days between each merged dialect conversion commit so that it will be easier to address issues (if any). So, landing this PR not before next week or so...

[zero9178]

Amazing, thank you so much!