[llvm] [AMDGPU] Track DS prefetch flush points on single-path loop. (PR #227274)
Vikash Gupta via llvm-commits
llvm-commits at lists.llvm.org
Tue Sep 29 04:17:38 PDT 2026
https://github.com/vg0204 created https://github.com/llvm/llvm-project/pull/227274
On `GFX12+`, the DS prefetch preheader flush tracked flush points only in a single-block loop. Walk the loop in issue order when each block has one in-loop successor, and use that order for the same tracking. A loop with an in-loop branch is unchanged.
>From 0e34fc2c98b2ae75ca43bf9fce2ddf9762743e3e Mon Sep 17 00:00:00 2001
From: vg0204 <Vikash.Gupta at amd.com>
Date: Tue, 29 Sep 2026 16:43:35 +0530
Subject: [PATCH] [AMDGPU] Track DS prefetch flush points on single-path loop.
On GFX12+, the DS prefetch preheader flush tracked flush points only in a
single-block loop. Walk the loop in issue order when each block has one
in-loop successor, and use that order for the same tracking. A loop with
an in-loop branch is unchanged.
---
llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp | 58 ++++++++---
.../waitcnt-loop-ds-prefetch-multiblock.mir | 98 +++++++++++++++++++
2 files changed, 144 insertions(+), 12 deletions(-)
create mode 100644 llvm/test/CodeGen/AMDGPU/waitcnt-loop-ds-prefetch-multiblock.mir
diff --git a/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp b/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
index 33fb3c7ca37a0..5b16d29671585 100644
--- a/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
+++ b/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
@@ -3261,6 +3261,37 @@ bool SIInsertWaitcnts::mayStoreIncrementingDSCNT(const MachineInstr &MI) const {
return MI.mayStore() && SIInstrInfo::isDS(MI);
}
+// Issue order when the body has no in-loop branch: header, then each block's
+// one in-loop successor. Returns false if that walk is not the whole loop.
+static bool
+appendSinglePathLoopBlocks(const MachineLoop &ML,
+ SmallVectorImpl<MachineBasicBlock *> &Order) {
+ MachineBasicBlock *Header = ML.getHeader();
+ if (!Header)
+ return false;
+
+ MachineBasicBlock *BB = Header;
+ SmallPtrSet<MachineBasicBlock *, 8> Seen;
+ do {
+ if (!Seen.insert(BB).second)
+ return false;
+ Order.push_back(BB);
+ MachineBasicBlock *Next = nullptr;
+ for (MachineBasicBlock *Succ : BB->successors()) {
+ if (!ML.contains(Succ) || Succ == Header)
+ continue;
+ if (Next)
+ return false;
+ Next = Succ;
+ }
+ if (!Next)
+ break;
+ BB = Next;
+ } while (true);
+
+ return Order.size() == ML.getNumBlocks();
+}
+
// Return flags indicating which counters should be flushed in the preheader of
// the given loop. We currently decide to flush in the following situations:
// For VMEM (FlushVmCnt):
@@ -3279,15 +3310,15 @@ bool SIInsertWaitcnts::mayStoreIncrementingDSCNT(const MachineInstr &MI) const {
// Flushing in preheader reduces wait overhead if the wait requirement in
// iteration 1 would otherwise be more strict (but unfortunately preheader
// flush decision is taken before knowing that).
-// 5. (Single-block loops only) The loop has DS prefetch reads with flush point
-// tracking. Some DS reads may be used in the same iteration (creating
-// "flush points"), but others remain unflushed at the backedge. When a DS
-// read is consumed in the same iteration, it and all prior reads are
-// "flushed" (FIFO order). No DS writes are allowed in the loop.
-// TODO: Find a way to extend to multi-block loops.
+// 5. The loop has DS prefetch reads with flush point tracking. Some DS reads
+// may be used in the same iteration (creating "flush points"), but others
+// remain unflushed at the backedge. When a DS read is consumed in the same
+// iteration, it and all prior reads are "flushed" (FIFO order). No DS
+// writes are allowed in the loop. Requires one issue order, so an in-loop
+// branch skips this case.
PreheaderFlushFlags
-SIInsertWaitcnts::getPreheaderFlushFlags(MachineLoop *ML,
- const WaitcntBrackets &Brackets) {
+SIInsertWaitcnts::LastDSReadPositionMap(MachineLoop *ML,
+ const WaitcntBrackets &Brackets) {
PreheaderFlushFlags Flags;
bool HasVMemLoad = false;
bool HasVMemStore = false;
@@ -3302,17 +3333,20 @@ SIInsertWaitcnts::getPreheaderFlushFlags(MachineLoop *ML,
DenseSet<MCRegUnit> VgprDefVMEM;
DenseSet<MCRegUnit> VgprDefDS;
- // Track DS reads for prefetch pattern with flush points (single-block only).
+ // Track DS reads for prefetch pattern with flush points.
// Keeps track of the last DS read (position counted from the top of the loop)
// to each VGPR. Read is considered consumed (and thus needs flushing) if
// the dest register has a use or is overwritten (by any later opertions).
DenseMap<MCRegUnit, unsigned> LastDSReadPositionMap;
unsigned DSReadPosition = 0;
- bool IsSingleBlock = ML->getNumBlocks() == 1;
- bool TrackDSFlushPoint = ST.hasExtendedWaitCounts() && IsSingleBlock;
+ SmallVector<MachineBasicBlock *, 8> BlockOrder;
+ bool SinglePath = appendSinglePathLoopBlocks(*ML, BlockOrder);
+ if (!SinglePath)
+ append_range(BlockOrder, ML->blocks());
+ bool TrackDSFlushPoint = ST.hasExtendedWaitCounts() && SinglePath;
unsigned LastDSFlushPosition = 0;
- for (MachineBasicBlock *MBB : ML->blocks()) {
+ for (MachineBasicBlock *MBB : BlockOrder) {
for (MachineInstr &MI : *MBB) {
if (isVMEMOrFlatVMEM(MI)) {
HasVMemLoad |= MI.mayLoad();
diff --git a/llvm/test/CodeGen/AMDGPU/waitcnt-loop-ds-prefetch-multiblock.mir b/llvm/test/CodeGen/AMDGPU/waitcnt-loop-ds-prefetch-multiblock.mir
new file mode 100644
index 0000000000000..9d334e954bb4f
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/waitcnt-loop-ds-prefetch-multiblock.mir
@@ -0,0 +1,98 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 6
+# RUN: llc -mtriple=amdgpu12.50-amd-amdhsa -run-pass=si-insert-waitcnts -o - %s | FileCheck %s
+
+---
+# DS prefetch flush points in a two-block loop with no in-loop branch.
+# Header DS reads are consumed in the latch; later latch reads stay unflushed.
+# Expected: s_wait_dscnt 0 in the preheader.
+name: ds_prefetch_flush_two_blocks
+tracksRegLiveness: true
+machineFunctionInfo:
+ isEntryFunction: true
+body: |
+ ; CHECK-LABEL: name: ds_prefetch_flush_two_blocks
+ ; CHECK: bb.0:
+ ; CHECK-NEXT: successors: %bb.1(0x80000000)
+ ; CHECK-NEXT: liveins: $sgpr0, $vgpr0
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: S_SETREG_IMM32_B32 1, 1601, implicit-def $mode, implicit $mode
+ ; CHECK-NEXT: $sgpr64_sgpr65 = S_MOV_B64 0
+ ; CHECK-NEXT: V_NOP_e32 implicit $exec
+ ; CHECK-NEXT: GLOBAL_PREFETCH_B8_SADDR $sgpr64_sgpr65, undef $vgpr0, 0, 8, implicit $exec
+ ; CHECK-NEXT: $vgpr10_vgpr11_vgpr12_vgpr13 = DS_READ_B128 $vgpr0, 0, 0, implicit $m0, implicit $exec
+ ; CHECK-NEXT: $vgpr28 = DS_READ_B32 $vgpr0, 0, 0, implicit $m0, implicit $exec
+ ; CHECK-NEXT: $vgpr32 = DS_READ_B32 $vgpr0, 0, 0, implicit $m0, implicit $exec
+ ; CHECK-NEXT: S_WAIT_DSCNT 0
+ ; CHECK-NEXT: S_BRANCH %bb.1
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: bb.1:
+ ; CHECK-NEXT: successors: %bb.2(0x80000000)
+ ; CHECK-NEXT: liveins: $sgpr0, $vgpr0, $vgpr10_vgpr11_vgpr12_vgpr13, $vgpr28, $vgpr32
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: $vgpr50 = V_ADD_F32_e32 $vgpr10, $vgpr11, implicit $mode, implicit $exec
+ ; CHECK-NEXT: S_WAIT_DSCNT 1
+ ; CHECK-NEXT: $vgpr51 = V_ADD_F32_e32 $vgpr28, $vgpr28, implicit $mode, implicit $exec
+ ; CHECK-NEXT: S_BARRIER
+ ; CHECK-NEXT: $vgpr20_vgpr21_vgpr22_vgpr23 = DS_READ_B128 $vgpr0, 64, 0, implicit $m0, implicit $exec
+ ; CHECK-NEXT: $vgpr24_vgpr25_vgpr26_vgpr27 = DS_READ_B128 $vgpr0, 80, 0, implicit $m0, implicit $exec
+ ; CHECK-NEXT: S_BRANCH %bb.2
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: bb.2:
+ ; CHECK-NEXT: successors: %bb.1(0x40000000), %bb.3(0x40000000)
+ ; CHECK-NEXT: liveins: $sgpr0, $vgpr0, $vgpr10_vgpr11_vgpr12_vgpr13, $vgpr20_vgpr21_vgpr22_vgpr23, $vgpr24_vgpr25_vgpr26_vgpr27, $vgpr28, $vgpr32
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: S_WAIT_DSCNT 2
+ ; CHECK-NEXT: $vgpr52 = V_ADD_F32_e32 $vgpr32, $vgpr32, implicit $mode, implicit $exec
+ ; CHECK-NEXT: $vgpr28_vgpr29_vgpr30_vgpr31 = DS_READ_B128 $vgpr0, 96, 0, implicit $m0, implicit $exec
+ ; CHECK-NEXT: $vgpr32_vgpr33_vgpr34_vgpr35 = DS_READ_B128 $vgpr0, 112, 0, implicit $m0, implicit $exec
+ ; CHECK-NEXT: S_WAIT_DSCNT 2
+ ; CHECK-NEXT: $vgpr53 = V_ADD_F32_e32 $vgpr24, $vgpr25, implicit $mode, implicit $exec
+ ; CHECK-NEXT: $sgpr0 = S_ADD_I32 $sgpr0, -1, implicit-def $scc
+ ; CHECK-NEXT: S_CBRANCH_SCC1 %bb.1, implicit $scc
+ ; CHECK-NEXT: S_BRANCH %bb.3
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: bb.3:
+ ; CHECK-NEXT: S_ENDPGM 0
+ bb.0:
+ successors: %bb.1
+ liveins: $sgpr0, $vgpr0
+
+ ; Preheader DS loads, used in the loop.
+ $vgpr10_vgpr11_vgpr12_vgpr13 = DS_READ_B128 $vgpr0, 0, 0, implicit $m0, implicit $exec
+ $vgpr28 = DS_READ_B32 $vgpr0, 0, 0, implicit $m0, implicit $exec
+ $vgpr32 = DS_READ_B32 $vgpr0, 0, 0, implicit $m0, implicit $exec
+ S_BRANCH %bb.1
+
+ bb.1:
+ successors: %bb.2
+ liveins: $sgpr0, $vgpr0, $vgpr10_vgpr11_vgpr12_vgpr13, $vgpr28, $vgpr32
+
+ $vgpr50 = V_ADD_F32_e32 $vgpr10, $vgpr11, implicit $mode, implicit $exec
+ $vgpr51 = V_ADD_F32_e32 $vgpr28, $vgpr28, implicit $mode, implicit $exec
+
+ S_BARRIER
+
+ ; Consumed in the latch. Flush points.
+ $vgpr20_vgpr21_vgpr22_vgpr23 = DS_READ_B128 $vgpr0, 64, 0, implicit $m0, implicit $exec
+ $vgpr24_vgpr25_vgpr26_vgpr27 = DS_READ_B128 $vgpr0, 80, 0, implicit $m0, implicit $exec
+ S_BRANCH %bb.2
+
+ bb.2:
+ successors: %bb.1, %bb.3
+ liveins: $sgpr0, $vgpr0, $vgpr10_vgpr11_vgpr12_vgpr13, $vgpr20_vgpr21_vgpr22_vgpr23, $vgpr24_vgpr25_vgpr26_vgpr27, $vgpr28, $vgpr32
+
+ $vgpr52 = V_ADD_F32_e32 $vgpr32, $vgpr32, implicit $mode, implicit $exec
+ ; Unflushed prefetches. Issued after the header reads, so the use below
+ ; does not flush them.
+ $vgpr28_vgpr29_vgpr30_vgpr31 = DS_READ_B128 $vgpr0, 96, 0, implicit $m0, implicit $exec
+ $vgpr32_vgpr33_vgpr34_vgpr35 = DS_READ_B128 $vgpr0, 112, 0, implicit $m0, implicit $exec
+ ; Flush point for the header reads.
+ $vgpr53 = V_ADD_F32_e32 $vgpr24, $vgpr25, implicit $mode, implicit $exec
+
+ $sgpr0 = S_ADD_I32 $sgpr0, -1, implicit-def $scc
+ S_CBRANCH_SCC1 %bb.1, implicit $scc
+ S_BRANCH %bb.3
+
+ bb.3:
+ S_ENDPGM 0
+...
More information about the llvm-commits
mailing list