[AMDGPU] Add second pass of the scheduler

[rampitec]

Summary

If during scheduling we have identified that we cannot keep optimistic
occupancy increase critical register pressure limit and try scheduling
of the whole function again. In this case blocks with smaller pressure
will have a chance for better scheduling.

[tstellar]

lib/Target/AMDGPU/GCNSchedStrategy.cpp

48–50 ↗

(On Diff #89972)

Can you store TargetOccupancy in the SIMachineFunctionInfo object? I think that would be a little cleaner.

[tstellar]

lib/Target/AMDGPU/GCNSchedStrategy.cpp

48–50 ↗

(On Diff #89972)

To clarify:

TargetOccupancy = MFI->getTargetOccupancy();

[rampitec]

lib/Target/AMDGPU/GCNSchedStrategy.cpp

48–50 ↗

(On Diff #89972)

It is the override for scheduler and it will not be always initialized to non-zero. I'm afraid if I expose this field in MFI it would be misleading.

[vpykhtin]

I'm just a bit confused: what gives scheduler more registry freedom on rescheduling run?

[rampitec]

In D30442#688729, @vpykhtin wrote:

I'm just a bit confused: what gives scheduler more registry freedom on rescheduling run?

With lower occupancy we have more registers available. If for any block we cannot maintain minimal occupancy we can use the same number of registers in other blocks. That is achieved by bumping critical limits.

[vpykhtin]

Ok, it's just not that easy to follow this in GCNMaxOccupancySchedStrategy::initialize.

[kzhuravl]

lib/Target/AMDGPU/GCNSchedStrategy.cpp

62 ↗

(On Diff #89972)

I think this should also respect the "amdgpu-waves-per-eu" attribute (https://clang.llvm.org/docs/AttributeReference.html#amdgpu-waves-per-eu)?

[arsenm]

How expensive is it to do this? The scheduler is already frequently the most expensive pass after RA, sometimes surpassing it

[rampitec]

In D30442#688792, @arsenm wrote:

How expensive is it to do this? The scheduler is already frequently the most expensive pass after RA, sometimes surpassing it

The algorithm is not very expensive itself. Liveins scanned for all defined registers once every region. This is probably most expensive part if there are a lot of registers. The main scan after that only touches those registers which are alive in the region. This is obviously more expensive than without it, but not terribly expensive given ready LIS we already have.

That would be possible to preserve LiveRegs between regions, but main scheduler loop can skip some regions.

lib/Target/AMDGPU/GCNSchedStrategy.cpp

62 ↗

(On Diff #89972)

It does when it calls getRegPressureSetLimit(). However if we are not limited or cannot keep within guessed optimistic limits we override TargetOccupancy and reschedule.

llvm/trunk/lib/Target/AMDGPU/GCNSchedStrategy.h

 
 #ifndef LLVM_LIB_TARGET_AMDGPU_GCNSCHEDSTRATEGY_H
 #define LLVM_LIB_TARGET_AMDGPU_GCNSCHEDSTRATEGY_H
 
 #include "llvm/CodeGen/MachineScheduler.h"
 
 namespace llvm {
 
+class SIMachineFunctionInfo;
 class SIRegisterInfo;
+class SISubtarget;
 
 /// This is a minimal scheduler strategy.  The main difference between this
 /// and the GenericScheduler is that GCNSchedStrategy uses different
 /// heuristics to determine excess/critical pressure sets.  Its goal is to
 /// maximize kernel occupancy (i.e. maximum number of waves per simd).
 class GCNMaxOccupancySchedStrategy : public GenericScheduler {
   friend class GCNScheduleDAGMILive;
 
   SUnit *pickNodeBidirectional(bool &IsTopNode);
 
   void pickNodeFromQueue(SchedBoundary &Zone, const CandPolicy &ZonePolicy,
                          const RegPressureTracker &RPTracker,
                          SchedCandidate &Cand);
 
   void initCandidate(SchedCandidate &Cand, SUnit *SU,
                      bool AtTop, const RegPressureTracker &RPTracker,
                      const SIRegisterInfo *SRI,
                      unsigned SGPRPressure, unsigned VGPRPressure);
 
   unsigned SGPRExcessLimit;
   unsigned VGPRExcessLimit;
   unsigned SGPRCriticalLimit;
   unsigned VGPRCriticalLimit;
 
+  unsigned TargetOccupancy;
+
+  MachineFunction *MF;
+
 public:
   GCNMaxOccupancySchedStrategy(const MachineSchedContext *C);
 
   SUnit *pickNode(bool &IsTopNode) override;
 
   void initialize(ScheduleDAGMI *DAG) override;
 };
 
 class GCNScheduleDAGMILive : public ScheduleDAGMILive {
 
+  const SISubtarget &ST;
+
+  const SIMachineFunctionInfo &MFI;
+
+  // Occupancy target at the begining of function scheduling cycle.
+  unsigned StartingOccupancy;
+
+  // Minimal real occupancy recorder for the function.
+  unsigned MinOccupancy;
+
+  // Scheduling stage number.
+  unsigned Stage;
+
+  // Vecor of regions recorder for later rescheduling
+  SmallVector<std::pair<const MachineBasicBlock::iterator,
+                        const MachineBasicBlock::iterator>, 32> Regions;
+
   // Region live-ins.
   DenseMap<unsigned, LaneBitmask> LiveIns;
 
   // Number of live-ins to the current region, first SGPR then VGPR.
   std::pair<unsigned, unsigned> LiveInPressure;
 
   // Collect current region live-ins.
   void discoverLiveIns();
 
   // Return current region pressure. First value is SGPR number, second is VGPR.
   std::pair<unsigned, unsigned> getRealRegPressure() const;
 
 public:
   GCNScheduleDAGMILive(MachineSchedContext *C,
-                       std::unique_ptr<MachineSchedStrategy> S) :
-    ScheduleDAGMILive(C, std::move(S)) {}
+                       std::unique_ptr<MachineSchedStrategy> S);
+
+  void enterRegion(MachineBasicBlock *bb,
+                   MachineBasicBlock::iterator begin,
+                   MachineBasicBlock::iterator end,
+                   unsigned regioninstrs) override;
 
   void schedule() override;
 
   void finalizeSchedule() override;
 };
 
 } // End namespace llvm
 
 #endif // GCNSCHEDSTRATEGY_H

llvm/trunk/lib/Target/AMDGPU/GCNSchedStrategy.cpp

 #include "llvm/Support/MathExtras.h"
 
 #define DEBUG_TYPE "misched"
 
 using namespace llvm;
 
 GCNMaxOccupancySchedStrategy::GCNMaxOccupancySchedStrategy(
     const MachineSchedContext *C) :
-    GenericScheduler(C) { }
+    GenericScheduler(C), TargetOccupancy(0), MF(nullptr) { }
 
 static unsigned getMaxWaves(unsigned SGPRs, unsigned VGPRs,
                             const MachineFunction &MF) {
 
   const SISubtarget &ST = MF.getSubtarget<SISubtarget>();
   const SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
   unsigned MinRegOccupancy = std::min(ST.getOccupancyWithNumSGPRs(SGPRs),
                                       ST.getOccupancyWithNumVGPRs(VGPRs));
   return std::min(MinRegOccupancy,
                   ST.getOccupancyWithLocalMemSize(MFI->getLDSSize(),
                                                   *MF.getFunction()));
 }
 
 void GCNMaxOccupancySchedStrategy::initialize(ScheduleDAGMI *DAG) {
   GenericScheduler::initialize(DAG);
 
   const SIRegisterInfo *SRI = static_cast<const SIRegisterInfo*>(TRI);
 
+  if (MF != &DAG->MF)
+    TargetOccupancy = 0;
+  MF = &DAG->MF;
+
+  const SISubtarget &ST = MF->getSubtarget<SISubtarget>();
+
   // FIXME: This is also necessary, because some passes that run after
   // scheduling and before regalloc increase register pressure.
   const int ErrorMargin = 3;
 
   SGPRExcessLimit = Context->RegClassInfo
     ->getNumAllocatableRegs(&AMDGPU::SGPR_32RegClass) - ErrorMargin;
   VGPRExcessLimit = Context->RegClassInfo
     ->getNumAllocatableRegs(&AMDGPU::VGPR_32RegClass) - ErrorMargin;
+  if (TargetOccupancy) {
+    SGPRCriticalLimit = ST.getMaxNumSGPRs(TargetOccupancy, true);
+    VGPRCriticalLimit = ST.getMaxNumVGPRs(TargetOccupancy);
+  } else {
     SGPRCriticalLimit = SRI->getRegPressureSetLimit(DAG->MF,
-                        SRI->getSGPRPressureSet()) - ErrorMargin;
+                                                    SRI->getSGPRPressureSet());
     VGPRCriticalLimit = SRI->getRegPressureSetLimit(DAG->MF,
-                        SRI->getVGPRPressureSet()) - ErrorMargin;
+                                                    SRI->getVGPRPressureSet());
+  }
+
+  SGPRCriticalLimit -= ErrorMargin;
+  VGPRCriticalLimit -= ErrorMargin;
 }
 
 void GCNMaxOccupancySchedStrategy::initCandidate(SchedCandidate &Cand, SUnit *SU,
                                      bool AtTop, const RegPressureTracker &RPTracker,
                                      const SIRegisterInfo *SRI,
                                      unsigned SGPRPressure,
                                      unsigned VGPRPressure) {
 
if (SU->isTopReady())
     Top.removeReady(SU);
   if (SU->isBottomReady())
     Bot.removeReady(SU);
 
   DEBUG(dbgs() << "Scheduling SU(" << SU->NodeNum << ") " << *SU->getInstr());
   return SU;
 }
 
+GCNScheduleDAGMILive::GCNScheduleDAGMILive(MachineSchedContext *C,
+                        std::unique_ptr<MachineSchedStrategy> S) :
+  ScheduleDAGMILive(C, std::move(S)),
+  ST(MF.getSubtarget<SISubtarget>()),
+  MFI(*MF.getInfo<SIMachineFunctionInfo>()),
+  StartingOccupancy(ST.getOccupancyWithLocalMemSize(MFI.getLDSSize(),
+                                                    *MF.getFunction())),
+  MinOccupancy(StartingOccupancy), Stage(0) {
+
+  DEBUG(dbgs() << "Starting occupancy is " << StartingOccupancy << ".\n");
+}
+
+void GCNScheduleDAGMILive::enterRegion(MachineBasicBlock *bb,
+                                       MachineBasicBlock::iterator begin,
+                                       MachineBasicBlock::iterator end,
+                                       unsigned regioninstrs) {
+  ScheduleDAGMILive::enterRegion(bb, begin, end, regioninstrs);
+
+  if (Stage == 0)
+    Regions.push_back(std::make_pair(begin, end));
+}
+
 void GCNScheduleDAGMILive::schedule() {
   std::vector<MachineInstr*> Unsched;
   Unsched.reserve(NumRegionInstrs);
   for (auto &I : *this)
     Unsched.push_back(&I);
 
   std::pair<unsigned, unsigned> PressureBefore;
   if (LIS) {
void GCNScheduleDAGMILive::schedule() {
   }
   unsigned WavesAfter = getMaxWaves(PressureAfter.first,
                                     PressureAfter.second, MF);
   unsigned WavesBefore = getMaxWaves(PressureBefore.first,
                                       PressureBefore.second, MF);
   DEBUG(dbgs() << "Occupancy before scheduling: " << WavesBefore <<
                   ", after " << WavesAfter << ".\n");
 
+  // We could not keep current target occupancy because of the just scheduled
+  // region. Record new occupancy for next scheduling cycle.
+  unsigned NewOccupancy = std::max(WavesAfter, WavesBefore);
+  if (NewOccupancy < MinOccupancy) {
+    MinOccupancy = NewOccupancy;
+    DEBUG(dbgs() << "Occupancy lowered for the function to "
+                 << MinOccupancy << ".\n");
+  }
+
   if (WavesAfter >= WavesBefore)
     return;
 
   DEBUG(dbgs() << "Attempting to revert scheduling.\n");
   RegionEnd = RegionBegin;
   for (MachineInstr *MI : Unsched) {
     if (MI->getIterator() != RegionEnd) {
       BB->remove(MI);
GCNScheduleDAGMILive::getRealRegPressure() const {
 
   DEBUG(dbgs() << "Real region's register pressure:\nSGPR = " << MaxSGPRs
                << "\nVGPR = " << MaxVGPRs << '\n');
 
   return std::make_pair(MaxSGPRs, MaxVGPRs);
 }
 
 void GCNScheduleDAGMILive::finalizeSchedule() {
+  // Retry function scheduling if we found resulting occupancy and it is
+  // lower than used for first pass scheduling. This will give more freedom
+  // to schedule low register pressure blocks.
+  // Code is partially copied from MachineSchedulerBase::scheduleRegions().
+
+  if (!LIS || StartingOccupancy <= MinOccupancy)
+    return;
+
+  DEBUG(dbgs() << "Retrying function scheduling with lowest recorded occupancy "
+               << MinOccupancy << ".\n");
+
+  Stage++;
+  GCNMaxOccupancySchedStrategy &S = (GCNMaxOccupancySchedStrategy&)*SchedImpl;
+  S.TargetOccupancy = MinOccupancy;
+
+  MachineBasicBlock *MBB = nullptr;
+  for (auto Region : Regions) {
+    RegionBegin = Region.first;
+    RegionEnd = Region.second;
+
+    if (RegionBegin->getParent() != MBB) {
+      if (MBB) finishBlock();
+      MBB = RegionBegin->getParent();
+      startBlock(MBB);
+    }
+
+    unsigned NumRegionInstrs = std::distance(begin(), end());
+    enterRegion(MBB, begin(), end(), NumRegionInstrs);
+
+    // Skip empty scheduling regions (0 or 1 schedulable instructions).
+    if (begin() == end() || begin() == std::prev(end())) {
+      exitRegion();
+      continue;
+    }
+    DEBUG(dbgs() << "********** MI Scheduling **********\n");
+    DEBUG(dbgs() << MF.getName()
+          << ":BB#" << MBB->getNumber() << " " << MBB->getName()
+          << "\n  From: " << *begin() << "    To: ";
+          if (RegionEnd != MBB->end()) dbgs() << *RegionEnd;
+          else dbgs() << "End";
+          dbgs() << " RegionInstrs: " << NumRegionInstrs << '\n');
+
+    schedule();
+
+    exitRegion();
+  }
+  finishBlock();
   LiveIns.shrink_and_clear();
 }