Skip to content

[AMDGPU] Restore SP correctly in functions with dynamic allocas #122743

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 4 commits into from
Closed

[AMDGPU] Restore SP correctly in functions with dynamic allocas #122743

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c34194a
[AMDGPU] Restore SP correctly in functions with dynamic allocas
easyonaadit Jan 13, 2025
14248c2
use BP path as fall-back
easyonaadit Jan 17, 2025
9f5b2fd
Update conditions for setting up base pointer
easyonaadit Jan 20, 2025
b5d25ea
Code cleanup
easyonaadit Jan 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions llvm/lib/Target/AMDGPU/SIFrameLowering.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1259,6 +1259,23 @@ void SIFrameLowering::emitEpilogue(MachineFunction &MF,
Register FramePtrRegScratchCopy;
Register SGPRForFPSaveRestoreCopy =
FuncInfo->getScratchSGPRCopyDstReg(FramePtrReg);

if (MFI.hasVarSizedObjects()) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's any kind of dynamic stack adjustment. I think this is the same conditions as hasFP (that's what I was using in my last attempt at this)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hasFP checks for variable-sized objects, stack realignment, disable-frame-pointer-elimination optimization, and these intrinsics: @llvm.frameaddress, @llvm.experimental.stackmap, @llvm.experimental.patchpoint.
Of these, I think we only need SP restoration in case of variable-sized objects and stack realignment.
If a function has only stack-realignment without variable-sized objects, then that is handled in the existing prolog/epilog.
So I think we only need to handle SP restoration for variable sized objects (unless there is something else which does a dynamic stack adjustment)

