[AMDGPU] Spill more than wavesize CSR SGPRs

[saiislam]

Summary

In case of more than wavesize CSR SGPR spills, lanes of reserved VGPR were getting
overwritten due to wrap around.

Reserve a VGPR (when NumVGPRSpillLanes = 0, WaveSize, 2*WaveSize, ..) and when one
of the two conditions is true:

1. One reserved VGPR being tracked by VGPRReservedForSGPRSpill is not yet reserved.
2. All spill lanes of reserved VGPR(s) are full and another spill lane is required.

[arsenm]

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

312

I think it might be clearer to split out the logic for the special reserved case.

Could you do something like

if (FuncInfo->VGPRReservedForSGPRSpill && NumVGPRSpillLanes <= WaveSize) {
  assert(VGPRResevedForSGPRSpill == SpillVGPRs.back().VGPR);
  LaneVGPR = VGPRResevedForSGPRSpill;
} else (VGPRIndex == 0) {
 ....
} else {
  LaneVGPR = SpillVGPRs.back().VGPR
}

llvm/test/CodeGen/AMDGPU/spill_more_than_wavesize_csr_sgprs.ll

11–12

I'm pretty sure this isn't broken if you do not have a stack object, but these should be unrelated so I think something is still wrong here

[saiislam]

1. Split the VGPR allocation logic as suggested

[arsenm]

llvm/test/CodeGen/AMDGPU/spill_more_than_wavesize_csr_sgprs.ll

1

Maybe also try with and without the reserved VGPR flag?

10

You removed the stack object, but my point was without your patch, this already works:

v_writelane_b32 v0, s98, 63
v_writelane_b32 v1, s99, 0
v_writelane_b32 v1, s100, 1

So that means this test should include with and without the stack object. We also need to solve the mystery of why this stack object is needed

[kerbowa]

llvm/test/CodeGen/AMDGPU/spill_more_than_wavesize_csr_sgprs.ll

10

Without the stack object FuncInfo->VGPRReservedForSGPRSpill is 0.

[kerbowa]

llvm/test/CodeGen/AMDGPU/spill_more_than_wavesize_csr_sgprs.ll

10

No VGPR is pre-reserved when there are no stack objects.
https://github.com/llvm/llvm-project/blob/master/llvm/lib/Target/AMDGPU/SIISelLowering.cpp#L11283

This patch will fix some current Navi issues as well where it's more common to run out of lanes in the pre-reserved register.

[saiislam]

Updated the test case to check for spilling in the second vgpr with and without stack object present, once all lanes of first vgpr has been used.

[kerbowa]

llvm/test/CodeGen/AMDGPU/spill_more_than_wavesize_csr_sgprs.ll

31

I think you need "store volatile i32 0, i32 addrspace(5)* %alloca" as well or there is no stack object.

[saiislam]

Corrected the test case

[kerbowa]

LGTM

[arsenm]

llvm/test/CodeGen/AMDGPU/spill_more_than_wavesize_csr_sgprs.ll

10

This doesn't really make sense though; why would this be a factor in the reservation logic?

[saiislam]

llvm/test/CodeGen/AMDGPU/spill_more_than_wavesize_csr_sgprs.ll

32

With this patch applied over recent trunk, I am seeing following behavior:

• without volatile in store instruction, correct behavior (using next vgpr for spilling). stack object instruction gets optimized out
• with volatile in store instruction, spills are still getting overwritten
• With -O0, with and without volatile spills are getting overwritten.

So, whenever stack object is present, wrong behavior is getting triggered somehow.

[arsenm]

llvm/test/CodeGen/AMDGPU/spill_more_than_wavesize_csr_sgprs.ll

32

Did you mean without the patch?

[saiislam]

llvm/test/CodeGen/AMDGPU/spill_more_than_wavesize_csr_sgprs.ll

32

No, I mean with the patch. Something changed in last 3-4 days.
Earlier, this patch used to work correctly for volatile store instruction as well.

[kerbowa]

llvm/test/CodeGen/AMDGPU/spill_more_than_wavesize_csr_sgprs.ll

32

It's because you reversed the logic. It should be 'NumVGPRSpillLanes < WaveSize'.

[saiislam]

Corrected condition on NumVGPRSpillLanes and updated the test case.

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

bool SIMachineFunctionInfo::allocateSGPRSpillToVGPR(MachineFunction &MF,
   const MCPhysReg *CSRegs = MRI.getCalleeSavedRegs();
 
   // Make sure to handle the case where a wide SGPR spill may span between two
   // VGPRs.
   for (unsigned I = 0; I < NumLanes; ++I, ++NumVGPRSpillLanes) {
     Register LaneVGPR;
     unsigned VGPRIndex = (NumVGPRSpillLanes % WaveSize);
 
-    if (VGPRIndex == 0 && !FuncInfo->VGPRReservedForSGPRSpill) {
+    // Reserve a VGPR (when NumVGPRSpillLanes = 0, WaveSize, 2*WaveSize, ..) and
+    // when one of the two conditions is true:
+    // 1. One reserved VGPR being tracked by VGPRReservedForSGPRSpill is not yet
+    // reserved.
+    // 2. All spill lanes of reserved VGPR(s) are full and another spill lane is
+    // required.
+    if (FuncInfo->VGPRReservedForSGPRSpill && NumVGPRSpillLanes <= WaveSize) {
+      assert(FuncInfo->VGPRReservedForSGPRSpill == SpillVGPRs.back().VGPR);
+      LaneVGPR = FuncInfo->VGPRReservedForSGPRSpill;
+    } else if (VGPRIndex == 0) {
       LaneVGPR = TRI->findUnusedRegister(MRI, &AMDGPU::VGPR_32RegClass, MF);
       if (LaneVGPR == AMDGPU::NoRegister) {
         // We have no VGPRs left for spilling SGPRs. Reset because we will not
         // partially spill the SGPR to VGPRs.
         SGPRToVGPRSpills.erase(FI);
         NumVGPRSpillLanes -= I;
         return false;
       }

llvm/test/CodeGen/AMDGPU/spill_more_than_wavesize_csr_sgprs.ll

+; RUN: llc -mtriple amdgcn-amd-amdhsa -mcpu=gfx803 -verify-machineinstrs < %s | FileCheck -enable-var-scope %s
+
+; CHECK-LABEL: {{^}}spill_more_than_wavesize_csr_sgprs:
+; CHECK:        v_writelane_b32 v0, s98, 63
+; CHECK-NEXT:   v_writelane_b32 v1, s99, 0
+; CHECK:        v_readlane_b32 s99, v1, 0
+; CHECK-NEXT:   v_readlane_b32 s98, v0, 63
+
+define void @spill_more_than_wavesize_csr_sgprs() {
+  call void asm sideeffect "",
+   "~{s35},~{s36},~{s37},~{s38},~{s39},~{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}"()
+  ret void
+}
+