[AMDGPU] Don't remove VGPR to AGPR dead spills from frame info

[bcahoon]

Summary

Removing dead frame indices for VGPR to AGPR spills is incorrect
when the frame index is shared by multiple objects, which may
occur due to stack slot coloring. The problem is that subsequent
code processes the other object will assert because the stack
frame index has been marked dead.

Removing dead frame indices is needed prior to stack slot
coloring, which is what happens with SGPR to VGPR spills. These
spills are lowered prior to stack slot coloring, but the VGPR
to AGPR spills are processed afterwards during the Prolog/Epilog
Inserter pass.

[foad]

Typo in summary "pills".

[arsenm]

Can we rip out this special "spill" handling since D109301 @cdevadas

[rampitec]

In D115996#3202932, @arsenm wrote:

Can we rip out this special "spill" handling since D109301 @cdevadas

I has unsolved problem with partial tuple spill. There is no partial copy.

[rampitec]

Will the slot remain in the final allocation?

[cdevadas]

In D115996#3203105, @rampitec wrote:

In D115996#3202932, @arsenm wrote:

Can we rip out this special "spill" handling since D109301 @cdevadas

I has unsolved problem with partial tuple spill. There is no partial copy.

I'm not sure how we are going to handle this partial copy differently.
The equivalent for partial-copy during regalloc would be possible only if the tryInstructionSplit inserts partial copy followed by the InlineSpiller handling the leftover tuple spills.
Currently, the former doesn't handle partial copy and the latter doesn't insert subrange tuple spills/restores.
Even if they are equipped to handle partial copy and spill, how do these independent functions communicate each other the exact tuple subrange that the other suppose to handle?

[bcahoon]

In D115996#3203119, @rampitec wrote:

Will the slot remain in the final allocation?

Yes. At least in the example I'm seeing.

[rampitec]

In D115996#3203263, @bcahoon wrote:

In D115996#3203119, @rampitec wrote:

Will the slot remain in the final allocation?

Yes. At least in the example I'm seeing.

Is there any way to get rid of them? Check they are actually dead, or move code later?

[bcahoon]

In D115996#3203344, @rampitec wrote:

In D115996#3203263, @bcahoon wrote:

In D115996#3203119, @rampitec wrote:

Will the slot remain in the final allocation?

Yes. At least in the example I'm seeing.

Is there any way to get rid of them? Check they are actually dead, or move code later?

Let me think about this a little and see if I can come up with a way to remove them if they are actually dead.

[bcahoon]

Changes so that the VGPR to AGPR spill slot can be removed if stack slot coloring has not allocated multiple objects to the same slot. Added a couple extra test cases to check for different cases when the slot can/can't be removed.

[rampitec]

Thanks! LGTM.

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

