[llvm-branch-commits] [mlir] [mlir][Transforms] Simplify `ConversionPatternRewriter::replaceOp` implementation (PR #158075)

[Matthias Springer via llvm-branch-commits]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/158075

Depends on #158067.


>From 8113b1d6c7600dec5ccf93d6c3fe356c08dbc067 Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
Date: Wed, 3 Sep 2025 07:35:47 +0000
Subject: [PATCH] proto

---
 .../Transforms/Utils/DialectConversion.cpp    | 52 +++++++------------
 1 file changed, 20 insertions(+), 32 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 4b483c32ecef9..52369c18faa61 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,
@@ -1709,8 +1711,9 @@ Value ConversionPatternRewriterImpl::findOrBuildReplacementValue(
   // mapping. This includes cached materializations. We try to reuse those
   // instead of generating duplicate IR.
   ValueVector repl = lookupOrNull(value, value.getType());
-  if (!repl.empty())
+  if (!repl.empty()) {
     return repl.front();
+  }
 
   // Check if the value is dead. No replacement value is needed in that case.
   // This is an approximate check that may have false negatives but does not
@@ -1718,22 +1721,14 @@ Value ConversionPatternRewriterImpl::findOrBuildReplacementValue(
   // building source materializations that are never used and that fold away.)
   if (llvm::all_of(value.getUsers(),
                    [&](Operation *op) { return replacedOps.contains(op); }) &&
-      !mapping.isMappedTo(value))
+      !mapping.isMappedTo(value)) {
     return Value();
+  }
 
   // No replacement value was found. Get the latest replacement value
   // (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 +1736,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 +1898,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.