Correct register pressure calculation in presence of subregs

[rampitec]

Summary

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.

[arsenm]

test/CodeGen/AMDGPU/schedule-regpressure.mir

62–68 ↗

(On Diff #88015)

Can you add this as a second test rather than modifying the existing test for m0

[rampitec]

Separated test per Matt's request.

[MatzeB]

I did not implement it this way yet because in general a bit in the lane mask doesn't necessarily represent 1 register pressure unit. Without a similar patch to {increase|decrease}RegPressure() pressure book keeping will go out of sync as well. I'd recomment running with -verify-misched when touching the regpressure code.

[rampitec]

Added -verify-misched to the test.

[rampitec]

In D29835#673763, @MatzeB wrote:

I did not implement it this way yet because in general a bit in the lane mask doesn't necessarily represent 1 register pressure unit.

You are right, yet this estimation seems more correct to me. I think if/when we have targets for which this is not true we could extend TRI interface to return the estimation.

Without a similar patch to {increase|decrease}RegPressure() pressure book keeping will go out of sync as well.

You are probably right, I was going to look into these places...

I'd recomment running with -verify-misched when touching the regpressure code.

Thanks, added.

[MatzeB]

In D29835#673808, @rampitec wrote:

In D29835#673763, @MatzeB wrote:

I did not implement it this way yet because in general a bit in the lane mask doesn't necessarily represent 1 register pressure unit.

You are right, yet this estimation seems more correct to me. I think if/when we have targets for which this is not true we could extend TRI interface to return the estimation.

The current code is correct, just not particularily exact.

We cannot apply this without a plan to deal with those cases, even if it is just aborting/falling back to a pessimistic behavior.

Without a similar patch to {increase|decrease}RegPressure() pressure book keeping will go out of sync as well.

You are probably right, I was going to look into these places...

We cannot apply a patch that just changes one part of the system.

I'd recomment running with -verify-misched when touching the regpressure code.

Thanks, added.

That was meant as a recommendation for running experiments (not necessarily for the test here) like llvm-lit -Dllc='llc -verify-misched'. I would expect the scheduler to fail in a number of cases when the pressure diffs and the decrease/increaseRegisterPressure() functions are out of sync.

[rampitec]

In D29835#673827, @MatzeB wrote:

In D29835#673808, @rampitec wrote:

In D29835#673763, @MatzeB wrote:

I did not implement it this way yet because in general a bit in the lane mask doesn't necessarily represent 1 register pressure unit.

You are right, yet this estimation seems more correct to me. I think if/when we have targets for which this is not true we could extend TRI interface to return the estimation.

The current code is correct, just not particularily exact.
We cannot apply this without a plan to deal with those cases, even if it is just aborting/falling back to a pessimistic behavior.

I guess if we really see a problem like that then TRI interface can be extended to return a weight given RegUnit and Lanemask. In fact I'm trying to patch other places like increase/decreaseRegPressure and I had to separate this logic into a function. The very same function can be exposed by targets. Do you want to create it right now?

Without a similar patch to {increase|decrease}RegPressure() pressure book keeping will go out of sync as well.

You are probably right, I was going to look into these places...

We cannot apply a patch that just changes one part of the system.

I'm adding this changes at the moment.

[rampitec]

In D29835#673828, @rampitec wrote:

I'm adding this changes at the moment.

In fact PressureDiff calculation and general RP tracking are independent. Now as I'm trying to extend the change for general RP tracking as well I really see many failures in lit tests...

[rampitec]

Added the same logic for all register pressure, not just PressureDiff.

[rampitec]

Ping.

[rampitec]

Added callback for a target to estimate register unit weight given a lane mask, since that is unclear if proportional weight is correct for all targets.

[alex-t]

lib/CodeGen/RegisterPressure.cpp

43

Given that getMaxLaneMaskForVReg returns all bit s that possibly can be set in the n=mask for given register class, what is the purpose of the AND? Unless you expect the bits not present in the register class to be set in LaneMask for the reg unit of this class.

[rampitec]

lib/CodeGen/RegisterPressure.cpp

43

In certain situations I just see -1 passed as a lane mask, mostly with Hexagon. In fact I can also see quite a few places where lane mask is initialized with default argument = LaneBitmask::getAll(), so until it is properly tracked at all places we need to take extra care. Without this I see a lot of assertions about pressure underflow from many lit tests.

[MatzeB]

The current implementation will miscalculate the PressureDiff objects. Could you revert this until we can further investigate and fix this.

llvm/trunk/include/llvm/Target/TargetRegisterInfo.h

723–724 ↗

(On Diff #89545)

This description makes no sense like that: Register units have no subregisters.

llvm/trunk/lib/CodeGen/MachineScheduler.cpp

1088 ↗

(On Diff #89545)

The LiveUses array here only contains register that have become live or completely dead, you do not get updates of some lanes getting added/remove at the moment. So the calculations for the pressure diffs will be out of sync when you do them at lane granularity like this.

[rampitec]

In D29835#686080, @MatzeB wrote:

The current implementation will miscalculate the PressureDiff objects. Could you revert this until we can further investigate and fix this.

I will revert it for now. There is yet another problem, for a target which does not override getRegUnitWeight that is possible to get pressure underflow because a full register weight is decreased with each subreg def.

[rampitec]

Reverted in r296182.

[MatzeB]

In D29835#686230, @rampitec wrote:

Reverted in r296182.

Thanks!

include/llvm/CodeGen/RegisterPressure.h

class PressureDiff {
   iterator nonconst_begin() { return &PressureChanges[0]; }
   iterator nonconst_end() { return &PressureChanges[MaxPSets]; }
 
 public:
   typedef const PressureChange* const_iterator;
   const_iterator begin() const { return &PressureChanges[0]; }
   const_iterator end() const { return &PressureChanges[MaxPSets]; }
 
-  void addPressureChange(unsigned RegUnit, bool IsDec,
+  void addPressureChange(RegisterMaskPair P, bool IsDec,
                          const MachineRegisterInfo *MRI);
 
   void dump(const TargetRegisterInfo &TRI) const;
 };
 
 /// List of registers defined and used by a machine instruction.
 class RegisterOperands {
 public:

lib/CodeGen/MachineScheduler.cpp

if (ShouldTrackLaneMasks) {
 
       for (const VReg2SUnit &V2SU
            : make_range(VRegUses.find(Reg), VRegUses.end())) {
         SUnit &SU = *V2SU.SU;
         if (SU.isScheduled || &SU == &ExitSU)
           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)
                  << ' ' << *SU.getInstr();
           dbgs() << "              to ";
           PDiff.dump(*TRI);
         );
       }
if (ShouldTrackLaneMasks) {
         SUnit *SU = V2SU.SU;
         // If this use comes before the reaching def, it cannot be a last use,
         // so decrease its pressure change.
         if (!SU->isScheduled && SU != &ExitSU) {
           LiveQueryResult LRQ =
               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();
               dbgs() << "              to ";
               PDiff.dump(*TRI);
             );
           }
         }

lib/CodeGen/RegisterPressure.cpp

 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/RegisterPressure.h"
 #include "llvm/CodeGen/LiveInterval.h"
 #include "llvm/CodeGen/LiveIntervalAnalysis.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace llvm;
 
+/// Compute register unit weight given the lane mask.
+static unsigned getRegUnitWeight(const MachineRegisterInfo &MRI,
+                                 unsigned RegUnit, LaneBitmask LaneMask) {
+  PSetIterator PSetI = MRI.getPressureSets(RegUnit);
+  unsigned Weight = PSetI.getWeight();
+
+  if (Weight > 1 && LaneMask.any() && !LaneMask.all() &&
+      TargetRegisterInfo::isVirtualRegister(RegUnit)) {
+    LaneBitmask Max = MRI.getMaxLaneMaskForVReg(RegUnit);
+    if (Max != LaneMask && !Max.all() && !Max.none())
+      Weight = (Weight * countPopulation(LaneMask.getAsInteger())) /
+                         countPopulation(Max.getAsInteger());
+  }
+  return Weight;
+}
+
+/// 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;
 
+  unsigned Weight = getRegUnitWeight(MRI, Reg, NewMask & ~PrevMask);
   PSetIterator PSetI = MRI.getPressureSets(Reg);
-  unsigned Weight = PSetI.getWeight();
   for (; PSetI.isValid(); ++PSetI)
     CurrSetPressure[*PSetI] += Weight;
 }
 
 /// Decrease pressure for each pressure set provided by TargetRegisterInfo.
 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;
 
+  unsigned Weight = 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;
   }
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 LLVM_DUMP_METHOD
void PressureDiff::dump(const TargetRegisterInfo &TRI) const {
   }
   dbgs() << '\n';
 }
 #endif
 
 void RegPressureTracker::increaseRegPressure(unsigned RegUnit,
                                              LaneBitmask PreviousMask,
                                              LaneBitmask NewMask) {
-  if (PreviousMask.any() || NewMask.none())
+  clampMasks(*MRI, RegUnit, PreviousMask, NewMask);
+  if ((NewMask & ~PreviousMask).none())
     return;
 
+  unsigned Weight = getRegUnitWeight(*MRI, RegUnit, NewMask & ~PreviousMask);
   PSetIterator PSetI = MRI->getPressureSets(RegUnit);
-  unsigned Weight = PSetI.getWeight();
   for (; PSetI.isValid(); ++PSetI) {
     CurrSetPressure[*PSetI] += Weight;
     P.MaxSetPressure[*PSetI] =
         std::max(P.MaxSetPressure[*PSetI], CurrSetPressure[*PSetI]);
   }
 }
 
 void RegPressureTracker::decreaseRegPressure(unsigned RegUnit,

 }
 
 void PressureDiffs::addInstruction(unsigned Idx,
                                    const RegisterOperands &RegOpers,
                                    const MachineRegisterInfo &MRI) {
   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();
+  int Weight = (int)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();
     for (; I != E && I->isValid(); ++I) {
       if (I->getPSet() >= *PSetI)
         break;
     }
     // If all pressure sets are more constrained, skip the remaining PSets.

test/CodeGen/AMDGPU/load-global-i32.ll

 ; GCN-NOHSA: buffer_store_dwordx4
 ; GCN-NOHSA: buffer_store_dwordx4
 
 ; GCN-NOHSA: buffer_store_dwordx4
 ; GCN-NOHSA: buffer_store_dwordx4
 ; GCN-NOHSA: buffer_store_dwordx4
 ; GCN-NOHSA: buffer_store_dwordx4
 
-; GCN-HSA: flat_store_dwordx4
-; GCN-HSA: flat_store_dwordx4
-; GCN-HSA: flat_store_dwordx4
-; GCN-HSA: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
 
-; GCN-HSA: flat_store_dwordx4
-; GCN-HSA: flat_store_dwordx4
-; GCN-HSA: flat_store_dwordx4
-; GCN-HSA: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
 
-; GCN-HSA: flat_store_dwordx4
-; GCN-HSA: flat_store_dwordx4
-; GCN-HSA: flat_store_dwordx4
-; GCN-HSA: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
 
-; GCN-HSA: flat_store_dwordx4
-; GCN-HSA: flat_store_dwordx4
-; GCN-HSA: flat_store_dwordx4
-; GCN-HSA: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
+; GCN-HSA-DAG: flat_store_dwordx4
 
 define void @global_sextload_v32i32_to_v32i64(<32 x i64> addrspace(1)* %out, <32 x i32> addrspace(1)* %in) #0 {
   %ld = load <32 x i32>, <32 x i32> addrspace(1)* %in
   %ext = sext <32 x i32> %ld to <32 x i64>
   store <32 x i64> %ext, <32 x i64> addrspace(1)* %out
   ret void
 }
 

