Inline: Fix handling of byval using non-alloca addrspace (llvm#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 llvm#97086

Author: arsenm

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, Arg->getType()->getPointerAddressSpace(),
+                     nullptr, Alignment, Arg->getName());
   NewAlloca->insertBefore(Caller->begin()->begin());
   IFI.StaticAllocas.push_back(NewAlloca);
 

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
+}