AMDGPU: Fix clobbering CSR VGPRs when spilling SGPR to it

[rampitec]

LGTM

[arsenm]

r309783

lib/Target/AMDGPU/SIFrameLowering.cpp

void SIFrameLowering::emitPrologue(MachineFunction &MF,
 
   uint32_t NumBytes = MFI.getStackSize();
   if (NumBytes != 0 && hasSP(MF)) {
     BuildMI(MBB, MBBI, DL, TII->get(AMDGPU::S_ADD_U32), StackPtrReg)
       .addReg(StackPtrReg)
       .addImm(NumBytes * ST.getWavefrontSize())
       .setMIFlag(MachineInstr::FrameSetup);
   }
+
+  for (const SIMachineFunctionInfo::SGPRSpillVGPRCSR &Reg
+         : FuncInfo->getSGPRSpillVGPRs()) {
+    if (!Reg.FI.hasValue())
+      continue;
+    TII->storeRegToStackSlot(MBB, MBBI, Reg.VGPR, true,
+                             Reg.FI.getValue(), &AMDGPU::VGPR_32RegClass,
+                             &TII->getRegisterInfo());
+  }
 }
 
 void SIFrameLowering::emitEpilogue(MachineFunction &MF,
                                    MachineBasicBlock &MBB) const {
   const SIMachineFunctionInfo *FuncInfo = MF.getInfo<SIMachineFunctionInfo>();
   if (FuncInfo->isEntryFunction())
     return;
 
+  const SISubtarget &ST = MF.getSubtarget<SISubtarget>();
+  const SIInstrInfo *TII = ST.getInstrInfo();
+  MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator();
+
+  for (const SIMachineFunctionInfo::SGPRSpillVGPRCSR &Reg
+         : FuncInfo->getSGPRSpillVGPRs()) {
+    if (!Reg.FI.hasValue())
+      continue;
+    TII->loadRegFromStackSlot(MBB, MBBI, Reg.VGPR,
+                              Reg.FI.getValue(), &AMDGPU::VGPR_32RegClass,
+                              &TII->getRegisterInfo());
+  }
+
   unsigned StackPtrReg = FuncInfo->getStackPtrOffsetReg();
   if (StackPtrReg == AMDGPU::NoRegister)
     return;
 
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   uint32_t NumBytes = MFI.getStackSize();
 
-  const SISubtarget &ST = MF.getSubtarget<SISubtarget>();
-  const SIInstrInfo *TII = ST.getInstrInfo();
-  MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator();
   DebugLoc DL;
 
   // FIXME: Clarify distinction between no set SP and SP. For callee functions,
   // it's really whether we need SP to be accurate or not.
 
   if (NumBytes != 0 && hasSP(MF)) {
     BuildMI(MBB, MBBI, DL, TII->get(AMDGPU::S_SUB_U32), StackPtrReg)
       .addReg(StackPtrReg)

lib/Target/AMDGPU/SIMachineFunctionInfo.h

struct SpilledReg {
 
     SpilledReg() = default;
     SpilledReg(unsigned R, int L) : VGPR (R), Lane (L) { }
 
     bool hasLane() { return Lane != -1;}
     bool hasReg() { return VGPR != AMDGPU::NoRegister;}
   };
 
+  struct SGPRSpillVGPRCSR {
+    // VGPR used for SGPR spills
+    unsigned VGPR;
+
+    // If the VGPR is a CSR, the stack slot used to save/restore it in the
+    // prolog/epilog.
+    Optional<int> FI;
+
+    SGPRSpillVGPRCSR(unsigned V, Optional<int> F) :
+      VGPR(V),
+      FI(F) {}
+  };
+
 private:
   // SGPR->VGPR spilling support.
   typedef std::pair<unsigned, unsigned> SpillRegMask;
 
   // Track VGPR + wave index for each subregister of the SGPR spilled to
   // frameindex key.
   DenseMap<int, std::vector<SpilledReg>> SGPRToVGPRSpills;
   unsigned NumVGPRSpillLanes = 0;
-  SmallVector<unsigned, 2> SpillVGPRs;
+  SmallVector<SGPRSpillVGPRCSR, 2> SpillVGPRs;
 
 public:
 
   SIMachineFunctionInfo(const MachineFunction &MF);
 
   ArrayRef<SpilledReg> getSGPRToVGPRSpills(int FrameIndex) const {
     auto I = SGPRToVGPRSpills.find(FrameIndex);
     return (I == SGPRToVGPRSpills.end()) ?
       ArrayRef<SpilledReg>() : makeArrayRef(I->second);
   }
 
+  ArrayRef<SGPRSpillVGPRCSR> getSGPRSpillVGPRs() const {
+    return SpillVGPRs;
+  }
+
   bool allocateSGPRSpillToVGPR(MachineFunction &MF, int FI);
   void removeSGPRToVGPRFrameIndices(MachineFrameInfo &MFI);
 
   bool hasCalculatedTID() const { return TIDReg != AMDGPU::NoRegister; };
   unsigned getTIDReg() const { return TIDReg; };
   void setTIDReg(unsigned Reg) { TIDReg = Reg; }
 
   // Add user SGPRs.

lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

 
 unsigned SIMachineFunctionInfo::addImplicitBufferPtr(const SIRegisterInfo &TRI) {
   ImplicitBufferPtrUserSGPR = TRI.getMatchingSuperReg(
     getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_64RegClass);
   NumUserSGPRs += 2;
   return ImplicitBufferPtrUserSGPR;
 }
 
+static bool isCalleeSavedReg(const MCPhysReg *CSRegs, MCPhysReg Reg) {
+  for (unsigned I = 0; CSRegs[I]; ++I) {
+    if (CSRegs[I] == Reg)
+      return true;
+  }
+
+  return false;
+}
+
 /// Reserve a slice of a VGPR to support spilling for FrameIndex \p FI.
 bool SIMachineFunctionInfo::allocateSGPRSpillToVGPR(MachineFunction &MF,
                                                     int FI) {
   std::vector<SpilledReg> &SpillLanes = SGPRToVGPRSpills[FI];
 
   // This has already been allocated.
   if (!SpillLanes.empty())
     return true;
 
   const SISubtarget &ST = MF.getSubtarget<SISubtarget>();
   const SIRegisterInfo *TRI = ST.getRegisterInfo();
   MachineFrameInfo &FrameInfo = MF.getFrameInfo();
   MachineRegisterInfo &MRI = MF.getRegInfo();
   unsigned WaveSize = ST.getWavefrontSize();
 
   unsigned Size = FrameInfo.getObjectSize(FI);
   assert(Size >= 4 && Size <= 64 && "invalid sgpr spill size");
   assert(TRI->spillSGPRToVGPR() && "not spilling SGPRs to VGPRs");
 
   int NumLanes = Size / 4;
 
+  const MCPhysReg *CSRegs = TRI->getCalleeSavedRegs(&MF);
+
   // Make sure to handle the case where a wide SGPR spill may span between two
   // VGPRs.
   for (int I = 0; I < NumLanes; ++I, ++NumVGPRSpillLanes) {
     unsigned LaneVGPR;
     unsigned VGPRIndex = (NumVGPRSpillLanes % WaveSize);
 
     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 won't
         // partially spill the SGPR to VGPRs.
         SGPRToVGPRSpills.erase(FI);
         NumVGPRSpillLanes -= I;
         return false;
       }
 
-      SpillVGPRs.push_back(LaneVGPR);
+      Optional<int> CSRSpillFI;
+      if (FrameInfo.hasCalls() && CSRegs && isCalleeSavedReg(CSRegs, LaneVGPR)) {
+        // TODO: Should this be a CreateSpillStackObject? This is technically a
+        // weird CSR spill.
+        CSRSpillFI = FrameInfo.CreateStackObject(4, 4, false);
+      }
+
+      SpillVGPRs.push_back(SGPRSpillVGPRCSR(LaneVGPR, CSRSpillFI));
 
       // Add this register as live-in to all blocks to avoid machine verifer
       // complaining about use of an undefined physical register.
       for (MachineBasicBlock &BB : MF)
         BB.addLiveIn(LaneVGPR);
     } else {
-      LaneVGPR = SpillVGPRs.back();
+      LaneVGPR = SpillVGPRs.back().VGPR;
     }
 
     SpillLanes.push_back(SpilledReg(LaneVGPR, VGPRIndex));
   }
 
   return true;
 }
 
 void SIMachineFunctionInfo::removeSGPRToVGPRFrameIndices(MachineFrameInfo &MFI) {
   for (auto &R : SGPRToVGPRSpills)
     MFI.RemoveStackObject(R.first);
 }

