[AMDGPU] Fix incorrect register usage tracking in GCNUpwardTracker

[vpykhtin]

Summary

This change fixes incorrect maximum register pressure calculation in GCNUpwardRPTracker: it reduced pressure of defs before incrementing pressure on uses losing the possible maximum pressure of defs + uses at the machine instruction.

After several attempts to fix it with fewer lines of code I decided it would be easier to introduce MachineInstrRegs class which collects register def/uses for a machine instruction with theirs linemasks. This class does job similar to standard llvm's RegisterOperands class but much smaller.

Trying to figure out a test for this.

[rampitec]

Where is GCNUpwardRPTracker::reset() now? What is the base revision this diff is taken against? The GCNUpwardRPTracker::reset() is right before recede() in the current revision, but I do not see it in the left pane as well. Same with some other functions like getDefRegMask.

lib/Target/AMDGPU/GCNRegPressure.cpp

314

You have GCNRPTracker::getDefRegMask() for this.

318

Do not you want to also erase it from LiveRegs if LaneMask.none()?

319

And there is getUsedRegMask() for this too.

330

Why not to do it right above, where you assign getAll()? It seems to be less work.

[rampitec]

lib/Target/AMDGPU/GCNRegPressure.cpp

318

You are now updating MaxPressure correctly, but your CurPressure is incorrect. At this point defs still contribute to the current pressure, they will be out only with the next recede. That may lead scheduler to wrong decisions about a current instruction. You can also run into a paradoxical situation that no single instruction has a pressure equal to max.

I still believe that in the situation where we want to account for both defs and uses contributing to the RP of an instruction we naturally have a two step advance/recede process - step before/past the instruction and actually step to it.

[vpykhtin]

lib/Target/AMDGPU/GCNRegPressure.cpp

314

This is another class, though getDefRegMask can be made nonmember

318

Why to erase?

Actually recede moves from the point after the instruction to the point before the instruction in top-down order accounting max pressure interim. recede never stops at the instruction.

319

I'm avoiding getLiveLaneMask call here before all registers collected

330

This avoids getLiveLaneMask call more than once for a register

[vpykhtin]

I inserted MachineInstrRegs between reset and recede, all functions are in old places. I should move MachineInstrRegs higher

[rampitec]

In D33289#758351, @vpykhtin wrote:

I inserted MachineInstrRegs between reset and recede, all functions are in old places. I should move MachineInstrRegs higher

I see it now, thanks. Can you move it higher please?

lib/Target/AMDGPU/GCNRegPressure.cpp

314

Probably static inline is just right for it. Even for getUsedRegMask if you pass LIS to it.

318

If you erase you will not iterate over it getRegPressure().

I understand your point that recede does not stop in steps, but I'm still concerned that you will not get a correct CurPressure, or even will not get CurPressure equal to max pressure anywhere in the region. How about that?

330

Do you see a common situation where the same register is used in the same instruction more than once? This sounds quite exotic to me, provided we are speaking about uses only, not defs.

[vpykhtin]

In D33289#758647, @rampitec wrote:

In D33289#758351, @vpykhtin wrote:

I inserted MachineInstrRegs between reset and recede, all functions are in old places. I should move MachineInstrRegs higher

I see it now, thanks. Can you move it higher please?

Yes I'll fix that.

lib/Target/AMDGPU/GCNRegPressure.cpp

318

Ok, now that live regs can be reused it may have the point to clear a register immediately. Previously I used stripEmpty on the set but only for debug printing purposes.

CurPressure isn't calculated for the at-the-instruction level, its calculated for the after recede point. I put an assert that CurPressure calculated correctly in the end of recede. CurPressure can never become MaxPressure, but I don't see a problem here. There is no at-the-instruction position in the tracker - it is always in between.

330

I agree its going to be very rare ocasion, but the main purpose of this class is to deduplicate register def/uses so it wouldn't be counted twice when calculating pressure. So if I deduplicated registers already then I can save some time on calculating mask.

[rampitec]

lib/Target/AMDGPU/GCNRegPressure.cpp

314

Actually getUsedRegMask is unused now, so you can just delete it.

330

I mean, if that is extremely rare you may lose more in the inexpensive but way more often called second loop.

[vpykhtin]

lib/Target/AMDGPU/GCNRegPressure.cpp

330

Ok, I'll make getDefRegMask and getUsedRegMask static and reuse it here removing bottom loop, thanks.

[rampitec]

lib/Target/AMDGPU/GCNRegPressure.cpp

