[AMDGPU][Scheduler] Make finalizeGCNRegion an overridable hook (NFC)

[lucas-rami]

This allows individual stages to make decisions after re-scheduling individual regions.

---

• Re-apply "[AMDGPU][Scheduler] Scoring system for rematerializations (#175050)" #177206
• [AMDGPU][Scheduler] Revert all regions when remat fails to increase occ. #177205
• [AMDGPU][Scheduler] Simplify scheduling revert logic #177203
• [AMDGPU][Scheduler] Make finalizeGCNRegion an overridable hook (NFC) #177199 👈
• main

[llvmbot]

@llvm/pr-subscribers-backend-amdgpu

Author: Lucas Ramirez (lucas-rami)

Changes

This allows individual stages to make decisions after re-scheduling individual regions.

---

Full diff: https://github.com/llvm/llvm-project/pull/177199.diff

2 Files Affected:

• (modified) llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp (+2-4)
• (modified) llvm/lib/Target/AMDGPU/GCNSchedStrategy.h (+3-3)

diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
index cabf759762a72..876adddcfbbaa 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
@@ -145,7 +145,6 @@ void GCNSchedStrategy::initialize(ScheduleDAGMI *DAG) {
   VGPRCriticalLimit -= std::min(VGPRLimitBias + ErrorMargin, VGPRCriticalLimit);
   SGPRExcessLimit -= std::min(SGPRLimitBias + ErrorMargin, SGPRExcessLimit);
   VGPRExcessLimit -= std::min(VGPRLimitBias + ErrorMargin, VGPRExcessLimit);
-
   LLVM_DEBUG(dbgs() << "VGPRCriticalLimit = " << VGPRCriticalLimit
                     << ", VGPRExcessLimit = " << VGPRExcessLimit
                     << ", SGPRCriticalLimit = " << SGPRCriticalLimit
@@ -1163,6 +1162,8 @@ void GCNScheduleDAGMILive::runSchedStages() {
 
       ScheduleDAGMILive::schedule();
       Stage->finalizeGCNRegion();
+      Stage->advanceRegion();
+      exitRegion();
     }
 
     Stage->finalizeGCNSchedStage();
@@ -1473,9 +1474,6 @@ void GCNSchedStage::finalizeGCNRegion() {
   if (DAG.RegionsWithIGLPInstrs[RegionIdx] &&
       StageID != GCNSchedStageID::UnclusteredHighRPReschedule)
     SavedMutations.swap(DAG.Mutations);
-
-  DAG.exitRegion();
-  advanceRegion();
 }
 
 void GCNSchedStage::checkScheduling() {
diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
index 95a931b9beb2a..6563c8d050ff9 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
@@ -367,12 +367,12 @@ class GCNSchedStage {
   // be skipped.
   virtual bool initGCNRegion();
 
+  // Finalize state after scheduling a region.
+  virtual void finalizeGCNRegion();
+
   // Track whether a new region is also a new MBB.
   void setupNewBlock();
 
-  // Finalize state after scheudling a region.
-  void finalizeGCNRegion();
-
   // Check result of scheduling.
   void checkScheduling();
 

[arsenm]

Isn't the point of finalize that it's done? Why would a stage need to do something else after it's done?

[lucas-rami]

"Finalize" may not be a super clear name indeed; even in the current implementation this is where re-scheduling is reverted if it turns out to be worse than before (stages can customize what "worse" means with the shouldRevertScheduling hook) so stages already have some power to control what happens after a region is re-scheduled. finalizeGCNSchedStage is another stage hook that is invoked after all regions have been re-scheduled.

Having the ability to do something after each region is re-scheduled will allow a future change to early exit and skip further re-schedulings if it deems them to be useless.