test/CodeGen/AMDGPU/byval-frame-setup.ll

entry:
   %add3 = add nsw i32 %tmp1, 2
   store volatile i32 %add3, i32* %arrayidx2, align 4
   store volatile i32 9, i32 addrspace(1)* null, align 4
   ret void
 }
 
 ; GCN-LABEL: {{^}}void_func_byval_struct_non_leaf:
 ; GCN: s_mov_b32 s5, s32
-; GCN: buffer_store_dword v32
+; GCN-DAG: buffer_store_dword v32
+; GCN-DAG: buffer_store_dword v33
 ; GCN: v_writelane_b32
 
-; GCN-DAG: s_add_u32 s32, s32, 0xa00{{$}}
+; GCN-DAG: s_add_u32 s32, s32, 0xb00{{$}}
 
 ; GCN-DAG: buffer_load_dword [[LOAD0:v[0-9]+]], off, s[0:3], s5 offset:4{{$}}
 ; GCN: v_add_i32_e32 [[ADD0:v[0-9]+]], vcc, 1, [[LOAD0]]
 ; GCN: buffer_store_dword [[ADD0]], off, s[0:3], s5 offset:4{{$}}
 
 ; GCN: buffer_load_dword [[LOAD1:v[0-9]+]], off, s[0:3], s5 offset:20{{$}}
 ; GCN: v_add_i32_e32 [[ADD1:v[0-9]+]], vcc, 2, [[LOAD1]]
 
 ; GCN: s_swappc_b64
 
 ; GCN: buffer_store_dword [[ADD1]], off, s[0:3], s5 offset:20{{$}}
 
 ; GCN: v_readlane_b32
 ; GCN: buffer_load_dword v32,