test/CodeGen/AMDGPU/schedule-regpressure-subregs.mir

+# RUN: llc -march=amdgcn -misched=converge -run-pass machine-scheduler -verify-misched %s -o - -debug-only=misched 2>&1 | FileCheck %s
+# REQUIRES: asserts
+
+# CHECK-LABEL: ScheduleDAGMILive::schedule starting
+
+# Check that def and use subregs count with the same weight
+# CHECK: %vreg9:sub1<def> = V_MUL_LO_I32 %vreg6:sub1, 3
+# CHECK: Pressure Diff : {{$}}
+
+# Check that a subreg does not count as a whole superreg
+# CHECK: %vreg9:sub0<def> = V_MUL_LO_I32 %vreg6:sub0, %vreg9:sub1
+# CHECK: Pressure Diff : VGPR_32 1{{$}}
+
+# Check that two subregs of the same register count as a whole register
+# CHECK: DS_WRITE2_B32 %vreg7, %vreg9:sub0, %vreg9:sub1
+# CHECK: Pressure Diff : VGPR_32 3{{$}}
+
+---
+name:            mo_pset
+alignment:       0
+exposesReturnsTwice: false
+legalized:       false
+regBankSelected: false
+selected:        false
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: sreg_128 }
+  - { id: 1, class: sgpr_64 }
+  - { id: 2, class: sreg_32_xm0 }
+  - { id: 3, class: sgpr_32 }
+  - { id: 4, class: vgpr_32 }
+  - { id: 5, class: sreg_32_xm0_xexec }
+  - { id: 6, class: vreg_64 }
+  - { id: 7, class: vgpr_32 }
+  - { id: 8, class: vgpr_32 }
+  - { id: 9, class: vreg_64 }
+liveins:
+  - { reg: '%sgpr4_sgpr5', virtual-reg: '%1' }
+frameInfo:
+  isFrameAddressTaken: false
+  isReturnAddressTaken: false
+  hasStackMap:     false
+  hasPatchPoint:   false
+  stackSize:       0
+  offsetAdjustment: 0
+  maxAlignment:    0
+  adjustsStack:    false
+  hasCalls:        false
+  maxCallFrameSize: 0
+  hasOpaqueSPAdjustment: false
+  hasVAStart:      false
+  hasMustTailInVarArgFunc: false
+body:             |
+  bb.0:
+    liveins: %sgpr4_sgpr5
+
+    %1 = COPY %sgpr4_sgpr5
+    %5 = S_LOAD_DWORD_IMM %1, 0, 0 :: (non-temporal dereferenceable invariant load 4 from `i32 addrspace(2)* undef`)
+    %m0 = S_MOV_B32 -1
+    %7 = COPY %5
+    %6 = DS_READ2_B32 %7, 0, 1, 0, implicit %m0, implicit %exec
+    undef %9.sub1 = V_MUL_LO_I32 %6.sub1, 3, implicit %exec
+    %9.sub0 = V_MUL_LO_I32 %6.sub0, %9.sub1, implicit %exec
+    DS_WRITE2_B32 %7, %9.sub0, %9.sub1, 4, 5, 0, implicit killed %m0, implicit %exec
+    S_ENDPGM
+
+...