288

It looks like defs are not really needed in this class. Uses needed because you walk them twice, but defs can be just directly processed. I.e. the code can be simplified and overhead somewhat reduced.

318

I see it as a sort of quantum tracker. It hides the intermediate step where pressure actually peaks from an observer. As long as we agree on that understanding I have no objection on submitting such a tracker where actual pressure "tunnels" through recede method as we have no actual interested observers currently. We might want to split the method in the future if we have them.

[vpykhtin]

fixed as per comments

[rampitec]

LGTM. Thanks!

[rampitec]

lib/Target/AMDGPU/GCNRegPressure.cpp

313

Are you sure it should always present? What if we have a dead def? I.e. an instruction defines a register which is never used. I guess it will not be reported by LIS.

If so this should be if (I != LiveRegs.end()) continue;

[rampitec]

lib/Target/AMDGPU/GCNRegPressure.cpp

313

Sorry, if (I == LiveRegs.end()) continue;

[vpykhtin]

lib/Target/AMDGPU/GCNRegPressure.cpp

313

Done, thanks!

[foad]

This change fixes incorrect maximum register pressure calculation in GCNUpwardRPTracker: it reduced pressure of defs before incrementing pressure on uses losing the possible maximum pressure of defs + uses at the machine instruction.

Hi @vpykhtin! I don't understand the reason for this change. Why should max pressure include both the uses and the defs of one instruction? The uses and defs are not live at the same time and can be allocated to overlapping physical registers (assuming the uses are killed by the instruction). There should be an exception for early-clobber def operands but they are not very common.

+ @piotr @critson

[piotr]

In D33289#4654033, @foad wrote:

This change fixes incorrect maximum register pressure calculation in GCNUpwardRPTracker: it reduced pressure of defs before incrementing pressure on uses losing the possible maximum pressure of defs + uses at the machine instruction.

Hi @vpykhtin! I don't understand the reason for this change. Why should max pressure include both the uses and the defs of one instruction? The uses and defs are not live at the same time and can be allocated to overlapping physical registers (assuming the uses are killed by the instruction). There should be an exception for early-clobber def operands but they are not very common.

+ @piotr @critson

Would it make sense to only do the AtMIPressure part for instructions with early clobber?

[vpykhtin]

In D33289#4654038, @piotr wrote:

In D33289#4654033, @foad wrote:

This change fixes incorrect maximum register pressure calculation in GCNUpwardRPTracker: it reduced pressure of defs before incrementing pressure on uses losing the possible maximum pressure of defs + uses at the machine instruction.

Hi @vpykhtin! I don't understand the reason for this change. Why should max pressure include both the uses and the defs of one instruction? The uses and defs are not live at the same time and can be allocated to overlapping physical registers (assuming the uses are killed by the instruction). There should be an exception for early-clobber def operands but they are not very common.

+ @piotr @critson

Would it make sense to only do the AtMIPressure part for instructions with early clobber?

You're right, we should not increment pressure for early-clobbers twice in AtMIPressure

[vpykhtin]

You're right, we should not increment pressure for early-clobbers twice in AtMIPressure

Sorry, disregard this comment. We should only account for early-clobbers this way.

[vpykhtin]

Is anyone working on this at the moment?

[foad]

In D33289#4654317, @vpykhtin wrote:

Is anyone working on this at the moment?

I'm not working on it but I have been thinking about it. The first problem is how to write regression tests for GCNRegPressure.

[vpykhtin]

In D33289#4654324, @foad wrote:

In D33289#4654317, @vpykhtin wrote:

Is anyone working on this at the moment?

I'm not working on it but I have been thinking about it. The first problem is how to write regression tests for GCNRegPressure.

I can probably come up with a unit test in the way similar to how we test LiveIntervals and LiveVariables.

[foad]

In D33289#4654325, @vpykhtin wrote:

In D33289#4654324, @foad wrote:

In D33289#4654317, @vpykhtin wrote:

Is anyone working on this at the moment?

I'm not working on it but I have been thinking about it. The first problem is how to write regression tests for GCNRegPressure.

I can probably come up with a unit test in the way similar to how we test LiveIntervals and LiveVariables.

I wonder if we can test it more directly by adding an analysis pass like this: https://github.com/GPUOpen-Drivers/llvm-project/commit/042be23e3d98963fb02833511a86f4e26378a04d
and then using something like opt -passes='print<amdgpu-reg-press>'.

[vpykhtin]

