[mlir][Transforms] Dialect conversion: add missing argument materialization. #121200
Conversation
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-core Author: Benjamin Chetioui (bchetioui) ChangesWhen replacing a block argument, previously to #117513, we would automatically insert a N->1 argument materialization. After #117513, this is no longer the case for 1->1 mappings. As a result, no materialization is added until Here is an example reproducer. Full diff: https://github.com/llvm/llvm-project/pull/121200.diff 1 Files Affected:
diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 255b0ba2559ee6..5229c0f8d7f2ce 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1375,12 +1375,18 @@ Block *ConversionPatternRewriterImpl::applySignatureConversion(
// used as a replacement.
auto replArgs =
newBlock->getArguments().slice(inputMap->inputNo, inputMap->size);
+ auto insertPoint = OpBuilder::InsertPoint(newBlock, newBlock->begin());
if (replArgs.size() == 1) {
- mapping.map(origArg, replArgs.front());
+ // We need an argument materialization to replace the block argument.
+ Value argMat = buildUnresolvedMaterialization(
+ MaterializationKind::Argument, insertPoint, origArg.getLoc(),
+ /*valueToMap=*/origArg, /*inputs=*/replArgs,
+ /*outputType=*/origArg.getType(), /*originalType=*/Type(), converter);
+ mapping.map(origArg, argMat);
} else {
insertNTo1Materialization(
- OpBuilder::InsertPoint(newBlock, newBlock->begin()), origArg.getLoc(),
- /*replacements=*/replArgs, /*outputValue=*/origArg, converter);
+ insertPoint, origArg.getLoc(), /*replacements=*/replArgs,
+ /*originalValue=*/origArg, converter);
}
appendRewrite<ReplaceBlockArgRewrite>(block, origArg, converter);
}
|
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
materialization. When replacing a block argument, previously to llvm#117513, we would automatically insert a N->1 argument materialization. After llvm#117513, this is no longer the case for 1->1 mappings. As a result, no materialization is added until `ReplaceBlockArgRewrite` is committed, where `findOrBuildReplacementValue` inserts a source materialization. The switch from an argument materialization to a source materialization causes legalization to fail.
bchetioui
force-pushed
the
broken-materializations
branch
from
bdeb48e to
dd88f99
Compare
There was a problem hiding this comment.
The switch from an argument materialization to a source materialization causes legalization to fail.
What kind of failure are you seeing? Is there no matching source materialization callback? It's possible that you have to turn an addArgumentMaterialization into an addSourceMaterialization.
| if (replArgs.size() == 1) { | ||
| mapping.map(origArg, replArgs.front()); | ||
| // We need an argument materialization to replace the block argument. | ||
| Value argMat = buildUnresolvedMaterialization( |
There was a problem hiding this comment.
This should not be necessary. Argument materialization are a workaround around missing 1:N support, but this is a 1:1 replacement.
This is the error I'm seeing: with IR I believe that the materialization is inserted here. But admittedly, I'm not that familiar with these internals---let me try to look into |
Could you try adjusting this line: https://github.com/openxla/xla/blob/eb9d08bae564680ff465d772ceb70f4d84542e8c/xla/mlir_hlo/mhlo/utils/type_conversion.cc#L112 As far as I can see in your IR, the only difference between before and now is that the |
Thanks---I did manage to fix that test by changing the converter! (Though what fixed it was instead adding a target materialization--- At this point, I'm debugging other related failures where this simple fix doesn't seem to work out (and for which I unfortunately don't have shareable reproducers). Nevertheless, I do expect that the issue is similar :) Thank you @matthias-springer and @zero9178 for the assist, I'll close this now---since it seems unnecessary! |
When replacing a block argument, previously to #117513, we would automatically insert a N->1 argument materialization. After #117513, this is no longer the case for 1->1 mappings.
As a result, no materialization is added until
ReplaceBlockArgRewriteis committed, wherefindOrBuildReplacementValueinserts a source materialization. The switch from an argument materialization to a source materialization causes legalization to fail.Here is an example reproducer.