[AMDGPU] Always create Stack Object for reserved VGPR

[ruiling]

Summary

As we may overwrite inactive lanes of a caller-save-vgpr, we should
always save/restore the reserved vgpr for sgpr spill.

llvm/lib/Target/AMDGPU/SILowerSGPRSpills.cpp

static bool lowerShiftReservedVGPR(MachineFunction &MF,
   // If there are no free lower VGPRs available, default to using the
   // pre-reserved register instead.
   const SIRegisterInfo *TRI = ST.getRegisterInfo();
   Register LowestAvailableVGPR =
       TRI->findUnusedRegister(MF.getRegInfo(), &AMDGPU::VGPR_32RegClass, MF);
   if (!LowestAvailableVGPR)
     LowestAvailableVGPR = PreReservedVGPR;
 
-  const MCPhysReg *CSRegs = MF.getRegInfo().getCalleeSavedRegs();
   MachineFrameInfo &FrameInfo = MF.getFrameInfo();
-  Optional<int> FI;
-  // Check if we are reserving a CSR. Create a stack object for a possible spill
-  // in the function prologue.
-  if (FuncInfo->isCalleeSavedReg(CSRegs, LowestAvailableVGPR))
-    FI = FrameInfo.CreateSpillStackObject(4, Align(4));
+  // Create a stack object for a possible spill in the function prologue.
+  // Note Non-CSR VGPR also need this as we may overwrite inactive lanes.
+  Optional<int> FI = FrameInfo.CreateSpillStackObject(4, Align(4));
 
   // Find saved info about the pre-reserved register.
   const auto *ReservedVGPRInfoItr =
       llvm::find_if(FuncInfo->getSGPRSpillVGPRs(),
                     [PreReservedVGPR](const auto &SpillRegInfo) {
                       return SpillRegInfo.VGPR == PreReservedVGPR;
                     });
 

llvm/test/CodeGen/AMDGPU/callee-frame-setup.ll

define void @callee_with_stack_no_fp_elim_csr_vgpr() #1 {
   store volatile i32 0, i32 addrspace(5)* %alloca
   call void asm sideeffect "; clobber v41", "~{v41}"()
   ret void
 }
 
 ; Use a copy to a free SGPR instead of introducing a second CSR VGPR.
 ; GCN-LABEL: {{^}}last_lane_vgpr_for_fp_csr:
 ; GCN: s_waitcnt
+; GCN-NEXT: s_or_saveexec_b64 [[COPY_EXEC0:s\[[0-9]+:[0-9]+\]]], -1{{$}}
+; MUBUF-NEXT:   buffer_store_dword [[CSR_VGPR:v[0-9]+]], off, s[0:3], s32 offset:12 ; 4-byte Folded Spill
+; FLATSCR-NEXT: scratch_store_dword off, [[CSR_VGPR:v[0-9]+]], s32 offset:12 ; 4-byte Folded Spill
+; GCN-NEXT: s_mov_b64 exec, [[COPY_EXEC0]]
 ; GCN-NEXT: v_writelane_b32 v1, s33, 63
 ; GCN-COUNT-60: v_writelane_b32 v1
 ; GCN: s_mov_b32 s33, s32
 ; GCN-COUNT-2: v_writelane_b32 v1
 ; MUBUF:   buffer_store_dword v41, off, s[0:3], s33 ; 4-byte Folded Spill
 ; FLATSCR: scratch_store_dword off, v41, s33 ; 4-byte Folded Spill
 ; MUBUF:   buffer_store_dword v{{[0-9]+}}, off, s[0:3], s33 offset:8
 ; FLATSCR: scratch_store_dword off, v{{[0-9]+}}, s33 offset:8
 ; GCN: ;;#ASMSTART
 ; GCN: v_writelane_b32 v1
 
-; MUBUF:        s_add_u32 s32, s32, 0x300
-; MUBUF:        s_sub_u32 s32, s32, 0x300
-; FLATSCR:      s_add_u32 s32, s32, 12
-; FLATSCR:      s_sub_u32 s32, s32, 12
+; MUBUF:        s_add_u32 s32, s32, 0x400
+; MUBUF:        s_sub_u32 s32, s32, 0x400
+; FLATSCR:      s_add_u32 s32, s32, 16
+; FLATSCR:      s_sub_u32 s32, s32, 16
 ; GCN-NEXT: v_readlane_b32 s33, v1, 63
