[llvm] Inline: Fix handling of byval using non-alloca addrspace (PR #97306)

[Matt Arsenault via llvm-commits]

https://github.com/arsenm created https://github.com/llvm/llvm-project/pull/97306

Use the address space of the original pointer argument instead
of querying the datalayout. This avoids producing a verifier error
since this was changing the address space for the user instructions.

Fixes #97086

>From 74453f3f277ad4716c2ad79989bd28e65541d2cc Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault at amd.com>
Date: Mon, 1 Jul 2024 16:44:34 +0200
Subject: [PATCH] Inline: Fix handling of byval using non-alloca addrspace

Use the address space of the original pointer argument instead
of querying the datalayout. This avoids producing a verifier error
since this was changing the address space for the user instructions.

Fixes #97086
---
 llvm/lib/Transforms/Utils/InlineFunction.cpp  |  5 ++-
 .../Inline/byval-with-non-alloca-addrspace.ll | 38 +++++++++++++++++++
 2 files changed, 41 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/Inline/byval-with-non-alloca-addrspace.ll

diff --git a/llvm/lib/Transforms/Utils/InlineFunction.cpp b/llvm/lib/Transforms/Utils/InlineFunction.cpp
index 0725addfbb90a..46605492490ea 100644
--- a/llvm/lib/Transforms/Utils/InlineFunction.cpp
+++ b/llvm/lib/Transforms/Utils/InlineFunction.cpp
@@ -1674,8 +1674,9 @@ static Value *HandleByValArgument(Type *ByValType, Value *Arg,
   if (ByValAlignment)
     Alignment = std::max(Alignment, *ByValAlignment);
 
-  AllocaInst *NewAlloca = new AllocaInst(ByValType, DL.getAllocaAddrSpace(),
-                                         nullptr, Alignment, Arg->getName());
+  AllocaInst *NewAlloca = new AllocaInst(
+      ByValType, cast<PointerType>(Arg->getType())->getAddressSpace(), nullptr,
+      Alignment, Arg->getName());
   NewAlloca->insertBefore(Caller->begin()->begin());
   IFI.StaticAllocas.push_back(NewAlloca);
 
diff --git a/llvm/test/Transforms/Inline/byval-with-non-alloca-addrspace.ll b/llvm/test/Transforms/Inline/byval-with-non-alloca-addrspace.ll
new file mode 100644
index 0000000000000..42ec0f2bf5699
--- /dev/null
+++ b/llvm/test/Transforms/Inline/byval-with-non-alloca-addrspace.ll
@@ -0,0 +1,38 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=always-inline %s | FileCheck %s
+
+; The byval parameters use a different address space from the alloca
+; address space. Make sure this is gracefully handled by using the
+; original byval pointer type's address space instead of what the
+; datalayout says to use.
+
+target datalayout = "A5"
+
+%struct = type { i64, i64 }
+
+define i64 @bar(ptr byval(%struct) %a) alwaysinline {
+; CHECK-LABEL: define i64 @bar(
+; CHECK-SAME: ptr byval([[STRUCT:%.*]]) [[A:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr [[STRUCT]], ptr [[A]], i64 0, i32 1
+; CHECK-NEXT:    [[TMP2:%.*]] = load i64, ptr [[TMP1]], align 4
+; CHECK-NEXT:    ret i64 0
+;
+  %1 = getelementptr %struct, ptr %a, i64 0, i32 1
+  %2 = load i64, ptr %1, align 4
+  ret i64 0
+}
+
+define i64 @foo(ptr %arg) {
+; CHECK-LABEL: define i64 @foo(
+; CHECK-SAME: ptr [[ARG:%.*]]) {
+; CHECK-NEXT:    [[ARG1:%.*]] = alloca [[STRUCT:%.*]], align 8
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 16, ptr [[ARG1]])
+; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr align 1 [[ARG1]], ptr align 1 [[ARG]], i64 16, i1 false)
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr [[STRUCT]], ptr [[ARG1]], i64 0, i32 1
+; CHECK-NEXT:    [[TMP2:%.*]] = load i64, ptr [[TMP1]], align 4
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 16, ptr [[ARG1]])
+; CHECK-NEXT:    ret i64 0
+;
+  %1 = call i64 @bar(ptr byval(%struct) align 8 %arg)
+  ret i64 0
+}