[llvm-branch-commits] [llvm] [AMDGPU] Apply occupancy-aware register allocation anti-hints (PR #218074)

[Lukas Sommer via llvm-branch-commits]

================
@@ -4173,6 +4177,126 @@ bool SIRegisterInfo::getRegAllocationHints(Register VirtReg,
   }
 }
 
+bool SIRegisterInfo::shouldApplyAntiHints(
+    const MachineFunction &MF, unsigned NumAllocatedVGPRs,
+    unsigned &MaxVGPRsForCurrentOccupancy) const {
+
+  const SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
+  unsigned DynamicVGPRBlockSize = MFI->getDynamicVGPRBlockSize();
+  unsigned TargetOccupancy = MFI->getOccupancy();
+  unsigned CurrentOccupancy =
+      ST.getOccupancyWithNumVGPRs(NumAllocatedVGPRs, DynamicVGPRBlockSize);
+  MaxVGPRsForCurrentOccupancy =
+      ST.getMaxNumVGPRs(CurrentOccupancy, DynamicVGPRBlockSize);
+
+  LLVM_DEBUG(dbgs() << "anti-hints: " << NumAllocatedVGPRs
+                    << " VGPRs allocated, target occupancy " << TargetOccupancy
+                    << ", current occupancy " << CurrentOccupancy << '\n');
+
+  // If we are already at lowest occupancy, then there is no need to protect
+  // against occupancy regression.
+  if (CurrentOccupancy == 1)
+    return true;
+
+  // Do not apply anti-hints if we are reaching close to the VGPR budget. For
+  // target occupancy, the 80% cutoff is a conservative: anti-hints are disabled
+  // early enough that later registers still have headroom to stay at target
+  // occupancy. For current occupancy, the 95% cutoff margin is used to apply
+  // anti-hints close to the limit of the current occupancy budget.
+  unsigned MaxVGPRsCutOffForTargetOccupancy =
+      (ST.getMaxNumVGPRs(TargetOccupancy, DynamicVGPRBlockSize) * 80) / 100;
+  unsigned MaxVGPRsCutOffForCurrentOccupancy =
+      (MaxVGPRsForCurrentOccupancy * 95) / 100;
+
+  if (NumAllocatedVGPRs >= MaxVGPRsCutOffForTargetOccupancy) {
+    LLVM_DEBUG(dbgs() << "anti-hints: not applied, at or above the "
+                      << MaxVGPRsCutOffForTargetOccupancy
+                      << " VGPR cutoff for target occupancy\n");
+    return false;
+  }
+
+  if (NumAllocatedVGPRs >= MaxVGPRsCutOffForCurrentOccupancy) {
+    LLVM_DEBUG(dbgs() << "anti-hints: not applied, at or above the "
+                      << MaxVGPRsCutOffForCurrentOccupancy
+                      << " VGPR cutoff for current occupancy\n");
+    return false;
+  }
+
+  return true;
+}
+
+void SIRegisterInfo::filterAndSortForAntiHintedRegs(
+    Register VirtReg, MutableArrayRef<MCPhysReg> CustomOrder,
+    const BitVector &AntiHintedRegUnits, const MachineFunction &MF,
+    const LiveRegMatrix *Matrix) const {
+
+  // Get total number of allocated VGPRs to determine the current occupancy.
+  unsigned NumVGPRs = 0;
+  unsigned NumAGPRs = 0;
+  if (Matrix) {
+    for (MCPhysReg Reg : AMDGPU::VGPR_32RegClass)
+      if (Matrix->isPhysRegUsed(Reg))
+        NumVGPRs = std::max(NumVGPRs, getHWRegIndex(Reg) + 1);
+    for (MCPhysReg Reg : AMDGPU::AGPR_32RegClass)
+      if (Matrix->isPhysRegUsed(Reg))
+        NumAGPRs = std::max(NumAGPRs, getHWRegIndex(Reg) + 1);
+  }
+  unsigned NumAllocatedVGPRs =
+      AMDGPU::getTotalNumVGPRs(ST.hasGFX90AInsts(), NumAGPRs, NumVGPRs);
+
+  // Early exit if we should not apply anti-hints.
+  unsigned MaxVGPRsForCurrentOccupancy = 0;
+  if (!shouldApplyAntiHints(MF, NumAllocatedVGPRs, MaxVGPRsForCurrentOccupancy))
+    return;
+
+  // Returns true if Reg fits within the current occupancy VGPR budget.
+  auto IsWithinBudget = [&](MCPhysReg Reg) -> bool {
+    const TargetRegisterClass *RC = getPhysRegBaseClass(Reg);
+
+    // No VGPR or AGPR usage.
+    if (!RC ||
+        (!isVGPRClass(RC) && !isAGPRClass(RC) && !isVectorSuperClass(RC)))
+      return true;
+
+    unsigned RegEndIndex =
+        getHWRegIndex(Reg) + divideCeil(getRegSizeInBits(*RC), 32);
+
+    unsigned MaxVGPR = NumVGPRs;
+    unsigned MaxAGPR = NumAGPRs;
+    if (isAGPRClass(RC))
+      MaxAGPR = std::max(MaxAGPR, RegEndIndex);
+    else
+      MaxVGPR = std::max(MaxVGPR, RegEndIndex);
+
+    return static_cast<unsigned>(AMDGPU::getTotalNumVGPRs(ST.hasGFX90AInsts(),
+                                                          MaxAGPR, MaxVGPR)) <=
+           MaxVGPRsForCurrentOccupancy;
+  };
+
+  // Find the cutoff point for the current occupancy VGPR budget.
+  auto *BeyondBudgetStart = llvm::find_if(
+      CustomOrder, [&](MCPhysReg Reg) { return !IsWithinBudget(Reg); });
+
+  // Only shuffle within the current occupancy VGPR budget.
+  auto *PartitionPoint = std::stable_partition(
+      CustomOrder.begin(), BeyondBudgetStart,
----------------
sommerlukas wrote:

Is using `BeyondBudgetStart` correct here for `gfx90a`? If one of the VGPRs is above budget, all AGPRs would be excluded from anti-hint partitioning, so anti-hints wouldn't be applied to AGPRs in that case. 

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