AMDGPU: Correct prolog SP initialization logic

[arsenm]

Summary

Having callees that will read SP is not the only reason we need to
reference the stack pointer.

[dstuttard]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

453

Is there a way to finesse this so it's only added when it is likely to be required?

[arsenm]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

453

This isn't really right either. hasNonSpillStackObjects isn't accurate, since we can access with an soffset of 0 in the entry point. This also will fail to catch cases with variable allocas. I think the logic here should probably be hasFP()?

[arsenm]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

453

I think this really is hasFP based on the implementation of getFrameRegister

[dstuttard]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

453

I tried hasFP and it didn't resolve the original issue I was fixing. I'm not 100% sure what's going on here now. The backend definitely generates an access using stackPtrOffsetReg when nonSpillStackObjects are present. Should the test include both?

[scott.linder]

Sorry for the delay in reviewing, I've tried to wrap my head around this and I'm still lost.

I think someone needs to define what all of these things mean. Some fundamental questions I don't think we have really adequately answered are:

• What is the SP, when is it defined, when can it be changed?
• What is the FP, when is it defined, when can it be changed?
• Can we decide ahead-of-time (ISel?) whether we will need to have an SGPR available for the immediate-overflow case in SIRegisterInfo::buildSpillLoadStore?
  • Should we be using the SP/FP to handle that case in the first place (temporarily munging the SP/FP has implications for e.g. trap handlers, signals, CFI, etc.)?
  • If not, should we add an additional SGPR argument to the spill or something similar to ensure there is an SGPR available so we don't just ICE in the case where one isn't? Is this possible with the way RA works?
  • Can we in some cases prove that the spare SGPR is not needed because we will be able to fit the offset into the immediate field of the buffer instruction? Can we actually leverage this to free up that SGPR to be used by RA, or is this a circular problem/impossible because of potential need for the SGPR much later? Is this part of the purpose of the emergency spill slot? Is this an issue because we can't actually spill SGPRs directly to VMEM, and the emergency slot is in VMEM?

It seems like both of these patches, while maybe fixing some specific cases, aren't really framing the solution in terms of what is fundamentally broken. I think I broke this further with my patch to add the scratch wave offset to the scratch SRD, but I think all of these issues existed before. I don't know what the right answer is, but I can't really review anything without at least that amount of context.

[arsenm]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

453

Do you have the testcase for the original issue?

[arsenm]

Correct the logic.

hasFP || hasCalls should be slightly too conservative. I have no idea what other problem you could have been seeing from using hasFP without a testcase

[arsenm]

ping

[scott.linder]

I don't know if this has anything to do with David's issue, but I think separately this makes sense. I'm just confused what the semantic difference between the two "do we need an SP" predicates is?

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

1299

How does this differ from frameRequiresSP?

[arsenm]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

1299

This is split mostly because of the distinctive handling of functions and kernels in hasFP. hasFP and requiresStackPointerReferences are almost the same, except the kernel case is a bit more specific since there's never a real reason to setup FP since it's known 0 (e.g. we never realign the kernel stack). requiresStackPointerReference is essentially hasFP for kernels, and frameRequiresSP is common to kernels and functions.

[scott.linder]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

1265

So this is the function for "the conditions which imply the need for setting up an SP which are shared for entry functions and non-entry functions; just checking this for a frame can give false-negatives"?

I think at first I missed that this is just a static function used to implement the other bits and was confused. I don't know if it actually would help, but maybe you could hint at this in the name with something like frameTriviallyRequiresSP, and note in the comment that this is just the conditions shared between all frames?

1299

Can you put a (possibly condensed) version of this in a comment on the function? I mistakenly thought the frameRequiresSP was public, so the comment doesn't need to actually reference it, but just stating what the function does would help with the ambiguity between all of the needs{S,F}P type functions.

1301

Can this just be an assert that it is called for an entry function, and then could the name of this function just include EntryFunction like e.g. emitEntryFunctionPrologue does?

[arsenm]

Add comments, name changes, revert to assert

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

1301

Yes, this is what I did originally but I decided a more wholistic function may be better for future uses

[scott.linder]

LGTM, thank you!

[arsenm]

ec8c172d01eb14eba890f36205da0613dda7f742

llvm/lib/Target/AMDGPU/SIFrameLowering.h

