[AMDGPU] New method to estimate register pressure

[rampitec]

Summary

This change introduces new method to estimate register pressure in
GCNScheduler. Standard RPTracker gives huge error due to the following
reasons:

1. It does not account for live-ins or live-outs if value is not used

in the region itself. That creates a huge error in a very common case
if there are a lot of live-thu registers.

2. It does not properly count subregs.
3. It assumes a register used as an input operand can be reused as an

output. This is not always possible by itself, this is not what RA
will finally do in many cases for various reasons not limited to RA's
inability to do so, and this is not so if the value is actually a
live-thu.

In addition we can now see clear separation between live-in pressure
which we cannot change with the scheduling and tentative pressure
which we can change.

[vpykhtin]

lib/Target/AMDGPU/GCNSchedStrategy.cpp

400

Should this specify slot type explicitly? Something like LIS->getInstructionIndex(*begin()).getBaseIndex() ? I just don't know which slot will be returned by default. For example if it return dead slot you will get ranges that actually live after the instruction.

419

When a map is indexed for non-existent key the value constructed with default constructor is inserted for this key. LaneBitmask::getNone ()is just a syntax sugar for LaneBitmask(). So you can just read LiveIns[Reg] here without initializing it.

476

Again this is redundant. If there is no such key in the map the LiveRegs[Reg] will return LaneBitmask() on read.

477

Please don't copy/paste map indexing. You can't assume compiler will always optimize this and the operation involves hash calculation on the key and array indexing (in the best case). I don't encorage such copy/paste indexing even for ordinary arrays.

[rampitec]

lib/Target/AMDGPU/GCNSchedStrategy.cpp

419

LaneBitmask does not have default initializer value. I guess it worth nothing that LaneBitmask() should result in zero initialization.

477

OK, I can get references. That will be more verbose but less work for compiler to optimize.

[vpykhtin]

lib/Target/AMDGPU/GCNSchedStrategy.cpp

419

It has, though using C++11 syntax:

LineBitmask.h: line 77

private:
  Type Mask = 0;
};

[rampitec]

lib/Target/AMDGPU/GCNSchedStrategy.cpp

419

Oh, thanks!

[rampitec]

lib/Target/AMDGPU/GCNSchedStrategy.cpp

400

It actually returns first MI in a bundle if any:

Mi2IndexMap::const_iterator itr = mi2iMap.find(getBundleStart(MI));

I guess that is what I'm looking here for.

[tstellar]

If you aren't going to use any of the RegisterPressure information, then you should disable register pressure tracking to improve compile time performance. Also, I think it would be nice if the queries could be abstracted so it was easy to switch between the standard ReigsterPressureTracker and the improved version you have implemented.

[rampitec]

In D30439#688636, @tstellar wrote:

If you aren't going to use any of the RegisterPressure information, then you should disable register pressure tracking to improve compile time performance. Also, I think it would be nice if the queries could be abstracted so it was easy to switch between the standard ReigsterPressureTracker and the improved version you have implemented.

It is not a tracker, just the estimation. It least not just yet. During scheduling standard RPTracker still used for pressure diffs.

[vpykhtin]

lib/Target/AMDGPU/GCNSchedStrategy.cpp

400

Right, but what is a slot type for this first bundle instruction?

[rampitec]

Removed redundant map references.
Added .getBaseIndex() for slot query.

[rampitec]

Added full context diff.

[vpykhtin]

LGTM.

lib/Target/AMDGPU/GCNSchedStrategy.h

public:
   GCNMaxOccupancySchedStrategy(const MachineSchedContext *C);
 
   SUnit *pickNode(bool &IsTopNode) override;
 
   void initialize(ScheduleDAGMI *DAG) override;
 };
 
 class GCNScheduleDAGMILive : public ScheduleDAGMILive {
+
+  // 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)) {}
 
   void schedule() override;
+
+  void finalizeSchedule() override;
 };
 
 } // End namespace llvm
 
 #endif // GCNSCHEDSTRATEGY_H

lib/Target/AMDGPU/GCNSchedStrategy.cpp

 //===----------------------------------------------------------------------===//
 
 #include "GCNSchedStrategy.h"
 #include "AMDGPUSubtarget.h"
 #include "SIInstrInfo.h"
 #include "SIMachineFunctionInfo.h"
 #include "SIRegisterInfo.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
+#include "llvm/Support/MathExtras.h"
 
 #define DEBUG_TYPE "misched"
 
 using namespace llvm;
 
 GCNMaxOccupancySchedStrategy::GCNMaxOccupancySchedStrategy(
     const MachineSchedContext *C) :
     GenericScheduler(C) { }
