AMDGPU: Preserve m0 value when spilling

[• tstellarAMD]

LGTM.

[arsenm]

r287844

lib/Target/AMDGPU/SIRegisterInfo.cpp

void SIRegisterInfo::spillSGPR(MachineBasicBlock::iterator MI,
 
   SIMachineFunctionInfo *MFI = MF->getInfo<SIMachineFunctionInfo>();
   MachineFrameInfo &FrameInfo = MF->getFrameInfo();
 
   bool SpillToSMEM = ST.hasScalarStores() && EnableSpillSGPRToSMEM;
 
   assert(SuperReg != AMDGPU::M0 && "m0 should never spill");
 
+  unsigned OffsetReg = AMDGPU::M0;
+  unsigned M0CopyReg = AMDGPU::NoRegister;
+
+  if (SpillToSMEM) {
+    if (RS->isRegUsed(AMDGPU::M0)) {
+      M0CopyReg = MRI.createVirtualRegister(&AMDGPU::SReg_32_XM0RegClass);
+      BuildMI(*MBB, MI, DL, TII->get(AMDGPU::COPY), M0CopyReg)
+        .addReg(AMDGPU::M0);
+    }
+  }
+
   // SubReg carries the "Kill" flag when SubReg == SuperReg.
   unsigned SubKillState = getKillRegState((NumSubRegs == 1) && IsKill);
   for (unsigned i = 0, e = NumSubRegs; i < e; ++i) {
     unsigned SubReg = NumSubRegs == 1 ?
       SuperReg : getSubReg(SuperReg, getSubRegFromChannel(i));
 
     if (SpillToSMEM) {
       int64_t FrOffset = FrameInfo.getObjectOffset(Index);
       unsigned Size = FrameInfo.getObjectSize(Index);
       unsigned Align = FrameInfo.getObjectAlignment(Index);
       MachinePointerInfo PtrInfo
         = MachinePointerInfo::getFixedStack(*MF, Index);
       MachineMemOperand *MMO
         = MF->getMachineMemOperand(PtrInfo, MachineMemOperand::MOStore,
                                    Size, Align);
 
-      unsigned OffsetReg = AMDGPU::M0;
       // Add i * 4 wave offset.
       //
       // SMEM instructions only support a single offset, so increment the wave
       // offset.
 
       int64_t Offset = ST.getWavefrontSize() * (FrOffset + 4 * i);
       if (Offset != 0) {
         BuildMI(*MBB, MI, DL, TII->get(AMDGPU::S_ADD_U32), OffsetReg)
           .addReg(MFI->getScratchWaveOffsetReg())
           .addImm(Offset);
       } else {
         BuildMI(*MBB, MI, DL, TII->get(AMDGPU::S_MOV_B32), OffsetReg)
           .addReg(MFI->getScratchWaveOffsetReg());
       }
 
       BuildMI(*MBB, MI, DL, TII->get(AMDGPU::S_BUFFER_STORE_DWORD_SGPR))
         .addReg(SubReg, getKillRegState(IsKill)) // sdata
         .addReg(MFI->getScratchRSrcReg())        // sbase
-        .addReg(OffsetReg)                       // soff
+        .addReg(OffsetReg, RegState::Kill)       // soff
         .addImm(0)                               // glc
         .addMemOperand(MMO);
 
       continue;
     }
 
     struct SIMachineFunctionInfo::SpilledReg Spill =
       MFI->getSpilledReg(MF, Index, i);
if (Spill.hasReg()) {
         .addFrameIndex(Index)                   // vaddr
         .addReg(MFI->getScratchRSrcReg())       // srrsrc
         .addReg(MFI->getScratchWaveOffsetReg()) // soffset
         .addImm(i * 4)                          // offset
         .addMemOperand(MMO);
     }
   }
 
+  if (M0CopyReg != AMDGPU::NoRegister) {
+    BuildMI(*MBB, MI, DL, TII->get(AMDGPU::COPY), AMDGPU::M0)
+      .addReg(M0CopyReg, RegState::Kill);
+  }
+
   MI->eraseFromParent();
   MFI->addToSpilledSGPRs(NumSubRegs);
 }
 
 void SIRegisterInfo::restoreSGPR(MachineBasicBlock::iterator MI,
                                  int Index,
                                  RegScavenger *RS) const {
   MachineFunction *MF = MI->getParent()->getParent();
   MachineRegisterInfo &MRI = MF->getRegInfo();
   MachineBasicBlock *MBB = MI->getParent();
   SIMachineFunctionInfo *MFI = MF->getInfo<SIMachineFunctionInfo>();
   MachineFrameInfo &FrameInfo = MF->getFrameInfo();
   const SISubtarget &ST =  MF->getSubtarget<SISubtarget>();
   const SIInstrInfo *TII = ST.getInstrInfo();
   const DebugLoc &DL = MI->getDebugLoc();
 
   unsigned NumSubRegs = getNumSubRegsForSpillOp(MI->getOpcode());
   unsigned SuperReg = MI->getOperand(0).getReg();
   bool SpillToSMEM = ST.hasScalarStores() && EnableSpillSGPRToSMEM;
 
   assert(SuperReg != AMDGPU::M0 && "m0 should never spill");
 
+  unsigned OffsetReg = AMDGPU::M0;
+  unsigned M0CopyReg = AMDGPU::NoRegister;
+
+  if (SpillToSMEM) {
+    if (RS->isRegUsed(AMDGPU::M0)) {
+      M0CopyReg = MRI.createVirtualRegister(&AMDGPU::SReg_32_XM0RegClass);
+      BuildMI(*MBB, MI, DL, TII->get(AMDGPU::COPY), M0CopyReg)
+        .addReg(AMDGPU::M0);
+    }
+  }
+
+  // SubReg carries the "Kill" flag when SubReg == SuperReg.
   int64_t FrOffset = FrameInfo.getObjectOffset(Index);
 
   for (unsigned i = 0, e = NumSubRegs; i < e; ++i) {
     unsigned SubReg = NumSubRegs == 1 ?
       SuperReg : getSubReg(SuperReg, getSubRegFromChannel(i));
 
     if (SpillToSMEM) {
       unsigned Size = FrameInfo.getObjectSize(Index);
       unsigned Align = FrameInfo.getObjectAlignment(Index);
       MachinePointerInfo PtrInfo
         = MachinePointerInfo::getFixedStack(*MF, Index);
       MachineMemOperand *MMO
         = MF->getMachineMemOperand(PtrInfo, MachineMemOperand::MOLoad,
                                    Size, Align);
 
-      unsigned OffsetReg = AMDGPU::M0;
-
       // Add i * 4 offset
       int64_t Offset = ST.getWavefrontSize() * (FrOffset + 4 * i);
       if (Offset != 0) {
         BuildMI(*MBB, MI, DL, TII->get(AMDGPU::S_ADD_U32), OffsetReg)
           .addReg(MFI->getScratchWaveOffsetReg())
           .addImm(Offset);
       } else {
         BuildMI(*MBB, MI, DL, TII->get(AMDGPU::S_MOV_B32), OffsetReg)
           .addReg(MFI->getScratchWaveOffsetReg());
       }
 
       BuildMI(*MBB, MI, DL, TII->get(AMDGPU::S_BUFFER_LOAD_DWORD_SGPR), SubReg)
-        .addReg(MFI->getScratchRSrcReg()) // sbase
-        .addReg(OffsetReg)                // soff
-        .addImm(0)                        // glc
+        .addReg(MFI->getScratchRSrcReg())  // sbase
+        .addReg(OffsetReg, RegState::Kill) // soff
+        .addImm(0)                         // glc
         .addMemOperand(MMO)
         .addReg(MI->getOperand(0).getReg(), RegState::ImplicitDefine);
 
       continue;
     }
 
     SIMachineFunctionInfo::SpilledReg Spill
       = MFI->getSpilledReg(MF, Index, i);
if (Spill.hasReg()) {
         .addMemOperand(MMO);
       BuildMI(*MBB, MI, DL,
               TII->get(AMDGPU::V_READFIRSTLANE_B32), SubReg)
         .addReg(TmpReg, RegState::Kill)
         .addReg(MI->getOperand(0).getReg(), RegState::ImplicitDefine);
     }
   }
 
+  if (M0CopyReg != AMDGPU::NoRegister) {
+    BuildMI(*MBB, MI, DL, TII->get(AMDGPU::COPY), AMDGPU::M0)
+      .addReg(M0CopyReg, RegState::Kill);
+  }
+
   MI->eraseFromParent();
 }
 
 void SIRegisterInfo::eliminateFrameIndex(MachineBasicBlock::iterator MI,
                                         int SPAdj, unsigned FIOperandNum,
                                         RegScavenger *RS) const {
   MachineFunction *MF = MI->getParent()->getParent();
   MachineRegisterInfo &MRI = MF->getRegInfo();

test/CodeGen/AMDGPU/spill-m0.ll

endif:
   %foo = call i32 asm sideeffect "s_add_i32 $0, $1, 1", "=s,{M0}"(i32 %m0) #0
   store i32 %foo, i32 addrspace(1)* %out
   ret void
 }
 
 @lds = internal addrspace(3) global [64 x float] undef
 
 ; GCN-LABEL: {{^}}spill_m0_lds:
