[llvm] [AMDGPU] Support partial and empty WWM pools for SGPR spills (PR #213491)

[Shilei Tian via llvm-commits]

================
@@ -372,42 +375,64 @@ void SILowerSGPRSpills::updateLaneVGPRDomInstr(
   }
 }
 
-void SILowerSGPRSpills::determineRegsForWWMAllocation(MachineFunction &MF,
-                                                      BitVector &RegMask) {
-  // Determine an optimal number of VGPRs for WWM allocation. The complement
-  // list will be available for allocating other VGPR virtual registers.
-  SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
+SmallVector<MCRegister>
+SILowerSGPRSpills::determineRegsForWWMAllocation(MachineFunction &MF) {
+  SmallVector<MCRegister> WWMRegCandidates;
+  if (!MaxNumVGPRsForWwmAllocation)
+    return WWMRegCandidates;
+
   MachineRegisterInfo &MRI = MF.getRegInfo();
   BitVector ReservedRegs = TRI->getReservedRegs(MF);
-  BitVector NonWwmAllocMask(TRI->getNumRegs());
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
+  unsigned MaxNumVGPRs = ST.getMaxNumVectorRegs(MF.getFunction()).first;
 
-  // FIXME: MaxNumVGPRsForWwmAllocation might need to be adjusted in the future
-  // to have a balanced allocation between WWM values and per-thread vector
-  // register operands.
-  unsigned NumRegs = MaxNumVGPRsForWwmAllocation;
-  NumRegs =
-      std::min(static_cast<unsigned>(MFI->getSGPRSpillVGPRs().size()), NumRegs);
-
-  auto [MaxNumVGPRs, MaxNumAGPRs] = ST.getMaxNumVectorRegs(MF.getFunction());
   // Try to use the highest available registers for now. Later after
   // vgpr-regalloc, they can be shifted to the lowest range.
-  unsigned I = 0;
   for (unsigned Reg = AMDGPU::VGPR0 + MaxNumVGPRs - 1;
-       (I < NumRegs) && (Reg >= AMDGPU::VGPR0); --Reg) {
+       WWMRegCandidates.size() < MaxNumVGPRsForWwmAllocation &&
+       Reg >= AMDGPU::VGPR0;
+       --Reg) {
     if (!ReservedRegs.test(Reg) &&
-        !MRI.isPhysRegUsed(Reg, /*SkipRegMaskTest=*/true)) {
-      TRI->markSuperRegs(RegMask, Reg);
-      ++I;
-    }
+        !MRI.isPhysRegUsed(Reg, /*SkipRegMaskTest=*/true))
+      WWMRegCandidates.push_back(Reg);
   }
 
-  if (I != NumRegs) {
+  return WWMRegCandidates;
+}
+
+void SILowerSGPRSpills::assignWWMRegs(MachineFunction &MF,
+                                      ArrayRef<MCRegister> WWMRegCandidates,
+                                      bool RequiresFullWWMPool) {
+  SIMachineFunctionInfo *FuncInfo = MF.getInfo<SIMachineFunctionInfo>();
+  if (FuncInfo->getSGPRSpillVGPRs().empty())
+    return;
+
+  BitVector WwmRegMask(TRI->getNumRegs());
+
+  unsigned DesiredPoolSize =
+      std::min(static_cast<unsigned>(FuncInfo->getSGPRSpillVGPRs().size()),
+               static_cast<unsigned>(MaxNumVGPRsForWwmAllocation));
+  unsigned SelectedPoolSize =
+      std::min<unsigned>(DesiredPoolSize, WWMRegCandidates.size());
+  // WWM register candidates are ordered high-to-low, so take the highest
+  // available registers when the desired pool is smaller than the candidate
+  // list.
+  for (MCRegister Reg : WWMRegCandidates.take_front(SelectedPoolSize))
+    TRI->markSuperRegs(WwmRegMask, Reg);
+
+  if (RequiresFullWWMPool && SelectedPoolSize != DesiredPoolSize) {
     // Reserve an arbitrary register and report the error.
-    TRI->markSuperRegs(RegMask, AMDGPU::VGPR0);
+    TRI->markSuperRegs(WwmRegMask, AMDGPU::VGPR0);
     MF.getFunction().getContext().emitError(
         "cannot find enough VGPRs for wwm-regalloc");
   }
+
+  BitVector NonWwmRegMask(WwmRegMask);
----------------
shiltian wrote:

https://github.com/llvm/llvm-project/pull/213827

https://github.com/llvm/llvm-project/pull/213491