[Mlir-commits] [mlir] [mlir][Transforms] Dialect conversion: add missing argument materialization. (PR #121200)

[Benjamin Chetioui]

https://github.com/bchetioui updated https://github.com/llvm/llvm-project/pull/121200

>From dd88f9959f4a7410eccd546b9308d216027377c9 Mon Sep 17 00:00:00 2001
From: Benjamin Chetioui <bchetioui at google.com>
Date: Fri, 27 Dec 2024 10:27:58 +0000
Subject: [PATCH] [mlir][Transforms] Dialect conversion: add missing argument
 materialization.

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 `ReplaceBlockArgRewrite`
is committed, where `findOrBuildReplacementValue` inserts a source
materialization. The switch from an argument materialization to a source
materialization causes legalization to fail.
---
 mlir/lib/Transforms/Utils/DialectConversion.cpp | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 255b0ba2559ee6..aa06fee998c8fa 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);
+      insertNTo1Materialization(insertPoint, origArg.getLoc(),
+                                /*replacements=*/replArgs,
+                                /*originalValue=*/origArg, converter);
     }
     appendRewrite<ReplaceBlockArgRewrite>(block, origArg, converter);
   }