private:
   void emitEntryFunctionScratchRsrcRegSetup(
       MachineFunction &MF, MachineBasicBlock &MBB,
       MachineBasicBlock::iterator I, const DebugLoc &DL,
       Register PreloadedPrivateBufferReg, Register ScratchRsrcReg,
       Register ScratchWaveOffsetReg) const;
 
 public:
   bool hasFP(const MachineFunction &MF) const override;
+
+  bool requiresStackPointerReference(const MachineFunction &MF) const;
 };
 
 } // end namespace llvm
 
 #endif // LLVM_LIB_TARGET_AMDGPU_SIFRAMELOWERING_H

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

for (MCPhysReg Reg : AllSGPRs) {
         break;
       }
     }
   } else {
     ScratchWaveOffsetReg = PreloadedScratchWaveOffsetReg;
   }
   assert(ScratchWaveOffsetReg);
 
-  if (MF.getFrameInfo().hasCalls()) {
+  if (requiresStackPointerReference(MF)) {
     Register SPReg = MFI->getStackPtrOffsetReg();
     assert(SPReg != AMDGPU::SP_REG);
     BuildMI(MBB, I, DL, TII->get(AMDGPU::S_MOV_B32), SPReg)
         .addImm(MF.getFrameInfo().getStackSize() * ST.getWavefrontSize());
   }
 
   if (hasFP(MF)) {
     Register FPReg = MFI->getFrameOffsetReg();
     assert(FPReg != AMDGPU::FP_REG);
     BuildMI(MBB, I, DL, TII->get(AMDGPU::S_MOV_B32), FPReg).addImm(0);
   }
BuildMI(MBB, I, DL, TII->get(Op), SPReg)
       .addImm(Amount * ST.getWavefrontSize());
   } else if (CalleePopAmount != 0) {
     llvm_unreachable("is this used?");
   }
 
   return MBB.erase(I);
 }
 
+/// Returns true if the frame will require a reference to the stack pointer.
+///
+/// This is the set of conditions common to setting up the stack pointer in a
+/// kernel, and for using a frame pointer in a callable function.
+///
+/// FIXME: Should also check hasOpaqueSPAdjustment and if any inline asm
+/// references SP.
+static bool frameTriviallyRequiresSP(const MachineFrameInfo &MFI) {
+  return MFI.hasVarSizedObjects() || MFI.hasStackMap() || MFI.hasPatchPoint();
+}
+
+// The FP for kernels is always known 0, so we never really need to setup an
+// explicit register for it. However, DisableFramePointerElim will force us to
+// use a register for it.
 bool SIFrameLowering::hasFP(const MachineFunction &MF) const {
   const MachineFrameInfo &MFI = MF.getFrameInfo();
 
   // For entry functions we can use an immediate offset in most cases, so the
   // presence of calls doesn't imply we need a distinct frame pointer.
   if (MFI.hasCalls() &&
       !MF.getInfo<SIMachineFunctionInfo>()->isEntryFunction()) {
     // All offsets are unsigned, so need to be addressed in the same direction
     // as stack growth.
 
     // FIXME: This function is pretty broken, since it can be called before the
     // frame layout is determined or CSR spills are inserted.
     return MFI.getStackSize() != 0;
   }
 
-  return MFI.hasVarSizedObjects() || MFI.isFrameAddressTaken() ||
-    MFI.hasStackMap() || MFI.hasPatchPoint() ||
+  return frameTriviallyRequiresSP(MFI) || MFI.isFrameAddressTaken() ||
     MF.getSubtarget<GCNSubtarget>().getRegisterInfo()->needsStackRealignment(MF) ||
     MF.getTarget().Options.DisableFramePointerElim(MF);
 }
+
+// This is essentially a reduced version of hasFP for entry functions. Since the
+// stack pointer is known 0 on entry to kernels, we never really need an FP
+// register. We may need to initialize the stack pointer depending on the frame
+// properties, which logically overlaps many of the cases where an ordinary
+// function would require an FP.
+bool SIFrameLowering::requiresStackPointerReference(
+    const MachineFunction &MF) const {
+  // Callable functions always require a stack pointer reference.
+  assert(MF.getInfo<SIMachineFunctionInfo>()->isEntryFunction() &&
+         "only expected to call this for entry points");
+
+  const MachineFrameInfo &MFI = MF.getFrameInfo();
+
+  // Entry points ordinarily don't need to initialize SP. We have to set it up
+  // for callees if there are any. Also note tail calls are impossible/don't
+  // make any sense for kernels.
+  if (MFI.hasCalls())
+    return true;
+
+  // We still need to initialize the SP if we're doing anything weird that
+  // references the SP, like variable sized stack objects.
+  return frameTriviallyRequiresSP(MFI);
+}

