[AMDGPU] Small cleanups in wait counter code

[stepthomas]

Summary

A small number of cleanups and refactors intended to enhance readability in
two passes that deal with s_waitcnt instructions.

Change-Id: Iff58371f659e8648b695868af009c5613de0e213

[foad]

Note that getScoreLB should always be <= getScoreUB, so code like getScoreLB < getScoreUB can be updated to getScoreRange != 0.

llvm/lib/Target/AMDGPU/GCNHazardRecognizer.cpp

1335

Not sure about the name. Maybe isWaitStoreCountZero would be better...?

1336

Don't need the outermost parens.

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

698

Can use getScoreRange here.

759

Can use getScoreRange here.

1204

Can use getScoreRange here.

1598

Can use getScoreRange here.

[stepthomas]

Address review comments

Rename isStoreCountWait() -> isStoreCountWaitZero(), use getScoreRange()
in more places.

[scott.linder]

Thank you for the cleanup! I left some suggestions, but I am also OK with everything as-is

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

251

If you did decide to add the hasPendingEvent overload I would suggest renaming this as well. All of these have the same semantics, just with different "filters", so I think they make sense as an overload set.

307–311

Some more thoughts while reviewing (although not applicable to your changes directly):

• Why is EXP_CNT special here?
• Is there a particular reason this relies on unsigned integer underflow? Does this map better onto the conceptual model of the scores/counts?

For either/both can anyone devise a comment to clarify? I suggested what I believe to be an equivalent version which I find easier to understand, but I am likely missing something.

699–716

Nit, but I'd also factor out the remaining common stuff

1226

I would suggest also adding a hasPendingEvent(InstCounterType T) const overload to use here, i.e.

bool hasPendingEvent(InstCounterType T) const {
  bool HasPending = PendingEvents & WaitEventMaskForInst[T];
  assert(HasPending == getScoreUB(T) > getScoreLB(T));
  return HasPending;
}

1241

1598

[stepthomas]

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

699–716

I'll decline this suggestion if I may - it actually makes it less clear to me.

1241

As I was typing this (and other cases) I was very much in two minds. I think I'll adopt this suggestion.

[stepthomas]

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

1226

I'll add this suggestion, including renaming the void parameter overloading, although I'll also change the return type to unsigned (and the void parameter overloading simply returns PendingEvents) because then I can use the new overloading in hasMixedPendingEvents() immediately blow this code (so getting an extra assertion check). It makes the assertion

assert((HasPending != 0) == (getScoreRange(T) != 0))

which is ugly, but still worth it.

[stepthomas]

Address more review comments

[foad]

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

307–311

I don't know why EXP_CNT is handled specially, but I agree your version of the code is simpler and easier to understand.

1204

You're just testing getScoreRange != 0 here, so you could use the new hasPendingEvent overload (I'm not really sure why we have both and which you should prefer)

1598

I would have thought that hasPendingEvent(SMEM_ACCESS) implies hasPendingEvent(LGKM_CNT) - doesn't it?

[scott.linder]

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

699–716

Works for me, I agree it is nice to see the whole "token" for each case in one line

[scott.linder]

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

1598

+1, the redundant check for a pending LGKM_CNT event should be removed (or I suppose moved to an assert if it makes the intent wrt. the bug clearer)

[stepthomas]

Rebase and address more review comments.

Apply a suggested change to setScoreUB() that I'd somehow missed before,
and remove a couple of redundant checks.

[foad]

LGTM, thanks!

[foad]

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

307–311

I did look into this a bit more, and it's to do with whether we still need to emit any wait for a counter that has overflowed. See D70418. But I still don't know why EXP_CNT is treated specially.

llvm/lib/Target/AMDGPU/GCNHazardRecognizer.cpp

for (auto &MI : MBB) {
           SIInstrInfo::isVMEM(MI) || SIInstrInfo::isSegmentSpecificFLAT(MI);
       if (HasLds && HasVmem)
         return true;
     }
   }
   return false;
 }
 
