[AMDGPU] Implement whole wave register spill

[cdevadas]

Summary

To reduce the register pressure during allocation,
when the allocator spills a virtual register that
corresponds to a whole wave mode operation, the
spill loads and restores should be activated for
all lanes by temporarily flipping all bits in exec
register to one just before the spills. It is not
implemented in the compiler as of today and this
patch enables the necessary support.

This is a pre-patch before the SGPR spill to virtual
VGPR lanes that would eventually causes the whole
wave register spills during allocation.

[yassingh]

• Rebase over new TARGET::PRED_COPY approach

[yassingh]

Rebase

[arsenm]

llvm/lib/Target/AMDGPU/SIDefines.h

934 ↗

(On Diff #528809)

I assumed the value would be 1. Change this to be the pre-shifted values rather than using 1 << flag in the getter?

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

12934–12940 ↗

(On Diff #528809)

Move this up with the other reserved register logic

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

549 ↗

(On Diff #528809)

Capitalization is weird, Capitalize clone

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

651 ↗

(On Diff #528809)

Why are the inbounds checks needed? The analagous MRI functions do not tolerate invalid register values

656–657 ↗

(On Diff #528809)

This should assert (like setFlag already does)

659 ↗

(On Diff #528809)

Ditto, why range check?

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

2151–2152

Braces and indent under if

2160

Indent off

[cdevadas]

llvm/lib/Target/AMDGPU/SIDefines.h

934 ↗

(On Diff #528809)

Will do.

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

549 ↗

(On Diff #528809)

Will do.

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

651 ↗

(On Diff #528809)

VRegFlags will be grown each time (if the size exceeds the initial size set) when a new virtual register is created/cloned. The MRI's 'VRegToType' has a similar bound check.

656–657 ↗

(On Diff #528809)

I guess it's already been discussed in the downstream review. There are instances when Physical registers are taken to spill - fastAlloc, CSR spills during PEI. To handle such instances, an early return was inserted intentionally for the PhysRegs instead of an assertion.

[arsenm]

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

651 ↗

(On Diff #528809)

VRegToType isn't really analagous. Untyped and typed registers can coexist, I'd expect it to be unnecessary here

[arsenm]

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

12934–12940 ↗

(On Diff #528809)

Actually, why is this not in getReservedRegs?

[yassingh]

Addressed Matt's comments:

• Updated flag
• Formatting

[yassingh]

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

651 ↗

(On Diff #528809)

Getting rid of inbounds check causes multiple test failures. (In spill192.mir, spill224.mir, spill288.mir, spill320.mir, etc).

[yassingh]

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

2151–2152

Missed indenting, will fix it.

[cdevadas]

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

2151–2152

Update the patch after clang-format.

[arsenm]

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

647

D149775 is adding another reserved register, but I'm pretty sure they cannot contextually overlap. After both patches are in they should be consolidated to reuse the same reserved register

[arsenm]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

1471 ↗

(On Diff #532851)

ST.getWaveMaskRegClass

1473 ↗

(On Diff #532851)

This misses VCC as a possibility, which is fixed by using getWaveMaskRegClass

llvm/lib/Target/AMDGPU/SILowerSGPRSpills.cpp

338 ↗

(On Diff #532851)

Same, use getWaveMaskRegClass

[yassingh]

• Review comments (use getWaveMaskRegClass(), getExec() functions)

[yassingh]

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

647

Will update once both lands. Added FIXME for now.

2151–2152

Oh right, clang-format reverts the indentation. Going to leave it as it is, else clang-format will keep trying.

[arsenm]

lgtm besides a few nits

llvm/include/llvm/CodeGen/MachineRegisterInfo.h

60 ↗

(On Diff #533300)

Didn't realize this already existed

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

4985 ↗

(On Diff #533300)

ST is already a member

4986 ↗

(On Diff #533300)

This is just this already, you don't need to query it

4987–4988 ↗

(On Diff #533300)

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

2218–2219

Indentation's off

2226

Indentation's off

[yassingh]

• review comments

[yassingh]

Reuploading diff

[yassingh]

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

2218–2219

clang format as well arc undos this indentation.

2226

same

[arsenm]

lgtm with nits

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

70 ↗

(On Diff #533868)

Documentation comment for what IncludeScratchCopy is

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

2218–2219

Well it's clearly wrong, so just manually fix it

2226

Well it's clearly wrong, so just manually fix it

[yassingh]

Review comments and fix formatting

[yassingh]

Rebase

[cdevadas]

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

2151

No need for braces in case of a single instruction.

2159

Wrong indentation.

2217

Avoid braces.

[yassingh]

rebase and formatting

[arsenm]

llvm/lib/Target/AMDGPU/SIInstructions.td

936 ↗

(On Diff #535427)

Should be marked convergent

[yassingh]

Mark SI_SPILL_WWM_V32 convergent, rebase.

[yassingh]

llvm/test/CodeGen/MIR/AMDGPU/long-branch-reg-all-sgpr-used.ll

42 ↗

(On Diff #536981)

We are using sgpr100, sgpr101 in our code below using inline asm calls why not sure what do they mean here?
cc @cdevadas

307 ↗

(On Diff #536981)

@crobeck I believe this test is now not serving the purpose you added it for. I tried freeing up s98 and s99(see below changes in the same file) expecting them to be used for 'longBranchReservedReg', but didn't work. Any pointers on why it might be?

[yassingh]

Rebase before merge

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

BitVector SIRegisterInfo::getReservedRegs(const MachineFunction &MF) const {
 
   if (hasBasePointer(MF)) {
     MCRegister BasePtrReg = getBaseRegister();
     reserveRegisterTuples(Reserved, BasePtrReg);
     assert(!isSubRegister(ScratchRSrcReg, BasePtrReg));
   }
 
   // FIXME: Use same reserved register introduced in D149775
   // SGPR used to preserve EXEC MASK around WWM spill/copy instructions.
   Register ExecCopyReg = MFI->getSGPRForEXECCopy();
   if (ExecCopyReg)
     reserveRegisterTuples(Reserved, ExecCopyReg);
 
   // Reserve VGPRs/AGPRs.
   //
   unsigned MaxNumVGPRs = ST.getMaxNumVGPRs(MF);
   unsigned MaxNumAGPRs = MaxNumVGPRs;
case AMDGPU::SI_SPILL_WWM_V32_SAVE: {
       assert(TII->getNamedOperand(*MI, AMDGPU::OpName::soffset)->getReg() ==
              MFI->getStackPtrOffsetReg());
 
       unsigned Opc = ST.enableFlatScratch() ? AMDGPU::SCRATCH_STORE_DWORD_SADDR
                                             : AMDGPU::BUFFER_STORE_DWORD_OFFSET;
       auto *MBB = MI->getParent();
       bool IsWWMRegSpill = TII->isWWMRegSpillOpcode(MI->getOpcode());
       if (IsWWMRegSpill) {
       TII->insertScratchExecCopy(*MF, *MBB, MI, DL, MFI->getSGPRForEXECCopy(),
                                  RS->isRegUsed(AMDGPU::SCC));
       }
       buildSpillLoadStore(
           *MBB, MI, DL, Opc, Index, VData->getReg(), VData->isKill(), FrameReg,
           TII->getNamedOperand(*MI, AMDGPU::OpName::offset)->getImm(),
           *MI->memoperands_begin(), RS);
       MFI->addToSpilledVGPRs(getNumSubRegsForSpillOp(MI->getOpcode()));
       if (IsWWMRegSpill)
       TII->restoreExec(*MF, *MBB, MI, DL, MFI->getSGPRForEXECCopy());
 
       MI->eraseFromParent();
       return true;
     }
     case AMDGPU::SI_SPILL_V32_RESTORE:
     case AMDGPU::SI_SPILL_V64_RESTORE:
     case AMDGPU::SI_SPILL_V96_RESTORE:
     case AMDGPU::SI_SPILL_V128_RESTORE:
case AMDGPU::SI_SPILL_WWM_V32_RESTORE: {
                                                          AMDGPU::OpName::vdata);
       assert(TII->getNamedOperand(*MI, AMDGPU::OpName::soffset)->getReg() ==
              MFI->getStackPtrOffsetReg());
 
       unsigned Opc = ST.enableFlatScratch() ? AMDGPU::SCRATCH_LOAD_DWORD_SADDR
                                             : AMDGPU::BUFFER_LOAD_DWORD_OFFSET;
       auto *MBB = MI->getParent();
       bool IsWWMRegSpill = TII->isWWMRegSpillOpcode(MI->getOpcode());
       if (IsWWMRegSpill) {
         TII->insertScratchExecCopy(*MF, *MBB, MI, DL, MFI->getSGPRForEXECCopy(),
                                  RS->isRegUsed(AMDGPU::SCC));
       }
       buildSpillLoadStore(
           *MBB, MI, DL, Opc, Index, VData->getReg(), VData->isKill(), FrameReg,
           TII->getNamedOperand(*MI, AMDGPU::OpName::offset)->getImm(),
           *MI->memoperands_begin(), RS);
-      if (IsWWMRegSpill) 
+      if (IsWWMRegSpill) {
         TII->restoreExec(*MF, *MBB, MI, DL, MFI->getSGPRForEXECCopy());
-  
+      }
       MI->eraseFromParent();
       return true;
     }
 
     default: {
       // Other access to frame index
       const DebugLoc &DL = MI->getDebugLoc();