[AMDGPU] Do not reassign spilled registers

[rampitec]

Summary

We cannot call LRM::unassign() if LRM::assign() was never called
before, these are symmetrical calls. There are two ways of
assigning a physical register to virtual, via LRM::assign() and
via VRM::assignVirt2Phys(). LRM::assign() will call the VRM to
assign the register and then update LiveIntervalUnion. Inline
spiller calls VRM directly and thus LiveIntervalUnion never gets
updated. A call to LRM::unassign() then asserts about inconsistent
liveness.

We have to note that not all callers of the InlineSpiller even
have LRM to pass, RegAllocPBQP does not have it, so we cannot
always pass LRM into the spiller.

The only way to get into that spiller LRE_DidCloneVirtReg() call
is from LiveRangeEdit::eliminateDeadDefs if we split an LI.

This patch refuses to reassign a LiveInterval created by a split
to workaround the problem. In fact we cannot reassign a spill
anyway as all registers of the needed class are occupied and we
are spilling.

Fixes: SWDEV-267996

[arsenm]

llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp

482

Can you tell if it's a spill from the VirtRegMap instead?

llvm/test/CodeGen/AMDGPU/nsa-reassign.mir

8–16

Don't need the registers section, should use the newer :class syntax

[rampitec]

llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp

482

I don't see how.

llvm/test/CodeGen/AMDGPU/nsa-reassign.mir

8–16

The point of using it is to have 'preffered-register'. I.e. after greedy but before rewriter.

[arsenm]

llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp

482

Does Register::isStackSlot return true here, or does VRM->getStackSlot() != NO_STACK_SLOT tell you?

[rampitec]

llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp

482

Both don't:

(gdb) p Register::isStackSlot((unsigned)Reg)
$12 = false
(gdb) p VRM->getStackSlot(Reg) == VirtRegMap::NO_STACK_SLOT
$11 = true

I.e. it is VRM::isAssignedReg() which is checked in the first place. If not PBQB this shall be considered an RA bug.

[rampitec]

I have done some more debugging on it. The only way to get into that spiller HoistSpillHelper::LRE_DidCloneVirtReg() call is from LiveRangeEdit::eliminateDeadDefs if we split an LI. That lives LRM in the inconsistent state so we cannot call unassign anymore. I will change the check to VRM->getPreSplitReg(Reg), it will work, but I have absolutely no idea how to write a reliable test for it.

[rampitec]

Refuse to reassign any registers which are result of an LI split.
I doubt I can create a proper test though.

[rampitec]

Added tests with forced split.

[arsenm]

llvm/lib/Target/AMDGPU/GCNNSAReassign.cpp

197

Typo incosistent

llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp

477

I don't understand if this is a bug with LRM::assign or not

478

Typo incosistent

[rampitec]

Fixed typo in comments.

[rampitec]

llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp

477

Not in the LRM. This is probably a bug, but in the implementation of the InlineSpiller which uses VRM directly and dies not call LRM::assign(). Nobody has hit it before because nobody tried to use unassign() after greedy I suppose. However, this is not really easy to fix it in the spiller.

[arsenm]

llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp

477

Can you add a fixme explaining this and file a bug for this

[rampitec]

Filed bug https://bugs.llvm.org/show_bug.cgi?id=48911 and updated comments.

llvm/lib/Target/AMDGPU/GCNNSAReassign.cpp

 /// in NSA image instructions. Later SIShrinkInstructions pass will relace NSA
 /// with sequential versions where possible.
 ///
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
 #include "GCNSubtarget.h"
 #include "SIMachineFunctionInfo.h"
+#include "SIInstrInfo.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/LiveRegMatrix.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/InitializePasses.h"
 
 using namespace llvm;
 
private:
   } NSA_Status;
 
   const GCNSubtarget *ST;
 
   const MachineRegisterInfo *MRI;
 
   const SIRegisterInfo *TRI;
 
