[llvm] [AMDGPU][MachineScheduler] Alternative way to control excess RP. (PR #68004)

[via llvm-commits]

https://github.com/alex-t created https://github.com/llvm/llvm-project/pull/68004

This pull request was created as a discussion place to illustrate the idea: to decide if the region has excess RP in finalizeGCNRegion based on the DAG.getRealRegPressure(RegionIdx). Decide if we should keep or revert the
result of the UnclusteredHighRP stage based on the RP after the stage: if the
RP is not less than before - revert.


>From 7d894c0b039e6804079c096f1bbeb2980cafe378 Mon Sep 17 00:00:00 2001
From: Alexander Timofeev <alexander.timofeev at amd.com>
Date: Mon, 2 Oct 2023 18:35:12 +0200
Subject: [PATCH] [AMDGPU][MachineScheduler] Alternative way to control excess
 RP.

---
 llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp | 52 ++++++++++++---------
 llvm/lib/Target/AMDGPU/GCNSchedStrategy.h   | 11 +++--
 2 files changed, 38 insertions(+), 25 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
index ce481e1f1a8bc48..793bbe90307efce 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
@@ -894,10 +894,24 @@ void GCNSchedStage::setupNewBlock() {
 
 void GCNSchedStage::finalizeGCNRegion() {
   DAG.Regions[RegionIdx] = std::pair(DAG.RegionBegin, DAG.RegionEnd);
-  DAG.RescheduleRegions[RegionIdx] = false;
-  if (S.HasHighPressure)
+  PressureAfter = DAG.getRealRegPressure(RegionIdx);
+
+  unsigned NewVGPRRP = PressureAfter.getVGPRNum(false);
+  unsigned NewAGPRRP = PressureAfter.getAGPRNum();
+  unsigned NewSGPRRP = PressureAfter.getSGPRNum();
+
+  if ((NewVGPRRP >= S.VGPRCriticalLimit - S.VGPRExcessMargin) ||
+      (NewAGPRRP >= S.VGPRCriticalLimit - S.VGPRExcessMargin) ||
+      (NewSGPRRP >= S.SGPRCriticalLimit - S.SGPRExcessMargin))
     DAG.RegionsWithHighRP[RegionIdx] = true;
 
+  if ((NewVGPRRP >= S.VGPRExcessLimit - S.VGPRExcessMargin) ||
+      (NewAGPRRP >= S.VGPRExcessLimit - S.SGPRExcessMargin) ||
+      (NewSGPRRP >= S.SGPRExcessLimit - S.VGPRExcessMargin)) {
+    DAG.RegionsWithExcessRP[RegionIdx] = true;
+    DAG.RescheduleRegions[RegionIdx] = true;
+  }
+
   // Revert scheduling if we have dropped occupancy or there is some other
   // reason that the original schedule is better.
   checkScheduling();
@@ -912,7 +926,6 @@ void GCNSchedStage::finalizeGCNRegion() {
 
 void GCNSchedStage::checkScheduling() {
   // Check the results of scheduling.
-  PressureAfter = DAG.getRealRegPressure(RegionIdx);
   LLVM_DEBUG(dbgs() << "Pressure after scheduling: " << print(PressureAfter));
   LLVM_DEBUG(dbgs() << "Region: " << RegionIdx << ".\n");
 
@@ -959,16 +972,6 @@ void GCNSchedStage::checkScheduling() {
                       << DAG.MinOccupancy << ".\n");
   }
 
-  unsigned MaxVGPRs = ST.getMaxNumVGPRs(MF);
-  unsigned MaxSGPRs = ST.getMaxNumSGPRs(MF);
-  if (PressureAfter.getVGPRNum(false) > MaxVGPRs ||
-      PressureAfter.getAGPRNum() > MaxVGPRs ||
-      PressureAfter.getSGPRNum() > MaxSGPRs) {
-    DAG.RescheduleRegions[RegionIdx] = true;
-    DAG.RegionsWithHighRP[RegionIdx] = true;
-    DAG.RegionsWithExcessRP[RegionIdx] = true;
-  }
-
   // Revert if this region's schedule would cause a drop in occupancy or
   // spilling.
   if (shouldRevertScheduling(WavesAfter)) {
@@ -1117,16 +1120,23 @@ bool OccInitialScheduleStage::shouldRevertScheduling(unsigned WavesAfter) {
 bool UnclusteredHighRPStage::shouldRevertScheduling(unsigned WavesAfter) {
   // If RP is not reduced in the unclustered reschedule stage, revert to the
   // old schedule.
-  if ((WavesAfter <= PressureBefore.getOccupancy(ST) &&
-       mayCauseSpilling(WavesAfter)) ||
-      GCNSchedStage::shouldRevertScheduling(WavesAfter)) {
-    LLVM_DEBUG(dbgs() << "Unclustered reschedule did not help.\n");
-    return true;
-  }
+  if (DAG.RegionsWithExcessRP[RegionIdx]) {
+    unsigned NewVGPRRP = PressureAfter.getVGPRNum(false);
+    unsigned NewAGPRRP = PressureAfter.getAGPRNum();
+    unsigned NewSGPRRP = PressureAfter.getSGPRNum();
 
-  // Do not attempt to relax schedule even more if we are already spilling.
-  if (isRegionWithExcessRP())
+    unsigned OldVGPRRP = PressureBefore.getVGPRNum(false);
+    unsigned OldAGPRRP = PressureBefore.getAGPRNum();
+    unsigned OldSGPRRP = PressureBefore.getSGPRNum();
+
+    if (NewVGPRRP > S.VGPRExcessLimit && NewVGPRRP >= OldVGPRRP)
+      return true;
+    if (NewAGPRRP > S.VGPRExcessLimit && NewAGPRRP >= OldAGPRRP)
+      return true;
+    if (NewSGPRRP > S.SGPRExcessLimit && NewSGPRRP >= OldSGPRRP)
+      return true;
     return false;
+  }
 
   LLVM_DEBUG(
       dbgs()
diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
index 7862ec1e894b62e..2119a6f3109bca8 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
@@ -56,10 +56,6 @@ class GCNSchedStrategy : public GenericScheduler {
 
   std::vector<unsigned> MaxPressure;
 
-  unsigned SGPRExcessLimit;
-
-  unsigned VGPRExcessLimit;
-
   unsigned TargetOccupancy;
 
   MachineFunction *MF;
@@ -94,10 +90,17 @@ class GCNSchedStrategy : public GenericScheduler {
 
   unsigned VGPRCriticalLimit;
 
+  unsigned SGPRExcessLimit;
+
+  unsigned VGPRExcessLimit;
+
   unsigned SGPRLimitBias = 0;
 
   unsigned VGPRLimitBias = 0;
 
+  unsigned VGPRExcessMargin = 1;
+  unsigned SGPRExcessMargin = 0;
+
   GCNSchedStrategy(const MachineSchedContext *C);
 
   SUnit *pickNode(bool &IsTopNode) override;