+; GCN-NEXT: s_or_saveexec_b64 [[COPY_EXEC1:s\[[0-9]+:[0-9]+\]]], -1{{$}}
+; MUBUF-NEXT:   buffer_load_dword [[CSR_VGPR]], off, s[0:3], s32 offset:12 ; 4-byte Folded Reload
+; FLATSCR-NEXT: scratch_load_dword [[CSR_VGPR]], off, s32 offset:12 ; 4-byte Folded Reload
+; GCN-NEXT: s_mov_b64 exec, [[COPY_EXEC1]]
 ; GCN-NEXT: s_waitcnt vmcnt(0)
 ; GCN-NEXT: s_setpc_b64
 define void @last_lane_vgpr_for_fp_csr() #1 {
   %alloca = alloca i32, addrspace(5)
   store volatile i32 0, i32 addrspace(5)* %alloca
   call void asm sideeffect "; clobber v41", "~{v41}"()
   call void asm sideeffect "",
     "~{s40},~{s41},~{s42},~{s43},~{s44},~{s45},~{s46},~{s47},~{s48},~{s49}
     ,~{s50},~{s51},~{s52},~{s53},~{s54},~{s55},~{s56},~{s57},~{s58},~{s59}
     ,~{s60},~{s61},~{s62},~{s63},~{s64},~{s65},~{s66},~{s67},~{s68},~{s69}
     ,~{s70},~{s71},~{s72},~{s73},~{s74},~{s75},~{s76},~{s77},~{s78},~{s79}
     ,~{s80},~{s81},~{s82},~{s83},~{s84},~{s85},~{s86},~{s87},~{s88},~{s89}
     ,~{s90},~{s91},~{s92},~{s93},~{s94},~{s95},~{s96},~{s97},~{s98},~{s99}
     ,~{s100},~{s101},~{s102}"() #1
 
   ret void
 }
 
 ; Use a copy to a free SGPR instead of introducing a second CSR VGPR.
 ; GCN-LABEL: {{^}}no_new_vgpr_for_fp_csr:
 ; GCN: s_waitcnt
+; GCN-NEXT: s_or_saveexec_b64 [[COPY_EXEC0:s\[[0-9]+:[0-9]+\]]], -1{{$}}
+; MUBUF-NEXT:   buffer_store_dword [[CSR_VGPR:v[0-9]+]], off, s[0:3], s32 offset:12 ; 4-byte Folded Spill
+; FLATSCR-NEXT: scratch_store_dword off, [[CSR_VGPR:v[0-9]+]], s32 offset:12 ; 4-byte Folded Spill
+; GCN-NEXT: s_mov_b64 exec, [[COPY_EXEC0]]
 ; GCN-COUNT-62: v_writelane_b32 v1,
 ; GCN: s_mov_b32 [[FP_COPY:s[0-9]+]], s33
 ; GCN-NEXT: s_mov_b32 s33, s32
 ; GCN: v_writelane_b32 v1,
 ; MUBUF:   buffer_store_dword v41, off, s[0:3], s33 ; 4-byte Folded Spill
 ; FLATSCR: scratch_store_dword off, v41, s33 ; 4-byte Folded Spill
 ; MUBUF:   buffer_store_dword
 ; FLATSCR: scratch_store_dword
 ; GCN: ;;#ASMSTART
 ; GCN: v_writelane_b32 v1,
 ; MUBUF:   buffer_load_dword v41, off, s[0:3], s33 ; 4-byte Folded Reload
 ; FLATSCR: scratch_load_dword v41, off, s33 ; 4-byte Folded Reload
-; MUBUF:        s_add_u32 s32, s32, 0x300
-; FLATSCR:      s_add_u32 s32, s32, 12
+; MUBUF:        s_add_u32 s32, s32, 0x400
+; FLATSCR:      s_add_u32 s32, s32, 16
 ; GCN-COUNT-64: v_readlane_b32 s{{[0-9]+}}, v1
-; MUBUF-NEXT:   s_sub_u32 s32, s32, 0x300
-; FLATSCR-NEXT: s_sub_u32 s32, s32, 12
+; MUBUF-NEXT:   s_sub_u32 s32, s32, 0x400
+; FLATSCR-NEXT: s_sub_u32 s32, s32, 16
 ; GCN-NEXT: s_mov_b32 s33, [[FP_COPY]]