+  const SIInstrInfo *TII;
+
   VirtRegMap *VRM;
 
   LiveRegMatrix *LRM;
 
   LiveIntervals *LIS;
 
   unsigned MaxNumVGPRs;
 
if (!Fast) {
       // cases a tuple is used because a vector variable contains different
       // parts of an address and it is either already consequitive or cannot
       // be reassigned if not. If needed it is better to rely on register
       // coalescer to process such address tuples.
       if (MRI->getRegClass(Reg) != &AMDGPU::VGPR_32RegClass || Op.getSubReg())
         return NSA_Status::FIXED;
 
       const MachineInstr *Def = MRI->getUniqueVRegDef(Reg);
 
-      if (Def && Def->isCopy() && Def->getOperand(1).getReg() == PhysReg)
+      if (Def && ((Def->isCopy() && Def->getOperand(1).getReg() == PhysReg) ||
+                  TII->isSpill(*Def)))
         return NSA_Status::FIXED;
 
       for (auto U : MRI->use_nodbg_operands(Reg)) {
         if (U.isImplicit())
           return NSA_Status::FIXED;
         const MachineInstr *UseInst = U.getParent();
         if (UseInst->isCopy() && UseInst->getOperand(0).getReg() == PhysReg)
           return NSA_Status::FIXED;

 
 bool GCNNSAReassign::runOnMachineFunction(MachineFunction &MF) {
   ST = &MF.getSubtarget<GCNSubtarget>();
   if (ST->getGeneration() < GCNSubtarget::GFX10)
     return false;
 
   MRI = &MF.getRegInfo();
   TRI = ST->getRegisterInfo();
+  TII = ST->getInstrInfo();
   VRM = &getAnalysis<VirtRegMap>();
   LRM = &getAnalysis<LiveRegMatrix>();
   LIS = &getAnalysis<LiveIntervals>();
 
   const SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
   MaxNumVGPRs = ST->getMaxNumVGPRs(MF);
   MaxNumVGPRs = std::min(ST->getMaxNumVGPRs(MFI->getOccupancy()), MaxNumVGPRs);
   CSRegs = MRI->getCalleeSavedRegs();

llvm/lib/Target/AMDGPU/GCNRegBankReassign.cpp

 /// that 4 has to be subtracted from an SGPR bank number to get the real value.
 /// This also corresponds to bit numbers in bank masks used in the pass.
 ///
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
 #include "GCNSubtarget.h"
 #include "SIMachineFunctionInfo.h"
+#include "SIInstrInfo.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/LiveRegMatrix.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/InitializePasses.h"
 

 
 private:
   const GCNSubtarget *ST;
 
   const MachineRegisterInfo *MRI;
 
   const SIRegisterInfo *TRI;
 
+  const SIInstrInfo *TII;
+
   MachineLoopInfo *MLI;
 
   VirtRegMap *VRM;
 
   LiveRegMatrix *LRM;
 
   LiveIntervals *LIS;
 
unsigned GCNRegBankReassign::getOperandGatherWeight(const MachineInstr& MI,
   }
   return countPopulation(Defs);
 }
 
 bool GCNRegBankReassign::isReassignable(Register Reg) const {
   if (Reg.isPhysical() || !VRM->isAssignedReg(Reg))
     return false;
 
   const MachineInstr *Def = MRI->getUniqueVRegDef(Reg);
 
   Register PhysReg = VRM->getPhys(Reg);
 
-  if (Def && Def->isCopy() && Def->getOperand(1).getReg() == PhysReg)
+  if (Def && ((Def->isCopy() && Def->getOperand(1).getReg() == PhysReg) ||
+              TII->isSpill(*Def)))
     return false;
 
   for (auto U : MRI->use_nodbg_operands(Reg)) {
     if (U.isImplicit())
       return false;
     const MachineInstr *UseInst = U.getParent();
     if (UseInst->isCopy() && UseInst->getOperand(0).getReg() == PhysReg)
       return false;

 
 bool GCNRegBankReassign::runOnMachineFunction(MachineFunction &MF) {
   ST = &MF.getSubtarget<GCNSubtarget>();
   if (!ST->hasRegisterBanking() || skipFunction(MF.getFunction()))
     return false;
 
   MRI = &MF.getRegInfo();
   TRI = ST->getRegisterInfo();
+  TII = ST->getInstrInfo();
   MLI = &getAnalysis<MachineLoopInfo>();
   VRM = &getAnalysis<VirtRegMap>();
   LRM = &getAnalysis<LiveRegMatrix>();
   LIS = &getAnalysis<LiveIntervals>();
 
   const SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
   unsigned Occupancy = MFI->getOccupancy();
   MaxNumVGPRs = ST->getMaxNumVGPRs(MF);

llvm/lib/Target/AMDGPU/SIInstrInfo.h

public:
   static bool isSGPRSpill(const MachineInstr &MI) {
     return MI.getDesc().TSFlags & SIInstrFlags::SGPRSpill;
   }
 
   bool isSGPRSpill(uint16_t Opcode) const {
     return get(Opcode).TSFlags & SIInstrFlags::SGPRSpill;
   }
 
+  static bool isSpill(const MachineInstr &MI) {
+    return MI.getDesc().TSFlags &
+           (SIInstrFlags::VGPRSpill | SIInstrFlags::SGPRSpill);
+  }
+
+  bool isSpill(uint16_t Opcode) const {
+    return get(Opcode).TSFlags &
+           (SIInstrFlags::VGPRSpill | SIInstrFlags::SGPRSpill);
+  }
+
   static bool isDPP(const MachineInstr &MI) {
     return MI.getDesc().TSFlags & SIInstrFlags::DPP;
   }
 
   bool isDPP(uint16_t Opcode) const {
     return get(Opcode).TSFlags & SIInstrFlags::DPP;
   }
 

llvm/test/CodeGen/AMDGPU/nsa-reassign.mir

+# RUN: llc -march=amdgcn -mcpu=gfx1010 -verify-machineinstrs -run-pass greedy,amdgpu-nsa-reassign,virtregrewriter -o - %s | FileCheck -check-prefix=GCN %s
+
+# GCN-LABEL: name: nsa_reassign
+# GCN: IMAGE_SAMPLE_C_L_V1_V5_nsa_gfx10 killed renamable $vgpr0, killed renamable $vgpr1, killed renamable $vgpr2, killed renamable $vgpr3, killed renamable $vgpr4,
+---
+name:            nsa_reassign
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: sgpr_256, preferred-register: '$sgpr0_sgpr1_sgpr2_sgpr3_sgpr4_sgpr5_sgpr6_sgpr7' }
+  - { id: 1, class: sgpr_128, preferred-register: '$sgpr8_sgpr9_sgpr10_sgpr11' }
+  - { id: 2, class: vgpr_32, preferred-register: '$vgpr2' }
+  - { id: 3, class: vgpr_32, preferred-register: '$vgpr4' }
+  - { id: 4, class: vgpr_32, preferred-register: '$vgpr6' }
+  - { id: 5, class: vgpr_32, preferred-register: '$vgpr8' }
+  - { id: 6, class: vgpr_32, preferred-register: '$vgpr10' }
+  - { id: 7, class: vgpr_32, preferred-register: '$vgpr0' }
+body:             |
+  bb.0:
+    %0 = IMPLICIT_DEF
+    %1 = IMPLICIT_DEF
+    %2 = IMPLICIT_DEF
+    %3 = IMPLICIT_DEF
+    %4 = IMPLICIT_DEF
+    %5 = IMPLICIT_DEF
+    %6 = IMPLICIT_DEF
+    %7:vgpr_32 = IMAGE_SAMPLE_C_L_V1_V5_nsa_gfx10 %2, %3, %4, %5, %6, %0, %1, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, implicit $exec :: (dereferenceable load 4 from custom "ImageResource")
+    S_ENDPGM 0, implicit %7
+...
+
+# GCN-LABEL: do_not_reassign_spill
+# GCN: IMAGE_SAMPLE_C_L_V1_V5_nsa_gfx10 killed renamable $vgpr2, killed renamable $vgpr4, killed renamable $vgpr6, killed renamable $vgpr8, killed renamable $vgpr10,
+---
+name:            do_not_reassign_spill
+tracksRegLiveness: true
+machineFunctionInfo:
+  stackPtrOffsetReg:  $sgpr32
+stack:
+  - { id: 0, type: default, offset: 0, size: 4, alignment: 4 }
+registers:
+  - { id: 0, class: sgpr_256, preferred-register: '$sgpr0_sgpr1_sgpr2_sgpr3_sgpr4_sgpr5_sgpr6_sgpr7' }
+  - { id: 1, class: sgpr_128, preferred-register: '$sgpr8_sgpr9_sgpr10_sgpr11' }
+  - { id: 2, class: vgpr_32, preferred-register: '$vgpr2' }
+  - { id: 3, class: vgpr_32, preferred-register: '$vgpr4' }
+  - { id: 4, class: vgpr_32, preferred-register: '$vgpr6' }
+  - { id: 5, class: vgpr_32, preferred-register: '$vgpr8' }
+  - { id: 6, class: vgpr_32, preferred-register: '$vgpr10' }
+  - { id: 7, class: vgpr_32, preferred-register: '$vgpr0' }
+body:             |
+  bb.0:
+    %0 = IMPLICIT_DEF
+    %1 = IMPLICIT_DEF
+    %2 = SI_SPILL_V32_RESTORE %stack.0, $sgpr32, 0, implicit $exec
+    %3 = IMPLICIT_DEF
+    %4 = IMPLICIT_DEF
+    %5 = IMPLICIT_DEF
+    %6 = IMPLICIT_DEF
+    %7:vgpr_32 = IMAGE_SAMPLE_C_L_V1_V5_nsa_gfx10 %2, %3, %4, %5, %6, %0, %1, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, implicit $exec :: (dereferenceable load 4 from custom "ImageResource")
+    S_ENDPGM 0, implicit %7
+...

llvm/test/CodeGen/AMDGPU/regbank-reassign.mir

bb.0:
     %12 = IMPLICIT_DEF
     %13 = IMPLICIT_DEF
     %14 = IMPLICIT_DEF
     %15 = IMPLICIT_DEF
     %3 = V_ADD_F64_e64 0, %0.sub2_sub3:vreg_128, 0, %1:vreg_64, 0, 0, implicit $mode, implicit $exec
     %4 = V_ADD_F64_e64 0, %0.sub0_sub1:vreg_128, 0, %2:vreg_64, 0, 0, implicit $mode, implicit $exec
     S_ENDPGM 0
 ...
+# GCN-LABEL: do_not_reassign_spill{{$}}
+# GCN: V_AND_B32_e32 killed $vgpr5, killed $vgpr1,
+---
+name:            do_not_reassign_spill
+tracksRegLiveness: true
+machineFunctionInfo:
+  stackPtrOffsetReg:  $sgpr32
+stack:
+  - { id: 0, type: default, offset: 0, size: 4, alignment: 4 }
+registers:
+  - { id: 0, class: vgpr_32, preferred-register: '$vgpr1' }
+  - { id: 1, class: vgpr_32, preferred-register: '$vgpr5' }
+  - { id: 2, class: vgpr_32 }
+body: |
+  bb.0:
+    %0 = SI_SPILL_V32_RESTORE %stack.0, $sgpr32, 0, implicit $exec
+    %1 = SI_SPILL_V32_RESTORE %stack.0, $sgpr32, 0, implicit $exec
+    %2 = V_AND_B32_e32 %1, %0, implicit $exec
+    S_ENDPGM 0
+...