[llvm] [AMDGPU] Add VGPR MSB-affinity pass to reduce s_set_vgpr_msb on gfx1250 (PR #216909)
Yuyang Zhang via llvm-commits
llvm-commits at lists.llvm.org
Sun Aug 23 21:25:56 PDT 2026
================
@@ -0,0 +1,1039 @@
+//===- AMDGPUVGPRMSBAffinity.cpp - VGPR MSB-group allocation hints --------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+/// \file
+/// On gfx1250 a wave may use 1024 VGPRs, but an instruction only addresses
+/// v0-v255; VGPRs 256-1023 are reached via per-slot MSB bits set by
+/// S_SET_VGPR_MSB, which AMDGPULowerVGPREncoding emits whenever a slot's MSB
+/// group changes between consecutive instructions.
+///
+/// This pre-RA pass (run after the scheduler fixes the order) records a desired
+/// MSB group per virtual register; SIRegisterInfo's allocation-hint hook then
+/// biases the greedy allocator toward it. The hint is soft, so it can never
+/// make allocation fail. Steps: build a schedule-driven affinity graph (edges
+/// between vregs that would cause a mode switch), cluster it under a per-group
+/// register-pressure cap, and pack the clusters into MSB groups.
+//
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPUVGPRMSBAffinity.h"
+#include "AMDGPU.h"
+#include "GCNSubtarget.h"
+#include "SIInstrInfo.h"
+#include "SIMachineFunctionInfo.h"
+#include "SIRegisterInfo.h"
+#include "Utils/AMDGPUBaseInfo.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
+#include "llvm/CodeGen/LiveIntervals.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineLoopInfo.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/SlotIndexes.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+#include <queue>
+
+using namespace llvm;
+
+#define DEBUG_TYPE "amdgpu-vgpr-msb-affinity-hints"
+
+static cl::opt<bool> EnableVGPRMSBAffinity(
+ "amdgpu-vgpr-msb-affinity", cl::Hidden, cl::init(false),
+ cl::desc("Bias VGPR allocation into 256-VGPR MSB groups to reduce "
+ "S_SET_VGPR_MSB insertions (gfx1250)"));
+
+static cl::opt<unsigned> BenefitPct(
+ "amdgpu-vgpr-msb-affinity-benefit-pct", cl::Hidden, cl::init(75),
+ cl::desc("Commit only if predicted plan switches < this % of the no-hint "
+ "switches (self-benefit gate; 0 disables it)"));
+
+static cl::opt<unsigned> MinBaseSwitch(
+ "amdgpu-vgpr-msb-affinity-min-base-switch", cl::Hidden, cl::init(500),
+ cl::desc("Skip when the naive baseline switch weight is below "
+ "this"));
+
+namespace {
+
+constexpr unsigned MSBGroupSize = 256;
+constexpr unsigned NumMSBGroups = 4;
+// Skip the plan when a group's planned load exceeds this percent of its cap.
+// Mild overflow is realizable (RA spills a few, most hints honored); severe
+// overflow is not.
+constexpr unsigned OverflowPct = 125;
+
+// Weighted, undirected affinity graph over virtual-register indices. An edge's
+// weight is the S_SET_VGPR_MSB cost paid if its two vregs land in different MSB
+// groups; FirstOrdinal (earliest program point) is a deterministic clustering
+// tie-break.
+class AffinityGraph {
+public:
+ // Canonical (order-independent) key for the edge between vreg indices A and
+ // B.
+ static uint64_t makeKey(unsigned A, unsigned B) {
+ if (A > B)
+ std::swap(A, B);
+ return (static_cast<uint64_t>(A) << 32) | B;
+ }
+ static unsigned lowEnd(uint64_t Key) {
+ return static_cast<unsigned>(Key >> 32);
+ }
+ static unsigned highEnd(uint64_t Key) {
+ return static_cast<unsigned>(Key & 0xffffffff);
+ }
+
+ // Accumulate Weight on the edge (A, B) and remember its earliest Ordinal.
+ void addEdge(unsigned A, unsigned B, uint64_t Weight, unsigned Ordinal) {
+ if (A == B)
+ return;
+ uint64_t Key = makeKey(A, B);
+ Weights[Key] += Weight;
+ FirstOrdinal.try_emplace(Key, Ordinal);
+ }
+
+ bool empty() const { return Weights.empty(); }
+ unsigned size() const { return Weights.size(); }
+ unsigned firstOrdinal(uint64_t Key) const { return FirstOrdinal.lookup(Key); }
+ const DenseMap<uint64_t, uint64_t> &edges() const { return Weights; }
+
+private:
+ DenseMap<uint64_t, uint64_t> Weights;
+ DenseMap<uint64_t, unsigned> FirstOrdinal;
+};
+
+// Union-find over vreg indices that greedily merges the heaviest affinity
+// edges, refusing a merge whose footprint would exceed a cap (one MSB group's
+// pressure limit). Footprints come from an injected functor so it stays
+// decoupled from LiveIntervals.
+class ClusterForest {
+public:
+ using FootprintFn = function_ref<int(ArrayRef<Register>)>;
+
+ ClusterForest(unsigned NumNodes, FootprintFn ComputeFootprint)
+ : Parent(NumNodes), Rank(NumNodes, 0), Epoch(NumNodes, 0),
+ Footprint(NumNodes, -1), Nodes(NumNodes),
+ ComputeFootprint(ComputeFootprint) {
+ for (unsigned I = 0; I < NumNodes; ++I)
+ Parent[I] = I;
+ }
+
+ // Seed the singleton cluster for node Idx with register Reg.
+ void addNode(unsigned Idx, Register Reg) { Nodes[Idx].push_back(Reg); }
+
+ unsigned find(unsigned X) {
+ while (Parent[X] != X) {
+ Parent[X] = Parent[Parent[X]];
+ X = Parent[X];
+ }
+ return X;
+ }
+
+ ArrayRef<Register> nodes(unsigned Root) const { return Nodes[Root]; }
+
+ // Cached simultaneously-live footprint of the cluster rooted at Root.
+ int footprintOf(unsigned Root) {
+ if (Footprint[Root] < 0)
+ Footprint[Root] = ComputeFootprint(Nodes[Root]);
+ return Footprint[Root];
+ }
+
+ // Greedily merge edges heaviest-first (ties broken by smallest footprint
+ // delta, then earliest ordinal, then key for determinism), refusing a merge
+ // that would push the merged footprint past MergeCap. Each refused edge is a
+ // cut.
+ void clusterByWeight(const AffinityGraph &Graph, unsigned MergeCap) {
+ struct Item {
+ uint64_t Weight;
+ int Delta;
+ unsigned Ordinal;
+ uint64_t Key;
+ unsigned RootA, RootB;
+ uint64_t EpochA, EpochB;
+ // Max-heap ordering: an Item that should be processed first must compare
+ // "greater" than the others.
+ bool operator<(const Item &O) const {
+ if (Weight != O.Weight)
+ return Weight < O.Weight; // higher weight first
+ if (Delta != O.Delta)
+ return Delta > O.Delta; // smaller footprint delta first
+ if (Ordinal != O.Ordinal)
+ return Ordinal > O.Ordinal; // earlier program order first
+ return Key > O.Key; // lower key first (determinism)
+ }
+ };
+ std::priority_queue<Item> Queue;
+ // (Re-)evaluate the merge for an edge against the current forest and
+ // enqueue it. Endpoints already in the same cluster are dropped.
+ auto PushEdge = [&](uint64_t Key, uint64_t Weight) {
+ unsigned RootA = find(AffinityGraph::lowEnd(Key));
+ unsigned RootB = find(AffinityGraph::highEnd(Key));
+ if (RootA == RootB)
+ return;
+ int Merged = unionFootprint(RootA, RootB);
+ int Delta = Merged - std::max(footprintOf(RootA), footprintOf(RootB));
+ Queue.push({Weight, Delta, Graph.firstOrdinal(Key), Key, RootA, RootB,
+ Epoch[RootA], Epoch[RootB]});
+ };
+ for (auto &[Key, Weight] : Graph.edges())
+ PushEdge(Key, Weight);
+ while (!Queue.empty()) {
+ Item Top = Queue.top();
+ Queue.pop();
+ unsigned RootA = find(AffinityGraph::lowEnd(Top.Key));
+ unsigned RootB = find(AffinityGraph::highEnd(Top.Key));
+ if (RootA == RootB)
+ continue;
+ // A touched cluster changed since this item was pushed -> its delta/roots
+ // are stale, so re-evaluate and re-enqueue rather than act on it.
+ if (RootA != Top.RootA || RootB != Top.RootB ||
+ Epoch[RootA] != Top.EpochA || Epoch[RootB] != Top.EpochB) {
+ PushEdge(Top.Key, Top.Weight);
+ continue;
+ }
+ int Merged = unionFootprint(RootA, RootB);
+ if (Merged > static_cast<int>(MergeCap))
+ continue; // Refuse: this edge becomes a cut.
+ mergeInto(RootA, RootB, Merged);
+ }
+ }
+
+private:
+ // Exact union footprint of two clusters (time-aware peak of their nodes).
+ int unionFootprint(unsigned RootA, unsigned RootB) {
+ SmallVector<Register, 16> Both(Nodes[RootA].begin(), Nodes[RootA].end());
+ Both.append(Nodes[RootB].begin(), Nodes[RootB].end());
+ return ComputeFootprint(Both);
+ }
+ // Rank-union RootB into RootA (the higher-rank root is kept), fold nodes and
+ // bump the kept root's epoch so stale queue items are detected.
+ void mergeInto(unsigned RootA, unsigned RootB, int MergedFootprint) {
+ if (Rank[RootA] < Rank[RootB])
+ std::swap(RootA, RootB);
+ Parent[RootB] = RootA;
+ if (Rank[RootA] == Rank[RootB])
+ ++Rank[RootA];
+ Nodes[RootA].append(Nodes[RootB].begin(), Nodes[RootB].end());
+ Nodes[RootB].clear();
+ Footprint[RootA] = MergedFootprint;
+ Footprint[RootB] = -1;
+ ++Epoch[RootA];
+ }
+
+ SmallVector<unsigned, 0> Parent;
+ SmallVector<unsigned, 0> Rank;
+ SmallVector<uint64_t, 0> Epoch;
+ SmallVector<int, 0> Footprint; // cached per-root footprint, -1 = stale.
+ SmallVector<SmallVector<Register, 4>, 0> Nodes;
+ FootprintFn ComputeFootprint;
+};
+
+class AMDGPUVGPRMSBAffinity {
+public:
+ bool run(MachineFunction &MF, LiveIntervals *LIS, MachineLoopInfo *MLI);
+
+private:
+ // Cluster packing strategy: pack hottest-first into the lowest group that
+ // fits (Compact) or into the least-loaded group so clusters spread and every
+ // used group keeps slack for the soft hints (Balanced).
+ enum class PackMode { Compact, Balanced };
+ // Scope the self-benefit gate is scored over: the whole function, or only the
+ // in-loop (recurring) switches.
+ enum class GateScope { WholeFunction, LoopOnly };
+
+ // Build the affinity graph, cluster, pack into MSB groups and commit hints
+ // for one region (a set of blocks). Vregs already in \p Assigned (hinted by a
+ // hotter region) are skipped; newly hinted vregs are added to it.
+ void processRegion(ArrayRef<MachineBasicBlock *> Blocks,
+ ArrayRef<Register> AllVGPRs, unsigned EffMSBGroups,
+ unsigned VGPRBudget, PackMode Mode, GateScope Scope,
+ DenseSet<unsigned> &Assigned, SIMachineFunctionInfo *MFI);
+
+ AffinityGraph buildAffinityGraph(ArrayRef<MachineBasicBlock *> Blocks) const;
+
+ SmallVector<unsigned, 0> collectHotRoots(ClusterForest &Forest,
+ const AffinityGraph &Graph) const;
+
+ void packClusters(ClusterForest &Forest, ArrayRef<unsigned> Roots,
+ unsigned EffMSBGroups, unsigned VGPRBudget, PackMode Mode,
+ MutableArrayRef<int> MSBLoad,
+ DenseMap<unsigned, int> &ClusterMSB) const;
+
+ // Per-group register cap. For power-of-two occupancy VGPRBudget is a whole
+ // number of 256-groups so every cap is 256; at a fractional occupancy the
+ // last group holds only VGPRBudget - (EffMSBGroups-1)*256 registers. Clamp to
+ // [1, 256] so a fractional group is never over-packed (which drops a wave).
+ static unsigned groupCap(unsigned Group, unsigned VGPRBudget) {
+ int Cap =
+ static_cast<int>(VGPRBudget) - static_cast<int>(Group * MSBGroupSize);
+ return static_cast<unsigned>(std::max(1, std::min<int>(MSBGroupSize, Cap)));
+ }
+
+ const SIInstrInfo *TII = nullptr;
+ const SIRegisterInfo *TRI = nullptr;
+ MachineRegisterInfo *MRI = nullptr;
+ LiveIntervals *LIS = nullptr;
+ MachineLoopInfo *MLI = nullptr;
+ const GCNSubtarget *STI = nullptr;
+
+ // Per-block edge weight: a loop-depth proxy for trip count, so an
+ // innermost-loop transition outweighs straight-line code by orders of
+ // magnitude.
+ uint64_t blockFreq(const MachineBasicBlock &MBB) const {
+ unsigned Depth = MLI ? MLI->getLoopDepth(&MBB) : 0;
+ return 1ull << std::min(4u * Depth, 40u);
+ }
+
+ // Value-group union-find (mutable for path compression): vregs that coalesce
+ // to one physreg must be counted once in the footprint. See buildValueGroups.
+ mutable SmallVector<unsigned, 0> VGParent;
+
+ bool isVGPRVirtReg(Register Reg) const {
+ return Reg.isVirtual() && TRI->isVGPRClass(MRI->getRegClass(Reg));
+ }
+
+ unsigned dwords(Register Reg) const {
+ // Integer-divide by 32 (a 16-bit vreg -> 0): rounding up over-counts
+ // lo16/hi16 pairs that share a dword and worsens True16 plans. The
+ // footprint only feeds soft hints, so the undercount is acceptable.
+ return TRI->getRegSizeInBits(*MRI->getRegClass(Reg)) / 32;
+ }
+
+ // Record the MSB-group affinity and also a concrete physreg hint in that
+ // group: the latter marks a known preference so greedy colors the vreg early
+ // and it claims its group before contention. Existing (copy) hints win.
+ void recordMSB(SIMachineFunctionInfo *MFI, Register Reg, unsigned MSB) {
+ MFI->setVGPRMSBAffinity(Reg, MSB);
+ if (MRI->getRegAllocationHint(Reg).second)
+ return;
+ const TargetRegisterClass *RC = MRI->getRegClass(Reg);
+ for (MCPhysReg P : *RC) {
+ if (!MRI->isReserved(P) && (TRI->getHWRegIndex(P) >> 8) == MSB) {
+ MRI->setRegAllocationHint(Reg, 0, P);
+ return;
+ }
+ }
+ }
+
+ unsigned vgFind(unsigned X) const {
+ while (VGParent[X] != X) {
+ VGParent[X] = VGParent[VGParent[X]];
+ X = VGParent[X];
+ }
+ return X;
+ }
+
+ void buildValueGroups(MachineFunction &MF) {
+ unsigned N = MRI->getNumVirtRegs();
+ VGParent.resize(N);
+ for (unsigned I = 0; I < N; ++I)
+ VGParent[I] = I;
+ auto UnionVGroup = [&](Register A, Register B) {
+ if (!isVGPRVirtReg(A) || !isVGPRVirtReg(B))
+ return;
+ unsigned RootA = vgFind(A.virtRegIndex()),
+ RootB = vgFind(B.virtRegIndex());
+ if (RootA != RootB)
+ VGParent[RootA] = RootB;
+ };
+ for (MachineBasicBlock &MBB : MF) {
+ for (MachineInstr &MI : MBB) {
+ // Coalesce tied def/use pairs (e.g. the WMMA accumulator src2 tied to
+ // dst). General COPYs are intentionally *not* unioned: they connect
+ // distinct values and would collapse unrelated footprints.
+ for (unsigned I = 0, E = MI.getNumOperands(); I < E; ++I) {
+ const MachineOperand &MO = MI.getOperand(I);
+ if (MO.isReg() && MO.isUse() && MO.isTied()) {
+ unsigned DefIdx = MI.findTiedOperandIdx(I);
+ const MachineOperand &Def = MI.getOperand(DefIdx);
+ if (Def.isReg())
+ UnionVGroup(MO.getReg(), Def.getReg());
+ }
+ }
+ // Coalesce the accumulator chain dst <- src2: across an unrolled K-loop
+ // this chains acc0->acc1->... into one value group so the footprint
+ // counts the accumulator once. Disjoint output tiles never merge.
+ if (SIInstrInfo::isWMMA(MI) || TII->isMAI(MI)) {
+ const MachineOperand *D =
+ TII->getNamedOperand(MI, AMDGPU::OpName::vdst);
+ const MachineOperand *S2 =
+ TII->getNamedOperand(MI, AMDGPU::OpName::src2);
+ if (D && D->isReg() && S2 && S2->isReg())
+ UnionVGroup(D->getReg(), S2->getReg());
+ }
+ }
+ }
+ }
+
+ // Peak simultaneously-live VGPR dwords in \p Regs, merging the live ranges of
+ // a value group so a coalescing value is counted once.
+ unsigned maxSimultaneousDwords(ArrayRef<Register> Regs) const {
+ DenseMap<unsigned, SmallVector<std::pair<SlotIndex, SlotIndex>, 2>> ByGroup;
+ DenseMap<unsigned, int> GroupSize;
+ for (Register Reg : Regs) {
+ if (!LIS->hasInterval(Reg))
+ continue;
+ unsigned G = vgFind(Reg.virtRegIndex());
----------------
yuyzhang512 wrote:
Done
https://github.com/llvm/llvm-project/pull/216909
More information about the llvm-commits
mailing list