+; GCN-NEXT: s_or_saveexec_b64 [[COPY_EXEC1:s\[[0-9]+:[0-9]+\]]], -1{{$}}
+; MUBUF-NEXT:   buffer_load_dword [[CSR_VGPR]], off, s[0:3], s32 offset:12 ; 4-byte Folded Reload
+; FLATSCR-NEXT: scratch_load_dword [[CSR_VGPR]], off, s32 offset:12 ; 4-byte Folded Reload
+; GCN-NEXT: s_mov_b64 exec, [[COPY_EXEC1]]
 ; GCN-NEXT: s_waitcnt vmcnt(0)
 ; GCN-NEXT: s_setpc_b64
 define void @no_new_vgpr_for_fp_csr() #1 {
   %alloca = alloca i32, addrspace(5)
   store volatile i32 0, i32 addrspace(5)* %alloca
   call void asm sideeffect "; clobber v41", "~{v41}"()
   call void asm sideeffect "",
     "~{s39},~{s40},~{s41},~{s42},~{s43},~{s44},~{s45},~{s46},~{s47},~{s48},~{s49}

 define void @realign_stack_no_fp_elim() #1 {
   %alloca = alloca i32, align 8192, addrspace(5)
   store volatile i32 0, i32 addrspace(5)* %alloca
   ret void
 }
 
 ; GCN-LABEL: {{^}}no_unused_non_csr_sgpr_for_fp:
 ; GCN: s_waitcnt
+; GCN-NEXT: s_or_saveexec_b64 [[COPY_EXEC0:s\[[0-9]+:[0-9]+\]]], -1{{$}}
+; MUBUF-NEXT:   buffer_store_dword [[CSR_VGPR:v[0-9]+]], off, s[0:3], s32 offset:8 ; 4-byte Folded Spill
+; FLATSCR-NEXT: scratch_store_dword off, [[CSR_VGPR:v[0-9]+]], s32 offset:8 ; 4-byte Folded Spill
+; GCN-NEXT: s_mov_b64 exec, [[COPY_EXEC0]]
 ; GCN-NEXT: v_writelane_b32 v1, s33, 2
 ; GCN-NEXT: v_writelane_b32 v1, s30, 0
 ; GCN-NEXT: s_mov_b32 s33, s32
 ; GCN: v_mov_b32_e32 [[ZERO:v[0-9]+]], 0
 ; GCN: v_writelane_b32 v1, s31, 1
 ; MUBUF:   buffer_store_dword [[ZERO]], off, s[0:3], s33 offset:4
 ; FLATSCR: scratch_store_dword off, [[ZERO]], s33 offset:4
 ; GCN-NEXT:     s_waitcnt vmcnt(0)
 ; GCN: ;;#ASMSTART
-; MUBUF:        v_readlane_b32 s4, v1, 0
-; MUBUF-NEXT:   s_add_u32 s32, s32, 0x200
+; MUBUF:        s_add_u32 s32, s32, 0x300
+; MUBUF-NEXT:   v_readlane_b32 s4, v1, 0
 ; MUBUF-NEXT:   v_readlane_b32 s5, v1, 1
-; FLATSCR:      v_readlane_b32 s0, v1, 0
-; FLATSCR-NEXT: s_add_u32 s32, s32, 8
+; FLATSCR:      s_add_u32 s32, s32, 12
+; FLATSCR-NEXT: v_readlane_b32 s0, v1, 0
 ; FLATSCR-NEXT: v_readlane_b32 s1, v1, 1
-; MUBUF-NEXT:   s_sub_u32 s32, s32, 0x200
-; FLATSCR-NEXT: s_sub_u32 s32, s32, 8
+; MUBUF-NEXT:   s_sub_u32 s32, s32, 0x300
+; FLATSCR-NEXT: s_sub_u32 s32, s32, 12
 ; GCN-NEXT:     v_readlane_b32 s33, v1, 2
+; GCN-NEXT: s_or_saveexec_b64 [[COPY_EXEC1:s\[[0-9]+:[0-9]+\]]], -1{{$}}
+; MUBUF-NEXT:   buffer_load_dword [[CSR_VGPR]], off, s[0:3], s32 offset:8 ; 4-byte Folded Reload
+; FLATSCR-NEXT: scratch_load_dword [[CSR_VGPR]], off, s32 offset:8 ; 4-byte Folded Reload
+; GCN-NEXT: s_mov_b64 exec, [[COPY_EXEC1]]
+; GCN-NEXT: s_waitcnt vmcnt(0)
 ; MUBUF-NEXT:   s_setpc_b64 s[4:5]
 ; FLATSCR-NEXT: s_setpc_b64 s[0:1]
 define void @no_unused_non_csr_sgpr_for_fp() #1 {
   %alloca = alloca i32, addrspace(5)
   store volatile i32 0, i32 addrspace(5)* %alloca
 
   ; Use all clobberable registers, so FP has to spill to a VGPR.
   call void asm sideeffect "",