SUnit *GCNMaxOccupancySchedStrategy::pickNode(bool &IsTopNode) {
   if (SU->isBottomReady())
     Bot.removeReady(SU);
 
   DEBUG(dbgs() << "Scheduling SU(" << SU->NodeNum << ") " << *SU->getInstr());
   return SU;
 }
 
 void GCNScheduleDAGMILive::schedule() {
-  const SIRegisterInfo *SRI = static_cast<const SIRegisterInfo*>(TRI);
-
   std::vector<MachineInstr*> Unsched;
   Unsched.reserve(NumRegionInstrs);
   for (auto &I : *this)
     Unsched.push_back(&I);
 
+  std::pair<unsigned, unsigned> PressureBefore;
+  if (LIS) {
+    DEBUG(dbgs() << "Pressure before scheduling:\n");
+    discoverLiveIns();
+    PressureBefore = getRealRegPressure();
+  }
+
   ScheduleDAGMILive::schedule();
+  if (!LIS)
+    return;
 
   // Check the results of scheduling.
   GCNMaxOccupancySchedStrategy &S = (GCNMaxOccupancySchedStrategy&)*SchedImpl;
-  std::vector<unsigned> UnschedPressure = getRegPressure().MaxSetPressure;
-  unsigned MaxSGPRs = std::max(
-    getTopRPTracker().getPressure().MaxSetPressure[SRI->getSGPRPressureSet()],
-    getBotRPTracker().getPressure().MaxSetPressure[SRI->getSGPRPressureSet()]);
-  unsigned MaxVGPRs = std::max(
-    getTopRPTracker().getPressure().MaxSetPressure[SRI->getVGPRPressureSet()],
-    getBotRPTracker().getPressure().MaxSetPressure[SRI->getVGPRPressureSet()]);
-  DEBUG(dbgs() << "Pressure after scheduling:\nSGPR = " << MaxSGPRs
-               << "\nVGPR = " << MaxVGPRs << '\n');
-  if (MaxSGPRs <= S.SGPRCriticalLimit &&
-      MaxVGPRs <= S.VGPRCriticalLimit) {
+  DEBUG(dbgs() << "Pressure after scheduling:\n");
+  auto PressureAfter = getRealRegPressure();
+  LiveIns.clear();
+
+  if (PressureAfter.first <= S.SGPRCriticalLimit &&
+      PressureAfter.second <= S.VGPRCriticalLimit) {
     DEBUG(dbgs() << "Pressure in desired limits, done.\n");
     return;
   }
-  unsigned WavesAfter = getMaxWaves(MaxSGPRs, MaxVGPRs, MF);
-  unsigned WavesUnsched = getMaxWaves(UnschedPressure[SRI->getSGPRPressureSet()],
-                            UnschedPressure[SRI->getVGPRPressureSet()], MF);
-  DEBUG(dbgs() << "Occupancy before scheduling: " << WavesUnsched <<
+  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");
-  if (WavesAfter >= WavesUnsched)
+
+  if (WavesAfter >= WavesBefore)
     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);
-      if (LIS)
-        LIS->handleMove(*MI, true);
     }
     // Reset read-undef flags and update them later.
     for (auto &Op : MI->operands())
       if (Op.isReg() && Op.isDef())
         Op.setIsUndef(false);
     RegisterOperands RegOpers;
     RegOpers.collect(*MI, *TRI, MRI, ShouldTrackLaneMasks, false);
     if (ShouldTrackLaneMasks) {
       // Adjust liveness and add missing dead+read-undef flags.
       SlotIndex SlotIdx = LIS->getInstructionIndex(*MI).getRegSlot();
       RegOpers.adjustLaneLiveness(*LIS, MRI, SlotIdx, MI);
     } else {
       // Adjust for missing dead-def flags.
       RegOpers.detectDeadDefs(*MI, *LIS);
     }
     RegionEnd = MI->getIterator();
     ++RegionEnd;
     DEBUG(dbgs() << "Scheduling " << *MI);
   }
   RegionBegin = Unsched.front()->getIterator();
 
   placeDebugValues();
 }