+static bool isStoreCountWait(const MachineInstr &I) {
+  return (I.getOpcode() == AMDGPU::S_WAITCNT_VSCNT &&
+          I.getOperand(0).getReg() == AMDGPU::SGPR_NULL &&
+          !I.getOperand(1).getImm());
+}
+
 bool GCNHazardRecognizer::fixLdsBranchVmemWARHazard(MachineInstr *MI) {
   if (!RunLdsBranchVmemWARHazardFixup)
     return false;
 
   assert(ST.hasLdsBranchVmemWARHazard());
 
   auto IsHazardInst = [](const MachineInstr &MI) {
     if (SIInstrInfo::isDS(MI))
       return 1;
     if (SIInstrInfo::isVMEM(MI) || SIInstrInfo::isSegmentSpecificFLAT(MI))
       return 2;
     return 0;
   };
 
   auto InstType = IsHazardInst(*MI);
   if (!InstType)
     return false;
 
   auto IsExpiredFn = [&IsHazardInst](const MachineInstr &I, int) {
-    return IsHazardInst(I) || (I.getOpcode() == AMDGPU::S_WAITCNT_VSCNT &&
-                               I.getOperand(0).getReg() == AMDGPU::SGPR_NULL &&
-                               !I.getOperand(1).getImm());
+    return IsHazardInst(I) || isStoreCountWait(I);
   };
 
   auto IsHazardFn = [InstType, &IsHazardInst](const MachineInstr &I) {
     if (!I.isBranch())
       return false;
 
     auto IsHazardFn = [InstType, IsHazardInst](const MachineInstr &I) {
       auto InstType2 = IsHazardInst(I);
       return InstType2 && InstType != InstType2;
     };
 
     auto IsExpiredFn = [InstType, &IsHazardInst](const MachineInstr &I, int) {
       auto InstType2 = IsHazardInst(I);
       if (InstType == InstType2)
         return true;
 
-      return I.getOpcode() == AMDGPU::S_WAITCNT_VSCNT &&
-             I.getOperand(0).getReg() == AMDGPU::SGPR_NULL &&
-             !I.getOperand(1).getImm();
+      return isStoreCountWait(I);
     };
 
     return ::getWaitStatesSince(IsHazardFn, &I, IsExpiredFn) !=
            std::numeric_limits<int>::max();
   };
 
   if (::getWaitStatesSince(IsHazardFn, MI, IsExpiredFn) ==
       std::numeric_limits<int>::max())

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

unsigned getScoreLB(InstCounterType T) const {
     return ScoreLBs[T];
   }
 
   unsigned getScoreUB(InstCounterType T) const {
     assert(T < NUM_INST_CNTS);
     return ScoreUBs[T];
   }
 
+  unsigned getScoreRange(InstCounterType T) const {
+    return getScoreUB(T) - getScoreLB(T);
+  }
+
   // Mapping from event to counter.
   InstCounterType eventCounter(WaitEventType E) {
-    if (WaitEventMaskForInst[VM_CNT] & (1 << E))
-      return VM_CNT;
-    if (WaitEventMaskForInst[LGKM_CNT] & (1 << E))
-      return LGKM_CNT;
-    if (WaitEventMaskForInst[VS_CNT] & (1 << E))
-      return VS_CNT;
-    assert(WaitEventMaskForInst[EXP_CNT] & (1 << E));
-    return EXP_CNT;
+    for (auto T : inst_counter_types()) {
+      if (WaitEventMaskForInst[T] & (1 << E))
+        return T;
+    }
+    llvm_unreachable("event type has no associated counter");
   }
 
-  unsigned getRegScore(int GprNo, InstCounterType T) {
+  unsigned getRegScore(int GprNo, InstCounterType T) const {
     if (GprNo < NUM_ALL_VGPRS) {
       return VgprScores[T][GprNo];
     }
     assert(T == LGKM_CNT);
     return SgprScores[GprNo - NUM_ALL_VGPRS];
   }
 
   bool merge(const WaitcntBrackets &Other);
 
   RegInterval getRegInterval(const MachineInstr *MI, const SIInstrInfo *TII,
                              const MachineRegisterInfo *MRI,
                              const SIRegisterInfo *TRI, unsigned OpNo) const;
 
   bool counterOutOfOrder(InstCounterType T) const;
   void simplifyWaitcnt(AMDGPU::Waitcnt &Wait) const;
   void simplifyWaitcnt(InstCounterType T, unsigned &Count) const;
-  void determineWait(InstCounterType T, unsigned ScoreToWait,
-                     AMDGPU::Waitcnt &Wait) const;
+  void determineWait(InstCounterType T, int RegNo, AMDGPU::Waitcnt &Wait) const;
   void applyWaitcnt(const AMDGPU::Waitcnt &Wait);
   void applyWaitcnt(InstCounterType T, unsigned Count);
   void updateByEvent(const SIInstrInfo *TII, const SIRegisterInfo *TRI,
                      const MachineRegisterInfo *MRI, WaitEventType E,
                      MachineInstr &MI);
 
   bool hasPending() const { return PendingEvents != 0; }
   bool hasPendingEvent(WaitEventType E) const {
     return PendingEvents & (1 << E);
   }
 
   bool hasMixedPendingEvents(InstCounterType T) const {
     unsigned Events = PendingEvents & WaitEventMaskForInst[T];
     // Return true if more than one bit is set in Events.
     return Events & (Events - 1);
private:
   void setScoreLB(InstCounterType T, unsigned Val) {
     assert(T < NUM_INST_CNTS);
     ScoreLBs[T] = Val;
   }
 
   void setScoreUB(InstCounterType T, unsigned Val) {
     assert(T < NUM_INST_CNTS);
     ScoreUBs[T] = Val;
     if (T == EXP_CNT) {
-      unsigned UB = ScoreUBs[T] - getWaitCountMax(EXP_CNT);
-      if (ScoreLBs[T] < UB && UB < ScoreUBs[T])
-        ScoreLBs[T] = UB;
+      unsigned UB = ScoreUBs[EXP_CNT] - getWaitCountMax(EXP_CNT);
+      if (ScoreLBs[EXP_CNT] < UB && UB < ScoreUBs[EXP_CNT])
+        ScoreLBs[EXP_CNT] = UB;
     }
   }
 
   void setRegScore(int GprNo, InstCounterType T, unsigned Val) {
     if (GprNo < NUM_ALL_VGPRS) {
       VgprUB = std::max(VgprUB, GprNo);
       VgprScores[T][GprNo] = Val;
     } else {
       assert(T == LGKM_CNT);
#endif
     }
   }
 }
 
 void WaitcntBrackets::print(raw_ostream &OS) {
   OS << '\n';
   for (auto T : inst_counter_types()) {
     unsigned LB = getScoreLB(T);
     unsigned UB = getScoreUB(T);
 
     switch (T) {
     case VM_CNT:
       OS << "    VM_CNT(" << UB - LB << "): ";
       break;
     case LGKM_CNT:
       OS << "    LGKM_CNT(" << UB - LB << "): ";
       break;
     case EXP_CNT:
       OS << "    EXP_CNT(" << UB - LB << "): ";
       break;
     case VS_CNT:
       OS << "    VS_CNT(" << UB - LB << "): ";
       break;
     default:
       OS << "    UNKNOWN(" << UB - LB << "): ";
       break;
     }
 
     if (LB < UB) {
       // Print vgpr scores.
       for (int J = 0; J <= VgprUB; J++) {
         unsigned RegScore = getRegScore(J, T);
         if (RegScore <= LB)
           continue;
         unsigned RelScore = RegScore - LB - 1;
void WaitcntBrackets::simplifyWaitcnt(AMDGPU::Waitcnt &Wait) const {
   simplifyWaitcnt(EXP_CNT, Wait.ExpCnt);
   simplifyWaitcnt(LGKM_CNT, Wait.LgkmCnt);
   simplifyWaitcnt(VS_CNT, Wait.VsCnt);
 }
 
 void WaitcntBrackets::simplifyWaitcnt(InstCounterType T,
                                       unsigned &Count) const {
   const unsigned LB = getScoreLB(T);
   const unsigned UB = getScoreUB(T);
 
   // The number of outstanding events for this type, T, can be calculated
   // as (UB - LB). If the current Count is greater than or equal to the number
   // of outstanding events, then the wait for this counter is redundant.
   if (Count >= UB - LB)
     Count = ~0u;
 }
 
-void WaitcntBrackets::determineWait(InstCounterType T, unsigned ScoreToWait,
+void WaitcntBrackets::determineWait(InstCounterType T, int RegNo,
                                     AMDGPU::Waitcnt &Wait) const {
+  unsigned ScoreToWait = getRegScore(RegNo, T);
+
   // If the score of src_operand falls within the bracket, we need an
   // s_waitcnt instruction.
   const unsigned LB = getScoreLB(T);
   const unsigned UB = getScoreUB(T);
   if ((UB >= ScoreToWait) && (ScoreToWait > LB)) {
     if ((T == VM_CNT || T == LGKM_CNT) &&
         hasPendingFlat() &&
         !ST->hasFlatLgkmVMemCountInOrder()) {
if (MI.isCall() && callWaitsOnFunctionEntry(MI)) {
           AMDGPU::getNamedOperandIdx(MI.getOpcode(), AMDGPU::OpName::src0);
 
       if (MI.getOperand(CallAddrOpIdx).isReg()) {
         RegInterval CallAddrOpInterval =
           ScoreBrackets.getRegInterval(&MI, TII, MRI, TRI, CallAddrOpIdx);
 
         for (int RegNo = CallAddrOpInterval.first;
              RegNo < CallAddrOpInterval.second; ++RegNo)
-          ScoreBrackets.determineWait(
-            LGKM_CNT, ScoreBrackets.getRegScore(RegNo, LGKM_CNT), Wait);
+          ScoreBrackets.determineWait(LGKM_CNT, RegNo, Wait);
 
         int RtnAddrOpIdx =
           AMDGPU::getNamedOperandIdx(MI.getOpcode(), AMDGPU::OpName::dst);
         if (RtnAddrOpIdx != -1) {
           RegInterval RtnAddrOpInterval =
             ScoreBrackets.getRegInterval(&MI, TII, MRI, TRI, RtnAddrOpIdx);
 
           for (int RegNo = RtnAddrOpInterval.first;
                RegNo < RtnAddrOpInterval.second; ++RegNo)
-            ScoreBrackets.determineWait(
-              LGKM_CNT, ScoreBrackets.getRegScore(RegNo, LGKM_CNT), Wait);
+            ScoreBrackets.determineWait(LGKM_CNT, RegNo, Wait);
         }
       }
     } else {
       // FIXME: Should not be relying on memoperands.
       // Look at the source operands of every instruction to see if
       // any of them results from a previous memory operation that affects
       // its current usage. If so, an s_waitcnt instruction needs to be
       // emitted.
if (MI.isCall() && callWaitsOnFunctionEntry(MI)) {
         unsigned AS = Memop->getAddrSpace();
         if (AS != AMDGPUAS::LOCAL_ADDRESS && AS != AMDGPUAS::FLAT_ADDRESS)
           continue;
         // No need to wait before load from VMEM to LDS.
         if (mayWriteLDSThroughDMA(MI))
           continue;
         unsigned RegNo = SQ_MAX_PGM_VGPRS + EXTRA_VGPR_LDS;
         // VM_CNT is only relevant to vgpr or LDS.
-        ScoreBrackets.determineWait(
-            VM_CNT, ScoreBrackets.getRegScore(RegNo, VM_CNT), Wait);
+        ScoreBrackets.determineWait(VM_CNT, RegNo, Wait);
         if (Memop->isStore()) {
-          ScoreBrackets.determineWait(
-              EXP_CNT, ScoreBrackets.getRegScore(RegNo, EXP_CNT), Wait);
+          ScoreBrackets.determineWait(EXP_CNT, RegNo, Wait);
         }
       }
 
       // Loop over use and def operands.
       for (unsigned I = 0, E = MI.getNumOperands(); I != E; ++I) {
         MachineOperand &Op = MI.getOperand(I);
         if (!Op.isReg())
           continue;
         RegInterval Interval =
             ScoreBrackets.getRegInterval(&MI, TII, MRI, TRI, I);
 
         const bool IsVGPR = TRI->isVectorRegister(*MRI, Op.getReg());
         for (int RegNo = Interval.first; RegNo < Interval.second; ++RegNo) {
           if (IsVGPR) {
             // RAW always needs an s_waitcnt. WAW needs an s_waitcnt unless the
             // previous write and this write are the same type of VMEM
             // instruction, in which case they're guaranteed to write their
             // results in order anyway.
             if (Op.isUse() || !SIInstrInfo::isVMEM(MI) ||
                 ScoreBrackets.hasOtherPendingVmemTypes(RegNo,
                                                        getVmemType(MI))) {
-              ScoreBrackets.determineWait(
-                  VM_CNT, ScoreBrackets.getRegScore(RegNo, VM_CNT), Wait);
+              ScoreBrackets.determineWait(VM_CNT, RegNo, Wait);
               ScoreBrackets.clearVgprVmemTypes(RegNo);
             }
             if (Op.isDef() || ScoreBrackets.hasPendingEvent(EXP_LDS_ACCESS)) {
-              ScoreBrackets.determineWait(
-                  EXP_CNT, ScoreBrackets.getRegScore(RegNo, EXP_CNT), Wait);
+              ScoreBrackets.determineWait(EXP_CNT, RegNo, Wait);
             }
           }
-          ScoreBrackets.determineWait(
-              LGKM_CNT, ScoreBrackets.getRegScore(RegNo, LGKM_CNT), Wait);
+          ScoreBrackets.determineWait(LGKM_CNT, RegNo, Wait);
         }
       }
     }
   }
 
   // The subtarget may have an implicit S_WAITCNT 0 before barriers. If it does
   // not, we need to ensure the subtarget is capable of backing off barrier
   // instructions in case there are any outstanding memory operations that may
   // cause an exception. Otherwise, insert an explicit S_WAITCNT 0 here.
   if (MI.getOpcode() == AMDGPU::S_BARRIER &&
       !ST->hasAutoWaitcntBeforeBarrier() && !ST->supportsBackOffBarrier()) {
     Wait = Wait.combined(AMDGPU::Waitcnt::allZero(ST->hasVscnt()));
   }
 
   // TODO: Remove this work-around, enable the assert for Bug 457939
   //       after fixing the scheduler. Also, the Shader Compiler code is
   //       independent of target.
   if (readsVCCZ(MI) && ST->hasReadVCCZBug()) {
     if (ScoreBrackets.getScoreLB(LGKM_CNT) <
             ScoreBrackets.getScoreUB(LGKM_CNT) &&
         ScoreBrackets.hasPendingEvent(SMEM_ACCESS)) {
       Wait.LgkmCnt = 0;
     }
   }
 
   // Verify that the wait is actually needed.
   ScoreBrackets.simplifyWaitcnt(Wait);
 
   if (ForceEmitZeroWaitcnts)
     Wait = AMDGPU::Waitcnt::allZero(ST->hasVscnt());
 
   if (ForceEmitWaitcnt[VM_CNT])
     Wait.VmCnt = 0;
   if (ForceEmitWaitcnt[EXP_CNT])
     Wait.ExpCnt = 0;
   if (ForceEmitWaitcnt[LGKM_CNT])
     Wait.LgkmCnt = 0;
   if (ForceEmitWaitcnt[VS_CNT])
     Wait.VsCnt = 0;
 
   if (FlushVmCnt) {
-    unsigned UB = ScoreBrackets.getScoreUB(VM_CNT);
-    unsigned LB = ScoreBrackets.getScoreLB(VM_CNT);
-    if (UB - LB != 0)
+    if (ScoreBrackets.getScoreRange(VM_CNT) != 0)
       Wait.VmCnt = 0;
   }
 
   return generateWaitcnt(Wait, MI.getIterator(), *MI.getParent(), ScoreBrackets,
                          OldWaitcntInstr);
 }
 
 // Add a waitcnt to flush the vmcnt counter at the end of the given block if
 // needed.
 bool SIInsertWaitcnts::generateWaitcntBlockEnd(MachineBasicBlock &Block,
                                                WaitcntBrackets &ScoreBrackets,
                                                MachineInstr *OldWaitcntInstr) {
   AMDGPU::Waitcnt Wait;
 
-  unsigned UB = ScoreBrackets.getScoreUB(VM_CNT);
-  unsigned LB = ScoreBrackets.getScoreLB(VM_CNT);
-  if (UB - LB == 0)
+  if (ScoreBrackets.getScoreRange(VM_CNT) == 0)
     return false;
 
   Wait.VmCnt = 0;
 
   return generateWaitcnt(Wait, Block.instr_end(), Block, ScoreBrackets,
                          OldWaitcntInstr);
 }
 
if (ST->hasReadVCCZBug() || !ST->partialVCCWritesUpdateVCCZ()) {
         // vccz bit, so when we detect that an instruction may read from a
         // corrupt vccz bit, we need to:
         // 1. Insert s_waitcnt lgkm(0) to wait for all outstanding SMRD
         //    operations to complete.
         // 2. Restore the correct value of vccz by writing the current value
         //    of vcc back to vcc.
         if (ST->hasReadVCCZBug() &&
             ScoreBrackets.getScoreLB(LGKM_CNT) <
                 ScoreBrackets.getScoreUB(LGKM_CNT) &&
             ScoreBrackets.hasPendingEvent(SMEM_ACCESS)) {
           // Writes to vcc while there's an outstanding smem read may get
           // clobbered as soon as any read completes.
           VCCZCorrect = false;
         } else {
           // Writes to vcc will fix any incorrect value in vccz.
           VCCZCorrect = true;
         }