[llvm] [AMDGPU] Track DS prefetch flush points on single-path loop. (PR #227274)
Vikash Gupta via llvm-commits
llvm-commits at lists.llvm.org
Thu Oct 1 04:42:18 PDT 2026
https://github.com/vg0204 updated https://github.com/llvm/llvm-project/pull/227274
>From 68690d84cce5bbefdcc86f680007409f0e6f0109 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 1/2] [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 | 54 ++++++++--
.../waitcnt-loop-ds-prefetch-multiblock.mir | 98 +++++++++++++++++++
2 files changed, 142 insertions(+), 10 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..bcf5403cb2e67 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,12 +3310,12 @@ 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) {
@@ -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
+...
>From 80acf34fd61630c6be9834d076e96cb770c4e551 Mon Sep 17 00:00:00 2001
From: vg0204 <Vikash.Gupta at amd.com>
Date: Thu, 1 Oct 2026 17:07:13 +0530
Subject: [PATCH 2/2] Moved CFG walk within loop into LoopInfo from the pass.
---
llvm/include/llvm/Support/GenericLoopInfo.h | 6 ++
.../llvm/Support/GenericLoopInfoImpl.h | 41 ++++++++++++
llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp | 33 +---------
llvm/unittests/Analysis/LoopInfoTest.cpp | 62 +++++++++++++++++++
4 files changed, 110 insertions(+), 32 deletions(-)
diff --git a/llvm/include/llvm/Support/GenericLoopInfo.h b/llvm/include/llvm/Support/GenericLoopInfo.h
index 9d3cc66028886..a7a03ff175cb0 100644
--- a/llvm/include/llvm/Support/GenericLoopInfo.h
+++ b/llvm/include/llvm/Support/GenericLoopInfo.h
@@ -329,6 +329,12 @@ template <class BlockT, class LoopT> class LoopBase {
LoopLatches.push_back(Pred);
}
+ /// Append this loop's blocks from the header to its single latch and return
+ /// true when each block before the latch has exactly one in-loop successor
+ /// and that path visits every block. Exit edges are ignored. Otherwise
+ /// return false and leave \p Order empty. blocks() is not this order.
+ bool getSinglePathBlocks(SmallVectorImpl<BlockT *> &Order) const;
+
/// Return all inner loops in the loop nest rooted by the loop in preorder,
/// with siblings in forward program order.
template <class Type>
diff --git a/llvm/include/llvm/Support/GenericLoopInfoImpl.h b/llvm/include/llvm/Support/GenericLoopInfoImpl.h
index f710c2b8f1615..20e7b20dfeca1 100644
--- a/llvm/include/llvm/Support/GenericLoopInfoImpl.h
+++ b/llvm/include/llvm/Support/GenericLoopInfoImpl.h
@@ -269,6 +269,47 @@ BlockT *LoopBase<BlockT, LoopT>::getLoopLatch() const {
return Latch;
}
+/// getSinglePathBlocks - Append this loop's blocks from the header to its
+/// single latch and return true when each block before the latch has exactly
+/// one in-loop successor and that path visits every block. Exit edges are
+/// ignored. Otherwise return false and leave \p Order empty.
+template <class BlockT, class LoopT>
+bool LoopBase<BlockT, LoopT>::getSinglePathBlocks(
+ SmallVectorImpl<BlockT *> &Order) const {
+ assert(!isInvalid() && "Loop not in a valid state!");
+ Order.clear();
+
+ // Null when the loop has more than one backedge.
+ BlockT *Latch = getLoopLatch();
+ if (!Latch)
+ return false;
+
+ // Null unless exactly one successor stays in the loop.
+ auto IsInLoopSucc = [&](BlockT *Succ, bool AllowRepeats) -> BlockT * {
+ assert(!AllowRepeats && "Unexpected parameter value.");
+ return contains(Succ) ? Succ : nullptr;
+ };
+
+ BlockT *BB = getHeader();
+ while (!is_contained(Order, BB)) {
+ Order.push_back(BB);
+ if (BB == Latch)
+ break;
+
+ BB = find_singleton<BlockT>(children<BlockT *>(BB), IsInLoopSucc);
+ if (!BB) {
+ Order.clear();
+ return false;
+ }
+ }
+
+ if (Order.size() != getNumBlocks()) {
+ Order.clear();
+ return false;
+ }
+ return true;
+}
+
//===----------------------------------------------------------------------===//
// APIs for updating loop information after changing the CFG
//
diff --git a/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp b/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
index bcf5403cb2e67..24b29aca0167c 100644
--- a/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
+++ b/llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp
@@ -3261,37 +3261,6 @@ 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):
@@ -3340,7 +3309,7 @@ SIInsertWaitcnts::getPreheaderFlushFlags(MachineLoop *ML,
DenseMap<MCRegUnit, unsigned> LastDSReadPositionMap;
unsigned DSReadPosition = 0;
SmallVector<MachineBasicBlock *, 8> BlockOrder;
- bool SinglePath = appendSinglePathLoopBlocks(*ML, BlockOrder);
+ bool SinglePath = ML->getSinglePathBlocks(BlockOrder);
if (!SinglePath)
append_range(BlockOrder, ML->blocks());
bool TrackDSFlushPoint = ST.hasExtendedWaitCounts() && SinglePath;
diff --git a/llvm/unittests/Analysis/LoopInfoTest.cpp b/llvm/unittests/Analysis/LoopInfoTest.cpp
index 5309d044d5805..55a3c60ecc36a 100644
--- a/llvm/unittests/Analysis/LoopInfoTest.cpp
+++ b/llvm/unittests/Analysis/LoopInfoTest.cpp
@@ -1787,3 +1787,65 @@ TEST(LoopInfoTest, RecomputeHoistsChild) {
LI.verify();
LI.destroy(OuterL);
}
+
+TEST(LoopInfoTest, SinglePathBlocks) {
+ const char *ModuleStr = "define void @single(i1 %c) {\n"
+ "entry:\n"
+ " br label %loop\n"
+ "loop:\n"
+ " br i1 %c, label %loop, label %exit\n"
+ "exit:\n"
+ " ret void\n"
+ "}\n"
+ "define void @chain(i1 %c) {\n"
+ "entry:\n"
+ " br label %header\n"
+ "header:\n"
+ " br label %latch\n"
+ "latch:\n"
+ " br i1 %c, label %header, label %exit\n"
+ "exit:\n"
+ " ret void\n"
+ "}\n"
+ "define void @diamond(i1 %c) {\n"
+ "entry:\n"
+ " br label %header\n"
+ "header:\n"
+ " br i1 %c, label %left, label %right\n"
+ "left:\n"
+ " br label %latch\n"
+ "right:\n"
+ " br label %latch\n"
+ "latch:\n"
+ " br i1 %c, label %header, label %exit\n"
+ "exit:\n"
+ " ret void\n"
+ "}\n";
+
+ LLVMContext Context;
+ std::unique_ptr<Module> M = makeLLVMModule(Context, ModuleStr);
+
+ runWithLoopInfo(*M, "single", [&](Function &F, LoopInfo &LI) {
+ Loop *L = LI.getLoopFor(getBlockByName(&F, "loop"));
+ SmallVector<BasicBlock *, 4> Order;
+ EXPECT_TRUE(L->getSinglePathBlocks(Order));
+ ASSERT_EQ(Order.size(), 1u);
+ EXPECT_EQ(Order[0], getBlockByName(&F, "loop"));
+ });
+
+ runWithLoopInfo(*M, "chain", [&](Function &F, LoopInfo &LI) {
+ Loop *L = LI.getLoopFor(getBlockByName(&F, "header"));
+ SmallVector<BasicBlock *, 4> Order;
+ EXPECT_TRUE(L->getSinglePathBlocks(Order));
+ ASSERT_EQ(Order.size(), 2u);
+ EXPECT_EQ(Order[0], getBlockByName(&F, "header"));
+ EXPECT_EQ(Order[1], getBlockByName(&F, "latch"));
+ });
+
+ runWithLoopInfo(*M, "diamond", [&](Function &F, LoopInfo &LI) {
+ Loop *L = LI.getLoopFor(getBlockByName(&F, "header"));
+ SmallVector<BasicBlock *, 4> Order;
+ EXPECT_FALSE(L->getSinglePathBlocks(Order));
+ EXPECT_TRUE(Order.empty());
+ });
+}
More information about the llvm-commits
mailing list