-; GCN: s_sub_u32 s32, s32, 0xa00{{$}}
+; GCN: buffer_load_dword v33,
+; GCN: s_sub_u32 s32, s32, 0xb00{{$}}
 ; GCN: s_setpc_b64
 define void  @void_func_byval_struct_non_leaf(%struct.ByValStruct* byval noalias nocapture align 4 %arg0, %struct.ByValStruct* byval noalias nocapture align 4 %arg1) #1 {
 entry:
   %arrayidx = getelementptr inbounds %struct.ByValStruct, %struct.ByValStruct* %arg0, i32 0, i32 0, i32 0
   %tmp = load volatile i32, i32* %arrayidx, align 4
   %add = add nsw i32 %tmp, 1
   store volatile i32 %add, i32* %arrayidx, align 4
   %arrayidx2 = getelementptr inbounds %struct.ByValStruct, %struct.ByValStruct* %arg1, i32 0, i32 0, i32 0
   %tmp1 = load volatile i32, i32* %arrayidx2, align 4
   %add3 = add nsw i32 %tmp1, 2
   call void @external_void_func_void()
   store volatile i32 %add3, i32* %arrayidx2, align 4
   store volatile i32 9, i32 addrspace(1)* null, align 4
   ret void
 }
 
 ; GCN-LABEL: {{^}}call_void_func_byval_struct_func:
 ; GCN: s_mov_b32 s5, s32
-; GCN: s_add_u32 s32, s32, 0xa00{{$}}
+; GCN: s_add_u32 s32, s32, 0xc00{{$}}
 ; GCN: v_writelane_b32
 
 ; GCN-DAG: s_add_u32 s32, s32, 0x800{{$}}
 ; GCN-DAG: v_mov_b32_e32 [[NINE:v[0-9]+]], 9
 ; GCN-DAG: v_mov_b32_e32 [[THIRTEEN:v[0-9]+]], 13
 
 ; GCN-DAG: buffer_store_dword [[NINE]], off, s[0:3], s5 offset:8
 ; GCN-DAG: buffer_store_dword [[THIRTEEN]], off, s[0:3], s5 offset:24

 ; GCN-DAG: buffer_store_dword [[LOAD6]], off, s[0:3], s32 offset:28
 ; GCN-DAG: buffer_store_dword [[LOAD7]], off, s[0:3], s32 offset:32
 
 ; GCN: s_swappc_b64
 ; GCN-NEXT: s_sub_u32 s32, s32, 0x800{{$}}
 
 ; GCN: v_readlane_b32
 
-; GCN: s_sub_u32 s32, s32, 0xa00{{$}}
+; GCN: s_sub_u32 s32, s32, 0xc00{{$}}
 ; GCN-NEXT: s_waitcnt
 ; GCN-NEXT: s_setpc_b64
 define void @call_void_func_byval_struct_func() #0 {
 entry:
   %arg0 = alloca %struct.ByValStruct, align 4
   %arg1 = alloca %struct.ByValStruct, align 4
   %tmp = bitcast %struct.ByValStruct* %arg0 to i8*
   call void @llvm.lifetime.start.p0i8(i64 32, i8* %tmp)

test/CodeGen/AMDGPU/call-graph-register-usage.ll

 ; GCN: ScratchSize: 404
 define void @use_stack1() #1 {
   %alloca = alloca [100 x i32], align 4
   call void asm sideeffect "; use $0", "v"([100 x i32]* %alloca) #0
   ret void
 }
 
 ; GCN-LABEL: {{^}}indirect_use_stack:
