[AMDGPU] Fix extra waitcnt being added with BUFFER_INVL2

[kerbowa]

Summary

The waitcnt pass would increment the number of vmem events for some buffer
invalidates that were not handled by the pass.

[rampitec]

LGTM with a nit: typo in comment.

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

950

Typo: affect.

[t-tye]

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

1323

Is this the right query? Or is the query really asking if the instruction requires the vmem counts to be tracked? Should the writeback also be handled here? Are invalidate/writeback considered to load/store memory?

llvm/lib/Target/AMDGPU/BUFInstructions.td

class MUBUF_Pseudo <string opName, dag outs, dag ins,
   bits<1> has_srsrc   = 1;
   bits<1> has_soffset = 1;
   bits<1> has_offset  = 1;
   bits<1> has_slc     = 1;
   bits<1> has_tfe     = 1;
   bits<4> elements    = 0;
   bits<1> has_sccb    = 1;
   bits<1> sccb_value  = 0;
+  bits<1> IsBufferInv = 0;
 }
 
 class MUBUF_Real <MUBUF_Pseudo ps> :
   InstSI <ps.OutOperandList, ps.InOperandList, ps.Mnemonic # ps.AsmOperands, []> {
 
   let isPseudo = 0;
   let isCodeGenOnly = 0;
 
