[Inliner] allocas created for byval now respect addrspace

[RossBrunton]

Summary

Some targets support allocas in non-default address spaces. These
allocas may be used as "byval" parameters for functions.

When inlining functions, the inliner needs to create an alloca to
hold each inlinee's local copy of those arguments. Before this
patch, those allocas would always be in the target's "alloca
address space". However, if the type of the byval parameter is a
pointer to some other address space, then the address space of the
new alloca will not match its user's in the function's body.

This change makes the byval alloca's address space match the
address space of the respective parameter's argument.

llvm/lib/Transforms/Utils/InlineFunction.cpp

static Value *HandleByValArgument(Type *ByValType, Value *Arg,
   Align Alignment(DL.getPrefTypeAlignment(ByValType));
 
   // If the byval had an alignment specified, we *must* use at least that
   // alignment, as it is required by the byval argument (and uses of the
   // pointer inside the callee).
   if (ByValAlignment > 0)
     Alignment = std::max(Alignment, Align(ByValAlignment));
 
+  // If we have a byval for a pointer type with a specific address space, we
+  // have to use that same address space when inlining the function, so the
+  // alloca should be in that address space.
+  auto AllocaAddressSpace = Arg->getType()->getPointerAddressSpace();
+
   Value *NewAlloca =
-      new AllocaInst(ByValType, DL.getAllocaAddrSpace(), nullptr, Alignment,
+      new AllocaInst(ByValType, AllocaAddressSpace, nullptr, Alignment,
                      Arg->getName(), &*Caller->begin()->begin());
   IFI.StaticAllocas.push_back(cast<AllocaInst>(NewAlloca));
 
   // Uses of the argument in the function should use our new alloca
   // instead.
   return NewAlloca;
 }
 

llvm/test/Transforms/Inline/byval-different-addrspace.ll

+; RUN: opt -always-inline %s -S | FileCheck %s
+
+declare void @consume_ptr(i32 addrspace(3)*);
+
+define void @inlinee(i32 addrspace(3)* byval(i32) %this) #0 {
+    call void @consume_ptr(i32 addrspace(3) *%this);
+    ret void
+}
+
+define void @inliner() {
+    %ptr = alloca i32, addrspace(3);
+    call void @inlinee(i32 addrspace(3) *%ptr);
+    ret void
+}
+
+attributes #0 = { alwaysinline }
+
+; CHECK: define void @inliner()
+
+; This call is inlined into @inliner from @inlinee
+; CHECK-DAG: call void @consume_ptr(i32 addrspace(3)* %[[ALLOCA:.*]])
+
+; Check that the alloca that was created by -always-inline is in address space 3
+; CHECK-DAG: %[[ALLOCA]] = alloca i32, align 4, addrspace(3)