[AMDGPU] Add WWM reserved VGPRs to WWMSpills

[cdevadas]

Summary

The custom VGPR spills inserted during frame lowering
maintain a separate list for WWM reserved registers.
Added them into WWMSpills that already tracks such
reserved registers. It unifies the spill insertion.

[arsenm]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

1245

If we're just going to add these to lane VGPRs, is there any real reason to distinguish the WWM registers? Could we just add them to LaneVGPRs to start with?

1245–1246

Theoretically we could want to call determineCalleeSaves multiple times, so I'm not sure accumulating state here is a good idea.

llvm/lib/Target/AMDGPU/SIPreAllocateWWMRegs.cpp

151 ↗

(On Diff #424259)

Drop this comment, it's not unconditionally true. The relevant problem is addressed in frame lowering

[cdevadas]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

1245

The separate list was maintained for Serialize info. We could directly add them into LaneVGPRs otherwise.

1245–1246

The best place to fill in the LaneVGPRs (both writelane and WWM reserved VGPRs) would be in processFunctionBeforeFrameFinalized. But that is invoked after determineCalleeSaves where we mask these spills off the CSR list so that the default spill insertion in PEI would defer them, and we could later custom insert them during FrameLowering.

At this moment I don't see a better place to move it. We could add a condition to enable this code only once if determineCalleeSaves needs to be called multiple times.

[cdevadas]

Removed the unwanted comment.

[arsenm]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

1245

The serialization is supposed to reflect whatever's here. It's not worth preserving for its own sake

[nhaehnle]

It makes sense to unify that handling, but I object to the term "lane VGPRs". In a way, it's the other way around: the registers that are used for SGPR->VGPR spills are de facto "WWM VGPRs", because they're used in a wave-wide manner, which is the opposite of lane-based usage.

[cdevadas]

Code rebase + review title and summary fixup.

[arsenm]

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

469–483

I'm confused by the relationship between "WWMReservedRegs" and WWMVGPRs.

[cdevadas]

Improved the comment for WWMReservedRegs to have more details especially to differentiate it from the superset WWMVGPRs.

[arsenm]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

1245

allocateWWMSpill would be a better name

[cdevadas]

Code rebase.

[arsenm]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

1245–1246

With the rename it's clearer what's going on here now. I'd prefer to keep the frame index creation in processFunctionBeforeFrameFinalized

[cdevadas]

Moved WWM reserved registers frame index creation into processFunctionBeforeFrameFinalized.

[cdevadas]

@sebastian-ne Is there a case when SIPreAllocateWWMRegs pass allocates a CSR VGPR to operands in WWM instructions?
If it is true, we may have a bug in the current handling that we failed to reset the WWMReservedRegs from the SavedVGPRs during determineCalleeSaves that allow us later to custom insert these spills in the emitPrologue/emitEpilogue functions.

Here is what we do for VGPRs used for SGPR spills.
https://github.com/llvm/llvm-project/blob/main/llvm/lib/Target/AMDGPU/SIFrameLowering.cpp#L1278

It doesn't come under the scope of this patch. But if this case can arise, we must fix it.

[sebastian-ne]

In D124193#3618553, @cdevadas wrote:

@sebastian-ne Is there a case when SIPreAllocateWWMRegs pass allocates a CSR VGPR to operands in WWM instructions?
If it is true, we may have a bug in the current handling that we failed to reset the WWMReservedRegs from the SavedVGPRs during determineCalleeSaves that allow us later to custom insert these spills in the emitPrologue/emitEpilogue functions.

Here is what we do for VGPRs used for SGPR spills.
https://github.com/llvm/llvm-project/blob/main/llvm/lib/Target/AMDGPU/SIFrameLowering.cpp#L1278

It doesn't come under the scope of this patch. But if this case can arise, we must fix it.

Sorry for getting back to you so late. You are correct, WWMReservedRegs should be removed from the SavedVGPRs during determineCalleeSaves.

Currently, if WWM registers fall on callee-saves, they are spilled twice, once as CSR, once as WWM register. At least it’s not incorrect.
Here’s a test demonstrating this: https://github.com/Flakebi/amdvlk-llvm-project/commit/0b0e056c5be110096a92c8175d6650d8535d4e33#diff-eddae1d72701ad027181cd89a7409a0b653da8debae655c3379c5b4d97934a73R1241-R1246

[cdevadas]

Rebase

[cdevadas]

Ping.

[cdevadas]

Ping

[nhaehnle]

LGTM

[cdevadas]

Code rebase

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

void SIFrameLowering::emitPrologue(MachineFunction &MF,
   uint32_t RoundedSize = NumBytes;
   // To avoid clobbering VGPRs in lanes that weren't active on function entry,
   // turn on all lanes before doing the spill to memory.
   Register ScratchExecCopy;
 
   Optional<int> FPSaveIndex = FuncInfo->FramePointerSaveIndex;
   Optional<int> BPSaveIndex = FuncInfo->BasePointerSaveIndex;
 
-  // VGPRs used for SGPR->VGPR spills
+  // Spill Whole-Wave Mode VGPRs.
   for (const auto &Reg : FuncInfo->getWWMSpills()) {
     Register VGPR = Reg.first;
     int FI = Reg.second;
     if (!ScratchExecCopy)
       ScratchExecCopy = buildScratchExecCopy(LiveRegs, MF, MBB, MBBI,
                                              /*IsProlog*/ true);
 
     buildPrologSpill(ST, TRI, *FuncInfo, LiveRegs, MF, MBB, MBBI, DL, VGPR, FI);
   }
 
-  for (auto ReservedWWM : FuncInfo->wwmAllocation()) {
-    if (!ScratchExecCopy)
-      ScratchExecCopy = buildScratchExecCopy(LiveRegs, MF, MBB, MBBI,
-                                             /*IsProlog*/ true);
-
-    buildPrologSpill(ST, TRI, *FuncInfo, LiveRegs, MF, MBB, MBBI, DL,
-                     std::get<0>(ReservedWWM), std::get<1>(ReservedWWM));
-  }
-
   if (ScratchExecCopy) {
     // FIXME: Split block and make terminator.
     unsigned ExecMov = ST.isWave32() ? AMDGPU::S_MOV_B32 : AMDGPU::S_MOV_B64;
     MCRegister Exec = ST.isWave32() ? AMDGPU::EXEC_LO : AMDGPU::EXEC;
     BuildMI(MBB, MBBI, DL, TII->get(ExecMov), Exec)
         .addReg(ScratchExecCopy, RegState::Kill);
     LiveRegs.addReg(ScratchExecCopy);
   }
for (const auto &Reg : FuncInfo->getWWMSpills()) {
     if (!ScratchExecCopy)
       ScratchExecCopy =
           buildScratchExecCopy(LiveRegs, MF, MBB, MBBI, /*IsProlog*/ false);
 
     buildEpilogRestore(ST, TRI, *FuncInfo, LiveRegs, MF, MBB, MBBI, DL, VGPR,
                        FI);
   }
 
-  for (auto ReservedWWM : FuncInfo->wwmAllocation()) {
-    if (!ScratchExecCopy)
-      ScratchExecCopy =
-          buildScratchExecCopy(LiveRegs, MF, MBB, MBBI, /*IsProlog*/ false);
-
-    buildEpilogRestore(ST, TRI, *FuncInfo, LiveRegs, MF, MBB, MBBI, DL,
-                       std::get<0>(ReservedWWM), std::get<1>(ReservedWWM));
-  }
-
   if (ScratchExecCopy) {
     // FIXME: Split block and make terminator.
     unsigned ExecMov = ST.isWave32() ? AMDGPU::S_MOV_B32 : AMDGPU::S_MOV_B64;
     MCRegister Exec = ST.isWave32() ? AMDGPU::EXEC_LO : AMDGPU::EXEC;
     BuildMI(MBB, MBBI, DL, TII->get(ExecMov), Exec)
         .addReg(ScratchExecCopy, RegState::Kill);
   }
 }
void SIFrameLowering::processFunctionBeforeFrameFinalized(
   MachineFrameInfo &MFI = MF.getFrameInfo();
 
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   const SIInstrInfo *TII = ST.getInstrInfo();
   const SIRegisterInfo *TRI = ST.getRegisterInfo();
   MachineRegisterInfo &MRI = MF.getRegInfo();
   SIMachineFunctionInfo *FuncInfo = MF.getInfo<SIMachineFunctionInfo>();
 
-  if (!FuncInfo->isEntryFunction()) {
-    // Spill VGPRs used for Whole Wave Mode
-    FuncInfo->allocateWWMReservedSpillSlots(MFI, *TRI);
-  }
-
   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);
 
for (MachineInstr &MI : MBB) {
       // FIXME: should exclude the general use of WRITELANE that doesn't really
       // cause such overwrites. At this point, it's hard to differentiate the
       // one used for spilling with the general usage.
       if (MI.getOpcode() == AMDGPU::V_WRITELANE_B32)
         MFI->allocateWWMSpill(MF, MI.getOperand(0).getReg());
     }
   }
 