llvm/test/CodeGen/AMDGPU/GlobalISel/dynamic-alloca-uniform.ll

 ; GFX9-NEXT:    s_load_dword s4, s[4:5], 0x0
 ; GFX9-NEXT:    s_add_u32 flat_scratch_lo, s6, s9
 ; GFX9-NEXT:    s_addc_u32 flat_scratch_hi, s7, 0
 ; GFX9-NEXT:    s_add_u32 s0, s0, s9
 ; GFX9-NEXT:    s_addc_u32 s1, s1, 0
 ; 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_movk_i32 s32, 0x400
 ; GFX9-NEXT:    s_lshl_b32 s4, s4, 6
 ; GFX9-NEXT:    s_add_u32 s4, s32, s4
 ; GFX9-NEXT:    v_mov_b32_e32 v0, 0
 ; GFX9-NEXT:    v_mov_b32_e32 v1, s4
 ; GFX9-NEXT:    s_mov_b32 s33, 0
 ; GFX9-NEXT:    buffer_store_dword v0, v1, s[0:3], 0 offen
 ; GFX9-NEXT:    s_endpgm
 ;
 ; GFX10-LABEL: kernel_dynamic_stackalloc_sgpr_align4:
 ; GFX10:       ; %bb.0:
 ; GFX10-NEXT:    s_add_u32 s6, s6, s9
+; GFX10-NEXT:    s_movk_i32 s32, 0x200
 ; GFX10-NEXT:    s_mov_b32 s33, 0
 ; GFX10-NEXT:    s_addc_u32 s7, s7, 0
 ; GFX10-NEXT:    s_setreg_b32 hwreg(HW_REG_FLAT_SCR_LO), s6
 ; GFX10-NEXT:    s_setreg_b32 hwreg(HW_REG_FLAT_SCR_HI), s7
 ; GFX10-NEXT:    s_load_dword s4, s[4:5], 0x0
 ; GFX10-NEXT:    s_add_u32 s0, s0, s9
 ; GFX10-NEXT:    s_addc_u32 s1, s1, 0
 ; GFX10-NEXT:    v_mov_b32_e32 v0, 0

 ; GFX9-NEXT:    s_load_dword s4, s[4:5], 0x0
 ; GFX9-NEXT:    s_add_u32 flat_scratch_lo, s6, s9
 ; GFX9-NEXT:    s_addc_u32 flat_scratch_hi, s7, 0
 ; GFX9-NEXT:    s_add_u32 s0, s0, s9
 ; GFX9-NEXT:    s_addc_u32 s1, s1, 0
 ; 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_movk_i32 s32, 0x400
 ; GFX9-NEXT:    s_lshl_b32 s4, s4, 6
 ; GFX9-NEXT:    s_add_u32 s4, s32, s4
 ; GFX9-NEXT:    v_mov_b32_e32 v0, 0
 ; GFX9-NEXT:    v_mov_b32_e32 v1, s4
 ; GFX9-NEXT:    s_mov_b32 s33, 0
 ; GFX9-NEXT:    buffer_store_dword v0, v1, s[0:3], 0 offen
 ; GFX9-NEXT:    s_endpgm
 ;
 ; GFX10-LABEL: kernel_dynamic_stackalloc_sgpr_align16:
 ; GFX10:       ; %bb.0:
 ; GFX10-NEXT:    s_add_u32 s6, s6, s9
