From a764c19d170d97888baa83ec8519427360d681f5 Mon Sep 17 00:00:00 2001
From: River Riddle <riverriddle@google.com>
Date: Sun, 14 Jul 2019 13:29:46 -0700
Subject: [PATCH] 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 | 3 ++-
 mlir/test/Transforms/test-legalizer.mlir  | 8 ++++++++
 mlir/test/lib/TestDialect/TestOps.td      | 5 +++++
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Transforms/DialectConversion.cpp b/mlir/lib/Transforms/DialectConversion.cpp
index fad5a60c453a..980bbd21e36e 100644
--- a/mlir/lib/Transforms/DialectConversion.cpp
+++ b/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.
diff --git a/mlir/test/Transforms/test-legalizer.mlir b/mlir/test/Transforms/test-legalizer.mlir
index 0e2a2d84b37b..6c38b6eba791 100644
--- a/mlir/test/Transforms/test-legalizer.mlir
+++ b/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
+}
diff --git a/mlir/test/lib/TestDialect/TestOps.td b/mlir/test/lib/TestDialect/TestOps.td
index 84f2f37db44e..e4ecc17cce12 100644
--- a/mlir/test/lib/TestDialect/TestOps.td
+++ b/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
 //===----------------------------------------------------------------------===//