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

[via llvm-commits]

Author: Jay Foad
Date: 2026-07-21T05:50:27+01:00
New Revision: ee98df545384f95feb72d39f88004075bd5329fc

URL: https://github.com/llvm/llvm-project/commit/ee98df545384f95feb72d39f88004075bd5329fc
DIFF: https://github.com/llvm/llvm-project/commit/ee98df545384f95feb72d39f88004075bd5329fc.diff

LOG: [AMDGPU] Fix simplification of VM_VSRC based on VMEM waits (#178711)

Teach simplifyVmVsrc that a wait for a VMEM counter (LOAD_CNT, DS_CNT
etc) only implies a corresponding wait for VM_VSRC if there are no
outstanding VMEM operations that use a different VMEM counter.

Added: 
    

Modified: 
    llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
    llvm/test/CodeGen/AMDGPU/expert_scheduling_gfx12.mir

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp b/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
index 54e30c7a80be4..7fa33458fa325 100644
--- a/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
+++ b/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
@@ -1309,15 +1309,26 @@ void WaitcntBrackets::simplifyXcnt(const AMDGPU::Waitcnt &CheckWait,
 
 void WaitcntBrackets::simplifyVmVsrc(const AMDGPU::Waitcnt &CheckWait,
                                      AMDGPU::Waitcnt &UpdateWait) const {
-  // 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);
+  // Waiting for a VMEM counter (like LOAD_CNT) implies an equivalent wait for
+  // VM_VSRC, because if the VMEM operation has completed then it must surely
+  // have read its VGPR sources, but only if there are no other outstanding VMEM
+  // operations that use a 
diff erent counter (like SAMPLE_CNT).
+  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;
+  for (AMDGPU::InstCounterType T : VmemCounters) {
+    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);
+  }
+
   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 4f15917679467..e50ce98284415 100644
--- a/llvm/test/CodeGen/AMDGPU/expert_scheduling_gfx12.mir
+++ b/llvm/test/CodeGen/AMDGPU/expert_scheduling_gfx12.mir
@@ -135,6 +135,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