[AMDGPU] Remove dead frame indices after sgpr spill.

[hsmhsm]

Summary

All those frame indices which are dead after sgpr spill should be removed from
the function frame. Othewise, there is a side effect such as re-mapping of free
frame index ids by the later pass(es) like "stack slot coloring" which in turn
could mess-up with the book keeping of "frame index to VGPR lane".

[hsmhsm]

Fix spell mistake in the comment and in assert line.

[foad]

llvm/test/CodeGen/AMDGPU/sgpr-spill-incorrect-fi-bookkeeping-bug.ll

9

This doesn't do anything unless you run FileCheck.

[hsmhsm]

Fix missing FileCheck in the lit test. Thanks Jay for noticing it.

[foad]

llvm/test/CodeGen/AMDGPU/sgpr-spill-incorrect-fi-bookkeeping-bug.ll

2

Does this work? It usually needs to be ... | FileCheck %s.

[hsmhsm]

llvm/test/CodeGen/AMDGPU/sgpr-spill-incorrect-fi-bookkeeping-bug.ll

2

No, it does not, it is again an incorrect check :( I have updated it. I am doing a clean build now, will run one round of lit test before pushing the updated fix.

[hsmhsm]

Fix FileCheck command line.

[cdevadas]

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

457

Isn't it possible to avoid this loop by handling the erase inside the original loop itself?
Maybe convert the range-based loop above into an iterator-based one?

llvm/test/CodeGen/AMDGPU/sgpr-spill-incorrect-fi-bookkeeping-bug.ll

12

Can you avoid all the unnamed variables?

[rampitec]

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

450

DenseSet, avoid std:: when possible.

[hsmhsm]

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

450

Will update it soon.

457

As I understand it, it does not help. The llvm DensMap<>::erase() function does not return next iterator. So, erasing DensMap entries while iterating it seems to be *not* possible. Please let me know if you have any idea.

llvm/test/CodeGen/AMDGPU/sgpr-spill-incorrect-fi-bookkeeping-bug.ll

12

I am coming up with altogether a different lit test. Will update it soon.

[cdevadas]

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

457

If erase doesn't return the iterator, better retain this loop. Thanks.

[hsmhsm]

Fix few review comments by CD and Stas.

[foad]

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

457

If you write the loop as for (auto &R : make_early_inc_range(SGPRToVGPRSpills)) then you can erase R inside the loop.

[hsmhsm]

Do make_early_inc_range traversal of SGPRToVGPRSpills container so that
the entries can be erased while iterating itself.

[hsmhsm]

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

457

Thanks Jay. This is a useful LLVM related information that I was not aware of.

[foad]

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

453

erase(R) should be slightly more efficient as it does not have to do another lookup in the map.

[hsmhsm]

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

453

DenseMap<>::erase() (overloaded) functions takes either key or iterator as arguments. Here R is neither key nor iterator. It is an item (key value pair).

https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/ADT/DenseMap.h#L302

[foad]

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

453

Oh yes, you're right. So it would be slightly more efficient to use an iterator-based for loop, but neater to use a range-based for loop! I don't mind which you choose.

[hsmhsm]

Rebase.

[cdevadas]

LGTM with a nitpick.

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

445–451

The extra comment line.

[hsmhsm]

Update comment.

llvm/lib/Target/AMDGPU/SILowerSGPRSpills.cpp

for (MachineBasicBlock &MBB : MF) {
       for (MachineInstr &MI : MBB) {
         if (MI.isDebugValue() && MI.getOperand(0).isFI() &&
             SpillFIs[MI.getOperand(0).getIndex()]) {
           MI.getOperand(0).ChangeToRegister(Register(), false /*isDef*/);
         }
       }
     }
 
+    // All those frame indices which are dead by now should be removed from the
+    // function frame. Othewise, there is a side effect such as re-mapping of
+    // free frame index ids by the later pass(es) like "stack slot coloring"
+    // which in turn could mess-up with the book keeping of "frame index to VGPR
+    // lane".
+    FuncInfo->removeDeadFrameIndices(MFI);
+
     MadeChange = true;
   } else if (FuncInfo->VGPRReservedForSGPRSpill) {
     FuncInfo->removeVGPRForSGPRSpill(FuncInfo->VGPRReservedForSGPRSpill, MF);
   }
 
   SaveBlocks.clear();
   RestoreBlocks.clear();
 
   // Updated the reserved registers with any VGPRs added for SGPR spills.
   if (NewReservedRegs)
     MRI.freezeReservedRegs(MF);
 
   return MadeChange;
 }

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