+; GFX10-NEXT:    s_movk_i32 s32, 0x200
 ; GFX10-NEXT:    s_mov_b32 s33, 0
 ; GFX10-NEXT:    s_addc_u32 s7, s7, 0
 ; GFX10-NEXT:    s_setreg_b32 hwreg(HW_REG_FLAT_SCR_LO), s6
 ; GFX10-NEXT:    s_setreg_b32 hwreg(HW_REG_FLAT_SCR_HI), s7
 ; GFX10-NEXT:    s_load_dword s4, s[4:5], 0x0
 ; GFX10-NEXT:    s_add_u32 s0, s0, s9
 ; GFX10-NEXT:    s_addc_u32 s1, s1, 0
 ; GFX10-NEXT:    v_mov_b32_e32 v0, 0

 ; GFX9-NEXT:    s_load_dword s4, s[4:5], 0x0
 ; GFX9-NEXT:    s_add_u32 flat_scratch_lo, s6, s9
 ; GFX9-NEXT:    s_addc_u32 flat_scratch_hi, s7, 0
 ; GFX9-NEXT:    s_add_u32 s0, s0, s9
 ; GFX9-NEXT:    s_addc_u32 s1, s1, 0
 ; 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_movk_i32 s32, 0x800
 ; GFX9-NEXT:    s_lshl_b32 s4, s4, 6
 ; GFX9-NEXT:    s_add_u32 s4, s32, s4
 ; GFX9-NEXT:    s_and_b32 s4, s4, 0xfffff800
 ; GFX9-NEXT:    v_mov_b32_e32 v0, 0
 ; GFX9-NEXT:    v_mov_b32_e32 v1, s4
 ; GFX9-NEXT:    s_mov_b32 s33, 0
 ; GFX9-NEXT:    buffer_store_dword v0, v1, s[0:3], 0 offen
 ; GFX9-NEXT:    s_endpgm
 ;
 ; GFX10-LABEL: kernel_dynamic_stackalloc_sgpr_align32:
 ; GFX10:       ; %bb.0:
 ; GFX10-NEXT:    s_add_u32 s6, s6, s9
+; GFX10-NEXT:    s_movk_i32 s32, 0x400
 ; GFX10-NEXT:    s_mov_b32 s33, 0
 ; GFX10-NEXT:    s_addc_u32 s7, s7, 0
 ; GFX10-NEXT:    s_setreg_b32 hwreg(HW_REG_FLAT_SCR_LO), s6
 ; GFX10-NEXT:    s_setreg_b32 hwreg(HW_REG_FLAT_SCR_HI), s7
 ; GFX10-NEXT:    s_load_dword s4, s[4:5], 0x0
 ; GFX10-NEXT:    s_add_u32 s0, s0, s9
 ; GFX10-NEXT:    s_addc_u32 s1, s1, 0
 ; GFX10-NEXT:    v_mov_b32_e32 v0, 0

