Fix a bug in DialectConversion when using RewritePattern.

When using a RewritePattern and replacing an operation with an existing value, that value may have already been replaced by something else. This cl ensures that only the final value is used when applying rewrites.

PiperOrigin-RevId: 258058488

mlir/lib/Transforms/DialectConversion.cpp

@@ -431,7 +431,8 @@ struct DialectConversionRewriter final : public PatternRewriter {
     // Apply all of the rewrites replacements requested during conversion.
     for (auto &repl : replacements) {
       for (unsigned i = 0, e = repl.newValues.size(); i != e; ++i)
-        repl.op->getResult(i)->replaceAllUsesWith(repl.newValues[i]);
+        repl.op->getResult(i)->replaceAllUsesWith(
+            mapping.lookupOrDefault(repl.newValues[i]));
 
       // if this operation defines any regions, drop any pending argument
       // rewrites.

mlir/test/Transforms/test-legalizer.mlir

@@ -91,3 +91,11 @@ func @dropped_input_in_use(%arg: i16, %arg2: i64) {
   // CHECK-NEXT: "work"{{.*}} : (i16)
   "work"(%arg) : (i16) -> ()
 }
+
+// CHECK-LABEL: func @up_to_date_replacement
+func @up_to_date_replacement(%arg: i8) -> i8 {
+  // CHECK-NEXT: return
+  %repl_1 = "test.rewrite"(%arg) : (i8) -> i8
+  %repl_2 = "test.rewrite"(%repl_1) : (i8) -> i8
+  return %repl_2 : i8
+}

mlir/test/lib/TestDialect/TestOps.td

@@ -425,6 +425,11 @@ def : Pat<(ILLegalOpD), (LegalOpA Test_LegalizerEnum_Failure)>;
 def : Pat<(ILLegalOpC), (ILLegalOpE), [], (addBenefit 10)>;
 def : Pat<(ILLegalOpE), (LegalOpA Test_LegalizerEnum_Success)>;
 
+// Check that patterns use the most up-to-date value when being replaced.
+def TestRewriteOp : TEST_Op<"rewrite">,
+  Arguments<(ins AnyType:$input)>, Results<(outs AnyType:$res)>;
+def : Pat<(TestRewriteOp $input), (replaceWithValue $input)>;
+
 //===----------------------------------------------------------------------===//
 // Test Type Legalization
 //===----------------------------------------------------------------------===//