[AMDGPU] Cache live-ins and register pressure in scheduler

[rampitec]

Summary

Using LIS can be quite expensive, so caching of calculated region
live-ins and pressure is implemented. It does two things:

1. Caches the info for the second stage when we schedule with decreased target occupancy.
2. Tracks the basic block from top to bottom thus eliminating the need to scan whole register file liveness at every region split in the middle of the block.

The scheduling is now done in 3 stages instead of two, with the first
one being really a no-op and only used to collect scheduling regions
as sent by the scheduler driver.

There is no functional change to the current behavior, only compilation
speed is affected. In general computeBlockPressure() could be simplified
if we switch to backward RP tracker, because scheduler sends regions
within a block starting from the last upward. We could use a natural
order of upward tracker to seamlessly change between regions of the same
block, since live reg set of a previous tracked region would become a
live-out of the next region. That however requires fixing upward tracker
to properly account defs and uses of the same instruction as both are
contributing to the current pressure. When we converge on the produced
pressure we should be able to switch between them back and forth. In
addition, backward tracker is less expensive as it uses LIS in recede
less often than forward uses it in advance.

At the moment the worst known case compilation time has improved from 26
minutes to 8.5.

[rampitec]

Added second very small cache to reuse calculated live set of a block as its only successor live-in set.
Simplified the loop inside computeBlockPressure() to catch end() of the block without code duplication.
Adjusted initial block iterator to the first instruction of the first (topologically) region to save few iterations in some blocks.
Diff is now created against it parent review D33105 to show only the real changes.

[rampitec]

Skip debug values when crossing basic block boundary.

[rampitec]

Switch to statefull tracker.

[vpykhtin]

LGTM.

lib/Target/AMDGPU/GCNSchedStrategy.cpp

340 ↗

