Correct register pressure calculation in presence of subregs

If a subreg is used in an instruction it counts as a whole superreg
for the purpose of register pressure calculation. This patch corrects
improper register pressure calculation by examining operand's lane mask.

Differential Revision: https://reviews.llvm.org/D29835

llvm-svn: 296009

5 files changed, 62 insertions, 13 deletions

diff --git a/llvm/lib/CodeGen/MachineScheduler.cpp b/llvm/lib/CodeGen/MachineScheduler.cpp
index 79b02c7de8c..066398358e8 100644
--- a/llvm/lib/CodeGen/MachineScheduler.cpp
+++ b/llvm/lib/CodeGen/MachineScheduler.cpp
@@ -1085,7 +1085,7 @@ void ScheduleDAGMILive::updatePressureDiffs(
           continue;
 
         PressureDiff &PDiff = getPressureDiff(&SU);
-        PDiff.addPressureChange(Reg, Decrement, &MRI);
+        PDiff.addPressureChange(P, Decrement, &MRI);
         DEBUG(
           dbgs() << "  UpdateRegP: SU(" << SU.NodeNum << ") "
                  << PrintReg(Reg, TRI) << ':' << PrintLaneMask(P.LaneMask)
@@ -1123,7 +1123,7 @@ void ScheduleDAGMILive::updatePressureDiffs(
               LI.Query(LIS->getInstructionIndex(*SU->getInstr()));
           if (LRQ.valueIn() == VNI) {
             PressureDiff &PDiff = getPressureDiff(SU);
-            PDiff.addPressureChange(Reg, true, &MRI);
+            PDiff.addPressureChange(P, true, &MRI);
             DEBUG(
               dbgs() << "  UpdateRegP: SU(" << SU->NodeNum << ") "
                      << *SU->getInstr();
diff --git a/llvm/lib/CodeGen/RegisterPressure.cpp b/llvm/lib/CodeGen/RegisterPressure.cpp
index c726edc88b4..b22b36a3607 100644
--- a/llvm/lib/CodeGen/RegisterPressure.cpp
+++ b/llvm/lib/CodeGen/RegisterPressure.cpp
@@ -46,16 +46,29 @@
 
 using namespace llvm;
 
+/// Clamp lane masks to maximum posible value.
+static void clampMasks(const MachineRegisterInfo &MRI, unsigned Reg,
+                       LaneBitmask& LaneMask1, LaneBitmask& LaneMask2) {
+  if (TargetRegisterInfo::isVirtualRegister(Reg)) {
+    LaneBitmask Max = MRI.getMaxLaneMaskForVReg(Reg);
+    LaneMask1 &= Max;
+    LaneMask2 &= Max;
+  }
+}
+
 /// Increase pressure for each pressure set provided by TargetRegisterInfo.
 static void increaseSetPressure(std::vector<unsigned> &CurrSetPressure,
                                 const MachineRegisterInfo &MRI, unsigned Reg,
                                 LaneBitmask PrevMask, LaneBitmask NewMask) {
   assert((PrevMask & ~NewMask).none() && "Must not remove bits");
-  if (PrevMask.any() || NewMask.none())
+
+  clampMasks(MRI, Reg, PrevMask, NewMask);
+  if ((NewMask & ~PrevMask).none())
     return;
 
+  const TargetRegisterInfo *TRI = MRI.getTargetRegisterInfo();
+  unsigned Weight = TRI->getRegUnitWeight(MRI, Reg, NewMask & ~PrevMask);
   PSetIterator PSetI = MRI.getPressureSets(Reg);
-  unsigned Weight = PSetI.getWeight();
   for (; PSetI.isValid(); ++PSetI)
     CurrSetPressure[*PSetI] += Weight;
 }
@@ -65,11 +78,13 @@ static void decreaseSetPressure(std::vector<unsigned> &CurrSetPressure,
                                 const MachineRegisterInfo &MRI, unsigned Reg,
                                 LaneBitmask PrevMask, LaneBitmask NewMask) {
   //assert((NewMask & !PrevMask) == 0 && "Must not add bits");
-  if (NewMask.any() || PrevMask.none())
+  clampMasks(MRI, Reg, PrevMask, NewMask);
+  if ((~NewMask & PrevMask).none())
     return;
 
+  const TargetRegisterInfo *TRI = MRI.getTargetRegisterInfo();
+  unsigned Weight = TRI->getRegUnitWeight(MRI, Reg, ~NewMask & PrevMask);
   PSetIterator PSetI = MRI.getPressureSets(Reg);
-  unsigned Weight = PSetI.getWeight();
   for (; PSetI.isValid(); ++PSetI) {
     assert(CurrSetPressure[*PSetI] >= Weight && "register pressure underflow");
     CurrSetPressure[*PSetI] -= Weight;
@@ -139,11 +154,14 @@ void PressureDiff::dump(const TargetRegisterInfo &TRI) const {
 void RegPressureTracker::increaseRegPressure(unsigned RegUnit,
                                              LaneBitmask PreviousMask,
                                              LaneBitmask NewMask) {
-  if (PreviousMask.any() || NewMask.none())
+  clampMasks(*MRI, RegUnit, PreviousMask, NewMask);
+  if ((NewMask & ~PreviousMask).none())
     return;
 
+  const TargetRegisterInfo *TRI = MRI->getTargetRegisterInfo();
+  unsigned Weight = TRI->getRegUnitWeight(*MRI, RegUnit,
+                                          NewMask & ~PreviousMask);
   PSetIterator PSetI = MRI->getPressureSets(RegUnit);
-  unsigned Weight = PSetI.getWeight();
   for (; PSetI.isValid(); ++PSetI) {
     CurrSetPressure[*PSetI] += Weight;
     P.MaxSetPressure[*PSetI] =
@@ -644,17 +662,19 @@ void PressureDiffs::addInstruction(unsigned Idx,
   PressureDiff &PDiff = (*this)[Idx];
   assert(!PDiff.begin()->isValid() && "stale PDiff");
   for (const RegisterMaskPair &P : RegOpers.Defs)
-    PDiff.addPressureChange(P.RegUnit, true, &MRI);
+    PDiff.addPressureChange(P, true, &MRI);
 
   for (const RegisterMaskPair &P : RegOpers.Uses)
-    PDiff.addPressureChange(P.RegUnit, false, &MRI);
+    PDiff.addPressureChange(P, false, &MRI);
 }
 
 /// Add a change in pressure to the pressure diff of a given instruction.
-void PressureDiff::addPressureChange(unsigned RegUnit, bool IsDec,
+void PressureDiff::addPressureChange(RegisterMaskPair P, bool IsDec,
                                      const MachineRegisterInfo *MRI) {
-  PSetIterator PSetI = MRI->getPressureSets(RegUnit);
-  int Weight = IsDec ? -PSetI.getWeight() : PSetI.getWeight();
+  const TargetRegisterInfo *TRI = MRI->getTargetRegisterInfo();
+  int Weight = (int)TRI->getRegUnitWeight(*MRI, P.RegUnit, P.LaneMask);
+  PSetIterator PSetI = MRI->getPressureSets(P.RegUnit);
+  if (IsDec) Weight = -Weight;
   for (; PSetI.isValid(); ++PSetI) {
     // Find an existing entry in the pressure diff for this PSet.
     PressureDiff::iterator I = nonconst_begin(), E = nonconst_end();
diff --git a/llvm/lib/CodeGen/TargetRegisterInfo.cpp b/llvm/lib/CodeGen/TargetRegisterInfo.cpp
index 66cdad278e8..0d3f8eb6fb5 100644
--- a/llvm/lib/CodeGen/TargetRegisterInfo.cpp
+++ b/llvm/lib/CodeGen/TargetRegisterInfo.cpp
@@ -412,6 +412,15 @@ bool TargetRegisterInfo::regmaskSubsetEqual(const uint32_t *mask0,
   return true;
 }
 
+/// Get the weight in units of pressure for a sub register of this register
+/// unit given a lane mask.
+unsigned TargetRegisterInfo::getRegUnitWeight(const MachineRegisterInfo &MRI,
+                                              unsigned RegUnit,
+                                              LaneBitmask LaneMask) const {
+  PSetIterator PSetI = MRI.getPressureSets(RegUnit);
+  return PSetI.getWeight();
+}
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 LLVM_DUMP_METHOD
 void TargetRegisterInfo::dumpReg(unsigned Reg, unsigned SubRegIndex,
diff --git a/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp b/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
index 39324cbbcc0..f1d037d1159 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
@@ -21,6 +21,7 @@
 #include "llvm/CodeGen/RegisterScavenging.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/Support/MathExtras.h"
 
 using namespace llvm;
 
@@ -1408,3 +1409,18 @@ const int *SIRegisterInfo::getRegUnitPressureSets(unsigned RegUnit) const {
     return Empty;
   return AMDGPURegisterInfo::getRegUnitPressureSets(RegUnit);
 }
+
+unsigned SIRegisterInfo::getRegUnitWeight(const MachineRegisterInfo &MRI,
+                                          unsigned RegUnit,
+                                          LaneBitmask LaneMask) const {
+  unsigned Weight = TargetRegisterInfo::getRegUnitWeight(MRI, RegUnit,
+                                                         LaneMask);
+  if (Weight > 1 && LaneMask.any() && !LaneMask.all() &&
+      isVirtualRegister(RegUnit)) {
+    LaneBitmask Max = MRI.getMaxLaneMaskForVReg(RegUnit);
+    if (Max != LaneMask && !Max.all() && !Max.none())
+      Weight = (Weight * countPopulation(LaneMask.getAsInteger())) /
+                         countPopulation(Max.getAsInteger());
+  }
+  return Weight;
+}
diff --git a/llvm/lib/Target/AMDGPU/SIRegisterInfo.h b/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
index 679ed229758..844f0317a1e 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
@@ -228,6 +228,10 @@ public:
 
   const int *getRegUnitPressureSets(unsigned RegUnit) const override;
 
+  unsigned getRegUnitWeight(const MachineRegisterInfo &MRI,
+                            unsigned RegUnit,
+                            LaneBitmask LaneMask) const override;
+
 private:
   void buildSpillLoadStore(MachineBasicBlock::iterator MI,
                            unsigned LoadStoreOp,