[llvm] [AMDGPU] Support alloca in AS0 (PR #136584)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Shilei Tian (shiltian)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/136584.diff


3 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp (+20) 
- (added) llvm/test/CodeGen/AMDGPU/amdgpu-codegenprepare-alloca-as0.ll (+17) 
- (removed) llvm/test/CodeGen/AMDGPU/assert-wrong-alloca-addrspace.ll (-16) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
index a37128b0d745a..a0ef7d9a7a4db 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
+++ b/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,23 @@ 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.getType(), DL.getAllocaAddrSpace(),
+                                           I.getArraySize());
+  NewAI->takeName(&I);
+  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;
diff --git a/llvm/test/CodeGen/AMDGPU/amdgpu-codegenprepare-alloca-as0.ll b/llvm/test/CodeGen/AMDGPU/amdgpu-codegenprepare-alloca-as0.ll
new file mode 100644
index 0000000000000..8a2b54c77ea5d
--- /dev/null
+++ b/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 ptr, align 8, 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
+}
diff --git a/llvm/test/CodeGen/AMDGPU/assert-wrong-alloca-addrspace.ll b/llvm/test/CodeGen/AMDGPU/assert-wrong-alloca-addrspace.ll
deleted file mode 100644
index 1e72e679e83c0..0000000000000
--- a/llvm/test/CodeGen/AMDGPU/assert-wrong-alloca-addrspace.ll
+++ /dev/null
@@ -1,16 +0,0 @@
-; RUN: not --crash llc -mtriple=amdgcn -mcpu=gfx900 -filetype=null %s 2>&1 | FileCheck %s
-
-; The alloca has the wrong address space and is passed to a call. The
-; FrameIndex was created with the natural 32-bit pointer type instead
-; of the declared 64-bit. Make sure we don't assert.
-
-; CHECK: LLVM ERROR: Cannot select: {{.*}}: i64 = FrameIndex<0>
-
-declare void @func(ptr)
-
-define void @main() {
-bb:
-  %alloca = alloca i32, align 4
-  call void @func(ptr %alloca)
-  ret void
-}

``````````

</details>


https://github.com/llvm/llvm-project/pull/136584