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" <[email protected]> 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 +... _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