call void asm sideeffect "; clobber all VGPRs except CSR v40",
     ,~{v30},~{v31},~{v32},~{v33},~{v34},~{v35},~{v36},~{v37},~{v38},~{v39}"()
   ret void
 }
 
 ; If the size of the offset exceeds the MUBUF offset field we need another
 ; scratch VGPR to hold the offset.
 ; GCN-LABEL: {{^}}spill_fp_to_memory_scratch_reg_needed_mubuf_offset
 ; MUBUF: s_or_saveexec_b64 s[4:5], -1
+; MUBUF: v_mov_b32_e32 v0, 0x1008
+; MUBUF-NEXT:  buffer_store_dword v39, v0, s[0:3], s32 offen ; 4-byte Folded Spill
 ; MUBUF: v_mov_b32_e32 v0, s33
-; GCN-NOT: v_mov_b32_e32 v0, 0x1008
-; MUBUF-NEXT: v_mov_b32_e32 v1, 0x1008
+; GCN-NOT: v_mov_b32_e32 v0, 0x100c
+; MUBUF-NEXT: v_mov_b32_e32 v1, 0x100c
 ; MUBUF-NEXT: buffer_store_dword v0, v1, s[0:3], s32 offen ; 4-byte Folded Spill
 ; FLATSCR: s_add_u32 [[SOFF:s[0-9]+]], s33, 0x1004
 ; FLATSCR: v_mov_b32_e32 v0, 0
 ; FLATSCR: scratch_store_dword off, v0, [[SOFF]]
 define void @spill_fp_to_memory_scratch_reg_needed_mubuf_offset([4096 x i8] addrspace(5)* byval([4096 x i8]) align 4 %arg) #3 {
   %alloca = alloca i32, addrspace(5)
   store volatile i32 0, i32 addrspace(5)* %alloca
 

llvm/test/CodeGen/AMDGPU/stack-realign.ll

 
 define void @spill_bp_to_memory_scratch_reg_needed_mubuf_offset(<32 x i32> %a, i32 %b, [4096 x i8] addrspace(5)* byval([4096 x i8]) align 4 %arg) #5 {
 ; If the size of the offset exceeds the MUBUF offset field we need another
 ; scratch VGPR to hold the offset.
 
 ; GCN-LABEL: spill_bp_to_memory_scratch_reg_needed_mubuf_offset
 ; GCN: s_or_saveexec_b64 s[4:5], -1
 ; GCN: v_mov_b32_e32 v0, s33
-; GCN-NOT: v_mov_b32_e32 v0, 0x1084
-; GCN-NEXT: v_mov_b32_e32 v1, 0x1084
-; GCN-NEXT: buffer_store_dword v0, v1, s[0:3], s32 offen
-; GCN: v_mov_b32_e32 v0, s34
 ; GCN-NOT: v_mov_b32_e32 v0, 0x1088
 ; GCN-NEXT: v_mov_b32_e32 v1, 0x1088
 ; GCN-NEXT: buffer_store_dword v0, v1, s[0:3], s32 offen
+; GCN: v_mov_b32_e32 v0, s34
+; GCN-NOT: v_mov_b32_e32 v0, 0x108c
+; GCN-NEXT: v_mov_b32_e32 v1, 0x108c
+; GCN-NEXT: buffer_store_dword v0, v1, s[0:3], s32 offen
   %local_val = alloca i32, align 128, addrspace(5)
   store volatile i32 %b, i32 addrspace(5)* %local_val, align 128
 
   call void asm sideeffect "; clobber nonpreserved SGPRs and 64 CSRs",
     "~{s4},~{s5},~{s6},~{s7},~{s8},~{s9}
     ,~{s10},~{s11},~{s12},~{s13},~{s14},~{s15},~{s16},~{s17},~{s18},~{s19}
     ,~{s20},~{s21},~{s22},~{s23},~{s24},~{s25},~{s26},~{s27},~{s28},~{s29}
     ,~{s40},~{s41},~{s42},~{s43},~{s44},~{s45},~{s46},~{s47},~{s48},~{s49}