[Mlir-commits] [mlir] [mlir][func] Fix incorrect API usage in `FuncOpConversion` (PR #113977)

[Matthias Springer]

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

This commit fixes a case of incorrect dialect conversion API usage during `FuncOpConversion`. `replaceAllUsesExcept` (same as `replaceAllUsesWith`) is currently not supported in a dialect conversion. `replaceUsesOfBlockArgument` should be used instead. It sometimes works anyway (like in this case), but that's just because of the way we insert materializations.

This commit is in preparation of merging the 1:1 and 1:N dialect conversion drivers. (At that point, the current use of `replaceAllUsesExcept` will no longer work.)

>From 9653e6bc77cbca3627e9ef2165d15bbf9a66ac94 Mon Sep 17 00:00:00 2001
From: Matthias Springer <mspringer at nvidia.com>
Date: Mon, 28 Oct 2024 23:40:10 +0100
Subject: [PATCH] [mlir][func] Fix incorrect API usage in `FuncOpConversion`

This commit fixes a case of incorrect dialect conversion API usage during `FuncOpConversion`. `replaceAllUsesExcept` (same as `replaceAllUsesWith`) is currently not supported in a dialect conversion. `replaceUsesOfBlockArgument` should be used instead. It sometimes works anyway (like in this case), but that's just because of the way we insert materializations.

This commit is in preparation of merging the 1:1 and 1:N dialect conversion drivers. (At that point, the current use of `replaceAllUsesExcept` will no longer work.)
---
 mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp b/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp
index 27c43e0daad072..c046ea1b824fc8 100644
--- a/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp
+++ b/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp
@@ -273,7 +273,7 @@ static void wrapExternalFunction(OpBuilder &builder, Location loc,
 static void restoreByValRefArgumentType(
     ConversionPatternRewriter &rewriter, const LLVMTypeConverter &typeConverter,
     ArrayRef<std::optional<NamedAttribute>> byValRefNonPtrAttrs,
-    LLVM::LLVMFuncOp funcOp) {
+    ArrayRef<BlockArgument> oldBlockArgs, LLVM::LLVMFuncOp funcOp) {
   // Nothing to do for function declarations.
   if (funcOp.isExternal())
     return;
@@ -281,8 +281,8 @@ static void restoreByValRefArgumentType(
   ConversionPatternRewriter::InsertionGuard guard(rewriter);
   rewriter.setInsertionPointToStart(&funcOp.getFunctionBody().front());
 
-  for (const auto &[arg, byValRefAttr] :
-       llvm::zip(funcOp.getArguments(), byValRefNonPtrAttrs)) {
+  for (const auto &[arg, oldArg, byValRefAttr] :
+       llvm::zip(funcOp.getArguments(), oldBlockArgs, byValRefNonPtrAttrs)) {
     // Skip argument if no `llvm.byval` or `llvm.byref` attribute.
     if (!byValRefAttr)
       continue;
@@ -295,7 +295,7 @@ static void restoreByValRefArgumentType(
         cast<TypeAttr>(byValRefAttr->getValue()).getValue());
 
     auto valueArg = rewriter.create<LLVM::LoadOp>(arg.getLoc(), resTy, arg);
-    rewriter.replaceAllUsesExcept(arg, valueArg, valueArg);
+    rewriter.replaceUsesOfBlockArgument(oldArg, valueArg);
   }
 }
 
@@ -309,6 +309,10 @@ mlir::convertFuncOpToLLVMFuncOp(FunctionOpInterface funcOp,
     return rewriter.notifyMatchFailure(
         funcOp, "Only support FunctionOpInterface with FunctionType");
 
+  // Keep track of the entry block arguments. They will be needed later.
+  SmallVector<BlockArgument> oldBlockArgs =
+      llvm::to_vector(funcOp.getArguments());
+
   // Convert the original function arguments. They are converted using the
   // LLVMTypeConverter provided to this legalization pattern.
   auto varargsAttr = funcOp->getAttrOfType<BoolAttr>(varargsAttrName);
@@ -438,7 +442,7 @@ mlir::convertFuncOpToLLVMFuncOp(FunctionOpInterface funcOp,
   // pointee type in the function body when converting `llvm.byval`/`llvm.byref`
   // function arguments.
   restoreByValRefArgumentType(rewriter, converter, byValRefNonPtrAttrs,
-                              newFuncOp);
+                              oldBlockArgs, newFuncOp);
 
   if (!shouldUseBarePtrCallConv(funcOp, &converter)) {
     if (funcOp->getAttrOfType<UnitAttr>(