+; GCN: s_mov_b32 m0, s6
+; GCN: v_interp_mov_f32
+
+; TOSMEM: s_mov_b32 vcc_hi, m0
+; TOSMEM: s_mov_b32 m0, s7
+; TOSMEM-NEXT: s_buffer_store_dword s{{[0-9]+}}, s{{\[[0-9]+:[0-9]+\]}}, m0 ; 4-byte Folded Spill
+; TOSMEM: s_mov_b32 m0, vcc_hi
+
+; TOSMEM: s_mov_b32 vcc_hi, m0
+; TOSMEM: s_add_u32 m0, s7, 0x100
+; TOSMEM: s_buffer_store_dword s{{[0-9]+}}, s{{\[[0-9]+:[0-9]+\]}}, m0 ; 8-byte Folded Spill
+; TOSMEM: s_add_u32 m0, s7, 0x200
+; TOSMEM: s_buffer_store_dword s{{[0-9]+}}, s{{\[[0-9]+:[0-9]+\]}}, m0 ; 8-byte Folded Spill
+; TOSMEM: s_mov_b32 m0, vcc_hi
+
+; TOSMEM: s_mov_b64 exec,
+; TOSMEM: s_cbranch_execz
+; TOSMEM: s_branch
+
+; TOSMEM: BB{{[0-9]+_[0-9]+}}:
+; TOSMEM-NEXT: s_add_u32 m0, s7, 0x100
+; TOSMEM-NEXT: s_buffer_load_dword s{{[0-9]+}}, s{{\[[0-9]+:[0-9]+\]}}, m0 ; 8-byte Folded Reload
+
+
 ; GCN-NOT: v_readlane_b32 m0
 ; GCN-NOT: s_buffer_store_dword m0
 ; GCN-NOT: s_buffer_load_dword m0
 define amdgpu_ps void @spill_m0_lds(<16 x i8> addrspace(2)* inreg, <16 x i8> addrspace(2)* inreg, <32 x i8> addrspace(2)* inreg, i32 inreg) #0 {
 main_body:
   %4 = call float @llvm.SI.fs.constant(i32 0, i32 0, i32 %3)
   %cmp = fcmp ueq float 0.0, %4
   br i1 %cmp, label %if, label %else

 endif:
   %export = phi float [%lds_data, %if], [%interp, %else]
   %5 = call i32 @llvm.SI.packf16(float %export, float %export)
   %6 = bitcast i32 %5 to float
   call void @llvm.SI.export(i32 15, i32 1, i32 1, i32 0, i32 1, float %6, float %6, float %6, float %6)
   ret void
 }
 
