[mlir][Transforms] Dialect conversion: Fix bug in UnresolvedMaterializationRewrite rollback

[matthias-springer]

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.

[llvmbot]

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/105949.diff

1 Files Affected:

• (modified) mlir/lib/Transforms/Utils/DialectConversion.cpp (+2-4)

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();
 }
 

[zero9178]

This seems fine to me, as in being a step in the right direction, but it doesn't seem guaranteed to me that the input operand was actually the operand that was originally mapped to the materilization:

llvm-project/mlir/lib/Transforms/Utils/DialectConversion.cpp

Line 1155 in 0b9a69b

mapping.map(mapping.lookupOrDefault(newOperand), castValue);

That said, this seems to be more of a pre-existing issue to me.

[matthias-springer]

This seems fine to me, as in being a step in the right direction, but it doesn't seem guaranteed to me that the input operand was actually the operand that was originally mapped to the materilization:

llvm-project/mlir/lib/Transforms/Utils/DialectConversion.cpp

Line 1155 in 0b9a69b

mapping.map(mapping.lookupOrDefault(newOperand), castValue);

That said, this seems to be more of a pre-existing issue to me.

You are right, it's not a perfect rollback and there are probably other issues that are not fixed by this PR.

It's actually not just that line of code that you pointed out. Here's another one:

llvm-project/mlir/lib/Transforms/Utils/DialectConversion.cpp

Line 2604 in 0b9a69b

rewriterImpl.mapping.map(result, castValue);

I think this is due to the fact that:

• An SSA value may be converted to multiple different types. E.g., different patterns with different type converters may create multiple conversions of the same SSA value to different types during remapValues. (And we want to convert the original value, not an already converted value.)
• The mapping data structure is a "flat" IRMapping. An SSA value can be mapped only to one other SSA value. We emulate the fact that there may be multiple conversions by chaining mappings. E.g., the first conversion will map to the second conversion. Kind of like a linked list. When we should really have a mapping of Value to SmallVector<Value>.

I'm thinking about fixing that. But not sure if it's worth it, because ultimately I want to get to a state where all IR changes are materialized immediately, and there is no need for a mapping anymore.

(The purpose of this change was mostly to cut down on code size. Bonus points if it fixes something. Anything that reduces complexity and the lines of code will make it easier to maintain the code base and/or build a new conversion driver.)

[matthias-springer]

@zero9178 I never merged this PR, but this issue was now triggered by a Google-internal test case when I asked @jpienaar to test #116470 internally before merging it. (I originally put these changes into #116470, but I think it's better to merge two smaller PRs with a day in-between.)

My original fix may not have been enough. (Or maybe it was, I don't want to take any chances...) This is now rolling back the mapping exactly to the previous state, addressing your comment here. I'm still not able to write a test case for this. (Because ReplaceOperationRewrite::rollback() and ReplaceOperationRewrite::rollback() also roll back the mapping. The dangling mapping is for target materialization that converts the result of an argument materialization.) I think this can only be triggered by accidental pointer reuse.

Can you take an other brief look at this before I merge it?

[zero9178]

LGTM, thank you! Greatly prefer this more principled approache.

One minor possible enhancement in the comment.