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

[Jay Foad via llvm-commits]

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

>From 4f29fd9af2e23cdb317c4e27fe7052c228f566db Mon Sep 17 00:00:00 2001
From: Jay Foad <jay.foad at amd.com>
Date: Thu, 29 Jan 2026 17:44:40 +0000
Subject: [PATCH] [AMDGPU] Fix simplification of VM_VSRC based on VMEM waits

---
 llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp   | 32 +++++++++++++++----
 .../AMDGPU/expert_scheduling_gfx12.mir        |  1 +
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp b/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
index 27f4ac389d437..df1a0f01a7c62 100644
--- a/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
+++ b/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
@@ -808,9 +808,12 @@ class WaitcntBrackets {
   }
 
   void setStateOnFunctionEntryOrReturn() {
-    setScoreUB(STORE_CNT, getScoreUB(STORE_CNT) +
-                              getWaitCountMax(Context->getLimits(), STORE_CNT));
-    PendingEvents |= Context->getWaitEventMask()[STORE_CNT];
+    if (Context->ST->hasVscnt()) {
+      setScoreUB(STORE_CNT,
+                 getScoreUB(STORE_CNT) +
+                     getWaitCountMax(Context->getLimits(), STORE_CNT));
+      PendingEvents |= Context->getWaitEventMask()[STORE_CNT];
+    }
   }
 
   ArrayRef<const MachineInstr *> getLDSDMAStores() const {
@@ -1351,10 +1354,25 @@ 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.VmVsrc >=
-      std::min({CheckWait.LoadCnt, CheckWait.StoreCnt, CheckWait.SampleCnt,
-                CheckWait.BvhCnt, CheckWait.DsCnt}))
-    UpdateWait.VmVsrc = ~0u;
+  static constexpr InstCounterType VmemCounters[] = {
+      LOAD_CNT, STORE_CNT, SAMPLE_CNT, BVH_CNT, DS_CNT};
+  unsigned VmemEvents =
+      llvm::accumulate(VmemCounters, 0u, [&](unsigned Acc, InstCounterType T) {
+        return Acc | Context->getWaitEventMask()[T];
+      });
+  unsigned PendingVmemEvents = PendingEvents & VmemEvents;
+  auto Simplify = [&](InstCounterType T, unsigned CheckCount) {
+    if (UpdateWait.VmVsrc >= CheckCount &&
+        (CheckCount == 0 || !counterOutOfOrder(T)) &&
+        (PendingVmemEvents & ~Context->getWaitEventMask()[T]) == 0)
+      UpdateWait.VmVsrc = ~0u;
+  };
+  Simplify(LOAD_CNT, CheckWait.LoadCnt);
+  Simplify(STORE_CNT, CheckWait.StoreCnt);
+  Simplify(SAMPLE_CNT, CheckWait.SampleCnt);
+  Simplify(BVH_CNT, CheckWait.BvhCnt);
+  Simplify(DS_CNT, CheckWait.DsCnt);
+
   simplifyWaitcnt(VM_VSRC, UpdateWait.VmVsrc);
 }
 
diff --git a/llvm/test/CodeGen/AMDGPU/expert_scheduling_gfx12.mir b/llvm/test/CodeGen/AMDGPU/expert_scheduling_gfx12.mir
index 932bd21ad4af7..77161924b5aaf 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 65411
     ; 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 0
+    ; GCN-NEXT: S_WAITCNT_DEPCTR 65411
     ; 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