[mlir][Transforms] Simplify ConversionPatternRewriter::replaceOp implementation

[matthias-springer]

Move the logic for building "out-of-thin-air" source materializations during op replacements from replaceOp to findOrBuildReplacementValue. That function already builds source materializations and can handle the case where an op result is dropped.

This commit is in preparation of turning replaceOp into a non-virtual function. (It is sufficient for replaceAllUsesWith and eraseOp to be virtual.)

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

Move the logic for building "out-of-thin-air" source materializations during op replacements from replaceOp to findOrBuildReplacementValue. That function already builds source materializations and can handle the case where an op result is dropped.

This commit is in preparation of turning replaceOp into a non-virtual function. (It is sufficient for replaceAllUsesWith and eraseOp to be virtual.)

---

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

1 Files Affected:

• (modified) mlir/lib/Transforms/Utils/DialectConversion.cpp (+16-30)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index f7565cfb0e45e..63b2917639cf6 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1618,6 +1618,8 @@ Block *ConversionPatternRewriterImpl::applySignatureConversion(
     if (!inputMap) {
       // This block argument was dropped and no replacement value was provided.
       // Materialize a replacement value "out of thin air".
+      // Note: Materialization must be built here because we cannot find a
+      // valid insertion point in the new block. (Will point to the old block.)
       Value mat =
           buildUnresolvedMaterialization(
               MaterializationKind::Source,
@@ -1725,15 +1727,6 @@ Value ConversionPatternRewriterImpl::findOrBuildReplacementValue(
   // (regardless of the type) and build a source materialization to the
   // original type.
   repl = lookupOrNull(value);
-  if (repl.empty()) {
-    // No replacement value is registered in the mapping. This means that the
-    // value is dropped and no longer needed. (If the value were still needed,
-    // a source materialization producing a replacement value "out of thin air"
-    // would have already been created during `replaceOp` or
-    // `applySignatureConversion`.)
-    return Value();
-  }
-
   // Note: `computeInsertPoint` computes the "earliest" insertion point at
   // which all values in `repl` are defined. It is important to emit the
   // materialization at that location because the same materialization may be
@@ -1741,13 +1734,19 @@ Value ConversionPatternRewriterImpl::findOrBuildReplacementValue(
   // in the conversion value mapping.) The insertion point of the
   // materialization must be valid for all future users that may be created
   // later in the conversion process.
-  Value castValue =
-      buildUnresolvedMaterialization(MaterializationKind::Source,
-                                     computeInsertPoint(repl), value.getLoc(),
-                                     /*valuesToMap=*/repl, /*inputs=*/repl,
-                                     /*outputTypes=*/value.getType(),
-                                     /*originalType=*/Type(), converter)
-          .front();
+  OpBuilder::InsertPoint ip;
+  if (repl.empty()) {
+    ip = computeInsertPoint(value);
+  } else {
+    ip = computeInsertPoint(repl);
+  }
+  Value castValue = buildUnresolvedMaterialization(
+                        MaterializationKind::Source, ip, value.getLoc(),
+                        /*valuesToMap=*/repl, /*inputs=*/repl,
+                        /*outputTypes=*/value.getType(),
+                        /*originalType=*/Type(), converter,
+                        /*isPureTypeConversion=*/!repl.empty())
+                        .front();
   return castValue;
 }
 
@@ -1897,21 +1896,8 @@ void ConversionPatternRewriterImpl::replaceOp(
   }
 
   // Create mappings for each of the new result values.
-  for (auto [repl, result] : llvm::zip_equal(newValues, op->getResults())) {
-    if (repl.empty()) {
-      // This result was dropped and no replacement value was provided.
-      // Materialize a replacement value "out of thin air".
-      buildUnresolvedMaterialization(
-          MaterializationKind::Source, computeInsertPoint(result),
-          result.getLoc(), /*valuesToMap=*/{result}, /*inputs=*/ValueRange(),
-          /*outputTypes=*/result.getType(), /*originalType=*/Type(),
-          currentTypeConverter, /*isPureTypeConversion=*/false);
-      continue;
-    }
-
-    // Remap result to replacement value.
+  for (auto [repl, result] : llvm::zip_equal(newValues, op->getResults()))
     mapping.map(static_cast<Value>(result), std::move(repl));
-  }
 
   appendRewrite<ReplaceOperationRewrite>(op, currentTypeConverter);
   // Mark this operation and all nested ops as replaced.

[llvmbot]

@llvm/pr-subscribers-mlir-core

Author: Matthias Springer (matthias-springer)

Changes

Move the logic for building "out-of-thin-air" source materializations during op replacements from replaceOp to findOrBuildReplacementValue. That function already builds source materializations and can handle the case where an op result is dropped.

This commit is in preparation of turning replaceOp into a non-virtual function. (It is sufficient for replaceAllUsesWith and eraseOp to be virtual.)

---

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

1 Files Affected:

• (modified) mlir/lib/Transforms/Utils/DialectConversion.cpp (+16-30)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index f7565cfb0e45e..63b2917639cf6 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1618,6 +1618,8 @@ Block *ConversionPatternRewriterImpl::applySignatureConversion(
     if (!inputMap) {
       // This block argument was dropped and no replacement value was provided.
       // Materialize a replacement value "out of thin air".
+      // Note: Materialization must be built here because we cannot find a
+      // valid insertion point in the new block. (Will point to the old block.)
       Value mat =
           buildUnresolvedMaterialization(
               MaterializationKind::Source,
@@ -1725,15 +1727,6 @@ Value ConversionPatternRewriterImpl::findOrBuildReplacementValue(
   // (regardless of the type) and build a source materialization to the
   // original type.
   repl = lookupOrNull(value);
-  if (repl.empty()) {
-    // No replacement value is registered in the mapping. This means that the
-    // value is dropped and no longer needed. (If the value were still needed,
-    // a source materialization producing a replacement value "out of thin air"
-    // would have already been created during `replaceOp` or
-    // `applySignatureConversion`.)
-    return Value();
-  }
-
   // Note: `computeInsertPoint` computes the "earliest" insertion point at
   // which all values in `repl` are defined. It is important to emit the
   // materialization at that location because the same materialization may be
@@ -1741,13 +1734,19 @@ Value ConversionPatternRewriterImpl::findOrBuildReplacementValue(
   // in the conversion value mapping.) The insertion point of the
   // materialization must be valid for all future users that may be created
   // later in the conversion process.
-  Value castValue =
-      buildUnresolvedMaterialization(MaterializationKind::Source,
-                                     computeInsertPoint(repl), value.getLoc(),
-                                     /*valuesToMap=*/repl, /*inputs=*/repl,
-                                     /*outputTypes=*/value.getType(),
-                                     /*originalType=*/Type(), converter)
-          .front();
+  OpBuilder::InsertPoint ip;
+  if (repl.empty()) {
+    ip = computeInsertPoint(value);
+  } else {
+    ip = computeInsertPoint(repl);
+  }
+  Value castValue = buildUnresolvedMaterialization(
+                        MaterializationKind::Source, ip, value.getLoc(),
+                        /*valuesToMap=*/repl, /*inputs=*/repl,
+                        /*outputTypes=*/value.getType(),
+                        /*originalType=*/Type(), converter,
+                        /*isPureTypeConversion=*/!repl.empty())
+                        .front();
   return castValue;
 }
 
@@ -1897,21 +1896,8 @@ void ConversionPatternRewriterImpl::replaceOp(
   }
 
   // Create mappings for each of the new result values.
-  for (auto [repl, result] : llvm::zip_equal(newValues, op->getResults())) {
-    if (repl.empty()) {
-      // This result was dropped and no replacement value was provided.
-      // Materialize a replacement value "out of thin air".
-      buildUnresolvedMaterialization(
-          MaterializationKind::Source, computeInsertPoint(result),
-          result.getLoc(), /*valuesToMap=*/{result}, /*inputs=*/ValueRange(),
-          /*outputTypes=*/result.getType(), /*originalType=*/Type(),
-          currentTypeConverter, /*isPureTypeConversion=*/false);
-      continue;
-    }
-
-    // Remap result to replacement value.
+  for (auto [repl, result] : llvm::zip_equal(newValues, op->getResults()))
     mapping.map(static_cast<Value>(result), std::move(repl));
-  }
 
   appendRewrite<ReplaceOperationRewrite>(op, currentTypeConverter);
   // Mark this operation and all nested ops as replaced.

[j2kun]

LGTM!