[AMDGPU] Make `AllocaInst` return AS5 in `getAssumedAddrSpace` #136798

shiltian · 2025-04-23T01:49:19Z

No description provided.

shiltian · 2025-04-23T01:49:34Z

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

This stack of pull requests is managed by Graphite. Learn more about stacking.

llvmbot · 2025-04-23T01:50:02Z

@llvm/pr-subscribers-backend-amdgpu

Author: Shilei Tian (shiltian)

Changes

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

2 Files Affected:

(modified) llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp (+3)
(added) llvm/test/Transforms/InferAddressSpaces/AMDGPU/alloca-as0.ll (+35)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
index b6cc5137d711a..2c4052a30b10f 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -951,6 +951,9 @@ bool AMDGPUTargetMachine::isNoopAddrSpaceCast(unsigned SrcAS,
 }
 
 unsigned AMDGPUTargetMachine::getAssumedAddrSpace(const Value *V) const {
+  if (isa<AllocaInst>(V))
+    return AMDGPUAS::PRIVATE_ADDRESS;
+
   const auto *LD = dyn_cast<LoadInst>(V);
   if (!LD) // TODO: Handle invariant load like constant.
     return AMDGPUAS::UNKNOWN_ADDRESS_SPACE;
diff --git a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/alloca-as0.ll b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/alloca-as0.ll
new file mode 100644
index 0000000000000..57dcd96594893
--- /dev/null
+++ b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/alloca-as0.ll
@@ -0,0 +1,35 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -passes=infer-address-spaces %s -o - | FileCheck %s
+
+declare void @bar(ptr)
+
+define i32 @static_alloca() {
+; CHECK-LABEL: define i32 @static_alloca() {
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast ptr [[ALLOCA]] to ptr addrspace(5)
+; CHECK-NEXT:    [[TMP2:%.*]] = addrspacecast ptr addrspace(5) [[TMP1]] to ptr
+; CHECK-NEXT:    call void @bar(ptr [[TMP2]])
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32, ptr addrspace(5) [[TMP1]], align 4
+; CHECK-NEXT:    ret i32 [[LOAD]]
+;
+  %alloca = alloca i32, align 4
+  call void @bar(ptr %alloca)
+  %load = load i32, ptr %alloca
+  ret i32 %load
+}
+
+define i32 @dynamic_alloca(i32 %n) {
+; CHECK-LABEL: define i32 @dynamic_alloca(
+; CHECK-SAME: i32 [[N:%.*]]) {
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca i32, i32 [[N]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast ptr [[ALLOCA]] to ptr addrspace(5)
+; CHECK-NEXT:    [[TMP2:%.*]] = addrspacecast ptr addrspace(5) [[TMP1]] to ptr
+; CHECK-NEXT:    call void @bar(ptr [[TMP2]])
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32, ptr addrspace(5) [[TMP1]], align 4
+; CHECK-NEXT:    ret i32 0
+;
+  %alloca = alloca i32, i32 %n, align 4
+  call void @bar(ptr %alloca)
+  %load = load i32, ptr %alloca
+  ret i32 0
+}

llvmbot · 2025-04-23T01:50:02Z

@llvm/pr-subscribers-llvm-transforms

Author: Shilei Tian (shiltian)

Changes

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

2 Files Affected:

(modified) llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp (+3)
(added) llvm/test/Transforms/InferAddressSpaces/AMDGPU/alloca-as0.ll (+35)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
index b6cc5137d711a..2c4052a30b10f 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -951,6 +951,9 @@ bool AMDGPUTargetMachine::isNoopAddrSpaceCast(unsigned SrcAS,
 }
 
 unsigned AMDGPUTargetMachine::getAssumedAddrSpace(const Value *V) const {
+  if (isa<AllocaInst>(V))
+    return AMDGPUAS::PRIVATE_ADDRESS;
+
   const auto *LD = dyn_cast<LoadInst>(V);
   if (!LD) // TODO: Handle invariant load like constant.
     return AMDGPUAS::UNKNOWN_ADDRESS_SPACE;
diff --git a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/alloca-as0.ll b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/alloca-as0.ll
new file mode 100644
index 0000000000000..57dcd96594893
--- /dev/null
+++ b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/alloca-as0.ll
@@ -0,0 +1,35 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -passes=infer-address-spaces %s -o - | FileCheck %s
+
+declare void @bar(ptr)
+
+define i32 @static_alloca() {
+; CHECK-LABEL: define i32 @static_alloca() {
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast ptr [[ALLOCA]] to ptr addrspace(5)
+; CHECK-NEXT:    [[TMP2:%.*]] = addrspacecast ptr addrspace(5) [[TMP1]] to ptr
+; CHECK-NEXT:    call void @bar(ptr [[TMP2]])
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32, ptr addrspace(5) [[TMP1]], align 4
+; CHECK-NEXT:    ret i32 [[LOAD]]
+;
+  %alloca = alloca i32, align 4
+  call void @bar(ptr %alloca)
+  %load = load i32, ptr %alloca
+  ret i32 %load
+}
+
+define i32 @dynamic_alloca(i32 %n) {
+; CHECK-LABEL: define i32 @dynamic_alloca(
+; CHECK-SAME: i32 [[N:%.*]]) {
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca i32, i32 [[N]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast ptr [[ALLOCA]] to ptr addrspace(5)
+; CHECK-NEXT:    [[TMP2:%.*]] = addrspacecast ptr addrspace(5) [[TMP1]] to ptr
+; CHECK-NEXT:    call void @bar(ptr [[TMP2]])
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32, ptr addrspace(5) [[TMP1]], align 4
+; CHECK-NEXT:    ret i32 0
+;
+  %alloca = alloca i32, i32 %n, align 4
+  call void @bar(ptr %alloca)
+  %load = load i32, ptr %alloca
+  ret i32 0
+}

shiltian · 2025-04-23T13:19:03Z

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

 ; GI-NEXT:    v_mov_b32_e32 v1, s15
 ; GI-NEXT:    s_mov_b32 s14, s16
 ; GI-NEXT:    s_movk_i32 s32, 0x400
-; GI-NEXT:    s_mov_b32 s36, 0


That's why I explicitly added -O0 in the test case before.

arsenm · 2025-04-23T13:37:17Z

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

 }

 unsigned AMDGPUTargetMachine::getAssumedAddrSpace(const Value *V) const {
+  if (isa<AllocaInst>(V))


Is this only reachable for flat typed values?

I don't think we need to differentiate them. It is either flat or private. It can't be something else, especially after #135820.

arsenm · 2025-04-23T13:37:52Z

llvm/test/Transforms/InferAddressSpaces/AMDGPU/alloca-as0.ll

+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast ptr [[ALLOCA]] to ptr addrspace(5)
+; CHECK-NEXT:    [[TMP2:%.*]] = addrspacecast ptr addrspace(5) [[TMP1]] to ptr


We probably don't manage to eliminate this cast pair...

No, this pair is for the function call. Not for the load.

By "eliminate this cast pair", do you mean to rewrite alloca with the right AS, or to use the original alloca directly for the function call?

I didn't say anything about a load. I don't think this pair is legally folded out by instcombine

Oh we don't run instcombine here. Let me add it.

No, shouldn't run instcombine here. It's unnecessary, we should see the raw pass output. It's a question of whether the full pass pipeline eliminates the cast. There was a recent DAG combine where IIRC the conclusion was instcombine folds these but it's not necessarily legal

K, does a follow-up work?

shiltian mentioned this pull request Apr 23, 2025

[AMDGPU] Support alloca in AS0 #136584

Closed

shiltian marked this pull request as ready for review April 23, 2025 01:49

shiltian requested a review from arsenm April 23, 2025 01:49

llvmbot added backend:AMDGPU llvm:transforms labels Apr 23, 2025

shiltian force-pushed the users/shiltian/alloca-as0-inferas branch from fb9ba07 to b5f7d3d Compare April 23, 2025 13:17

shiltian force-pushed the users/shiltian/amdgpu-alloca-as0 branch from ae09399 to c77ed29 Compare April 23, 2025 13:17

shiltian commented Apr 23, 2025

View reviewed changes

arsenm reviewed Apr 23, 2025

View reviewed changes

shiltian mentioned this pull request Apr 23, 2025

[Attributor] Use getAssumedAddrSpace to get address space for AllocaInst #136865

Closed

shiltian force-pushed the users/shiltian/amdgpu-alloca-as0 branch from c77ed29 to d04b738 Compare April 23, 2025 19:19

shiltian force-pushed the users/shiltian/alloca-as0-inferas branch from b5f7d3d to 8b975d2 Compare April 23, 2025 19:19

shiltian force-pushed the users/shiltian/amdgpu-alloca-as0 branch from d04b738 to 1f1cbe2 Compare April 23, 2025 20:48

shiltian force-pushed the users/shiltian/alloca-as0-inferas branch from 8b975d2 to 9d2612c Compare April 23, 2025 20:48

[AMDGPU] Make AllocaInst return AS5 in getAssumedAddrSpace

2d75ec2

shiltian force-pushed the users/shiltian/alloca-as0-inferas branch from 9d2612c to 2d75ec2 Compare April 23, 2025 20:51

shiltian force-pushed the users/shiltian/amdgpu-alloca-as0 branch from 1f1cbe2 to 01e6991 Compare April 23, 2025 20:51

shiltian closed this Apr 25, 2025

shiltian deleted the users/shiltian/alloca-as0-inferas branch April 25, 2025 14:33

[AMDGPU] Make AllocaInst return AS5 in getAssumedAddrSpace #136798

[AMDGPU] Make AllocaInst return AS5 in getAssumedAddrSpace #136798

Uh oh!

Conversation

shiltian commented Apr 23, 2025

Uh oh!

shiltian commented Apr 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Apr 23, 2025

Uh oh!

llvmbot commented Apr 23, 2025

Uh oh!

shiltian Apr 23, 2025

Choose a reason for hiding this comment

Uh oh!

arsenm Apr 23, 2025

Choose a reason for hiding this comment

Uh oh!

shiltian Apr 23, 2025

Choose a reason for hiding this comment

Uh oh!

arsenm Apr 23, 2025

Choose a reason for hiding this comment

Uh oh!

shiltian Apr 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

arsenm Apr 23, 2025

Choose a reason for hiding this comment

Uh oh!

shiltian Apr 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

arsenm Apr 23, 2025

Choose a reason for hiding this comment

Uh oh!

shiltian Apr 23, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

[AMDGPU] Make `AllocaInst` return AS5 in `getAssumedAddrSpace` #136798

[AMDGPU] Make `AllocaInst` return AS5 in `getAssumedAddrSpace` #136798

shiltian commented Apr 23, 2025 •

edited

Loading

shiltian Apr 23, 2025 •

edited

Loading

shiltian Apr 23, 2025 •

edited

Loading