[Mlir-commits] [mlir] [mlir][Transforms] Dialect conversion: Fix bug in `UnresolvedMaterializationRewrite` rollback (PR #105949)

[Matthias Springer]

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

When an unresolved materialization (`unrealized_conversion_cast` op) is rolled back, the mapping should be rolled back as well, regardless of whether it is a source, target or argument materialization. Otherwise, we accumulate pointers to erased IR in the `mapping`. This is harmless in most cases, but can cause issues when a new operation is allocated at the same memory location and the pointer is "reused".

It is not possible to write a test case for this because I cannot trigger the pointer reuse programmatically.

>From 0b9a69b3beb03f0de9b3c699e9b32a7b61b3ca57 Mon Sep 17 00:00:00 2001
From: Matthias Springer <mspringer at nvidia.com>
Date: Sat, 24 Aug 2024 18:34:45 +0200
Subject: [PATCH] [mlir][Transforms] Dialect conversion: Fix bug in
 `UnresolvedMaterializationRewrite` rollback

When an unresolved materializations (`unrealized_conversion_cast` op) is rolled back, the mapping should be rolled back as well, regardless of whether it is a source, target or argument materialization. Otherwise, we accumulate pointers to erased IR in the `mapping`. This is harmless in most cases, but can cause issues when a new operation is allocated at the same memory location and the pointer is "reused".

It is not possible to write a test case for this because I cannot trigger the pointer reuse programmatically.
---
 mlir/lib/Transforms/Utils/DialectConversion.cpp | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 4058ed39621198..11e40b74b00424 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1062,10 +1062,8 @@ void CreateOperationRewrite::rollback() {
 }
 
 void UnresolvedMaterializationRewrite::rollback() {
-  if (getMaterializationKind() == MaterializationKind::Target) {
-    for (Value input : op->getOperands())
-      rewriterImpl.mapping.erase(input);
-  }
+  for (Value input : op->getOperands())
+    rewriterImpl.mapping.erase(input);
   op->erase();
 }