+; GCN-LABEL: {{^}}restore_m0_lds:
+; TOSMEM: s_cmp_eq_u32
+; TOSMEM: s_mov_b32 vcc_hi, m0
+; TOSMEM: s_mov_b32 m0, s3
+; TOSMEM: s_buffer_store_dword s4, s[84:87], m0 ; 4-byte Folded Spill
+; TOSMEM: s_mov_b32 m0, vcc_hi
+; TOSMEM: s_cbranch_scc1
+
+; TOSMEM: s_mov_b32 m0, -1
+
+; TOSMEM: s_mov_b32 vcc_hi, m0
+; TOSMEM: s_mov_b32 m0, s3
+; TOSMEM: s_buffer_load_dword s4, s[84:87], m0 ; 8-byte Folded Reload
+; TOSMEM: s_add_u32 m0, s3, 0x100
+; TOSMEM: s_waitcnt lgkmcnt(0)
+; TOSMEM: s_buffer_load_dword s5, s[84:87], m0 ; 8-byte Folded Reload
+; TOSMEM: s_mov_b32 m0, vcc_hi
+; TOSMEM: s_waitcnt lgkmcnt(0)
+
+; TOSMEM: ds_write_b64
+
+; TOSMEM: s_mov_b32 vcc_hi, m0
+; TOSMEM: s_add_u32 m0, s3, 0x200
+; TOSMEM: s_buffer_load_dword s0, s[84:87], m0 ; 4-byte Folded Reload
+; TOSMEM: s_mov_b32 m0, vcc_hi
+; TOSMEM: s_waitcnt lgkmcnt(0)
+; TOSMEM: s_mov_b32 m0, s0
+; TOSMEM: ; use m0
+
+; TOSMEM: s_dcache_wb
+; TOSMEM: s_endpgm
+define void @restore_m0_lds(i32 %arg) {
+  %m0 = call i32 asm sideeffect "s_mov_b32 m0, 0", "={M0}"() #0
+  %sval = load volatile i64, i64 addrspace(2)* undef
+  %cmp = icmp eq i32 %arg, 0
+  br i1 %cmp, label %ret, label %bb
+
+bb:
+  store volatile i64 %sval, i64 addrspace(3)* undef
+  call void asm sideeffect "; use $0", "{M0}"(i32 %m0) #0
+  br label %ret
+
+ret:
+  ret void
+}
+
 declare float @llvm.SI.fs.constant(i32, i32, i32) readnone
 
 declare i32 @llvm.SI.packf16(float, float) readnone
 
 declare void @llvm.SI.export(i32, i32, i32, i32, i32, float, float, float, float)
 
 attributes #0 = { nounwind }