I wonder if we can test it more directly by adding an analysis pass like this: https://github.com/GPUOpen-Drivers/llvm-project/commit/042be23e3d98963fb02833511a86f4e26378a04d
and then using something like opt -passes='print<amdgpu-reg-press>'.

We can try to print reg pressure at every instruction

[vpykhtin]

Something like this? (don't forget to expand *.mir file diff, it's not shown by default)

https://github.com/llvm/llvm-project/compare/main...vpykhtin:llvm-project:rp_printer

lib/Target/AMDGPU/GCNRegPressure.h

 
 protected:
   const LiveIntervals &LIS;
   LiveRegSet LiveRegs;
   GCNRegPressure CurPressure, MaxPressure;
   const MachineInstr *LastTrackedMI = nullptr;
   mutable const MachineRegisterInfo *MRI = nullptr;
   GCNRPTracker(const LiveIntervals &LIS_) : LIS(LIS_) {}
-  LaneBitmask getDefRegMask(const MachineOperand &MO) const;
-  LaneBitmask getUsedRegMask(const MachineOperand &MO) const;
 public:
   // live regs for the current state
   const decltype(LiveRegs) &getLiveRegs() const { return LiveRegs; }
   const MachineInstr *getLastTrackedMI() const { return LastTrackedMI; }
 
   void clearMaxPressure() { MaxPressure.clear(); }
 
   // returns MaxPressure, resetting it

lib/Target/AMDGPU/GCNRegPressure.cpp

 //===------------------------- GCNRegPressure.cpp - -----------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 /// \file
 //
 //===----------------------------------------------------------------------===//
 
 #include "GCNRegPressure.h"