+
+static inline void setMask(const MachineRegisterInfo &MRI,
+                           const SIRegisterInfo *SRI, unsigned Reg,
+                           LaneBitmask &PrevMask, LaneBitmask NewMask,
+                           unsigned &SGPRs, unsigned &VGPRs) {
+  int NewRegs = countPopulation(NewMask.getAsInteger()) -
+                countPopulation(PrevMask.getAsInteger());
+  if (SRI->isSGPRReg(MRI, Reg))
+    SGPRs += NewRegs;
+  if (SRI->isVGPR(MRI, Reg))
+    VGPRs += NewRegs;
+  assert ((int)SGPRs >= 0 && (int)VGPRs >= 0);
+  PrevMask = NewMask;
+}
+
+void GCNScheduleDAGMILive::discoverLiveIns() {
+  unsigned SGPRs = 0;
+  unsigned VGPRs = 0;
+
+  const SIRegisterInfo *SRI = static_cast<const SIRegisterInfo*>(TRI);
+  SlotIndex SI = LIS->getInstructionIndex(*begin()).getBaseIndex();
+  assert (SI.isValid());
+
+  DEBUG(dbgs() << "Region live-ins:");
+  for (unsigned I = 0, E = MRI.getNumVirtRegs(); I != E; ++I) {
+    unsigned Reg = TargetRegisterInfo::index2VirtReg(I);
+    if (MRI.reg_nodbg_empty(Reg))
+      continue;
+    const LiveInterval &LI = LIS->getInterval(Reg);
+    LaneBitmask LaneMask = LaneBitmask::getNone();
+    if (LI.hasSubRanges()) {
+      for (const auto &S : LI.subranges())
+        if (S.liveAt(SI))
+          LaneMask |= S.LaneMask;
+    } else if (LI.liveAt(SI)) {
+      LaneMask = MRI.getMaxLaneMaskForVReg(Reg);
+    }
+
+    if (LaneMask.any()) {
+      setMask(MRI, SRI, Reg, LiveIns[Reg], LaneMask, SGPRs, VGPRs);
+
+      DEBUG(dbgs() << ' ' << PrintVRegOrUnit(Reg, SRI) << ':'
+                   << PrintLaneMask(LiveIns[Reg]));
+    }
+  }
+
+  LiveInPressure = std::make_pair(SGPRs, VGPRs);
+
+  DEBUG(dbgs() << "\nLive-in pressure:\nSGPR = " << SGPRs
+               << "\nVGPR = " << VGPRs << '\n');
+}
+
+std::pair<unsigned, unsigned>
+GCNScheduleDAGMILive::getRealRegPressure() const {
+  unsigned SGPRs, MaxSGPRs, VGPRs, MaxVGPRs;
+  SGPRs = MaxSGPRs = LiveInPressure.first;
+  VGPRs = MaxVGPRs = LiveInPressure.second;
+
+  const SIRegisterInfo *SRI = static_cast<const SIRegisterInfo*>(TRI);
+  DenseMap<unsigned, LaneBitmask> LiveRegs(LiveIns);
+
+  for (const MachineInstr &MI : *this) {
+    if (MI.isDebugValue())
+      continue;
+    SlotIndex SI = LIS->getInstructionIndex(MI).getBaseIndex();
+    assert (SI.isValid());
+
+    // Remove dead registers or mask bits.
+    for (auto &It : LiveRegs) {
+      if (It.second.none())
+        continue;
+      const LiveInterval &LI = LIS->getInterval(It.first);
+      if (LI.hasSubRanges()) {
+        for (const auto &S : LI.subranges())
+          if (!S.liveAt(SI))
+            setMask(MRI, SRI, It.first, It.second, It.second & ~S.LaneMask,
+                    SGPRs, VGPRs);
+      } else if (!LI.liveAt(SI)) {
+        setMask(MRI, SRI, It.first, It.second, LaneBitmask::getNone(),
+                SGPRs, VGPRs);
+      }
+    }
+
+    // Add new registers or mask bits.
+    for (const auto &MO : MI.defs()) {
+      if (!MO.isReg())
+        continue;
+      unsigned Reg = MO.getReg();
+      if (!TargetRegisterInfo::isVirtualRegister(Reg))
+        continue;
+      unsigned SubRegIdx = MO.getSubReg();
+      LaneBitmask LaneMask = SubRegIdx != 0
+                             ? TRI->getSubRegIndexLaneMask(SubRegIdx)
+                             : MRI.getMaxLaneMaskForVReg(Reg);
+      LaneBitmask &LM = LiveRegs[Reg];
+      setMask(MRI, SRI, Reg, LM, LM | LaneMask, SGPRs, VGPRs);
+    }
+    MaxSGPRs = std::max(MaxSGPRs, SGPRs);
+    MaxVGPRs = std::max(MaxVGPRs, VGPRs);
+  }
+
+  DEBUG(dbgs() << "Real region's register pressure:\nSGPR = " << MaxSGPRs
+               << "\nVGPR = " << MaxVGPRs << '\n');
+
+  return std::make_pair(MaxSGPRs, MaxVGPRs);
+}
+
+void GCNScheduleDAGMILive::finalizeSchedule() {
+  LiveIns.shrink_and_clear();
+}