[llvm] [AMDGPU][Scheduler] Refactor VGPR rematerialization during scheduling (PR #118722)

[Jeffrey Byrnes via llvm-commits]

================
@@ -1467,284 +1488,460 @@ void GCNSchedStage::revertScheduling() {
   DAG.Regions[RegionIdx] = std::pair(DAG.RegionBegin, DAG.RegionEnd);
 }
 
-void PreRARematStage::collectRematerializableInstructions() {
+bool PreRARematStage::canIncreaseOccupancyOrReduceSpill() {
   const SIRegisterInfo *SRI = static_cast<const SIRegisterInfo *>(DAG.TRI);
-  for (unsigned I = 0, E = DAG.MRI.getNumVirtRegs(); I != E; ++I) {
-    Register Reg = Register::index2VirtReg(I);
-    if (!DAG.LIS->hasInterval(Reg))
-      continue;
-
-    // TODO: Handle AGPR and SGPR rematerialization
-    if (!SRI->isVGPRClass(DAG.MRI.getRegClass(Reg)) ||
-        !DAG.MRI.hasOneDef(Reg) || !DAG.MRI.hasOneNonDBGUse(Reg))
-      continue;
 
-    MachineOperand *Op = DAG.MRI.getOneDef(Reg);
-    MachineInstr *Def = Op->getParent();
-    if (Op->getSubReg() != 0 || !isTriviallyReMaterializable(*Def))
-      continue;
-
-    MachineInstr *UseI = &*DAG.MRI.use_instr_nodbg_begin(Reg);
-    if (Def->getParent() == UseI->getParent())
+  RA_DEBUG(dbgs() << "Collecting rematerializable instructions in "
+                  << MF.getFunction().getName() << '\n');
+
+  // Models excess register pressure in a region.
+  struct ExcessRP {
+    // Amount of VGPR spilling, if any.
+    unsigned Spilling = 0;
+    // Number of VGPRs to save to increase occupancy. std::nullopt when
+    // occupancy is SGPR-limited.
+    std::optional<unsigned> Occupancy = std::nullopt;
+  };
+  // Maps optimizable regions (i.e., regions at minimum and VGPR-limited
+  // occupancy, or regions with VGPR spilling) to their excess RP.
+  DenseMap<unsigned, ExcessRP> OptRegions;
+
+  // Collect optimizable regions.
+  unsigned NumRegionsSpilling = 0;
+  bool CanIncreaseOccupancy = true;
+  for (unsigned I = 0, E = DAG.Regions.size(); I != E; ++I) {
+    if (!DAG.RegionsWithMinOcc[I])
       continue;
+    GCNRegPressure &RP = DAG.Pressure[I];
+
+    unsigned NumVGPRs = RP.getVGPRNum(ST.hasGFX90AInsts());
+    ExcessRP Excess;
+    if ((Excess.Spilling = ST.getNumVGPRsToEliminateSpilling(NumVGPRs))) {
+      // Region has VGPR spilling, we may not be able to increase occupancy but
+      // we can at least try to reduce spilling as much as possible.
+      ++NumRegionsSpilling;
+      RA_DEBUG(dbgs() << "Region " << I << " is spilling VGPRs, save "
+                      << Excess.Spilling << " VGPR(s) to eliminate spilling\n");
+    }
+    if (ST.getOccupancyWithNumSGPRs(RP.getSGPRNum()) != DAG.MinOccupancy) {
+      // Occupancy is VGPR-limited. If occupancy is minimal and there is
+      // spilling, the number of registers to save to increase occupancy
+      // includes the number of spilled registers to save; deduct the latter
+      // from the former to only get the number of registers to save to
+      // increase occupancy as if there was no spilling.
+      Excess.Occupancy =
+          ST.getNumVGPRsToIncreaseOccupancy(NumVGPRs) - Excess.Spilling;
+      RA_DEBUG(dbgs() << "Region " << I << " has min. occupancy: save "
+                      << Excess.Occupancy << " VGPR(s) to improve occupancy\n");
+    } else {
+      CanIncreaseOccupancy = false;
+    }
+    if (Excess.Spilling || Excess.Occupancy)
+      OptRegions.insert({I, Excess});
+  }
+  if (OptRegions.empty())
+    return false;
 
-    // We are only collecting defs that are defined in another block and are
-    // live-through or used inside regions at MinOccupancy. This means that the
-    // register must be in the live-in set for the region.
-    bool AddedToRematList = false;
-    for (unsigned I = 0, E = DAG.Regions.size(); I != E; ++I) {
-      auto It = DAG.LiveIns[I].find(Reg);
-      if (It != DAG.LiveIns[I].end() && !It->second.none()) {
-        if (DAG.RegionsWithMinOcc[I]) {
-          RematerializableInsts[I][Def] = UseI;
-          AddedToRematList = true;
+  bool FuncHasSpilling = NumRegionsSpilling != 0;
+  unsigned RematSpillingCutoff = 0;
+  TargetOcc = DAG.MinOccupancy + (CanIncreaseOccupancy ? 1 : 0);
+
+  // Accounts for a reduction in RP in an optimizable region. Returns whether we
+  // estimate that we have identified enough rematerialization opportunities to
+  // increase function occupancy.
+  auto ReduceRPInRegion = [&](auto OptIt, LaneBitmask Mask) -> bool {
+    auto NumRegs = SIRegisterInfo::getNumCoveredRegs(Mask);
+    unsigned I = OptIt->getFirst();
+    ExcessRP &Excess = OptIt->getSecond();
+
+    // While there is spilling, saved registers only serve to reduce it.
+    if (Excess.Spilling) {
+      unsigned Reduction = std::min(NumRegs, Excess.Spilling);
+      Excess.Spilling -= Reduction;
+      NumRegs -= Reduction;
+      if (!Excess.Spilling) {
+        // We have eliminated spilling in the region.
+        LLVM_DEBUG(dbgs() << "sinking eliminates spilling in region " << I
+                          << '\n');
+        if (--NumRegionsSpilling)
----------------
jrbyrnes wrote:

Should it be `if (!--NumRegionsSpilling)` -- that is, the limit is set when we eliminate all regions with spilling?

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