class MUBUF_Invalidate <string opName, SDPatternOperator node = null_frag> :
   MUBUF_Pseudo<opName, (outs), (ins), "", [(node)]> {
 
   let AsmMatchConverter = "";
 
   let hasSideEffects = 1;
   let mayLoad = 0;
   let mayStore = 0;
 
-  // Set everything to 0.
+  let IsBufferInv = 1;
+  // Set everything else to 0.
   let offen       = 0;
   let idxen       = 0;
   let addr64      = 0;
   let has_vdata   = 0;
   let has_vaddr   = 0;
   let has_glc     = 0;
   let has_dlc     = 0;
   let glc_value   = 0;
let SubtargetPredicate = HasPackedD16VMem in {
   defm TBUFFER_STORE_FORMAT_D16_XY   : MTBUF_Real_AllAddr_vi <0x0d>;
   defm TBUFFER_STORE_FORMAT_D16_XYZ  : MTBUF_Real_AllAddr_vi <0x0e>;
   defm TBUFFER_STORE_FORMAT_D16_XYZW : MTBUF_Real_AllAddr_vi <0x0f>;
 } // End HasUnpackedD16VMem.
 
 def MUBUFInfoTable : GenericTable {
   let FilterClass = "MUBUF_Pseudo";
   let CppTypeName = "MUBUFInfo";
-  let Fields = ["Opcode", "BaseOpcode", "elements", "has_vaddr", "has_srsrc", "has_soffset"];
+  let Fields = [
+    "Opcode", "BaseOpcode", "elements", "has_vaddr", "has_srsrc", "has_soffset",
+    "IsBufferInv"
+  ];
 
   let PrimaryKey = ["Opcode"];
   let PrimaryKeyName = "getMUBUFOpcodeHelper";
 }
 
 def getMUBUFInfoFromOpcode : SearchIndex {
   let Table = MUBUFInfoTable;
   let Key = ["Opcode"];

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

 /// multiple event types are in flight.
 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
 #include "GCNSubtarget.h"
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
 #include "SIMachineFunctionInfo.h"
+#include "Utils/AMDGPUBaseInfo.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/CodeGen/MachinePostDominators.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/DebugCounter.h"
 #include "llvm/Support/TargetParser.h"
 using namespace llvm;
 
bool SIInsertWaitcnts::generateWaitcntInstBefore(
   setForceEmitWaitcnt();
 
   if (MI.isMetaInstruction())
     return false;
 
   AMDGPU::Waitcnt Wait;
   bool Modified = false;
 
-  // See if this instruction has a forced S_WAITCNT VM.
-  // TODO: Handle other cases of NeedsWaitcntVmBefore()
+  // FIXME: This should have already been handled by the memory legalizer.
+  // Removing this currently doesn't affect any lit tests, but we need to
+  // verify that nothing was relying on this. The number of buffer invalidates
+  // being handled here should not be expanded.
   if (MI.getOpcode() == AMDGPU::BUFFER_WBINVL1 ||
       MI.getOpcode() == AMDGPU::BUFFER_WBINVL1_SC ||
       MI.getOpcode() == AMDGPU::BUFFER_WBINVL1_VOL ||
       MI.getOpcode() == AMDGPU::BUFFER_GL0_INV ||
       MI.getOpcode() == AMDGPU::BUFFER_GL1_INV) {
     Wait.VmCnt = 0;
   }
 
if (TII->isDS(Inst) && TII->usesLGKM_CNT(Inst)) {
     assert(FlatASCount);
 
     // This is a flat memory operation that access both VMEM and LDS, so note it
     // - it will require that both the VM and LGKM be flushed to zero if it is
     // pending when a VM or LGKM dependency occurs.
     if (FlatASCount > 1)
       ScoreBrackets->setPendingFlat();
   } else if (SIInstrInfo::isVMEM(Inst) &&
-             // TODO: get a better carve out.
-             Inst.getOpcode() != AMDGPU::BUFFER_WBINVL1 &&
-             Inst.getOpcode() != AMDGPU::BUFFER_WBINVL1_SC &&
-             Inst.getOpcode() != AMDGPU::BUFFER_WBINVL1_VOL &&
-             Inst.getOpcode() != AMDGPU::BUFFER_GL0_INV &&
-             Inst.getOpcode() != AMDGPU::BUFFER_GL1_INV) {
+             !llvm::AMDGPU::getMUBUFIsBufferInv(Inst.getOpcode())) {
     if (!ST->hasVscnt())
       ScoreBrackets->updateByEvent(TII, TRI, MRI, VMEM_ACCESS, Inst);
     else if ((Inst.mayLoad() && !SIInstrInfo::isAtomicNoRet(Inst)) ||
              /* IMAGE_GET_RESINFO / IMAGE_GET_LOD */
              (TII->isMIMG(Inst) && !Inst.mayLoad() && !Inst.mayStore()))
       ScoreBrackets->updateByEvent(TII, TRI, MRI, VMEM_READ_ACCESS, Inst);
     else if (Inst.mayStore())
       ScoreBrackets->updateByEvent(TII, TRI, MRI, VMEM_WRITE_ACCESS, Inst);

llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.h

 
 LLVM_READONLY
 bool getMUBUFHasSrsrc(unsigned Opc);
 
 LLVM_READONLY
 bool getMUBUFHasSoffset(unsigned Opc);
 
 LLVM_READONLY
+bool getMUBUFIsBufferInv(unsigned Opc);
+
+LLVM_READONLY
 bool getSMEMIsBuffer(unsigned Opc);
 
 LLVM_READONLY
 bool getVOP1IsSingle(unsigned Opc);
 
 LLVM_READONLY
 bool getVOP2IsSingle(unsigned Opc);
 

llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp

 
 struct MUBUFInfo {
   uint16_t Opcode;
   uint16_t BaseOpcode;
   uint8_t elements;
   bool has_vaddr;
   bool has_srsrc;
   bool has_soffset;
+  bool IsBufferInv;
 };
 
 struct MTBUFInfo {
   uint16_t Opcode;
   uint16_t BaseOpcode;
   uint8_t elements;
   bool has_vaddr;
   bool has_srsrc;
bool getMUBUFHasSrsrc(unsigned Opc) {
   return Info ? Info->has_srsrc : false;
 }
 
 bool getMUBUFHasSoffset(unsigned Opc) {
   const MUBUFInfo *Info = getMUBUFOpcodeHelper(Opc);
   return Info ? Info->has_soffset : false;
 }
 
+bool getMUBUFIsBufferInv(unsigned Opc) {
+  const MUBUFInfo *Info = getMUBUFOpcodeHelper(Opc);
+  return Info ? Info->IsBufferInv : false;
+}
+
 bool getSMEMIsBuffer(unsigned Opc) {
   const SMInfo *Info = getSMEMOpcodeHelper(Opc);
   return Info ? Info->IsBuffer : false;
 }
 
 bool getVOP1IsSingle(unsigned Opc) {
   const VOPInfo *Info = getVOP1OpcodeHelper(Opc);
   return Info ? Info->IsSingle : false;

llvm/test/CodeGen/AMDGPU/waitcnt-preexisting.mir

bb.0:
     ; GFX9-LABEL: name: test_waitcnt_preexisting_ignore_kill
     ; GFX9: S_WAITCNT 0
     ; GFX9: $vgpr0 = FLAT_LOAD_DWORD $vgpr0_vgpr1, 0, 0, implicit $exec, implicit $flat_scr
     ; GFX9: S_WAITCNT 3952
     ; GFX9: KILL $vgpr0
     $vgpr0 = FLAT_LOAD_DWORD $vgpr0_vgpr1, 0, 0, implicit $exec, implicit $flat_scr
     S_WAITCNT 3952
     KILL $vgpr0
+...
 
 # Combine preexisting waitcnt with wait added to the start of a non-entry function.
 
 ---
 name:            test_waitcnt_preexisting_func_start
 body:             |
   bb.0:
     ; GFX9-LABEL: name: test_waitcnt_preexisting_func_start
     ; GFX9: S_WAITCNT 0
     ; GFX9-NOT: S_WAITCNT 0
     ; GFX9: S_ENDPGM 0
     S_WAITCNT 0
     S_ENDPGM 0
 ...
+
+# Verify that extra waitcnt are not added after buffer invalidate instructions.
+
+---
+name:            test_waitcnt_preexisting_buffer_inv
+body:             |
+  bb.0:
+    ; GFX9-LABEL: name: test_waitcnt_preexisting_buffer_inv
+    ; GFX9: S_WAITCNT 0
+    ; GFX9: $vgpr0_vgpr1 = GLOBAL_LOAD_DWORDX2 $vgpr0_vgpr1, 0, 0, implicit $exec
+    ; GFX9: S_WAITCNT 3952
+    ; GFX9: BUFFER_INVL2 implicit $exec
+    ; GFX9: BUFFER_WBINVL1_VOL implicit $exec
+    ; GFX9: $vgpr0 = FLAT_LOAD_DWORD $vgpr0_vgpr1, 0, 0, implicit $exec, implicit $flat_scr
+    ; GFX9: S_WAITCNT 112
+    ; GFX9: FLAT_STORE_DWORD $vgpr0_vgpr1, $vgpr0, 0, 0, implicit $exec, implicit $flat_scr
+    ; GFX9: S_ENDPGM 0
+    $vgpr0_vgpr1 = GLOBAL_LOAD_DWORDX2 $vgpr0_vgpr1, 0, 0, implicit $exec
+    S_WAITCNT 3952
+    BUFFER_INVL2 implicit $exec
+    BUFFER_WBINVL1_VOL implicit $exec
+    $vgpr0 = FLAT_LOAD_DWORD $vgpr0_vgpr1, 0, 0, implicit $exec, implicit $flat_scr
+    FLAT_STORE_DWORD $vgpr0_vgpr1, $vgpr0, 0, 0, implicit $exec, implicit $flat_scr
+    S_ENDPGM 0
+...