// If we have over-aligned dynamic-sized objects, then restore SP with
// saved-BP, else restore it with saved-FP.
if (FuncInfo->isStackRealigned()) {
Register BasePtrReg = TRI.getBaseRegister();
BuildMI(MBB, MBBI, DL, TII->get(AMDGPU::S_ADD_I32), StackPtrReg)
.addReg(BasePtrReg)
.addImm(RoundedSize * getScratchScaleFactor(ST))
.setMIFlag(MachineInstr::FrameDestroy);
} else {
BuildMI(MBB, MBBI, DL, TII->get(AMDGPU::S_ADD_I32), StackPtrReg)
.addReg(FramePtrReg)
.addImm(RoundedSize * getScratchScaleFactor(ST))
.setMIFlag(MachineInstr::FrameDestroy);
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should SP restoration be marked as FrameDestroy ?
Since SP is actually modified during ISel, and there is no counter FrameSetup action in the prolog for this.

}
}
if (FPSaved) {
// CSR spill restores should use FP as base register. If
// SGPRForFPSaveRestoreCopy is not true, restore the previous value of FP
Expand Down
5 changes: 4 additions & 1 deletion llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -525,8 +525,11 @@ Register SIRegisterInfo::getFrameRegister(const MachineFunction &MF) const {
bool SIRegisterInfo::hasBasePointer(const MachineFunction &MF) const {
// When we need stack realignment, we can't reference off of the
// stack pointer, so we reserve a base pointer.
// For functions with dynamically sized stack objects, we need to reference
// off the base pointer in the epilog to restore the stack frame.
const MachineFrameInfo &MFI = MF.getFrameInfo();
return MFI.getNumFixedObjects() && shouldRealignStack(MF);
return (MFI.getNumFixedObjects() || MFI.hasVarSizedObjects()) &&
shouldRealignStack(MF);
}

Register SIRegisterInfo::getBaseRegister() const { return AMDGPU::SGPR34; }
Expand Down
47 changes: 32 additions & 15 deletions llvm/test/CodeGen/AMDGPU/GlobalISel/dynamic-alloca-uniform.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,13 +80,14 @@ define void @func_dynamic_stackalloc_sgpr_align4() {
; GFX9-NEXT: buffer_store_dword v0, v1, s[0:3], 0 offen
; GFX9-NEXT: s_waitcnt lgkmcnt(0)
; GFX9-NEXT: s_load_dword s4, s[4:5], 0x0
; GFX9-NEXT: s_mov_b32 s33, s7
; GFX9-NEXT: s_waitcnt lgkmcnt(0)
; GFX9-NEXT: s_lshl2_add_u32 s4, s4, 15
; GFX9-NEXT: s_and_b32 s4, s4, -16
; GFX9-NEXT: s_lshl_b32 s4, s4, 6
; GFX9-NEXT: s_add_u32 s32, s6, s4
; GFX9-NEXT: s_add_i32 s32, s33, 0x400
; GFX9-NEXT: s_addk_i32 s32, 0xfc00
; GFX9-NEXT: s_mov_b32 s33, s7
; GFX9-NEXT: s_waitcnt vmcnt(0)
; GFX9-NEXT: s_setpc_b64 s[30:31]
;
Expand All @@ -103,7 +104,6 @@ define void @func_dynamic_stackalloc_sgpr_align4() {
; GFX10-NEXT: s_load_dwordx2 s[4:5], s[4:5], 0x0
; GFX10-NEXT: v_mov_b32_e32 v0, 0
; GFX10-NEXT: v_mov_b32_e32 v1, s6
; GFX10-NEXT: s_mov_b32 s33, s7
; GFX10-NEXT: buffer_store_dword v0, v1, s[0:3], 0 offen
; GFX10-NEXT: s_waitcnt lgkmcnt(0)
; GFX10-NEXT: s_load_dword s4, s[4:5], 0x0
Expand All @@ -112,6 +112,8 @@ define void @func_dynamic_stackalloc_sgpr_align4() {
; GFX10-NEXT: s_and_b32 s4, s4, -16
; GFX10-NEXT: s_lshl_b32 s4, s4, 5
; GFX10-NEXT: s_add_u32 s32, s6, s4
; GFX10-NEXT: s_add_i32 s32, s33, 0x200
; GFX10-NEXT: s_mov_b32 s33, s7
; GFX10-NEXT: s_addk_i32 s32, 0xfe00
; GFX10-NEXT: s_setpc_b64 s[30:31]
;
Expand All @@ -124,10 +126,9 @@ define void @func_dynamic_stackalloc_sgpr_align4() {
; GFX11-NEXT: s_getpc_b64 s[0:1]
; GFX11-NEXT: s_add_u32 s0, s0, gv@gotpcrel32@lo+4
; GFX11-NEXT: s_addc_u32 s1, s1, gv@gotpcrel32@hi+12
; GFX11-NEXT: v_mov_b32_e32 v0, 0
; GFX11-NEXT: s_load_b64 s[0:1], s[0:1], 0x0
; GFX11-NEXT: s_mov_b32 s2, s32
; GFX11-NEXT: s_mov_b32 s33, s3
; GFX11-NEXT: s_load_b64 s[0:1], s[0:1], 0x0
; GFX11-NEXT: v_mov_b32_e32 v0, 0
; GFX11-NEXT: scratch_store_b32 off, v0, s2
; GFX11-NEXT: s_waitcnt lgkmcnt(0)
; GFX11-NEXT: s_load_b32 s0, s[0:1], 0x0
Expand All @@ -136,8 +137,10 @@ define void @func_dynamic_stackalloc_sgpr_align4() {
; GFX11-NEXT: s_delay_alu instid0(SALU_CYCLE_1) | instskip(NEXT) | instid1(SALU_CYCLE_1)
; GFX11-NEXT: s_and_b32 s0, s0, -16
; GFX11-NEXT: s_lshl_b32 s0, s0, 5
; GFX11-NEXT: s_delay_alu instid0(SALU_CYCLE_1) | instskip(NEXT) | instid1(SALU_CYCLE_1)
; GFX11-NEXT: s_delay_alu instid0(SALU_CYCLE_1)
; GFX11-NEXT: s_add_u32 s32, s2, s0
; GFX11-NEXT: s_add_i32 s32, s33, 16
; GFX11-NEXT: s_mov_b32 s33, s3
; GFX11-NEXT: s_add_i32 s32, s32, -16
; GFX11-NEXT: s_setpc_b64 s[30:31]
%n = load i32, ptr addrspace(4) @gv, align 4
Expand Down Expand Up @@ -221,13 +224,14 @@ define void @func_dynamic_stackalloc_sgpr_align16() {
; GFX9-NEXT: buffer_store_dword v0, v1, s[0:3], 0 offen
; GFX9-NEXT: s_waitcnt lgkmcnt(0)
; GFX9-NEXT: s_load_dword s4, s[4:5], 0x0
; GFX9-NEXT: s_mov_b32 s33, s7
; GFX9-NEXT: s_waitcnt lgkmcnt(0)
; GFX9-NEXT: s_lshl2_add_u32 s4, s4, 15
; GFX9-NEXT: s_and_b32 s4, s4, -16
; GFX9-NEXT: s_lshl_b32 s4, s4, 6
; GFX9-NEXT: s_add_u32 s32, s6, s4
; GFX9-NEXT: s_add_i32 s32, s33, 0x400
; GFX9-NEXT: s_addk_i32 s32, 0xfc00
; GFX9-NEXT: s_mov_b32 s33, s7
; GFX9-NEXT: s_waitcnt vmcnt(0)
; GFX9-NEXT: s_setpc_b64 s[30:31]
;
Expand All @@ -244,7 +248,6 @@ define void @func_dynamic_stackalloc_sgpr_align16() {
; GFX10-NEXT: s_load_dwordx2 s[4:5], s[4:5], 0x0
; GFX10-NEXT: v_mov_b32_e32 v0, 0
; GFX10-NEXT: v_mov_b32_e32 v1, s6
; GFX10-NEXT: s_mov_b32 s33, s7
; GFX10-NEXT: buffer_store_dword v0, v1, s[0:3], 0 offen
; GFX10-NEXT: s_waitcnt lgkmcnt(0)
; GFX10-NEXT: s_load_dword s4, s[4:5], 0x0
Expand All @@ -253,6 +256,8 @@ define void @func_dynamic_stackalloc_sgpr_align16() {
; GFX10-NEXT: s_and_b32 s4, s4, -16
; GFX10-NEXT: s_lshl_b32 s4, s4, 5
; GFX10-NEXT: s_add_u32 s32, s6, s4
; GFX10-NEXT: s_add_i32 s32, s33, 0x200
; GFX10-NEXT: s_mov_b32 s33, s7
; GFX10-NEXT: s_addk_i32 s32, 0xfe00
; GFX10-NEXT: s_setpc_b64 s[30:31]
;
Expand All @@ -265,10 +270,9 @@ define void @func_dynamic_stackalloc_sgpr_align16() {
; GFX11-NEXT: s_getpc_b64 s[0:1]
; GFX11-NEXT: s_add_u32 s0, s0, gv@gotpcrel32@lo+4
; GFX11-NEXT: s_addc_u32 s1, s1, gv@gotpcrel32@hi+12
; GFX11-NEXT: v_mov_b32_e32 v0, 0
; GFX11-NEXT: s_load_b64 s[0:1], s[0:1], 0x0
; GFX11-NEXT: s_mov_b32 s2, s32
; GFX11-NEXT: s_mov_b32 s33, s3
; GFX11-NEXT: s_load_b64 s[0:1], s[0:1], 0x0
; GFX11-NEXT: v_mov_b32_e32 v0, 0
; GFX11-NEXT: scratch_store_b32 off, v0, s2
; GFX11-NEXT: s_waitcnt lgkmcnt(0)
; GFX11-NEXT: s_load_b32 s0, s[0:1], 0x0
Expand All @@ -277,8 +281,10 @@ define void @func_dynamic_stackalloc_sgpr_align16() {
; GFX11-NEXT: s_delay_alu instid0(SALU_CYCLE_1) | instskip(NEXT) | instid1(SALU_CYCLE_1)
; GFX11-NEXT: s_and_b32 s0, s0, -16
; GFX11-NEXT: s_lshl_b32 s0, s0, 5
; GFX11-NEXT: s_delay_alu instid0(SALU_CYCLE_1) | instskip(NEXT) | instid1(SALU_CYCLE_1)
; GFX11-NEXT: s_delay_alu instid0(SALU_CYCLE_1)
; GFX11-NEXT: s_add_u32 s32, s2, s0
; GFX11-NEXT: s_add_i32 s32, s33, 16
; GFX11-NEXT: s_mov_b32 s33, s3
; GFX11-NEXT: s_add_i32 s32, s32, -16
; GFX11-NEXT: s_setpc_b64 s[30:31]
%n = load i32, ptr addrspace(4) @gv, align 16
Expand Down Expand Up @@ -355,6 +361,8 @@ define void @func_dynamic_stackalloc_sgpr_align32(ptr addrspace(1) %out) {
; GFX9-NEXT: s_mov_b32 s6, s33
; GFX9-NEXT: s_add_i32 s33, s32, 0x7c0
; GFX9-NEXT: s_and_b32 s33, s33, 0xfffff800
; GFX9-NEXT: s_mov_b32 s7, s34
; GFX9-NEXT: s_mov_b32 s34, s32
; GFX9-NEXT: s_addk_i32 s32, 0x1000
; GFX9-NEXT: s_getpc_b64 s[4:5]
; GFX9-NEXT: s_add_u32 s4, s4, gv@gotpcrel32@lo+4
Expand All @@ -373,6 +381,8 @@ define void @func_dynamic_stackalloc_sgpr_align32(ptr addrspace(1) %out) {
; GFX9-NEXT: s_and_b32 s4, s4, -16
; GFX9-NEXT: s_lshl_b32 s4, s4, 6
; GFX9-NEXT: s_add_u32 s32, s5, s4
; GFX9-NEXT: s_add_i32 s32, s34, 0x1000
; GFX9-NEXT: s_mov_b32 s34, s7
; GFX9-NEXT: s_addk_i32 s32, 0xf000
; GFX9-NEXT: s_waitcnt vmcnt(0)
; GFX9-NEXT: s_setpc_b64 s[30:31]
Expand All @@ -382,8 +392,10 @@ define void @func_dynamic_stackalloc_sgpr_align32(ptr addrspace(1) %out) {
; GFX10-NEXT: s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
; GFX10-NEXT: s_mov_b32 s6, s33
; GFX10-NEXT: s_add_i32 s33, s32, 0x3e0
; GFX10-NEXT: s_addk_i32 s32, 0x800
; GFX10-NEXT: s_mov_b32 s7, s34
; GFX10-NEXT: s_and_b32 s33, s33, 0xfffffc00
; GFX10-NEXT: s_mov_b32 s34, s32
; GFX10-NEXT: s_addk_i32 s32, 0x800
; GFX10-NEXT: s_getpc_b64 s[4:5]
; GFX10-NEXT: s_add_u32 s4, s4, gv@gotpcrel32@lo+4
; GFX10-NEXT: s_addc_u32 s5, s5, gv@gotpcrel32@hi+12
Expand All @@ -401,6 +413,8 @@ define void @func_dynamic_stackalloc_sgpr_align32(ptr addrspace(1) %out) {
; GFX10-NEXT: s_and_b32 s4, s4, -16
; GFX10-NEXT: s_lshl_b32 s4, s4, 5
; GFX10-NEXT: s_add_u32 s32, s5, s4
; GFX10-NEXT: s_add_i32 s32, s34, 0x800
; GFX10-NEXT: s_mov_b32 s34, s7
; GFX10-NEXT: s_addk_i32 s32, 0xf800
; GFX10-NEXT: s_setpc_b64 s[30:31]
;
Expand All @@ -409,8 +423,10 @@ define void @func_dynamic_stackalloc_sgpr_align32(ptr addrspace(1) %out) {
; GFX11-NEXT: s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
; GFX11-NEXT: s_mov_b32 s2, s33
; GFX11-NEXT: s_add_i32 s33, s32, 31
; GFX11-NEXT: s_add_i32 s32, s32, 64
; GFX11-NEXT: s_mov_b32 s3, s34
; GFX11-NEXT: s_and_not1_b32 s33, s33, 31
; GFX11-NEXT: s_mov_b32 s34, s32
; GFX11-NEXT: s_add_i32 s32, s32, 64
; GFX11-NEXT: s_getpc_b64 s[0:1]
; GFX11-NEXT: s_add_u32 s0, s0, gv@gotpcrel32@lo+4
; GFX11-NEXT: s_addc_u32 s1, s1, gv@gotpcrel32@hi+12
Expand All @@ -429,7 +445,8 @@ define void @func_dynamic_stackalloc_sgpr_align32(ptr addrspace(1) %out) {
; GFX11-NEXT: s_delay_alu instid0(SALU_CYCLE_1) | instskip(NEXT) | instid1(SALU_CYCLE_1)
; GFX11-NEXT: s_lshl_b32 s0, s0, 5
; GFX11-NEXT: s_add_u32 s32, s1, s0
; GFX11-NEXT: s_delay_alu instid0(SALU_CYCLE_1)
; GFX11-NEXT: s_add_i32 s32, s34, 64
; GFX11-NEXT: s_mov_b32 s34, s3
; GFX11-NEXT: s_addk_i32 s32, 0xffc0
; GFX11-NEXT: s_setpc_b64 s[30:31]
%n = load i32, ptr addrspace(4) @gv
Expand Down
5 changes: 5 additions & 0 deletions llvm/test/CodeGen/AMDGPU/GlobalISel/non-entry-alloca.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,7 @@ define void @func_non_entry_block_static_alloca_align4(ptr addrspace(1) %out, i3
; GCN-NEXT: .LBB2_3: ; %bb.2
; GCN-NEXT: s_or_b64 exec, exec, s[4:5]
; GCN-NEXT: v_mov_b32_e32 v0, 0
; GCN-NEXT: s_add_i32 s32, s33, 0x400
; GCN-NEXT: global_store_dword v[0:1], v0, off
; GCN-NEXT: s_waitcnt vmcnt(0)
; GCN-NEXT: s_addk_i32 s32, 0xfc00
Expand Down Expand Up @@ -216,8 +217,10 @@ define void @func_non_entry_block_static_alloca_align64(ptr addrspace(1) %out, i
; GCN-NEXT: s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
; GCN-NEXT: s_mov_b32 s7, s33
; GCN-NEXT: s_add_i32 s33, s32, 0xfc0
; GCN-NEXT: s_mov_b32 s8, s34
; GCN-NEXT: v_cmp_eq_u32_e32 vcc, 0, v2
; GCN-NEXT: s_and_b32 s33, s33, 0xfffff000
; GCN-NEXT: s_mov_b32 s34, s32
; GCN-NEXT: s_addk_i32 s32, 0x2000
; GCN-NEXT: s_and_saveexec_b64 s[4:5], vcc
; GCN-NEXT: s_cbranch_execz .LBB3_2
Expand All @@ -240,8 +243,10 @@ define void @func_non_entry_block_static_alloca_align64(ptr addrspace(1) %out, i
; GCN-NEXT: .LBB3_2: ; %bb.1
; GCN-NEXT: s_or_b64 exec, exec, s[4:5]
; GCN-NEXT: v_mov_b32_e32 v0, 0
; GCN-NEXT: s_add_i32 s32, s34, 0x2000
; GCN-NEXT: global_store_dword v[0:1], v0, off
; GCN-NEXT: s_waitcnt vmcnt(0)
; GCN-NEXT: s_mov_b32 s34, s8
; GCN-NEXT: s_addk_i32 s32, 0xe000
; GCN-NEXT: s_mov_b32 s33, s7
; GCN-NEXT: s_setpc_b64 s[30:31]
Expand Down
Loading
Loading