[AMDGPU][Waitcnt] Fix handling of loops with many bottom blocks

[msearles]

Summary

In terms of waitcnt insertion/if necessary, the waitcnt pass forces convergence for a loop. Previously, that kicked if greater than 2 passes over a loop, which doesn't account for loop with many bottom blocks. So, increase the threshold to (n+1), where n is the number of bottom blocks. This gives the pass an opportunity to consider the contribution of each bottom block, to the overall loop, before the forced convergence potentially kicks in.

[rampitec]

LGTM

lib/Target/AMDGPU/SIInsertWaitcnts.cpp

 // at the end of the loop footer.
 class LoopWaitcntData {
 public:
   LoopWaitcntData() = default;
   ~LoopWaitcntData() = default;
 
   void incIterCnt() { IterCnt++; }
   void resetIterCnt() { IterCnt = 0; }
-  int32_t getIterCnt() { return IterCnt; }
+  unsigned getIterCnt() { return IterCnt; }
 
   void setWaitcnt(MachineInstr *WaitcntIn) { LfWaitcnt = WaitcntIn; }
   MachineInstr *getWaitcnt() const { return LfWaitcnt; }
 
   void print() { LLVM_DEBUG(dbgs() << "  iteration " << IterCnt << '\n';); }
 
 private:
   // s_waitcnt added at the end of loop footer to stablize wait scores
if (EmitWaitcnt != 0 || IsForceEmitWaitcnt) {
           if (!ScoreBracket) {
             assert(!BlockVisitedSet.count(TBB));
             BlockWaitcntBracketsMap[TBB] =
                 llvm::make_unique<BlockWaitcntBrackets>();
             ScoreBracket = BlockWaitcntBracketsMap[TBB].get();
           }
           ScoreBracket->setRevisitLoop(true);
           LLVM_DEBUG(dbgs()
-                         << "set-revisit: Block"
+                         << "set-revisit2: Block"
                          << ContainingLoop->getHeader()->getNumber() << '\n';);
         }
       }
 
       // Update an existing waitcount, or make a new one.
       unsigned Enc = AMDGPU::encodeWaitcnt(IV,
                       ForceEmitWaitcnt[VM_CNT] ? 0 : CntVal[VM_CNT],
                       ForceEmitWaitcnt[EXP_CNT] ? 0 : CntVal[EXP_CNT],
#endif
 
   // 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 true if the given basic block is a "bottom" block of a loop. This
-/// differs from MachineLoop::getBottomBlock in that it works even if the loop
-/// is discontiguous. This also handles multiple back-edges for the same
-/// "header" block of a loop.
+/// Return true if the given basic block is a "bottom" block of a loop.
+/// This works even if the loop is discontiguous. This also handles
+/// multiple back-edges for the same "header" block of a loop.
 bool SIInsertWaitcnts::isLoopBottom(const MachineLoop *Loop,
                                     const MachineBasicBlock *Block) {
   for (MachineBasicBlock *MBB : Loop->blocks()) {
     if (MBB == Block && MBB->isSuccessor(Loop->getHeader())) {
       return true;
     }
   }
   return false;
#endif
   // Check if we need to force convergence at loop footer.
   MachineLoop *ContainingLoop = MLI->getLoopFor(&Block);
   if (ContainingLoop && isLoopBottom(ContainingLoop, &Block)) {
     LoopWaitcntData *WaitcntData = LoopWaitcntDataMap[ContainingLoop].get();
     WaitcntData->print();
     LLVM_DEBUG(dbgs() << '\n';);
 
     // The iterative waitcnt insertion algorithm aims for optimal waitcnt
-    // placement and doesn't always guarantee convergence for a loop. Each
-    // loop should take at most 2 iterations for it to converge naturally.
-    // When this max is reached and result doesn't converge, we force
-    // convergence by inserting a s_waitcnt at the end of loop footer.
-    if (WaitcntData->getIterCnt() > 2) {
+    // placement, but doesn't guarantee convergence for a loop. Each
+    // loop should take at most (n+1) iterations for it to converge naturally,
+    // where n is the number of bottom blocks. If this threshold is reached and
+    // the result hasn't converged, then we force convergence by inserting
+    // a s_waitcnt at the end of loop footer.
+    if (WaitcntData->getIterCnt() > (countNumBottomBlocks(ContainingLoop) + 1)) {
       // To ensure convergence, need to make wait events at loop footer be no
       // more than those from the previous iteration.
       // As a simplification, instead of tracking individual scores and
       // generating the precise wait count, just wait on 0.
       bool HasPending = false;
       MachineInstr *SWaitInst = WaitcntData->getWaitcnt();
       for (enum InstCounterType T = VM_CNT; T < NUM_INST_CNTS;
            T = (enum InstCounterType)(T + 1)) {
         if (ScoreBrackets->getScoreUB(T) > ScoreBrackets->getScoreLB(T)) {
           ScoreBrackets->setScoreLB(T, ScoreBrackets->getScoreUB(T));
           HasPending = true;
+          break;
         }
       }
 
       if (HasPending) {
         if (!SWaitInst) {
-          SWaitInst = Block.getParent()->CreateMachineInstr(
-              TII->get(AMDGPU::S_WAITCNT), DebugLoc());
+          SWaitInst = BuildMI(Block, Block.getFirstNonPHI(),
+                              DebugLoc(), TII->get(AMDGPU::S_WAITCNT))
+                              .addImm(0);
           TrackedWaitcntSet.insert(SWaitInst);
-          const MachineOperand &Op = MachineOperand::CreateImm(0);
-          SWaitInst->addOperand(MF, Op);
 #if 0 // TODO: Format the debug output
           OutputTransformBanner("insertWaitcntInBlock",0,"Create:",context);
           OutputTransformAdd(SWaitInst, context);
 #endif
         }
 #if 0 // TODO: ??
         _DEV( REPORTED_STATS->force_waitcnt_converge = 1; )
 #endif
for (ReversePostOrderTraversal<MachineFunction *>::rpo_iterator
     if (ContainingLoop && ContainingLoop->getHeader() == &MBB) {
       unsigned Count = countNumBottomBlocks(ContainingLoop);
 
       // If the loop has multiple back-edges, and so more than one "bottom"
       // basic block, we have to guarantee a re-walk over every blocks.
       if ((std::count(BlockWaitcntProcessedSet.begin(),
                       BlockWaitcntProcessedSet.end(), &MBB) < Count)) {
         BlockWaitcntBracketsMap[&MBB]->setRevisitLoop(true);
-        LLVM_DEBUG(dbgs() << "set-revisit: Block"
+        LLVM_DEBUG(dbgs() << "set-revisit1: Block"
                           << ContainingLoop->getHeader()->getNumber() << '\n';);
       }
     }
 
     // Walk over the instructions.
     insertWaitcntInBlock(MF, MBB);
 
-    // Flag that waitcnts have been processed at least once.
+    // Record that waitcnts have been processed at least once for this block.
     BlockWaitcntProcessedSet.push_back(&MBB);
 
     // See if we want to revisit the loop. If a loop has multiple back-edges,
     // we shouldn't revisit the same "bottom" basic block.
     if (ContainingLoop && isLoopBottom(ContainingLoop, &MBB) &&
         std::count(BlockWaitcntProcessedSet.begin(),
                    BlockWaitcntProcessedSet.end(), &MBB) == 1) {
       MachineBasicBlock *EntryBB = ContainingLoop->getHeader();
bool SIInsertWaitcnts::runOnMachineFunction(MachineFunction &MF) {
   if (!MFI->isEntryFunction()) {
     // Wait for any outstanding memory operations that the input registers may
     // depend on. We can't track them and it's better to the wait after the
     // costly call sequence.
 
     // TODO: Could insert earlier and schedule more liberally with operations
     // that only use caller preserved registers.
     MachineBasicBlock &EntryBB = MF.front();
-    BuildMI(EntryBB, EntryBB.getFirstNonPHI(), DebugLoc(), TII->get(AMDGPU::S_WAITCNT))
+    auto SWaitInst = BuildMI(EntryBB, EntryBB.getFirstNonPHI(),
+                             DebugLoc(), TII->get(AMDGPU::S_WAITCNT))
                              .addImm(0);
 
+    LLVM_DEBUG(dbgs() << "insertWaitcntInBlock\n"
+               << "New Instr: " << *SWaitInst << '\n');
+
     Modified = true;
   }
 
   return Modified;
 }

test/CodeGen/AMDGPU/waitcnt-back-edge-loop.mir

bb.4:
     S_BRANCH %bb.3
 
   bb.5:
 
     $vgpr4 = V_MAC_F32_e32 killed $vgpr0, killed $vgpr3, killed $vgpr4, implicit $exec
     EXP_DONE 12, killed $vgpr4, undef $vgpr0, undef $vgpr0, undef $vgpr0, 0, 0, 15, implicit $exec
     S_ENDPGM
 ...
+---
+
+# GCN-LABEL: name: waitcnt-multiple-back-edges{{$}}
+# GCN: bb.0:
+# GCN: S_WAITCNT 0
+# GCN-NEXT: S_BRANCH %bb.2
+
+name: waitcnt-multiple-back-edges
+body: |
+  bb.0:
+    S_BRANCH %bb.2
+
+  bb.1:
+    S_BRANCH %bb.2
+
+  bb.2:
+    S_CBRANCH_VCCZ %bb.1, implicit $vcc
+
+  bb.3:
+    S_CBRANCH_VCCNZ %bb.5, implicit $vcc
+
+  bb.4:
+    BUFFER_ATOMIC_ADD_OFFSET renamable $vgpr0, renamable $sgpr12_sgpr13_sgpr14_sgpr15, 0, 4, 0, implicit $exec
+    S_CBRANCH_SCC0 %bb.2, implicit $scc
+    S_BRANCH %bb.6
+
+  bb.5:
+    S_CBRANCH_SCC0 %bb.2, implicit $scc
+    S_BRANCH %bb.6
+
+  bb.6:
+    S_CBRANCH_SCC1 %bb.0, implicit $scc
+    S_ENDPGM
+...