void SIFrameLowering::processFunctionBeforeFrameFinalized(
   SIMachineFunctionInfo *FuncInfo = MF.getInfo<SIMachineFunctionInfo>();
 
   const bool SpillVGPRToAGPR = ST.hasMAIInsts() && FuncInfo->hasSpilledVGPRs()
                                && EnableSpillVGPRToAGPR;
 
   if (SpillVGPRToAGPR) {
     // To track the spill frame indices handled in this pass.
     BitVector SpillFIs(MFI.getObjectIndexEnd(), false);
+    BitVector NonVGPRSpillFIs(MFI.getObjectIndexEnd(), false);
 
     bool SeenDbgInstr = false;
 
     for (MachineBasicBlock &MBB : MF) {
       for (MachineInstr &MI : llvm::make_early_inc_range(MBB)) {
+        int FrameIndex;
         if (MI.isDebugInstr())
           SeenDbgInstr = true;
 
         if (TII->isVGPRSpill(MI)) {
           // Try to eliminate stack used by VGPR spills before frame
           // finalization.
           unsigned FIOp = AMDGPU::getNamedOperandIdx(MI.getOpcode(),
                                                      AMDGPU::OpName::vaddr);
           int FI = MI.getOperand(FIOp).getIndex();
           Register VReg =
             TII->getNamedOperand(MI, AMDGPU::OpName::vdata)->getReg();
           if (FuncInfo->allocateVGPRSpillToAGPR(MF, FI,
                                                 TRI->isAGPR(MRI, VReg))) {
             // FIXME: change to enterBasicBlockEnd()
             RS->enterBasicBlock(MBB);
             TRI->eliminateFrameIndex(MI, 0, FIOp, RS);
             SpillFIs.set(FI);
             continue;
           }
-        }
-      }
-    }
+        } else if (TII->isStoreToStackSlot(MI, FrameIndex) ||
+                   TII->isLoadFromStackSlot(MI, FrameIndex))
+          NonVGPRSpillFIs.set(FrameIndex);
+
+    // Stack slot coloring may assign different objets to the same stack slot.
+    // If not, then the VGPR to AGPR spill slot is dead.
+    for (unsigned FI : SpillFIs.set_bits())
+      if (!NonVGPRSpillFIs.test(FI))
+        FuncInfo->setVGPRToAGPRSpillDead(FI);
 
     for (MachineBasicBlock &MBB : MF) {
       for (MCPhysReg Reg : FuncInfo->getVGPRSpillAGPRs())
         MBB.addLiveIn(Reg);
 
       for (MCPhysReg Reg : FuncInfo->getAGPRSpillVGPRs())
         MBB.addLiveIn(Reg);
 

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

struct SGPRSpillVGPR {
     Optional<int> FI;
 
     SGPRSpillVGPR(Register V, Optional<int> F) : VGPR(V), FI(F) {}
   };
 
   struct VGPRSpillToAGPR {
     SmallVector<MCPhysReg, 32> Lanes;
     bool FullyAllocated = false;
+    bool IsDead = false;
   };
 
   // Map WWM VGPR to a stack slot that is used to save/restore it in the
   // prolog/epilog.
   MapVector<Register, Optional<int>> WWMReservedRegs;
 
 private:
   // Track VGPR + wave index for each subregister of the SGPR spilled to
public:
   }
 
   MCPhysReg getVGPRToAGPRSpill(int FrameIndex, unsigned Lane) const {
     auto I = VGPRToAGPRSpills.find(FrameIndex);
     return (I == VGPRToAGPRSpills.end()) ? (MCPhysReg)AMDGPU::NoRegister
                                          : I->second.Lanes[Lane];
   }
 
+  void setVGPRToAGPRSpillDead(int FrameIndex) {
+    auto I = VGPRToAGPRSpills.find(FrameIndex);
+    if (I != VGPRToAGPRSpills.end())
+      I->second.IsDead = true;
+  }
+
   bool haveFreeLanesForSGPRSpill(const MachineFunction &MF,
                                  unsigned NumLane) const;
   bool allocateSGPRSpillToVGPR(MachineFunction &MF, int FI);
   bool reserveVGPRforSGPRSpills(MachineFunction &MF);
   bool allocateVGPRSpillToAGPR(MachineFunction &MF, int FI, bool isAGPRtoVGPR);
   void removeDeadFrameIndices(MachineFrameInfo &MFI);
 
   int getScavengeFI(MachineFrameInfo &MFI, const SIRegisterInfo &TRI);

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

void SIMachineFunctionInfo::removeDeadFrameIndices(MachineFrameInfo &MFI) {
   // 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)
+    if (R.second.IsDead)
       MFI.RemoveStackObject(R.first);
   }
 }
 
 int SIMachineFunctionInfo::getScavengeFI(MachineFrameInfo &MFI,
                                          const SIRegisterInfo &TRI) {
   if (ScavengeFI)
     return *ScavengeFI;

llvm/test/CodeGen/AMDGPU/same-slot-agpr-sgpr.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx90a -verify-machineinstrs -run-pass=prologepilog -o - %s | FileCheck %s
+
+---
+name: same_slot_agpr_sgpr
+tracksRegLiveness: true
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: 0, size: 8, alignment: 4 }
+machineFunctionInfo:
+  hasSpilledVGPRs: true
+  scratchRSrcReg: $sgpr0_sgpr1_sgpr2_sgpr3
+  stackPtrOffsetReg: $sgpr32
+body: |
+  bb.0:
+    ; CHECK-LABEL: name: same_slot_agpr_sgpr
+    ; CHECK: liveins: $agpr0, $agpr1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: $vgpr0 = IMPLICIT_DEF
+    ; CHECK-NEXT: $agpr1 = V_ACCVGPR_WRITE_B32_e64 killed $vgpr0, implicit $exec
+    ; CHECK-NEXT: $sgpr4_sgpr5 = IMPLICIT_DEF
+    ; CHECK-NEXT: $sgpr6_sgpr7 = S_MOV_B64 $exec
+    ; CHECK-NEXT: $exec = S_MOV_B64 3, implicit-def $vgpr0
+    ; CHECK-NEXT: BUFFER_STORE_DWORD_OFFSET killed $vgpr0, $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 8, 0, 0, 0, implicit $exec :: (store (s32) into %stack.1, addrspace 5)
+    ; CHECK-NEXT: $vgpr0 = V_WRITELANE_B32 $sgpr4, 0, undef $vgpr0, implicit $sgpr4_sgpr5
+    ; CHECK-NEXT: $vgpr0 = V_WRITELANE_B32 $sgpr5, 1, $vgpr0, implicit killed $sgpr4_sgpr5
+    ; CHECK-NEXT: $agpr1 = V_ACCVGPR_WRITE_B32_e64 killed $vgpr0, implicit $exec
+    ; CHECK-NEXT: $vgpr0 = BUFFER_LOAD_DWORD_OFFSET $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 8, 0, 0, 0, implicit $exec :: (load (s32) from %stack.1, addrspace 5)
+    ; CHECK-NEXT: $exec = S_MOV_B64 killed $sgpr6_sgpr7, implicit killed $vgpr0
+    ; CHECK-NEXT: S_ENDPGM 0
+    $vgpr0 = IMPLICIT_DEF
+    SI_SPILL_AV32_SAVE killed $vgpr0, %stack.0, $sgpr32, 0, implicit $exec :: (store (s32) into %stack.0, addrspace 5)
+    $sgpr4_sgpr5 = IMPLICIT_DEF
+    SI_SPILL_S64_SAVE killed renamable $sgpr4_sgpr5, %stack.0, implicit $exec, implicit $sgpr32 :: (store (s64) into %stack.0, align 4, addrspace 5)
+    S_ENDPGM 0
+...
+---
+name: diff_slot_agpr_sgpr
+tracksRegLiveness: true
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: 0, size: 4, alignment: 4 }
+  - { id: 1, name: '', type: spill-slot, offset: 0, size: 8, alignment: 4 }
+machineFunctionInfo:
+  hasSpilledVGPRs: true
+  scratchRSrcReg: $sgpr0_sgpr1_sgpr2_sgpr3
+  stackPtrOffsetReg: $sgpr32
+body: |
+  bb.0:
+    ; CHECK-LABEL: name: diff_slot_agpr_sgpr
+    ; CHECK: liveins: $agpr0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: $vgpr0 = IMPLICIT_DEF
+    ; CHECK-NEXT: $agpr0 = V_ACCVGPR_WRITE_B32_e64 killed $vgpr0, implicit $exec
+    ; CHECK-NEXT: $sgpr4_sgpr5 = IMPLICIT_DEF
+    ; CHECK-NEXT: $sgpr6_sgpr7 = S_MOV_B64 $exec
+    ; CHECK-NEXT: $exec = S_MOV_B64 3, implicit-def $vgpr0
+    ; CHECK-NEXT: BUFFER_STORE_DWORD_OFFSET killed $vgpr0, $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 8, 0, 0, 0, implicit $exec :: (store (s32) into %stack.2, addrspace 5)
+    ; CHECK-NEXT: $vgpr0 = V_WRITELANE_B32 $sgpr4, 0, undef $vgpr0, implicit $sgpr4_sgpr5
+    ; CHECK-NEXT: $vgpr0 = V_WRITELANE_B32 $sgpr5, 1, $vgpr0, implicit killed $sgpr4_sgpr5
+    ; CHECK-NEXT: BUFFER_STORE_DWORD_OFFSET killed $vgpr0, $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 0, 0, 0, 0, implicit $exec :: (store (s32) into %stack.1, addrspace 5)
+    ; CHECK-NEXT: $vgpr0 = BUFFER_LOAD_DWORD_OFFSET $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 8, 0, 0, 0, implicit $exec :: (load (s32) from %stack.2, addrspace 5)
+    ; CHECK-NEXT: $exec = S_MOV_B64 killed $sgpr6_sgpr7, implicit killed $vgpr0
+    ; CHECK-NEXT: S_ENDPGM 0
+    $vgpr0 = IMPLICIT_DEF
+    SI_SPILL_AV32_SAVE killed $vgpr0, %stack.0, $sgpr32, 0, implicit $exec :: (store (s32) into %stack.0, addrspace 5)
+    $sgpr4_sgpr5 = IMPLICIT_DEF
+    SI_SPILL_S64_SAVE killed renamable $sgpr4_sgpr5, %stack.1, implicit $exec, implicit $sgpr32 :: (store (s64) into %stack.0, align 4, addrspace 5)
+    S_ENDPGM 0
+...
+---
+name: dead_vgpr_slot
+tracksRegLiveness: true
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: 0, size: 4, alignment: 4 }
+machineFunctionInfo:
+  hasSpilledVGPRs: true
+  stackPtrOffsetReg: $sgpr32
+body: |
+  bb.0:
+    ; CHECK-LABEL: name: dead_vgpr_slot
+    ; CHECK: liveins: $agpr0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: $vgpr0 = IMPLICIT_DEF
+    ; CHECK-NEXT: $agpr0 = V_ACCVGPR_WRITE_B32_e64 killed $vgpr0, implicit $exec
+    ; CHECK-NEXT: S_ENDPGM 0
+    $vgpr0 = IMPLICIT_DEF
+    SI_SPILL_AV32_SAVE killed $vgpr0, %stack.0, $sgpr32, 0, implicit $exec :: (store (s32) into %stack.0, addrspace 5)
+    S_ENDPGM 0
+...