[AMDGPU] Support alloca in AS0

This PR lowers an alloca in AS0 to an alloca in AS5 followed by an addrspacecast
back to AS0.

Author: shiltian

llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp

@@ -330,6 +330,9 @@ class AMDGPUCodeGenPrepareImpl
   bool visitBitreverseIntrinsicInst(IntrinsicInst &I);
   bool visitMinNum(IntrinsicInst &I);
   bool visitSqrt(IntrinsicInst &I);
+
+  bool visitAllocaInst(AllocaInst &I);
+
   bool run();
 };
 
@@ -2355,6 +2358,24 @@ bool AMDGPUCodeGenPrepareImpl::visitSqrt(IntrinsicInst &Sqrt) {
   return true;
 }
 
+// Rewrite alloca with AS0 to alloca with AS5 followed by a addrspace cast.
+bool AMDGPUCodeGenPrepareImpl::visitAllocaInst(AllocaInst &I) {
+  if (I.getAddressSpace() == DL.getAllocaAddrSpace())
+    return false;
+  assert(I.getAddressSpace() == 0 && "An alloca can't be in random AS");
+  IRBuilder<> Builder(&I);
+  AllocaInst *NewAI = Builder.CreateAlloca(
+      I.getAllocatedType(), DL.getAllocaAddrSpace(), I.getArraySize());
+  NewAI->takeName(&I);
+  NewAI->setAlignment(I.getAlign());
+  NewAI->copyMetadata(I);
+  Value *CastI = Builder.CreateAddrSpaceCast(NewAI, I.getType(),
+                                             NewAI->getName() + ".cast");
+  I.replaceAllUsesWith(CastI);
+  I.eraseFromParent();
+  return true;
+}
+
 bool AMDGPUCodeGenPrepare::runOnFunction(Function &F) {
   if (skipFunction(F))
     return false;

llvm/test/CodeGen/AMDGPU/amdgpu-codegenprepare-alloca-as0.ll

@@ -0,0 +1,17 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -mcpu=hawaii -passes=amdgpu-codegenprepare %s | FileCheck %s
+
+declare void @foo(ptr)
+
+define void @bar() {
+; CHECK-LABEL: define void @bar
+; CHECK-SAME: () #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca i32, align 4, addrspace(5)
+; CHECK-NEXT:    [[ALLOCA_CAST:%.*]] = addrspacecast ptr addrspace(5) [[ALLOCA]] to ptr
+; CHECK-NEXT:    call void @foo(ptr [[ALLOCA_CAST]])
+; CHECK-NEXT:    ret void
+;
+  %alloca = alloca i32, align 4
+  call void @foo(ptr %alloca)
+  ret void
+}