+#include "llvm/CodeGen/RegisterPressure.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "misched"
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 LLVM_DUMP_METHOD
 void llvm::printLivesAt(SlotIndex SI,
static bool isEqual(const GCNRPTracker::LiveRegSet &S1,
   for (const auto &P : S1) {
     auto I = S2.find(P.first);
     if (I == S2.end() || I->second != P.second)
       return false;
   }
   return true;
 }
 
-static GCNRPTracker::LiveRegSet
-stripEmpty(const GCNRPTracker::LiveRegSet &LR) {
-  GCNRPTracker::LiveRegSet Res;
-  for (const auto &P : LR) {
-    if (P.second.any())
-      Res.insert(P);
-  }
-  return Res;
-}
 #endif
 
 ///////////////////////////////////////////////////////////////////////////////
 // GCNRegPressure
 
 unsigned GCNRegPressure::getRegKind(unsigned Reg,
                                     const MachineRegisterInfo &MRI) {
   assert(TargetRegisterInfo::isVirtualRegister(Reg));
void GCNRegPressure::print(raw_ostream &OS, const SISubtarget *ST) const {
   if (ST) OS << "(O" << ST->getOccupancyWithNumSGPRs(getSGPRNum()) << ')';
   OS << ", LVGPR WT: " << getVGPRTuplesWeight()
      << ", LSGPR WT: " << getSGPRTuplesWeight();
   if (ST) OS << " -> Occ: " << getOccupancy(*ST);
   OS << '\n';
 }
 #endif
 
+
+static LaneBitmask getDefRegMask(const MachineOperand &MO,
+                                 const MachineRegisterInfo &MRI) {
+  assert(MO.isDef() && MO.isReg() &&
+    TargetRegisterInfo::isVirtualRegister(MO.getReg()));
+
+  // We don't rely on read-undef flag because in case of tentative schedule
+  // tracking it isn't set correctly yet. This works correctly however since
+  // use mask has been tracked before using LIS.
+  return MO.getSubReg() == 0 ?
+    MRI.getMaxLaneMaskForVReg(MO.getReg()) :
+    MRI.getTargetRegisterInfo()->getSubRegIndexLaneMask(MO.getSubReg());
+}
+
+static LaneBitmask getUsedRegMask(const MachineOperand &MO,
+                                  const MachineRegisterInfo &MRI,
+                                  const LiveIntervals &LIS) {
+  assert(MO.isUse() && MO.isReg() &&
+         TargetRegisterInfo::isVirtualRegister(MO.getReg()));
+
+  if (auto SubReg = MO.getSubReg())
+    return MRI.getTargetRegisterInfo()->getSubRegIndexLaneMask(SubReg);
+
+  auto MaxMask = MRI.getMaxLaneMaskForVReg(MO.getReg());
+  if (MaxMask.getAsInteger() == 1) // cannot have subregs
+    return MaxMask;
+
+  // For a tentative schedule LIS isn't updated yet but livemask should remain
+  // the same on any schedule. Subreg defs can be reordered but they all must
+  // dominate uses anyway.
+  auto SI = LIS.getInstructionIndex(*MO.getParent()).getBaseIndex();
+  return getLiveLaneMask(MO.getReg(), SI, LIS, MRI);
+}
+
+SmallVector<RegisterMaskPair, 8> collectVirtualRegUses(const MachineInstr &MI,
+                                              const LiveIntervals &LIS,
+                                              const MachineRegisterInfo &MRI) {
+  SmallVector<RegisterMaskPair, 8> Res;
+  for (const auto &MO : MI.operands()) {
+    if (!MO.isReg() || !TargetRegisterInfo::isVirtualRegister(MO.getReg()))
+      continue;
+    if (!MO.isUse() || !MO.readsReg())
+      continue;
+
+    auto const UsedMask = getUsedRegMask(MO, MRI, LIS);
+
+    auto Reg = MO.getReg();
+    auto I = std::find_if(Res.begin(), Res.end(), [Reg](const RegisterMaskPair &RM) {
+      return RM.RegUnit == Reg;
+    });
+    if (I != Res.end())
+      I->LaneMask |= UsedMask;
+    else
+      Res.push_back(RegisterMaskPair(Reg, UsedMask));
+  }
+  return Res;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 // GCNRPTracker
 
 LaneBitmask llvm::getLiveLaneMask(unsigned Reg,
                                   SlotIndex SI,
                                   const LiveIntervals &LIS,
                                   const MachineRegisterInfo &MRI) {
   LaneBitmask LiveMask;
if (!LIS.hasInterval(Reg))
       continue;
     auto LiveMask = getLiveLaneMask(Reg, SI, LIS, MRI);
     if (LiveMask.any())
       LiveRegs[Reg] = LiveMask;
   }
   return LiveRegs;
 }
 
-LaneBitmask GCNRPTracker::getDefRegMask(const MachineOperand &MO) const {
-  assert(MO.isDef() && MO.isReg() &&
-    TargetRegisterInfo::isVirtualRegister(MO.getReg()));
-
-  // We don't rely on read-undef flag because in case of tentative schedule
-  // tracking it isn't set correctly yet. This works correctly however since
-  // use mask has been tracked before using LIS.
-  return MO.getSubReg() == 0 ?
-    MRI->getMaxLaneMaskForVReg(MO.getReg()) :
-    MRI->getTargetRegisterInfo()->getSubRegIndexLaneMask(MO.getSubReg());
-}
-
-LaneBitmask GCNRPTracker::getUsedRegMask(const MachineOperand &MO) const {
-  assert(MO.isUse() && MO.isReg() &&
-         TargetRegisterInfo::isVirtualRegister(MO.getReg()));
-
-  if (auto SubReg = MO.getSubReg())
-    return MRI->getTargetRegisterInfo()->getSubRegIndexLaneMask(SubReg);
-
-  auto MaxMask = MRI->getMaxLaneMaskForVReg(MO.getReg());
-  if (MaxMask.getAsInteger() == 1) // cannot have subregs
-    return MaxMask;
-
-  // For a tentative schedule LIS isn't updated yet but livemask should remain
-  // the same on any schedule. Subreg defs can be reordered but they all must
-  // dominate uses anyway.
-  auto SI = LIS.getInstructionIndex(*MO.getParent()).getBaseIndex();
-  return getLiveLaneMask(MO.getReg(), SI, LIS, *MRI);
-}
-
 void GCNUpwardRPTracker::reset(const MachineInstr &MI,
                                const LiveRegSet *LiveRegsCopy) {
   MRI = &MI.getParent()->getParent()->getRegInfo();
   if (LiveRegsCopy) {
     if (&LiveRegs != LiveRegsCopy)
       LiveRegs = *LiveRegsCopy;
   } else {
     LiveRegs = getLiveRegsAfter(MI, LIS);
   }
   MaxPressure = CurPressure = getRegPressure(*MRI, LiveRegs);
 }
 
 void GCNUpwardRPTracker::recede(const MachineInstr &MI) {
   assert(MRI && "call reset first");
 
   LastTrackedMI = &MI;
 
   if (MI.isDebugValue())
     return;
 
-  // process all defs first to ensure early clobbers are handled correctly
-  // iterating over operands() to catch implicit defs
-  for (const auto &MO : MI.operands()) {
-    if (!MO.isReg() || !MO.isDef() ||
-      !TargetRegisterInfo::isVirtualRegister(MO.getReg()))
+  auto const RegUses = collectVirtualRegUses(MI, LIS, *MRI);
+
+  // calc pressure at the MI (defs + uses)
+  auto AtMIPressure = CurPressure;
+  for (const auto &U : RegUses) {
+    auto LiveMask = LiveRegs[U.RegUnit];
+    AtMIPressure.inc(U.RegUnit, LiveMask, LiveMask | U.LaneMask, *MRI);
+  }
+  // update max pressure
+  MaxPressure = max(AtMIPressure, MaxPressure);
+
+  for (const auto &MO : MI.defs()) {
+    if (!MO.isReg() || !TargetRegisterInfo::isVirtualRegister(MO.getReg()) ||
+         MO.isDead())
       continue;
 
     auto Reg = MO.getReg();
-    auto &LiveMask = LiveRegs[Reg];
+    auto I = LiveRegs.find(Reg);
+    assert(I != LiveRegs.end());
+    auto &LiveMask = I->second;
     auto PrevMask = LiveMask;
-    LiveMask &= ~getDefRegMask(MO);
+    LiveMask &= ~ getDefRegMask(MO, *MRI);
     CurPressure.inc(Reg, PrevMask, LiveMask, *MRI);
+    if (LiveMask.none())
+      LiveRegs.erase(I);
   }
-
-  // then all uses
-  for (const auto &MO : MI.uses()) {
-    if (!MO.isReg() || !MO.readsReg() ||
-      !TargetRegisterInfo::isVirtualRegister(MO.getReg()))
-      continue;
-
-    auto Reg = MO.getReg();
-    auto &LiveMask = LiveRegs[Reg];
+  for (const auto &U : RegUses) {
+    auto &LiveMask = LiveRegs[U.RegUnit];
     auto PrevMask = LiveMask;
-    LiveMask |= getUsedRegMask(MO);
-    CurPressure.inc(Reg, PrevMask, LiveMask, *MRI);
+    LiveMask |= U.LaneMask;
+    CurPressure.inc(U.RegUnit, PrevMask, LiveMask, *MRI);
   }
-
-  MaxPressure = max(MaxPressure, CurPressure);
+  assert(CurPressure == getRegPressure(*MRI, LiveRegs));
 }
 
 bool GCNDownwardRPTracker::reset(const MachineInstr &MI,
                                  const LiveRegSet *LiveRegsCopy) {
   MRI = &MI.getParent()->getParent()->getRegInfo();
   LastTrackedMI = nullptr;
   MBBEnd = MI.getParent()->end();
   NextMI = &MI;
   NextMI = skipDebugInstructionsForward(NextMI, MBBEnd);
   if (NextMI == MBBEnd)
     return false;
void GCNDownwardRPTracker::advanceToNext() {
   for (const auto &MO : LastTrackedMI->defs()) {
     if (!MO.isReg())
       continue;
     unsigned Reg = MO.getReg();
     if (!TargetRegisterInfo::isVirtualRegister(Reg))
       continue;
     auto &LiveMask = LiveRegs[Reg];
     auto PrevMask = LiveMask;
-    LiveMask |= getDefRegMask(MO);
+    LiveMask |= getDefRegMask(MO, *MRI);
     CurPressure.inc(Reg, PrevMask, LiveMask, *MRI);
   }
 
   MaxPressure = max(MaxPressure, CurPressure);
 }
 
 bool GCNDownwardRPTracker::advance() {
   // If we have just called reset live set is actual.
if (I == TrackedLR.end()) {
              << " isn't found in tracked set\n";
     }
   }
 }
 
 bool GCNUpwardRPTracker::isValid() const {
   const auto &SI = LIS.getInstructionIndex(*LastTrackedMI).getBaseIndex();
   const auto LISLR = llvm::getLiveRegs(SI, LIS, *MRI);
-  const auto TrackedLR = stripEmpty(LiveRegs);
+  const auto &TrackedLR = LiveRegs;
 
   if (!isEqual(LISLR, TrackedLR)) {
     dbgs() << "\nGCNUpwardRPTracker error: Tracked and"
               " LIS reported livesets mismatch:\n";
     printLivesAt(SI, LIS, *MRI);
     reportMismatch(LISLR, TrackedLR, MRI->getTargetRegisterInfo());
     return false;
   }