[AMDGPU] Waitcnt pass: propagate info in the case of a single basic block loop

[msearles]

Summary

Modify the waitcnt pass to propagate info in the case of a single basic block loop. mergeInputScoreBrackets() does this for us; update it so that it processes the single bb's score bracket when processing the single bb's preds. It is, after all, a pred of itself, so it's score bracket is needed.

[rampitec]

lib/Target/AMDGPU/SIInsertWaitcnts.cpp

1313 ↗

(On Diff #138201)

I think you can use std::find().

1328 ↗

(On Diff #138201)

Do you need to make a copy if block is not a self predecessor? You are not going to reuse it then as far as I understand.

[msearles]

Adjust per reviewer comments.

[rampitec]

LGTM, just fix formatting before submit.

lib/Target/AMDGPU/SIInsertWaitcnts.cpp

1320 ↗

(On Diff #138391)

Please add spaces after commas.

[msearles]

Fix formatting per reviewer comments

llvm/trunk/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

 };
 
 #define ForAllWaitEventType(w)                                                 \
   for (enum WaitEventType w = (enum WaitEventType)0;                           \
        (w) < (enum WaitEventType)NUM_WAIT_EVENTS;                              \
        (w) = (enum WaitEventType)((w) + 1))
 
 // This is a per-basic-block object that maintains current score brackets
-// of each wait-counter, and a per-register scoreboard for each wait-couner.
+// of each wait counter, and a per-register scoreboard for each wait counter.
 // We also maintain the latest score for every event type that can change the
 // waitcnt in order to know if there are multiple types of events within
 // the brackets. When multiple types of event happen in the bracket,
-// wait-count may get decreased out of order, therefore we need to put in
+// wait count may get decreased out of order, therefore we need to put in
 // "s_waitcnt 0" before use.
 class BlockWaitcntBrackets {
 public:
   BlockWaitcntBrackets() {
     for (enum InstCounterType T = VM_CNT; T < NUM_INST_CNTS;
          T = (enum InstCounterType)(T + 1)) {
       memset(VgprScores[T], 0, sizeof(VgprScores[T]));
     }
unsigned int BlockWaitcntBrackets::updateByWait(InstCounterType T,
   if ((UB >= ScoreToWait) && (ScoreToWait > LB)) {
     if (T == VM_CNT && hasPendingFlat()) {
       // If there is a pending FLAT operation, and this is a VM waitcnt,
       // then we need to force a waitcnt 0 for VM.
       NeedWait = CNT_MASK(T);
       setScoreLB(T, getScoreUB(T));
     } else if (counterOutOfOrder(T)) {
       // Counter can get decremented out-of-order when there
-      // are multiple types event in the brack. Also emit an s_wait counter
+      // are multiple types event in the bracket. Also emit an s_wait counter
       // with a conservative value of 0 for the counter.
       NeedWait = CNT_MASK(T);
       setScoreLB(T, getScoreUB(T));
     } else {
       NeedWait = CNT_MASK(T);
       setScoreLB(T, ScoreToWait);
     }
   }
case AMDGPU::S_MEMREALTIME:
       ScoreBrackets->updateByEvent(TII, TRI, MRI, SMEM_ACCESS, Inst);
       break;
     default:
       break;
     }
   }
 }
 
+// Merge the score brackets of the Block's predecessors;
+// this merged score bracket is used when adding waitcnts to the Block
 void SIInsertWaitcnts::mergeInputScoreBrackets(MachineBasicBlock &Block) {
   BlockWaitcntBrackets *ScoreBrackets = BlockWaitcntBracketsMap[&Block].get();
   int32_t MaxPending[NUM_INST_CNTS] = {0};
   int32_t MaxFlat[NUM_INST_CNTS] = {0};
   bool MixedExpTypes = false;
 
-  // Clear the score bracket state.
-  ScoreBrackets->clear();
-
-  // Compute the number of pending elements on block entry.
+  // For single basic block loops, we need to retain the Block's
+  // score bracket to have accurate Pred info. So, make a copy of Block's
+  // score bracket, clear() it (which retains several important bits of info),
+  // populate, and then replace en masse. For non-single basic block loops,
+  // just clear Block's current score bracket and repopulate in-place.
+  bool IsSelfPred;
+  std::unique_ptr<BlockWaitcntBrackets> S;
+
+  IsSelfPred = (std::find(Block.pred_begin(), Block.pred_end(), &Block))
+    != Block.pred_end();
+  if (IsSelfPred) {
+    S = llvm::make_unique<BlockWaitcntBrackets>(*ScoreBrackets);
+    ScoreBrackets = S.get();
+  }
 
-  // IMPORTANT NOTE: If iterative handling of loops is added, the code will
-  // need to handle single BBs with backedges to themselves. This means that
-  // they will need to retain and not clear their initial state.
+  ScoreBrackets->clear();
 
   // See if there are any uninitialized predecessors. If so, emit an
   // s_waitcnt 0 at the beginning of the block.
-  for (MachineBasicBlock *pred : Block.predecessors()) {
+  for (MachineBasicBlock *Pred : Block.predecessors()) {
     BlockWaitcntBrackets *PredScoreBrackets =
-        BlockWaitcntBracketsMap[pred].get();
-    bool Visited = BlockVisitedSet.count(pred);
+        BlockWaitcntBracketsMap[Pred].get();
+    bool Visited = BlockVisitedSet.count(Pred);
     if (!Visited || PredScoreBrackets->getWaitAtBeginning()) {
       continue;
     }
     for (enum InstCounterType T = VM_CNT; T < NUM_INST_CNTS;
          T = (enum InstCounterType)(T + 1)) {
       int span =
           PredScoreBrackets->getScoreUB(T) - PredScoreBrackets->getScoreLB(T);
       MaxPending[T] = std::max(MaxPending[T], span);
if (pred_exp_ub > PredScoreBrackets->getScoreLB(EXP_CNT)) {
                        PredScoreBrackets->getScoreUB(EXP_CNT);
       if (new_exp_ub > 0) {
         ScoreBrackets->setEventUB(
             EXP_GPR_LOCK,
             std::max(ScoreBrackets->getEventUB(EXP_GPR_LOCK), new_exp_ub));
       }
     }
   }
+
+  // if a single block loop, update the score brackets. Not needed for other
+  // blocks, as we did this in-place
+  if (IsSelfPred) {
+    BlockWaitcntBracketsMap[&Block] = llvm::make_unique<BlockWaitcntBrackets>(*ScoreBrackets);
+  }
 }
 
 /// Return the "bottom" block of a loop. This differs from
 /// MachineLoop::getBottomBlock in that it works even if the loop is
 /// discontiguous.
 MachineBasicBlock *SIInsertWaitcnts::loopBottom(const MachineLoop *Loop) {
   MachineBasicBlock *Bottom = Loop->getHeader();
   for (MachineBasicBlock *MBB : Loop->blocks())

llvm/trunk/test/CodeGen/AMDGPU/waitcnt-loop-single-basic-block.mir

+# RUN: llc -o - %s -march=amdgcn -run-pass=si-insert-waitcnts -verify-machineinstrs | FileCheck -check-prefix=GCN %s
+
+# Check that the waitcnt propogates info in the case of a single basic block loop
+
+# GCN-LABEL: waitcnt-loop-single-basic-block
+# GCN: bb.0
+# GCN: S_WAITCNT 3952
+# GCN-NEXT: GLOBAL_STORE_DWORD
+# GCN: S_WAITCNT 3953
+# GCN-NEXT: GLOBAL_STORE_DWORD
+
+...
+name: waitcnt-loop-single-basic-block
+body: |
+  bb.0:
+    S_BRANCH %bb.1
+  bb.1:
+    GLOBAL_STORE_DWORD $vgpr7_vgpr8, $vgpr11, 0, 0, 0, implicit $exec
+    $vgpr21 = GLOBAL_LOAD_DWORD $vgpr4_vgpr5, 0, 0, 0, implicit $exec
+    $vgpr10 = GLOBAL_LOAD_DWORD $vgpr10_vgpr11, 0, 0, 0, implicit $exec
+    GLOBAL_STORE_DWORD $vgpr14_vgpr15, $vgpr21, 0, 0, 0, implicit $exec
+    $vgpr11 = GLOBAL_LOAD_DWORD $vgpr11_vgpr12, 0, 0, 0, implicit $exec
+    S_CBRANCH_SCC1 %bb.1, implicit $scc
+  bb.2:
+    S_ENDPGM
+...