llvm/test/CodeGen/AMDGPU/GlobalISel/non-entry-alloca.ll

 define amdgpu_kernel void @kernel_non_entry_block_static_alloca_uniformly_reached_align4(i32 addrspace(1)* %out, i32 %arg.cond0, i32 %arg.cond1, i32 %in) {
 ; GCN-LABEL: kernel_non_entry_block_static_alloca_uniformly_reached_align4:
 ; GCN:       ; %bb.0: ; %entry
 ; GCN-NEXT:    s_add_u32 flat_scratch_lo, s6, s9
 ; GCN-NEXT:    s_load_dword s6, s[4:5], 0x8
 ; GCN-NEXT:    s_addc_u32 flat_scratch_hi, s7, 0
 ; GCN-NEXT:    s_add_u32 s0, s0, s9
 ; GCN-NEXT:    s_addc_u32 s1, s1, 0
-; GCN-NEXT:    s_mov_b32 s33, 0
+; GCN-NEXT:    s_movk_i32 s32, 0x400
 ; GCN-NEXT:    s_waitcnt lgkmcnt(0)
 ; GCN-NEXT:    s_cmp_lg_u32 s6, 0
 ; GCN-NEXT:    s_cselect_b32 s6, 1, 0
 ; GCN-NEXT:    s_and_b32 s6, s6, 1
 ; GCN-NEXT:    s_cmp_lg_u32 s6, 0
+; GCN-NEXT:    s_mov_b32 s33, 0
 ; GCN-NEXT:    s_cbranch_scc1 BB0_3
 ; GCN-NEXT:  ; %bb.1: ; %bb.0
 ; GCN-NEXT:    s_load_dword s6, s[4:5], 0xc
 ; GCN-NEXT:    s_waitcnt lgkmcnt(0)
 ; GCN-NEXT:    s_cmp_lg_u32 s6, 0
 ; GCN-NEXT:    s_cselect_b32 s6, 1, 0
 ; GCN-NEXT:    s_and_b32 s6, s6, 1
 ; GCN-NEXT:    s_cmp_lg_u32 s6, 0

 define amdgpu_kernel void @kernel_non_entry_block_static_alloca_uniformly_reached_align64(i32 addrspace(1)* %out, i32 %arg.cond, i32 %in) {
 ; GCN-LABEL: kernel_non_entry_block_static_alloca_uniformly_reached_align64:
 ; GCN:       ; %bb.0: ; %entry
 ; GCN-NEXT:    s_add_u32 flat_scratch_lo, s6, s9
 ; GCN-NEXT:    s_load_dword s6, s[4:5], 0x8
 ; GCN-NEXT:    s_addc_u32 flat_scratch_hi, s7, 0
 ; GCN-NEXT:    s_add_u32 s0, s0, s9
 ; GCN-NEXT:    s_addc_u32 s1, s1, 0
-; GCN-NEXT:    s_mov_b32 s33, 0
+; GCN-NEXT:    s_movk_i32 s32, 0x1000
 ; GCN-NEXT:    s_waitcnt lgkmcnt(0)
 ; GCN-NEXT:    s_cmp_lg_u32 s6, 0
 ; GCN-NEXT:    s_cselect_b32 s6, 1, 0
 ; GCN-NEXT:    s_and_b32 s6, s6, 1
 ; GCN-NEXT:    s_cmp_lg_u32 s6, 0
+; GCN-NEXT:    s_mov_b32 s33, 0
 ; GCN-NEXT:    s_cbranch_scc1 BB1_2
 ; GCN-NEXT:  ; %bb.1: ; %bb.0
 ; GCN-NEXT:    s_load_dwordx2 s[6:7], s[4:5], 0x0
 ; GCN-NEXT:    s_load_dword s4, s[4:5], 0xc
 ; GCN-NEXT:    s_add_u32 s5, s32, 0x1000
 ; GCN-NEXT:    s_and_b32 s5, s5, 0xfffff000
 ; GCN-NEXT:    s_add_u32 s8, s5, 4
 ; GCN-NEXT:    v_mov_b32_e32 v1, 0

llvm/test/CodeGen/AMDGPU/non-entry-alloca.ll

 define amdgpu_kernel void @kernel_non_entry_block_static_alloca_uniformly_reached_align4(i32 addrspace(1)* %out, i32 %arg.cond0, i32 %arg.cond1, i32 %in) {
 ; GCN-LABEL: kernel_non_entry_block_static_alloca_uniformly_reached_align4:
 ; GCN:       ; %bb.0: ; %entry
 ; GCN-NEXT:    s_add_u32 flat_scratch_lo, s6, s9
 ; GCN-NEXT:    s_addc_u32 flat_scratch_hi, s7, 0
 ; GCN-NEXT:    s_add_u32 s0, s0, s9
 ; GCN-NEXT:    s_load_dwordx4 s[8:11], s[4:5], 0x8
 ; GCN-NEXT:    s_addc_u32 s1, s1, 0
+; GCN-NEXT:    s_movk_i32 s32, 0x400
 ; GCN-NEXT:    s_mov_b32 s33, 0
 ; GCN-NEXT:    s_waitcnt lgkmcnt(0)
 ; GCN-NEXT:    s_cmp_lg_u32 s8, 0
 ; GCN-NEXT:    s_cbranch_scc1 BB0_3
 ; GCN-NEXT:  ; %bb.1: ; %bb.0
 ; GCN-NEXT:    s_cmp_lg_u32 s9, 0
 ; GCN-NEXT:    s_cbranch_scc1 BB0_3
 ; GCN-NEXT:  ; %bb.2: ; %bb.1

 define amdgpu_kernel void @kernel_non_entry_block_static_alloca_uniformly_reached_align64(i32 addrspace(1)* %out, i32 %arg.cond, i32 %in) {
 ; GCN-LABEL: kernel_non_entry_block_static_alloca_uniformly_reached_align64:
 ; GCN:       ; %bb.0: ; %entry
 ; GCN-NEXT:    s_add_u32 flat_scratch_lo, s6, s9
 ; GCN-NEXT:    s_addc_u32 flat_scratch_hi, s7, 0
 ; GCN-NEXT:    s_load_dwordx2 s[6:7], s[4:5], 0x8
 ; GCN-NEXT:    s_add_u32 s0, s0, s9
 ; GCN-NEXT:    s_addc_u32 s1, s1, 0
+; GCN-NEXT:    s_movk_i32 s32, 0x1000
 ; GCN-NEXT:    s_mov_b32 s33, 0
 ; GCN-NEXT:    s_waitcnt lgkmcnt(0)
 ; GCN-NEXT:    s_cmp_lg_u32 s6, 0
 ; GCN-NEXT:    s_cbranch_scc1 BB1_2
 ; GCN-NEXT:  ; %bb.1: ; %bb.0
 ; GCN-NEXT:    s_add_i32 s6, s32, 0x1000
 ; GCN-NEXT:    s_and_b32 s6, s6, 0xfffff000
 ; GCN-NEXT:    s_lshl_b32 s7, s7, 2