-; GCN: ScratchSize: 2120
+; GCN: ScratchSize: 2124
 define void @indirect_use_stack() #1 {
   %alloca = alloca [16 x i32], align 4
   call void asm sideeffect "; use $0", "v"([16 x i32]* %alloca) #0
   call void @use_stack0()
   ret void
 }
 
 ; GCN-LABEL: {{^}}indirect_2_level_use_stack:
 ; GCN: is_dynamic_callstack = 0
-; GCN: ScratchSize: 2120
+; GCN: ScratchSize: 2124
 define amdgpu_kernel void @indirect_2_level_use_stack() #0 {
   call void @indirect_use_stack()
   ret void
 }
 
 
 ; Should be maximum of callee usage
 ; GCN-LABEL: {{^}}multi_call_use_use_stack:

 ; NumVgprs: 24
 ; GCN: ScratchSize: 16384
 define amdgpu_kernel void @usage_external_recurse() #0 {
   call void @external_recurse()
   ret void
 }
 
 ; GCN-LABEL: {{^}}direct_recursion_use_stack:
-; GCN: ScratchSize: 2052
+; GCN: ScratchSize: 2056
 define void @direct_recursion_use_stack(i32 %val) #2 {
   %alloca = alloca [512 x i32], align 4
   call void asm sideeffect "; use $0", "v"([512 x i32]* %alloca) #0
   %cmp = icmp eq i32 %val, 0
   br i1 %cmp, label %ret, label %call
 
 call:
   %val.sub1 = sub i32 %val, 1
   call void @direct_recursion_use_stack(i32 %val.sub1)
   br label %ret
 
 ret:
   ret void
 }
 
 ; GCN-LABEL: {{^}}usage_direct_recursion:
 ; GCN: is_ptr64 = 1
 ; GCN: is_dynamic_callstack = 1
-; GCN: workitem_private_segment_byte_size = 2052
+; GCN: workitem_private_segment_byte_size = 2056
 define amdgpu_kernel void @usage_direct_recursion(i32 %n) #0 {
   call void @direct_recursion_use_stack(i32 %n)
   ret void
 }
 
 
 attributes #0 = { nounwind norecurse }
 attributes #1 = { nounwind noinline norecurse }
 attributes #2 = { nounwind noinline }

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

define void @callee_with_stack() #0 {
   %alloca = alloca i32
   store volatile i32 0, i32* %alloca
   ret void
 }
 
 ; GCN-LABEL: {{^}}callee_with_stack_and_call:
 ; GCN: ; BB#0:
 ; GCN-NEXT: s_waitcnt
+; GCN: s_mov_b32 s5, s32
+; GCN: buffer_store_dword v32, off, s[0:3], s5 offset:8
 
-; GCN-DAG: s_mov_b32 s5, s32
 ; GCN-DAG: v_writelane_b32 v32, s33,
 ; GCN-DAG: v_writelane_b32 v32, s34,
 ; GCN-DAG: v_writelane_b32 v32, s35,
-; GCN-DAG: buffer_store_dword v0, off, s[0:3], s5 offset:4{{$}}
-; GCN-DAG: s_add_u32 s32, s32, 0x200{{$}}
+; GCN-DAG: s_add_u32 s32, s32, 0x300{{$}}
 ; GCN-DAG: v_mov_b32_e32 v0, 0{{$}}
+; GCN: buffer_store_dword v0, off, s[0:3], s5 offset:4{{$}}
 ; GCN-DAG: s_mov_b32 s33, s5
 
 
 ; GCN: s_swappc_b64
 ; GCN: s_mov_b32 s5, s33
 ; GCN-DAG: v_readlane_b32 s35,
 ; GCN-DAG: v_readlane_b32 s34,
 ; GCN-DAG: v_readlane_b32 s33,
+; GCN: buffer_load_dword v32, off, s[0:3], s5 offset:8
 ; GCN: s_waitcnt
 ; GCN-NEXT: s_setpc_b64
 define void @callee_with_stack_and_call() #0 {
   %alloca = alloca i32
   store volatile i32 0, i32* %alloca
   call void @external_void_func_void()
   ret void
 }
 
 ; Should be able to copy incoming stack pointer directly to inner
 ; call's stack pointer argument.
 
+; There is stack usage only because of the need to evict a VGPR for
+; spilling CSR SGPRs.
+
 ; GCN-LABEL: {{^}}callee_no_stack_with_call:
 ; GCN: s_waitcnt
