[llvm] [AMDGPU] Fix simplification of VM_VSRC based on VMEM waits (PR #178711)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Jay Foad (jayfoad)

<details>
<summary>Changes</summary>



---
Full diff: https://github.com/llvm/llvm-project/pull/178711.diff


2 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp (+26-9) 
- (modified) llvm/test/CodeGen/AMDGPU/expert_scheduling_gfx12.mir (+1) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp b/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
index 87e204b3d29e8..c1e5fbf2b6927 100644
--- a/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
+++ b/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
@@ -625,9 +625,11 @@ class WaitcntBrackets {
   }
 
   void setStateOnFunctionEntryOrReturn() {
-    setScoreUB(AMDGPU::STORE_CNT,
-               getScoreUB(AMDGPU::STORE_CNT) + getLimit(AMDGPU::STORE_CNT));
-    PendingEvents |= Context->getWaitEvents(AMDGPU::STORE_CNT);
+    if (Context->ST.hasVscnt()) {
+      setScoreUB(AMDGPU::STORE_CNT,
+                 getScoreUB(AMDGPU::STORE_CNT) + getLimit(AMDGPU::STORE_CNT));
+      PendingEvents |= Context->getWaitEvents(AMDGPU::STORE_CNT);
+    }
   }
 
   ArrayRef<const MachineInstr *> getLDSDMAStores() const {
@@ -1295,12 +1297,27 @@ void WaitcntBrackets::simplifyVmVsrc(const AMDGPU::Waitcnt &CheckWait,
   // Waiting for some counters implies waiting for VM_VSRC, since an
   // instruction that decrements a counter on completion would have
   // decremented VM_VSRC once its VGPR operands had been read.
-  if (CheckWait.get(AMDGPU::VM_VSRC) >=
-      std::min({CheckWait.get(AMDGPU::LOAD_CNT),
-                CheckWait.get(AMDGPU::STORE_CNT),
-                CheckWait.get(AMDGPU::SAMPLE_CNT),
-                CheckWait.get(AMDGPU::BVH_CNT), CheckWait.get(AMDGPU::DS_CNT)}))
-    UpdateWait.set(AMDGPU::VM_VSRC, ~0u);
+  static constexpr AMDGPU::InstCounterType VmemCounters[] = {
+      AMDGPU::LOAD_CNT, AMDGPU::STORE_CNT, AMDGPU::SAMPLE_CNT, AMDGPU::BVH_CNT,
+      AMDGPU::DS_CNT};
+  HWEvents VmemEvents = llvm::accumulate(
+      VmemCounters, HWEvents(), [&](HWEvents Acc, AMDGPU::InstCounterType T) {
+        return Acc | Context->getWaitEvents(T);
+      });
+  HWEvents PendingVmemEvents = PendingEvents & VmemEvents;
+  auto Simplify = [&](AMDGPU::InstCounterType T) {
+    unsigned CheckCount = CheckWait.get(T);
+    if (UpdateWait.get(AMDGPU::VM_VSRC) >= CheckCount &&
+        (CheckCount == 0 || !counterOutOfOrder(T)) &&
+        (PendingVmemEvents & ~Context->getWaitEvents(T)) == 0)
+      UpdateWait.set(AMDGPU::VM_VSRC, ~0u);
+  };
+  Simplify(AMDGPU::LOAD_CNT);
+  Simplify(AMDGPU::STORE_CNT);
+  Simplify(AMDGPU::SAMPLE_CNT);
+  Simplify(AMDGPU::BVH_CNT);
+  Simplify(AMDGPU::DS_CNT);
+
   simplifyWaitcnt(UpdateWait, AMDGPU::VM_VSRC);
 }
 
diff --git a/llvm/test/CodeGen/AMDGPU/expert_scheduling_gfx12.mir b/llvm/test/CodeGen/AMDGPU/expert_scheduling_gfx12.mir
index 6f320f538440b..0f548a9c21f26 100644
--- a/llvm/test/CodeGen/AMDGPU/expert_scheduling_gfx12.mir
+++ b/llvm/test/CodeGen/AMDGPU/expert_scheduling_gfx12.mir
@@ -133,6 +133,7 @@ body:             |
     ; GCN-NEXT: S_WAITCNT_DEPCTR .VmVsrc_0
     ; GCN-NEXT: $vgpr0_vgpr1_vgpr2_vgpr3 = IMAGE_LOAD_V4_V1_gfx12 $vgpr4, $sgpr0_sgpr1_sgpr2_sgpr3_sgpr4_sgpr5_sgpr6_sgpr7, 15, 0, 0, 0, 0, 0, 0, implicit $exec :: (dereferenceable load (s128), addrspace 8)
     ; GCN-NEXT: S_WAIT_LOADCNT_DSCNT .Loadcnt_0_Dscnt_0
+    ; GCN-NEXT: S_WAITCNT_DEPCTR .VmVsrc_0
     ; GCN-NEXT: S_SETREG_IMM32_B32 0, 2074, implicit-def $mode, implicit $mode
     ; GCN-NEXT: SI_RETURN_TO_EPILOG $vgpr0, $vgpr1, $vgpr2, $vgpr3
     $vgpr4 = V_MOV_B32_e32 $vgpr0, implicit $exec, implicit $exec

``````````

</details>


https://github.com/llvm/llvm-project/pull/178711