+  // Allocate spill slots for WWM reserved VGPRs.
+  for (Register Reg : MFI->getWWMReservedRegs()) {
+    const TargetRegisterClass *RC = TRI->getPhysRegClass(Reg);
+    MFI->allocateWWMSpill(MF, Reg, TRI->getSpillSize(*RC),
+                          TRI->getSpillAlign(*RC));
+  }
+
   for (MachineBasicBlock &MBB : MF) {
     for (auto &Reg : MFI->getWWMSpills())
       MBB.addLiveIn(Reg.first);
 
     MBB.sortUniqueLiveIns();
   }
 
   // Ignore the SGPRs the default implementation found.
void SIFrameLowering::determineCalleeSaves(MachineFunction &MF,
   // Note a new VGPR CSR may be introduced if one is used for the spill, but we
   // don't want to report it here.
   //
   // FIXME: Is this really hasReservedCallFrame?
   const bool WillHaveFP =
       FrameInfo.hasCalls() &&
       (SavedVGPRs.any() || !allStackObjectsAreDead(FrameInfo));
 
-  // VGPRs used for SGPR spilling need to be specially inserted in the prolog,
-  // so don't allow the default insertion to handle them.
+  // The Whole-Wave VGPRs need to be specially inserted in the prolog, so don't
+  // allow the default insertion to handle them.
   for (auto &Reg : MFI->getWWMSpills())
     SavedVGPRs.reset(Reg.first);
 
   LivePhysRegs LiveRegs;
   LiveRegs.init(*TRI);
 
   if (WillHaveFP || hasFP(MF)) {
     assert(!MFI->SGPRForFPSaveRestoreCopy && !MFI->FramePointerSaveIndex &&

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

 
 public:
   struct VGPRSpillToAGPR {
     SmallVector<MCPhysReg, 32> Lanes;
     bool FullyAllocated = false;
     bool IsDead = false;
   };
 
-  // Track VGPRs reserved for WWM.
-  SmallSetVector<Register, 8> WWMReservedRegs;
-
-  /// Track stack slots used for save/restore of reserved WWM VGPRs in the
-  /// prolog/epilog.
-
-  /// FIXME: This is temporary state only needed in PrologEpilogInserter, and
-  /// doesn't really belong here. It does not require serialization
-  SmallVector<int, 8> WWMReservedFrameIndexes;
-
-  void allocateWWMReservedSpillSlots(MachineFrameInfo &MFI,
-                                     const SIRegisterInfo &TRI);
-
-  auto wwmAllocation() const {
-    assert(WWMReservedRegs.size() == WWMReservedFrameIndexes.size());
-    return zip(WWMReservedRegs, WWMReservedFrameIndexes);
-  }
-
 private:
   // Track VGPR + wave index for each subregister of the SGPR spilled to
   // frameindex key.
   DenseMap<int, std::vector<SIRegisterInfo::SpilledReg>> SGPRToVGPRSpills;
   unsigned NumVGPRSpillLanes = 0;
   SmallVector<Register, 2> SpillVGPRs;
   using WWMSpillsMap = MapVector<Register, int>;
   // To track the registers used in instructions that can potentially modify the
   // inactive lanes. The WWM instructions and the writelane instructions for
   // spilling SGPRs to VGPRs fall under such category of operations. The VGPRs
   // modified by them should be spilled/restored at function prolog/epilog to
   // avoid any undesired outcome. Each entry in this map holds a pair of values,
   // the VGPR and its stack slot index.
   WWMSpillsMap WWMSpills;
 
+  using ReservedRegSet = SmallSetVector<Register, 8>;
+  // To track the VGPRs reserved for WWM instructions. They get stack slots
+  // later during PrologEpilogInserter and get added into the superset WWMSpills
+  // for actual spilling. A separate set makes the register reserved part and
+  // the serialization easier.
+  ReservedRegSet WWMReservedRegs;
+
   DenseMap<int, VGPRSpillToAGPR> VGPRToAGPRSpills;
 
   // AGPRs used for VGPR spills.
   SmallVector<MCPhysReg, 32> SpillAGPR;
 
   // VGPRs used for AGPR spills.
   SmallVector<MCPhysReg, 32> SpillVGPR;
 
   // Emergency stack slot. Sometimes, we create this before finalizing the stack
clone(BumpPtrAllocator &Allocator, MachineFunction &DestMF,
         const DenseMap<MachineBasicBlock *, MachineBasicBlock *> &Src2DstMBB)
       const override;
 
   bool initializeBaseYamlFields(const yaml::SIMachineFunctionInfo &YamlMFI,
                                 const MachineFunction &MF,
                                 PerFunctionMIParsingState &PFS,
                                 SMDiagnostic &Error, SMRange &SourceRange);
 
-  void reserveWWMRegister(Register Reg) {
-    WWMReservedRegs.insert(Reg);
-  }
+  void reserveWWMRegister(Register Reg) { WWMReservedRegs.insert(Reg); }
 
   ArrayRef<SIRegisterInfo::SpilledReg>
   getSGPRToVGPRSpills(int FrameIndex) const {
     auto I = SGPRToVGPRSpills.find(FrameIndex);
     return (I == SGPRToVGPRSpills.end())
                ? ArrayRef<SIRegisterInfo::SpilledReg>()
                : makeArrayRef(I->second);
   }
 
   ArrayRef<Register> getSGPRSpillVGPRs() const { return SpillVGPRs; }
   const WWMSpillsMap &getWWMSpills() const { return WWMSpills; }
+  const ReservedRegSet &getWWMReservedRegs() const { return WWMReservedRegs; }
 
   void allocateWWMSpill(MachineFunction &MF, Register VGPR, uint64_t Size = 4,
                         Align Alignment = Align(4));
 
   ArrayRef<MCPhysReg> getAGPRSpillVGPRs() const {
     return SpillAGPR;
   }
 

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

bool SIMachineFunctionInfo::removeDeadFrameIndices(
   for (auto &R : VGPRToAGPRSpills) {
     if (R.second.IsDead)
       MFI.RemoveStackObject(R.first);
   }
 
   return HaveSGPRToMemory;
 }
 
-void SIMachineFunctionInfo::allocateWWMReservedSpillSlots(
-    MachineFrameInfo &MFI, const SIRegisterInfo &TRI) {
-  assert(WWMReservedFrameIndexes.empty());
-
-  WWMReservedFrameIndexes.resize(WWMReservedRegs.size());
-
-  int I = 0;
-  for (Register VGPR : WWMReservedRegs) {
-    const TargetRegisterClass *RC = TRI.getPhysRegClass(VGPR);
-    WWMReservedFrameIndexes[I++] = MFI.CreateSpillStackObject(
-        TRI.getSpillSize(*RC), TRI.getSpillAlign(*RC));
-  }
-}
-
 int SIMachineFunctionInfo::getScavengeFI(MachineFrameInfo &MFI,
                                          const SIRegisterInfo &TRI) {
   if (ScavengeFI)
     return *ScavengeFI;
   if (isEntryFunction()) {
     ScavengeFI = MFI.CreateFixedObject(
         TRI.getSpillSize(AMDGPU::SGPR_32RegClass), 0, false);
   } else {
: ExplicitKernArgSize(MFI.getExplicitKernArgSize()),
       HighBitsOf32BitAddress(MFI.get32BitAddressHighBits()),
       Occupancy(MFI.getOccupancy()),
       ScratchRSrcReg(regToString(MFI.getScratchRSrcReg(), TRI)),
       FrameOffsetReg(regToString(MFI.getFrameOffsetReg(), TRI)),
       StackPtrOffsetReg(regToString(MFI.getStackPtrOffsetReg(), TRI)),
       BytesInStackArgArea(MFI.getBytesInStackArgArea()),
       ReturnsVoid(MFI.returnsVoid()),
       ArgInfo(convertArgumentInfo(MFI.getArgInfo(), TRI)), Mode(MFI.getMode()) {
-  for (Register Reg : MFI.WWMReservedRegs)
+  for (Register Reg : MFI.getWWMReservedRegs())
     WWMReservedRegs.push_back(regToString(Reg, TRI));
 
   if (MFI.getVGPRForAGPRCopy())
     VGPRForAGPRCopy = regToString(MFI.getVGPRForAGPRCopy(), TRI);
   auto SFI = MFI.getOptionalScavengeFI();
   if (SFI)
     ScavengeFI = yaml::FrameIndex(*SFI, MF.getFrameInfo());
 }

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

BitVector SIRegisterInfo::getReservedRegs(const MachineFunction &MF) const {
   }
 
   // On GFX908, in order to guarantee copying between AGPRs, we need a scratch
   // VGPR available at all times.
   if (ST.hasMAIInsts() && !ST.hasGFX90AInsts()) {
     reserveRegisterTuples(Reserved, MFI->getVGPRForAGPRCopy());
   }
 
-  for (Register Reg : MFI->WWMReservedRegs)
+  for (Register Reg : MFI->getWWMReservedRegs())
     reserveRegisterTuples(Reserved, Reg);
 
   // FIXME: Stop using reserved registers for this.
   for (MCPhysReg Reg : MFI->getAGPRSpillVGPRs())
     reserveRegisterTuples(Reserved, Reg);
 
   for (MCPhysReg Reg : MFI->getVGPRSpillAGPRs())
     reserveRegisterTuples(Reserved, Reg);