for (int I = NumLanes - 1; I >= 0; --I) {
     SpillRegs.push_back(*NextSpillReg);
     Spill.Lanes[I] = *NextSpillReg++;
   }
 
   return Spill.FullyAllocated;
 }
 
 void SIMachineFunctionInfo::removeDeadFrameIndices(MachineFrameInfo &MFI) {
-  // The FP & BP spills haven't been inserted yet, so keep them around.
-  for (auto &R : SGPRToVGPRSpills) {
-    if (R.first != FramePointerSaveIndex && R.first != BasePointerSaveIndex)
+  // Remove dead frame indices from function frame, however keep FP & BP since
+  // spills for them haven't been inserted yet. And also make sure to remove the
+  // frame indices from `SGPRToVGPRSpills` data structure, otherwise, it could
+  // result in an unexpected side effect and bug, in case of any re-mapping of
+  // freed frame indices by later pass(es) like "stack slot coloring".
+  for (auto &R : make_early_inc_range(SGPRToVGPRSpills)) {
+    if (R.first != FramePointerSaveIndex && R.first != BasePointerSaveIndex) {
       MFI.RemoveStackObject(R.first);
+      SGPRToVGPRSpills.erase(R.first);
+    }
   }
 
   // All other SPGRs must be allocated on the default stack, so reset the stack
   // ID.
   for (int i = MFI.getObjectIndexBegin(), e = MFI.getObjectIndexEnd(); i != e;
        ++i)
     if (i != FramePointerSaveIndex && i != BasePointerSaveIndex)
       MFI.setStackID(i, TargetStackID::Default);
 
   for (auto &R : VGPRToAGPRSpills) {
     if (R.second.FullyAllocated)

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

bool SIRegisterInfo::spillSGPR(MachineBasicBlock::iterator MI,
   bool SpillToVGPR = !VGPRSpills.empty();
   if (OnlyToVGPR && !SpillToVGPR)
     return false;
 
   assert(SpillToVGPR || (SB.SuperReg != SB.MFI.getStackPtrOffsetReg() &&
                          SB.SuperReg != SB.MFI.getFrameOffsetReg()));
 
   if (SpillToVGPR) {
+
+    assert(SB.NumSubRegs == VGPRSpills.size() &&
+           "Num of VGPR lanes should be equal to num of SGPRs spilled");
+
     for (unsigned i = 0, e = SB.NumSubRegs; i < e; ++i) {
       Register SubReg =
           SB.NumSubRegs == 1
               ? SB.SuperReg
               : Register(getSubReg(SB.SuperReg, SB.SplitParts[i]));
       SIMachineFunctionInfo::SpilledReg Spill = VGPRSpills[i];
 
       bool UseKill = SB.IsKill && i == SB.NumSubRegs - 1;

llvm/test/CodeGen/AMDGPU/sgpr-spill-incorrect-fi-bookkeeping-bug.ll

+; RUN: llc -mtriple=amdgcn--amdhsa -mcpu=gfx900 -verify-machineinstrs < %s | FileCheck %s
+
+; This tests for a bug that caused a crash in SIRegisterInfo::spillSGPR()
+; which was due to incorrect book-keeping of removed dead frame indices.
+
+; CHECK-LABEL: {{^}}kernel0:
+define amdgpu_kernel void @kernel0(i32 addrspace(1)* %out, i32 %in) #1 {
+  call void asm sideeffect "", "~{v[0:7]}" () #0
+  call void asm sideeffect "", "~{v[8:15]}" () #0
+  call void asm sideeffect "", "~{v[16:19]}"() #0
+  call void asm sideeffect "", "~{v[20:21]}"() #0
+  call void asm sideeffect "", "~{v22}"() #0
+
+  %val0 = call <2 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val1 = call <4 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val2 = call <8 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val3 = call <16 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val4 = call <2 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val5 = call <4 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val6 = call <8 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val7 = call <16 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val8 = call <2 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val9 = call <4 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val10 = call <8 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val11 = call <16 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val12 = call <2 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val13 = call <4 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val14 = call <8 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val15 = call <16 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val16 = call <2 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val17 = call <4 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val18 = call <8 x i32> asm sideeffect "; def $0", "=s" () #0
+  %val19 = call <16 x i32> asm sideeffect "; def $0", "=s" () #0
+  %cmp = icmp eq i32 %in, 0
+  br i1 %cmp, label %bb0, label %ret
+
+bb0:
+  call void asm sideeffect "; use $0", "s"(<2 x i32> %val0) #0
+  call void asm sideeffect "; use $0", "s"(<4 x i32> %val1) #0
+  call void asm sideeffect "; use $0", "s"(<8 x i32> %val2) #0
+  call void asm sideeffect "; use $0", "s"(<16 x i32> %val3) #0
+  call void asm sideeffect "; use $0", "s"(<2 x i32> %val4) #0
+  call void asm sideeffect "; use $0", "s"(<4 x i32> %val5) #0
+  call void asm sideeffect "; use $0", "s"(<8 x i32> %val6) #0
+  call void asm sideeffect "; use $0", "s"(<16 x i32> %val7) #0
+  call void asm sideeffect "; use $0", "s"(<2 x i32> %val8) #0
+  call void asm sideeffect "; use $0", "s"(<4 x i32> %val9) #0
+  call void asm sideeffect "; use $0", "s"(<8 x i32> %val10) #0
+  call void asm sideeffect "; use $0", "s"(<16 x i32> %val11) #0
+  call void asm sideeffect "; use $0", "s"(<2 x i32> %val12) #0
+  call void asm sideeffect "; use $0", "s"(<4 x i32> %val13) #0
+  call void asm sideeffect "; use $0", "s"(<8 x i32> %val14) #0
+  call void asm sideeffect "; use $0", "s"(<16 x i32> %val15) #0
+  call void asm sideeffect "; use $0", "s"(<2 x i32> %val16) #0
+  call void asm sideeffect "; use $0", "s"(<4 x i32> %val17) #0
+  call void asm sideeffect "; use $0", "s"(<8 x i32> %val18) #0
+  call void asm sideeffect "; use $0", "s"(<16 x i32> %val19) #0
+  br label %ret
+
+ret:
+  ret void
+}
+
+attributes #0 = { nounwind }
+attributes #1 = { nounwind "amdgpu-waves-per-eu"="10,10" }