[AMDGPU][Waitcnt] Fix handling of flat instrs

[msearles]

Summary

On GFX9 and earlier, flat memory ops may decrement VMCNT out-of-order as well as LGKMCNT out-of-order.

[t-tye]

lib/Target/AMDGPU/SIInsertWaitcnts.cpp

740–741

Update comment to:

// If there is a pending FLAT operation, this is a VM or
// LGKM waitcnt, and the target can report early
// completion, then we need to force a waitcnt 0.}

[msearles]

Update comment as suggested by reviewer.

[arsenm]

test/CodeGen/AMDGPU/waitcnt.mir

33–37

Since this is a change for a specific subtarget, should this test have multiple sets of check lines instead of just changing these

[msearles]

Add subtarget-specific checks

[msearles]

ping

[msearles]

ping

[rampitec]

LGTM

lib/Target/AMDGPU/AMDGPUSubtarget.h

public:
   bool hasFlatGlobalInsts() const {
     return FlatGlobalInsts;
   }
 
   bool hasFlatScratchInsts() const {
     return FlatScratchInsts;
   }
 
+  bool hasFlatLgkmVMemCountInOrder() const {
+    return getGeneration() > GFX9;
+  }
+
   bool hasD16LoadStore() const {
     return getGeneration() >= GFX9;
   }
 
   /// Return if most LDS instructions have an m0 use that require m0 to be
   /// iniitalized.
   bool ldsRequiresM0Init() const {
     return getGeneration() < GFX9;

lib/Target/AMDGPU/SIInsertWaitcnts.cpp

 // 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
 // "s_waitcnt 0" before use.
 class BlockWaitcntBrackets {
 public:
-  BlockWaitcntBrackets() {
+  BlockWaitcntBrackets(const SISubtarget *SubTarget) : ST(SubTarget) {
     for (enum InstCounterType T = VM_CNT; T < NUM_INST_CNTS;
          T = (enum InstCounterType)(T + 1)) {
       memset(VgprScores[T], 0, sizeof(VgprScores[T]));
     }
   }
 
   ~BlockWaitcntBrackets() = default;
 
public:
   void setMixedExpTypes(bool MixedExpTypesIn) {
     MixedExpTypes = MixedExpTypesIn;
   }
 
   void print(raw_ostream &);
   void dump() { print(dbgs()); }
 
 private:
+  const SISubtarget *ST = nullptr;
   bool WaitAtBeginning = false;
   bool RevisitLoop = false;
   bool MixedExpTypes = false;
   int32_t PostOrder = 0;
   MachineInstr *Waitcnt = nullptr;
   int32_t ScoreLBs[NUM_INST_CNTS] = {0};
   int32_t ScoreUBs[NUM_INST_CNTS] = {0};
   int32_t EventUBs[NUM_WAIT_EVENTS] = {0};
if (ScoreToWait == -1) {
     setScoreLB(T, getScoreUB(T));
     return NeedWait;
   }
 
   // If the score of src_operand falls within the bracket, we need an
   // s_waitcnt instruction.
   const int32_t LB = getScoreLB(T);
   const int32_t UB = getScoreUB(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.
+    if ((T == VM_CNT || T == LGKM_CNT) &&
+        hasPendingFlat() &&
+        !ST->hasFlatLgkmVMemCountInOrder()) {
+      // If there is a pending FLAT operation, and this is a VMem or LGKM
+      // waitcnt and the target can report early completion, then we need
+      // to force a waitcnt 0.
       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 bracket. Also emit an s_wait counter
       // with a conservative value of 0 for the counter.
       NeedWait = CNT_MASK(T);
       setScoreLB(T, getScoreUB(T));
if (EmitWaitcnt != 0 || IsForceEmitWaitcnt) {
         MachineLoop *ContainingLoop = MLI->getLoopFor(MI.getParent());
         if (ContainingLoop) {
           MachineBasicBlock *TBB = ContainingLoop->getHeader();
           BlockWaitcntBrackets *ScoreBracket =
               BlockWaitcntBracketsMap[TBB].get();
           if (!ScoreBracket) {
             assert(!BlockVisitedSet.count(TBB));
             BlockWaitcntBracketsMap[TBB] =
-                llvm::make_unique<BlockWaitcntBrackets>();
+                llvm::make_unique<BlockWaitcntBrackets>(ST);
             ScoreBracket = BlockWaitcntBracketsMap[TBB].get();
           }
           ScoreBracket->setRevisitLoop(true);
           DEBUG(dbgs() << "set-revisit: Block"
                        << ContainingLoop->getHeader()->getNumber() << '\n';);
         }
       }
 
for (ReversePostOrderTraversal<MachineFunction *>::rpo_iterator
            E = RPOT.end(), J = RPOT.begin();
        I != E;) {
     MachineBasicBlock &MBB = **I;
 
     BlockVisitedSet.insert(&MBB);
 
     BlockWaitcntBrackets *ScoreBrackets = BlockWaitcntBracketsMap[&MBB].get();
     if (!ScoreBrackets) {
-      BlockWaitcntBracketsMap[&MBB] = llvm::make_unique<BlockWaitcntBrackets>();
+      BlockWaitcntBracketsMap[&MBB] = llvm::make_unique<BlockWaitcntBrackets>(ST);
       ScoreBrackets = BlockWaitcntBracketsMap[&MBB].get();
     }
     ScoreBrackets->setPostOrder(MBB.getNumber());
     MachineLoop *ContainingLoop = MLI->getLoopFor(&MBB);
     if (ContainingLoop && LoopWaitcntDataMap[ContainingLoop] == nullptr)
       LoopWaitcntDataMap[ContainingLoop] = llvm::make_unique<LoopWaitcntData>();
 
     // If we are walking into the block from before the loop, then guarantee

test/CodeGen/AMDGPU/waitcnt.mir

 # CHECK: FLAT_LOAD_DWORD
 # CHECK: FLAT_LOAD_DWORDX4
 # Global loads will return in order so we should:
 # s_waitcnt vmcnt(1) lgkmcnt(1)
 # CHECK-NEXT: S_WAITCNT 369
 
 # CHECK-LABEL: bb.1:
 # CHECK: FLAT_LOAD_DWORD
-# CHECK: S_WAITCNT 368
+# CHECK: S_WAITCNT 112
 # CHECK: FLAT_LOAD_DWORDX4
-# The first load has no mem operand, so we should assume it accesses the flat
-# address space.
-# s_waitcnt lgkmcnt(1)
-# CHECK-NEXT: S_WAITCNT 383
 
 # CHECK-LABEL: bb.2:
 # CHECK: FLAT_LOAD_DWORD
-# CHECK: S_WAITCNT 368
+# CHECK: S_WAITCNT 112
 # CHECK: FLAT_LOAD_DWORDX4
 
-# One outstanding load accesses the flat address space.
-# s_waitcnt lgkmcnt(1)
-# CHECK-NEXT: S_WAITCNT 383
-
 name: flat_zero_waitcnt
 
 body: |
   bb.0:
     successors: %bb.1
     $vgpr0 = FLAT_LOAD_DWORD $vgpr1_vgpr2, 0, 0, 0, implicit $exec, implicit $flat_scr :: (load 4 from %ir.global4)
     $vgpr3_vgpr4_vgpr5_vgpr6 = FLAT_LOAD_DWORDX4 $vgpr7_vgpr8, 0, 0, 0, implicit $exec, implicit $flat_scr :: (load 16 from %ir.global16)
     $vgpr0 = V_MOV_B32_e32 $vgpr1, implicit $exec