[Mlir-commits] [mlir] [MLIR][Transforms] Fix dialect conversion inverse mapping (PR #104648)

[Billy Zhu]

https://github.com/zyx-billy updated https://github.com/llvm/llvm-project/pull/104648

>From 76334c7a735276269f2539f2a8ee949edf380743 Mon Sep 17 00:00:00 2001
From: Billy Zhu <billyzhu at modular.com>
Date: Fri, 16 Aug 2024 16:09:50 -0700
Subject: [PATCH 1/3] update inverse mapping

---
 mlir/lib/Transforms/Utils/DialectConversion.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 8a4c7463a69a95..bf2990c257bad2 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -2527,6 +2527,7 @@ LogicalResult OperationConverter::legalizeConvertedOpResultTypes(
           opReplacement->getConverter());
       rewriterImpl.mapping.map(result, castValue);
       inverseMapping[castValue].push_back(result);
+      llvm::erase(inverseMapping[newValue], result);
     }
   }
 

>From 5763c6ab173e40d2e4b4b91138ac340bcfeef6f6 Mon Sep 17 00:00:00 2001
From: Billy Zhu <billyzhu at modular.com>
Date: Fri, 16 Aug 2024 17:04:08 -0700
Subject: [PATCH 2/3] add test case

---
 .../test-legalize-target-materialization-no-uses.mlir      | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/mlir/test/Transforms/test-legalize-target-materialization-no-uses.mlir b/mlir/test/Transforms/test-legalize-target-materialization-no-uses.mlir
index 19dd2b2ee4b2f2..04db95ba7a9bac 100644
--- a/mlir/test/Transforms/test-legalize-target-materialization-no-uses.mlir
+++ b/mlir/test/Transforms/test-legalize-target-materialization-no-uses.mlir
@@ -25,3 +25,10 @@ func.func @foo() {
   "test.type_consumer"(%2) : (i64) -> ()
   return
 }
+
+// CHECK-LABEL: @direct_forward
+func.func @direct_forward(%arg0 : i16) -> i16 {
+  // CHECK-NEXT: return
+  %0 = "test.type_changer"(%arg0) : (i16) -> i16
+  return %0 : i16
+}

>From 74ba45caec5cd069c8f96b812a0c6df05022ead2 Mon Sep 17 00:00:00 2001
From: Billy Zhu <billyzhu at modular.com>
Date: Mon, 19 Aug 2024 12:45:11 -0700
Subject: [PATCH 3/3] split test

---
 .../test-legalize-target-materialization-no-uses.mlir         | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/mlir/test/Transforms/test-legalize-target-materialization-no-uses.mlir b/mlir/test/Transforms/test-legalize-target-materialization-no-uses.mlir
index 04db95ba7a9bac..5e694a04af49a8 100644
--- a/mlir/test/Transforms/test-legalize-target-materialization-no-uses.mlir
+++ b/mlir/test/Transforms/test-legalize-target-materialization-no-uses.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt -test-target-materialization-with-no-uses %s | FileCheck %s
+// RUN: mlir-opt -test-target-materialization-with-no-uses -split-input-file %s | FileCheck %s
 
 // The conversion is set up as follows:
 // - type_changer ops are illegal;
@@ -26,6 +26,8 @@ func.func @foo() {
   return
 }
 
+// -----
+
 // CHECK-LABEL: @direct_forward
 func.func @direct_forward(%arg0 : i16) -> i16 {
   // CHECK-NEXT: return