(On Diff #99097)

Turn this into a helper debug func?

[rampitec]

lib/Target/AMDGPU/GCNSchedStrategy.cpp

340 ↗

(On Diff #99097)

I want to switch to your print() function from tracker in a short while. It is now this way to compare logs to old implementation.

[rampitec]

lib/Target/AMDGPU/GCNSchedStrategy.cpp

340 ↗

(On Diff #99097)

Actually printLivesAt() is a heavy one, recomputing live set. How about factoring this into GCNRPTracker::printLiveRegs()?

[vpykhtin]

lib/Target/AMDGPU/GCNSchedStrategy.cpp

340 ↗

(On Diff #99097)

Later commit :)

[rampitec]

Rebase to master.

[rampitec]

lib/Target/AMDGPU/GCNSchedStrategy.cpp

340 ↗

(On Diff #99097)

D33244

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

class GCNScheduleDAGMILive : public ScheduleDAGMILive {
   unsigned StartingOccupancy;
 
   // Minimal real occupancy recorder for the function.
   unsigned MinOccupancy;
 
   // Scheduling stage number.
   unsigned Stage;
 
+  // Current region index.
+  size_t RegionIdx;
+
   // Vecor of regions recorder for later rescheduling
   SmallVector<std::pair<MachineBasicBlock::iterator,
                         MachineBasicBlock::iterator>, 32> Regions;
 
-  // Region live-ins.
-  GCNRPTracker::LiveRegSet LiveIns;
+  // Region live-in cache.
+  SmallVector<GCNRPTracker::LiveRegSet, 32> LiveIns;
+
+  // Region pressure cache.
+  SmallVector<GCNRegPressure, 32> Pressure;
 
-  // Collect current region live-ins.
-  void discoverLiveIns();
+  // Temporary basic block live-in cache.
+  DenseMap<const MachineBasicBlock*, GCNRPTracker::LiveRegSet> MBBLiveIns;
 
   // Return current region pressure.
   GCNRegPressure getRealRegPressure() const;
 
+  // Compute and cache live-ins and pressure for all regions in block.
+  void computeBlockPressure(const MachineBasicBlock *MBB);
+
+
 public:
   GCNScheduleDAGMILive(MachineSchedContext *C,
                        std::unique_ptr<MachineSchedStrategy> S);
 
   void schedule() override;
 
   void finalizeSchedule() override;
 };
 
 } // End namespace llvm
 
 #endif // GCNSCHEDSTRATEGY_H

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

 
 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) {
+  MinOccupancy(StartingOccupancy), Stage(0), RegionIdx(0) {
 
   DEBUG(dbgs() << "Starting occupancy is " << StartingOccupancy << ".\n");
 }
 
 void GCNScheduleDAGMILive::schedule() {
+  if (Stage == 0) {
+    // Just record regions at the first pass.
+    Regions.push_back(std::make_pair(RegionBegin, RegionEnd));
+    return;
+  }
+
   std::vector<MachineInstr*> Unsched;
   Unsched.reserve(NumRegionInstrs);
   for (auto &I : *this)
     Unsched.push_back(&I);
 
   GCNRegPressure PressureBefore;
   if (LIS) {
-    discoverLiveIns();
-    PressureBefore = getRealRegPressure();
+    PressureBefore = Pressure[RegionIdx];
 
-    DEBUG(dbgs() << "Pressure before scheduling:\nSGPR = "
+    DEBUG(const SIRegisterInfo *SRI = static_cast<const SIRegisterInfo*>(TRI);
+          dbgs() << "Pressure before scheduling:\nRegion live-ins:";
+          for (unsigned I = 0, E = MRI.getNumVirtRegs(); I != E; ++I) {
+            unsigned Reg = TargetRegisterInfo::index2VirtReg(I);
+            auto It = LiveIns[RegionIdx].find(Reg);
+            if (It != LiveIns[RegionIdx].end() && It->second.any())
+              dbgs() << ' ' << PrintVRegOrUnit(Reg, SRI) << ':'
+                     << PrintLaneMask(It->second);
+          }
+          auto P = llvm::getRegPressure(MRI, LiveIns[RegionIdx]);
+          dbgs() << "\nLive-in pressure:\nSGPR = " << P.getSGPRNum()
+                 << "\nVGPR = " << P.getVGPRNum()
+                 << "\nReal region's register pressure:\nSGPR = "
                  << PressureBefore.getSGPRNum()
                  << "\nVGPR = " << PressureBefore.getVGPRNum() << '\n');
-
   }
 
   ScheduleDAGMILive::schedule();
-  if (Stage == 0)
-    Regions.push_back(std::make_pair(RegionBegin, RegionEnd));
+  Regions[RegionIdx] = std::make_pair(RegionBegin, RegionEnd);
 
   if (!LIS)
     return;
 
   // Check the results of scheduling.
   GCNMaxOccupancySchedStrategy &S = (GCNMaxOccupancySchedStrategy&)*SchedImpl;
   auto PressureAfter = getRealRegPressure();
 
   DEBUG(dbgs() << "Pressure after scheduling:\nSGPR = "
                << PressureAfter.getSGPRNum()
                << "\nVGPR = " << PressureAfter.getVGPRNum() << '\n');
 
-  LiveIns.clear();
-
   if (PressureAfter.getSGPRNum() <= S.SGPRCriticalLimit &&
       PressureAfter.getVGPRNum() <= S.VGPRCriticalLimit) {
+    Pressure[RegionIdx] = PressureAfter;
     DEBUG(dbgs() << "Pressure in desired limits, done.\n");
     return;
   }
   unsigned WavesAfter = getMaxWaves(PressureAfter.getSGPRNum(),
                                     PressureAfter.getVGPRNum(), MF);
   unsigned WavesBefore = getMaxWaves(PressureBefore.getSGPRNum(),
                                      PressureBefore.getVGPRNum(), 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)
+  if (WavesAfter >= WavesBefore) {
+    Pressure[RegionIdx] = PressureAfter;
     return;
+  }
 
   DEBUG(dbgs() << "Attempting to revert scheduling.\n");
   RegionEnd = RegionBegin;
   for (MachineInstr *MI : Unsched) {
     if (MI->getIterator() != RegionEnd) {
       BB->remove(MI);
       BB->insert(RegionEnd, MI);
       LIS->handleMove(*MI, true);
if (ShouldTrackLaneMasks) {
       // Adjust for missing dead-def flags.
       RegOpers.detectDeadDefs(*MI, *LIS);
     }
     RegionEnd = MI->getIterator();
     ++RegionEnd;
     DEBUG(dbgs() << "Scheduling " << *MI);
   }
   RegionBegin = Unsched.front()->getIterator();
-  if (Stage == 0)
-    Regions.back() = std::make_pair(RegionBegin, RegionEnd);
+  Regions[RegionIdx] = std::make_pair(RegionBegin, RegionEnd);
 
   placeDebugValues();
 }
 
-void GCNScheduleDAGMILive::discoverLiveIns() {
+GCNRegPressure GCNScheduleDAGMILive::getRealRegPressure() const {
   GCNDownwardRPTracker RPTracker(*LIS);
-  RPTracker.reset(*begin());
+  RPTracker.advance(begin(), end(), &LiveIns[RegionIdx]);
+  return RPTracker.moveMaxPressure();
+}
 
-  LiveIns = RPTracker.moveLiveRegs();
+void GCNScheduleDAGMILive::computeBlockPressure(const MachineBasicBlock *MBB) {
+  GCNDownwardRPTracker RPTracker(*LIS);
 
-  DEBUG(GCNRegPressure LiveInPressure = RPTracker.moveMaxPressure();
-        const SIRegisterInfo *SRI = static_cast<const SIRegisterInfo*>(TRI);
-        dbgs() << "Region live-ins:";
-        for (unsigned I = 0, E = MRI.getNumVirtRegs(); I != E; ++I) {
-          unsigned Reg = TargetRegisterInfo::index2VirtReg(I);
-          auto It = LiveIns.find(Reg);
-          if (It != LiveIns.end())
-            dbgs() << ' ' << PrintVRegOrUnit(Reg, SRI) << ':'
-                   << PrintLaneMask(It->second);
-        }
-        dbgs() << "\nLive-in pressure:\nSGPR = "
-               << LiveInPressure.getSGPRNum()
-               << "\nVGPR = " << LiveInPressure.getVGPRNum() << '\n');
+  // If the block has the only successor then live-ins of that successor are
+  // live-outs of the current block. We can reuse calculated live set if the
+  // successor will be sent to scheduling past current block.
+  const MachineBasicBlock *OnlySucc = nullptr;
+  if (MBB->succ_size() == 1 && !(*MBB->succ_begin())->empty()) {
+    SlotIndexes *Ind = LIS->getSlotIndexes();
+    if (Ind->getMBBStartIdx(MBB) < Ind->getMBBStartIdx(*MBB->succ_begin()))
+      OnlySucc = *MBB->succ_begin();
+  }
+
+  // Scheduler sends regions from the end of the block upwards.
+  size_t CurRegion = RegionIdx;
+  for (size_t E = Regions.size(); CurRegion != E; ++CurRegion)
+    if (Regions[CurRegion].first->getParent() != MBB)
+      break;
+  --CurRegion;
+
+  auto I = MBB->begin();
+  auto LiveInIt = MBBLiveIns.find(MBB);
+  if (LiveInIt != MBBLiveIns.end()) {
+    auto LiveIn = std::move(LiveInIt->second);
+    RPTracker.reset(*MBB->begin(), &LiveIn);
+    MBBLiveIns.erase(LiveInIt);
+  } else {
+    I = Regions[CurRegion].first;
+    RPTracker.reset(*I);
+
+  for ( ; ; ) {
+    I = RPTracker.getNext();
+
+    if (Regions[CurRegion].first == I) {
+      LiveIns[CurRegion] = RPTracker.getLiveRegs();
+      RPTracker.clearMaxPressure();
     }
 
-GCNRegPressure GCNScheduleDAGMILive::getRealRegPressure() const {
-  GCNDownwardRPTracker RPTracker(*LIS);
-  RPTracker.advance(begin(), end(), &LiveIns);
-  return RPTracker.moveMaxPressure();
+    if (Regions[CurRegion].second == I) {
+      Pressure[CurRegion] = RPTracker.moveMaxPressure();
+      if (CurRegion-- == RegionIdx)
+        break;
+    }
+    RPTracker.advanceToNext();
+    RPTracker.advanceBeforeNext();
+  }
+
+  if (OnlySucc) {
+    if (I != MBB->end()) {
+      RPTracker.advanceToNext();
+      RPTracker.advance(MBB->end());
+    }
+    RPTracker.reset(*OnlySucc->begin(), &RPTracker.getLiveRegs());
+    RPTracker.advanceBeforeNext();
+    MBBLiveIns[OnlySucc] = RPTracker.moveLiveRegs();
+  }
 }
 
 void GCNScheduleDAGMILive::finalizeSchedule() {
+  GCNMaxOccupancySchedStrategy &S = (GCNMaxOccupancySchedStrategy&)*SchedImpl;
+  DEBUG(dbgs() << "All regions recorded, starting actual scheduling.\n");
+
+  LiveIns.resize(Regions.size());
+  Pressure.resize(Regions.size());
+
+  do {
+    Stage++;
+    RegionIdx = 0;
+    MachineBasicBlock *MBB = nullptr;
+
+    if (Stage > 1) {
       // 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;
+        break;
 
-  DEBUG(dbgs() << "Retrying function scheduling with lowest recorded occupancy "
-               << MinOccupancy << ".\n");
+      DEBUG(dbgs()
+              << "Retrying function scheduling with lowest recorded occupancy "
 
-  Stage++;
-  GCNMaxOccupancySchedStrategy &S = (GCNMaxOccupancySchedStrategy&)*SchedImpl;
       S.setTargetOccupancy(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);
+        if (Stage == 1)
+          computeBlockPressure(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();
+      ++RegionIdx;
     }
     finishBlock();
-  LiveIns.shrink_and_clear();
+
+  } while (Stage < 2);
 }