[llvm-branch-commits] [llvm] [CodeGen] Use validated wave counts for AMDGPU spill costs (PR #226215)
Yaxun Liu via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Thu Sep 24 11:04:50 PDT 2026
https://github.com/yxsamliu updated https://github.com/llvm/llvm-project/pull/226215
>From 55ce386fb935bf07bac79a876c402e214480f3ba Mon Sep 17 00:00:00 2001
From: "Yaxun (Sam) Liu" <yaxun.liu at amd.com>
Date: Thu, 24 Sep 2026 14:03:11 -0400
Subject: [PATCH] [CodeGen] Use validated wave counts for AMDGPU spill costs
Lane-based block frequencies can understate the cost of a spill in
divergent GPU code: a wave still executes a block with only some lanes
active. Use measured block-wave counts to weight SpillPlacement's costs
relative to the original entry-wave count.
Opt in on AMDGPU only. Accept a measured count only when its IR block
maps uniquely to a machine block with matching predecessors and
successors. Keep the existing MBFI cost for unmeasured or rejected
blocks, and leave branch probabilities and general BFI unchanged.
---
llvm/lib/CodeGen/SpillPlacement.cpp | 98 ++++++-
.../CodeGen/AMDGPU/wave-profile-spill.mir | 251 ++++++++++++++++++
2 files changed, 348 insertions(+), 1 deletion(-)
create mode 100644 llvm/test/CodeGen/AMDGPU/wave-profile-spill.mir
diff --git a/llvm/lib/CodeGen/SpillPlacement.cpp b/llvm/lib/CodeGen/SpillPlacement.cpp
index 55a96a22a00ec..de39ad94ada53 100644
--- a/llvm/lib/CodeGen/SpillPlacement.cpp
+++ b/llvm/lib/CodeGen/SpillPlacement.cpp
@@ -28,13 +28,23 @@
#include "llvm/CodeGen/SpillPlacement.h"
#include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
#include "llvm/CodeGen/EdgeBundles.h"
#include "llvm/CodeGen/MachineBasicBlock.h"
#include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/Passes.h"
+#include "llvm/IR/CFG.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/ProfDataUtils.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/ScaledNumber.h"
#include <algorithm>
#include <cassert>
#include <cstdint>
@@ -44,6 +54,13 @@ using namespace llvm;
#define DEBUG_TYPE "spill-code-placement"
+static cl::opt<bool> EnableWaveProfiledSpill(
+ "enable-wave-profiled-spill", cl::Hidden, cl::init(false),
+ cl::desc("Use validated AMDGPU wave counts for spill placement"));
+static cl::opt<bool> ReportWaveProfiledSpill(
+ "report-wave-profiled-spill", cl::Hidden, cl::init(false),
+ cl::desc("Report wave-count mapping and spill-cost coverage"));
+
char SpillPlacementWrapperLegacy::ID = 0;
char &llvm::SpillPlacementID = SpillPlacementWrapperLegacy::ID;
@@ -229,6 +246,26 @@ void SpillPlacement::releaseMemory() {
TodoList.clear();
}
+static bool hasMatchingWaveProfileBlock(
+ const MachineBasicBlock &MBB,
+ const DenseMap<const BasicBlock *, const MachineBasicBlock *> &UniqueMBBs) {
+ const BasicBlock *BB = MBB.getBasicBlock();
+ auto HasMatchingEdges = [&](auto IRBlocks, auto MachineBlocks) {
+ SmallPtrSet<const BasicBlock *, 4> Expected;
+ for (const BasicBlock *IRBlock : IRBlocks)
+ Expected.insert(IRBlock);
+ for (const MachineBasicBlock *MachineBlock : MachineBlocks) {
+ const BasicBlock *IRBlock = MachineBlock->getBasicBlock();
+ if (!IRBlock || UniqueMBBs.lookup(IRBlock) != MachineBlock ||
+ !Expected.erase(IRBlock))
+ return false;
+ }
+ return Expected.empty();
+ };
+ return HasMatchingEdges(predecessors(BB), MBB.predecessors()) &&
+ HasMatchingEdges(successors(BB), MBB.successors());
+}
+
void SpillPlacement::run(MachineFunction &mf, EdgeBundles *Bundles,
MachineBlockFrequencyInfo *MBFI) {
MF = &mf;
@@ -240,13 +277,72 @@ void SpillPlacement::run(MachineFunction &mf, EdgeBundles *Bundles,
TodoList.clear();
TodoList.setUniverse(bundles->getNumBundles());
+ SmallVector<uint64_t> WaveCounts;
+ BitVector HasWaveCount;
+ uint64_t EntryWaveCount = 0;
+ DenseMap<const BasicBlock *, BlockFrequency> WaveFrequencies;
+ unsigned RejectedWaveBlocks = 0;
+ const bool HasWaveProfile =
+ EnableWaveProfiledSpill &&
+ mf.getFunction().getParent()->getTargetTriple().isAMDGPU() &&
+ extractMappedBlockWaveCounts(mf.getFunction(), WaveCounts, HasWaveCount,
+ EntryWaveCount) &&
+ EntryWaveCount != 0;
+ if (HasWaveProfile) {
+ DenseMap<const BasicBlock *, const MachineBasicBlock *> UniqueMBBs;
+ for (const MachineBasicBlock &MBB : mf) {
+ const BasicBlock *BB = MBB.getBasicBlock();
+ if (!BB || BB->getParent() != &mf.getFunction())
+ continue;
+ auto [It, Inserted] = UniqueMBBs.try_emplace(BB, &MBB);
+ if (!Inserted)
+ It->second = nullptr;
+ }
+
+ unsigned Index = 0;
+ for (const BasicBlock &BB : mf.getFunction()) {
+ if (!HasWaveCount[Index]) {
+ ++Index;
+ continue;
+ }
+ const MachineBasicBlock *MBB = UniqueMBBs.lookup(&BB);
+ if (!MBB || !hasMatchingWaveProfileBlock(*MBB, UniqueMBBs)) {
+ if (ReportWaveProfiledSpill)
+ ++RejectedWaveBlocks;
+ ++Index;
+ continue;
+ }
+ uint64_t Count = WaveCounts[Index++];
+ uint64_t Frequency =
+ ScaledNumber<uint64_t>::getFraction(Count, EntryWaveCount)
+ .scale(MBFI->getEntryFreq().getFrequency());
+ WaveFrequencies[&BB] =
+ BlockFrequency(Count ? std::max(Frequency, uint64_t(1)) : 0);
+ }
+ }
+
// Compute total ingoing and outgoing block frequencies for all bundles.
BlockFrequencies.resize(mf.getNumBlockIDs());
setThreshold(MBFI->getEntryFreq());
+ unsigned ChangedWaveCosts = 0;
for (auto &I : mf) {
unsigned Num = I.getNumber();
- BlockFrequencies[Num] = MBFI->getBlockFreq(&I);
+ auto Wave = WaveFrequencies.find(I.getBasicBlock());
+ BlockFrequencies[Num] =
+ Wave == WaveFrequencies.end() ? MBFI->getBlockFreq(&I) : Wave->second;
+ if (Wave != WaveFrequencies.end()) {
+ if (ReportWaveProfiledSpill)
+ ChangedWaveCosts += Wave->second != MBFI->getBlockFreq(&I);
+ LLVM_DEBUG(dbgs() << "Wave spill frequency " << mf.getName() << " "
+ << printMBBReference(I) << " = "
+ << Wave->second.getFrequency() << " (entry "
+ << MBFI->getEntryFreq().getFrequency() << ")\n");
+ }
}
+ if (ReportWaveProfiledSpill && EnableWaveProfiledSpill)
+ errs() << "WAVE_PROFILED_SPILL\t" << mf.getName() << "\t" << HasWaveProfile
+ << "\t" << WaveFrequencies.size() << "\t" << RejectedWaveBlocks
+ << "\t" << ChangedWaveCosts << "\n";
}
/// activate - mark node n as active if it wasn't already.
diff --git a/llvm/test/CodeGen/AMDGPU/wave-profile-spill.mir b/llvm/test/CodeGen/AMDGPU/wave-profile-spill.mir
new file mode 100644
index 0000000000000..6d4e0ec5b4387
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/wave-profile-spill.mir
@@ -0,0 +1,251 @@
+# REQUIRES: amdgpu-registered-target, x86-registered-target
+# RUN: split-file %s %t
+# RUN: opt -passes=pgo-instr-gen -pgo-instrument-entry -S %t/diamond.ll -o %t/gen.ll
+# RUN: %python %t/raw.py %t/gen.ll normal > %t/profile.raw
+# RUN: llvm-profdata merge %t/profile.raw -o %t/profile.data
+# RUN: opt -passes=pgo-instr-use -pgo-instrument-entry -pgo-test-profile-file=%t/profile.data -S %t/diamond.ll -o %t/use.ll
+# RUN: FileCheck %s --check-prefix=IR < %t/use.ll
+# RUN: llvm-as %t/use.ll -o %t/use.bc
+# RUN: llvm-dis %t/use.bc -o %t/roundtrip.ll
+# RUN: %python %t/mir.py %t/roundtrip.ll %t/body.mir > %t/use.mir
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=spill-code-placement -enable-wave-profiled-spill -report-wave-profiled-spill -filetype=null %t/use.mir 2>&1 | FileCheck %s --check-prefix=WAVES
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=spill-code-placement -report-wave-profiled-spill -filetype=null %t/use.mir 2>&1 | FileCheck %s --check-prefix=OFF
+# RUN: %python %t/mir.py %t/diamond.ll %t/body.mir > %t/no-profile.mir
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=spill-code-placement -enable-wave-profiled-spill -report-wave-profiled-spill -filetype=null %t/no-profile.mir 2>&1 | FileCheck %s --check-prefix=FALLBACK
+# RUN: %python %t/mir.py %t/use.ll %t/x86.mir x86 > %t/x86-input.mir
+# RUN: llc -mtriple=x86_64-unknown-linux-gnu -run-pass=spill-code-placement -enable-wave-profiled-spill -report-wave-profiled-spill -filetype=null %t/x86-input.mir 2>&1 | FileCheck %s --check-prefix=FALLBACK
+# RUN: %python %t/raw.py %t/gen.ll zero-entry > %t/zero.raw
+# RUN: llvm-profdata merge %t/zero.raw -o %t/zero.data
+# RUN: opt -passes=pgo-instr-use -pgo-instrument-entry -pgo-test-profile-file=%t/zero.data -S %t/diamond.ll -o %t/zero.ll
+# RUN: %python %t/mir.py %t/zero.ll %t/body.mir > %t/zero.mir
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=spill-code-placement -enable-wave-profiled-spill -report-wave-profiled-spill -filetype=null %t/zero.mir 2>&1 | FileCheck %s --check-prefix=FALLBACK
+# RUN: %python %t/mir.py %t/use.ll %t/duplicate.mir > %t/duplicate-input.mir
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=spill-code-placement -enable-wave-profiled-spill -report-wave-profiled-spill -filetype=null %t/duplicate-input.mir 2>&1 | FileCheck %s --check-prefix=DUPLICATE
+# RUN: %python %t/mir.py %t/use.ll %t/redirected.mir > %t/redirected-input.mir
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=spill-code-placement -enable-wave-profiled-spill -report-wave-profiled-spill -filetype=null %t/redirected-input.mir 2>&1 | FileCheck %s --check-prefix=REDIRECTED
+# RUN: %python %t/mir.py %t/use.ll %t/unmapped.mir > %t/unmapped-input.mir
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=spill-code-placement -enable-wave-profiled-spill -report-wave-profiled-spill -filetype=null %t/unmapped-input.mir 2>&1 | FileCheck %s --check-prefix=UNMAPPED
+# RUN: %python %t/mir.py %t/use.ll %t/choice.mir > %t/choice-input.mir
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=greedy -stress-regalloc=2 -verify-regalloc -verify-machineinstrs -o - %t/choice-input.mir | FileCheck %s --check-prefix=ALLOC-OFF
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=greedy -stress-regalloc=2 -verify-regalloc -verify-machineinstrs -enable-wave-profiled-spill -o - %t/choice-input.mir | FileCheck %s --check-prefix=ALLOC-ON
+# RUN: %python %t/mir.py %t/diamond.ll %t/choice.mir > %t/choice-no-profile.mir
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=greedy -stress-regalloc=2 -verify-regalloc -verify-machineinstrs -enable-wave-profiled-spill -o - %t/choice-no-profile.mir | FileCheck %s --check-prefix=ALLOC-OFF
+
+# The measured right arm has half the lane frequency but the same wave
+# frequency as entry. Direct costs require matching machine control flow.
+# IR: define void @diamond{{.*}}!wave.profile
+# IR: !{i64 2, i64 {{-?[0-9]+}}, i64 100, i64 0, i64 100, i64 0}
+# WAVES: WAVE_PROFILED_SPILL diamond 1 2 0 1
+# OFF-NOT: WAVE_PROFILED_SPILL
+# FALLBACK: WAVE_PROFILED_SPILL diamond 0 0 0 0
+# DUPLICATE: WAVE_PROFILED_SPILL diamond 1 1 1 1
+# REDIRECTED: WAVE_PROFILED_SPILL diamond 1 1 1 0
+# A valid function profile is insufficient when no machine blocks map to it.
+# UNMAPPED: WAVE_PROFILED_SPILL diamond 1 0 2 0
+# ALLOC-OFF-LABEL: body:
+# ALLOC-OFF: bb.1.left:
+# ALLOC-OFF-NOT: lr-split COPY
+# ALLOC-OFF: bb.2.right:
+# ALLOC-OFF: = lr-split COPY
+# ALLOC-OFF: S_NOP 0, implicit-def dead $vgpr0
+# ALLOC-OFF: = lr-split COPY
+# ALLOC-ON-LABEL: body:
+# ALLOC-ON: bb.1.left:
+# ALLOC-ON: = lr-split COPY
+# ALLOC-ON: S_NOP 0, implicit-def dead $vgpr1
+# ALLOC-ON: = lr-split COPY
+# ALLOC-ON: bb.2.right:
+# ALLOC-ON-NOT: lr-split COPY
+# ALLOC-ON: bb.3.exit:
+
+#--- diamond.ll
+target triple = "amdgcn-amd-amdhsa"
+define void @diamond(i1 %condition, ptr addrspace(1) %out) {
+entry:
+ br i1 %condition, label %left, label %right
+left:
+ store volatile i32 1, ptr addrspace(1) %out
+ br label %exit
+right:
+ store volatile i32 2, ptr addrspace(1) %out
+ br label %exit
+exit:
+ ret void
+}
+
+#--- raw.py
+import hashlib
+import re
+import struct
+import sys
+from pathlib import Path
+
+ir = Path(sys.argv[1]).read_text()
+lane_counts = {"entry": 6400, "left": 3200, "right": 3200, "exit": 6400}
+block = None
+counts = {}
+for line in ir.splitlines():
+ label = re.match(r"^([\w.]+):", line)
+ if label:
+ block = label[1]
+ inc = re.search(r"@llvm.instrprof.increment\(.*i64 (-?\d+), i32 (\d+), i32 (\d+)\)", line)
+ if inc:
+ cfg_hash, num_counters, index = map(int, inc.groups())
+ cfg_hash %= 1 << 64
+ counts[index] = lane_counts[block]
+assert len(counts) == num_counters
+waves = [100] * num_counters
+if sys.argv[2] == "zero-entry":
+ waves[0] = 0
+values = [counts[i] for i in range(num_counters)] + waves
+name = b"diamond"
+name_ref = int.from_bytes(hashlib.md5(name).digest()[:8], "little")
+names = bytes([len(name), 0]) + name
+counter_delta = 80
+names_delta = counter_delta + len(values) * 8
+header = [0xff6c70726f667281, 12 | (1 << 56), 0, 1, 0, len(values),
+ 0, 0, 0, 0, 0, 0, len(names), counter_delta, 0, names_delta, 0, 0, 2]
+record = struct.pack("<7QI4HII4x", name_ref, cfg_hash, counter_delta,
+ 0, 0, 0, 0, len(values), 0, 0, 0, 64, 0, len(waves))
+sys.stdout.buffer.write(struct.pack("<19Q", *header) + record +
+ struct.pack("<" + "Q" * len(values), *values) +
+ names + bytes((-len(names)) % 8))
+
+#--- mir.py
+import sys
+import textwrap
+from pathlib import Path
+
+print("--- |")
+ir = Path(sys.argv[1]).read_text()
+if len(sys.argv) > 3 and sys.argv[3] == "x86":
+ ir = ir.replace('target triple = "amdgcn-amd-amdhsa"',
+ 'target triple = "x86_64-unknown-linux-gnu"')
+print(textwrap.indent(ir, " "))
+print("...")
+print(Path(sys.argv[2]).read_text())
+
+#--- body.mir
+---
+name: diamond
+body: |
+ bb.0.entry:
+ successors: %bb.1(0x40000000), %bb.2(0x40000000)
+ S_CBRANCH_SCC1 %bb.1, implicit $scc
+ S_BRANCH %bb.2
+ bb.1.left:
+ successors: %bb.3(0x80000000)
+ S_BRANCH %bb.3
+ bb.2.right:
+ successors: %bb.3(0x80000000)
+ S_BRANCH %bb.3
+ bb.3.exit:
+ S_ENDPGM 0
+...
+
+#--- x86.mir
+---
+name: diamond
+body: |
+ bb.0.entry:
+ successors: %bb.1(0x40000000), %bb.2(0x40000000)
+ JCC_1 %bb.1, 4, implicit $eflags
+ JMP_1 %bb.2
+ bb.1.left:
+ successors: %bb.3(0x80000000)
+ JMP_1 %bb.3
+ bb.2.right:
+ successors: %bb.3(0x80000000)
+ JMP_1 %bb.3
+ bb.3.exit:
+ RET64
+...
+
+#--- duplicate.mir
+---
+name: diamond
+body: |
+ bb.0.entry:
+ successors: %bb.1(0x40000000), %bb.2(0x40000000)
+ S_CBRANCH_SCC1 %bb.1, implicit $scc
+ S_BRANCH %bb.2
+ bb.1.left:
+ successors: %bb.4(0x80000000)
+ S_BRANCH %bb.4
+ bb.2.right:
+ successors: %bb.3(0x80000000)
+ S_BRANCH %bb.3
+ bb.3.exit:
+ S_ENDPGM 0
+ bb.4.left:
+ successors: %bb.3(0x80000000)
+ S_BRANCH %bb.3
+...
+
+#--- redirected.mir
+---
+name: diamond
+body: |
+ bb.0.entry:
+ successors: %bb.1(0x40000000), %bb.2(0x40000000)
+ S_CBRANCH_SCC1 %bb.1, implicit $scc
+ S_BRANCH %bb.2
+ bb.1.left:
+ successors: %bb.2(0x80000000)
+ S_BRANCH %bb.2
+ bb.2.right:
+ successors: %bb.3(0x80000000)
+ S_BRANCH %bb.3
+ bb.3.exit:
+ S_ENDPGM 0
+...
+
+#--- unmapped.mir
+---
+name: diamond
+body: |
+ bb.0.entry:
+ successors: %bb.1(0x40000000), %bb.2(0x40000000)
+ S_CBRANCH_SCC1 %bb.1, implicit $scc
+ S_BRANCH %bb.2
+ bb.1.left:
+ successors: %bb.3(0x80000000)
+ S_BRANCH %bb.3
+ bb.2:
+ successors: %bb.3(0x80000000)
+ S_BRANCH %bb.3
+ bb.3.exit:
+ S_ENDPGM 0
+...
+
+#--- choice.mir
+---
+name: diamond
+tracksRegLiveness: true
+machineFunctionInfo:
+ scratchRSrcReg: '$sgpr0_sgpr1_sgpr2_sgpr3'
+ stackPtrOffsetReg: '$sgpr32'
+body: |
+ bb.0.entry:
+ successors: %bb.1(0x40000000), %bb.2(0x40000000)
+ liveins: $sgpr4
+ S_NOP 0, implicit-def %0:vgpr_32
+ S_CMP_EQ_U32 $sgpr4, 0, implicit-def $scc
+ S_CBRANCH_SCC1 %bb.1, implicit $scc
+ S_BRANCH %bb.2
+ bb.1.left:
+ successors: %bb.3(0x80000000)
+ S_NOP 0, implicit %0
+ S_NOP 0, implicit-def dead $vgpr1
+ S_NOP 0, implicit %0
+ S_BRANCH %bb.3
+ bb.2.right:
+ successors: %bb.3(0x80000000)
+ S_NOP 0, implicit %0
+ S_NOP 0, implicit-def dead $vgpr0
+ S_NOP 0, implicit %0
+ S_BRANCH %bb.3
+ bb.3.exit:
+ S_NOP 0, implicit %0
+ S_ENDPGM 0
+...
More information about the llvm-branch-commits
mailing list