-; GCN-NOT: s32
+; GCN: s_mov_b32 s5, s32
+; GCN: buffer_store_dword v32, off, s[0:3], s5 offset:4
+; GCN-DAG: v_writelane_b32 v32, s33, 0
+; GCN-DAG: v_writelane_b32 v32, s34, 1
 ; GCN: s_mov_b32 s33, s5
 ; GCN: s_swappc_b64
 ; GCN: s_mov_b32 s5, s33
-; GCN-NOT: s32
+
+; GCN-DAG: v_readlane_b32 s34, v32, 1
+; GCN-DAG: v_readlane_b32 s33, v32, 0
+; GCN: buffer_load_dword v32, off, s[0:3], s5 offset:4
+; GCN: s_sub_u32 s32, s32, 0x200
+
 ; GCN: s_setpc_b64
 define void @callee_no_stack_with_call() #0 {
   call void @external_void_func_void()
   ret void
 }
 
 declare void @external_void_func_void() #0
 
 attributes #0 = { nounwind }
 attributes #1 = { nounwind "no-frame-pointer-elim"="true" }

test/CodeGen/AMDGPU/nested-calls.ll

 ; RUN: llc -march=amdgcn -mcpu=fiji -mattr=-flat-for-global -amdgpu-sroa=0 -amdgpu-function-calls -verify-machineinstrs < %s | FileCheck -enable-var-scope -check-prefix=GCN -check-prefix=VI %s
 ; RUN: llc -march=amdgcn -mcpu=hawaii -amdgpu-function-calls -amdgpu-sroa=0 -verify-machineinstrs < %s | FileCheck -enable-var-scope -check-prefix=GCN -check-prefix=CI %s
 ; RUN: llc -march=amdgcn -mcpu=gfx900 -mattr=-flat-for-global -amdgpu-sroa=0 -amdgpu-function-calls -verify-machineinstrs < %s | FileCheck -enable-var-scope -check-prefix=GCN -check-prefix=GFX9 -check-prefix=VI %s
 
 ; Test calls when called by other callable functions rather than
 ; kernels.
 
 declare void @external_void_func_i32(i32) #0
 
 ; GCN-LABEL: {{^}}test_func_call_external_void_func_i32_imm:
 ; GCN: s_waitcnt
-; GCN-NOT: s32
+; GCN: s_mov_b32 s5, s32
+; Spill CSR VGPR used for SGPR spilling
+; GCN: buffer_store_dword v32, off, s[0:3], s5 offset:4
+; GCN-DAG: s_add_u32 s32, s32, 0x200
+; GCN-DAG: v_writelane_b32 v32, s33, 0
+; GCN-DAG: v_writelane_b32 v32, s34, 1
+; GCN-DAG: v_writelane_b32 v32, s35, 2
+
 ; GCN: s_swappc_b64
-; GCN-NOT: s32
+
+; GCN: v_readlane_b32 s35, v32, 2
+; GCN: v_readlane_b32 s34, v32, 1
+; GCN: v_readlane_b32 s33, v32, 0
+; GCN: buffer_load_dword v32, off, s[0:3], s5 offset:4
+; GCN: s_sub_u32 s32, s32, 0x200
 ; GCN: s_setpc_b64
 define void @test_func_call_external_void_func_i32_imm() #0 {
   call void @external_void_func_i32(i32 42)
   ret void
 }
 
 ; GCN-LABEL: {{^}}test_func_call_external_void_func_i32_imm_stack_use:
 ; GCN: s_waitcnt
 ; GCN: s_mov_b32 s5, s32
-; GCN: s_add_u32 s32, s32, 0x1100{{$}}
+; GCN: s_add_u32 s32, s32, 0x1200{{$}}
 ; GCN: buffer_store_dword v{{[0-9]+}}, off, s[0:3], s5 offset
 ; GCN: s_swappc_b64
-; GCN: s_sub_u32 s32, s32, 0x1100{{$}}
+; GCN: s_sub_u32 s32, s32, 0x1200{{$}}
 ; GCN: s_setpc_b64
 define void @test_func_call_external_void_func_i32_imm_stack_use() #0 {
   %alloca = alloca [16 x i32], align 4
   %gep0 = getelementptr inbounds [16 x i32], [16 x i32]* %alloca, i32 0, i32 0
   %gep15 = getelementptr inbounds [16 x i32], [16 x i32]* %alloca, i32 0, i32 16
   store volatile i32 0, i32* %gep0
   store volatile i32 0, i32* %gep15
   call void @external_void_func_i32(i32 42)
   ret void
 }
 
 attributes #0 = { nounwind }
 attributes #1 = { nounwind readnone }
 attributes #2 = { nounwind noinline }