[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:33:02 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());
+      GroupSize[G] = std::max<int>(GroupSize[G], dwords(Reg));
+      auto &Segs = ByGroup[G];
+      for (const LiveRange::Segment &S : LIS->getInterval(Reg))
+        Segs.emplace_back(S.start, S.end);
+    }
+    SmallVector<std::pair<SlotIndex, int>, 64> Events;
+    for (auto &[G, Segs] : ByGroup) {
+      llvm::sort(Segs);
+      int Sz = GroupSize[G];
+      SlotIndex CurS, CurE;
+      bool Open = false;
+      auto Flush = [&] {
+        Events.emplace_back(CurS, Sz);
+        Events.emplace_back(CurE, -Sz);
+      };
+      for (auto &[S, E] : Segs) {
+        if (Open && S <= CurE) {
+          CurE = std::max(CurE, E); // overlaps: extend the open interval
+        } else {
+          if (Open) // gap: close the previous interval
+            Flush();
+          CurS = S;
+          CurE = E;
+          Open = true;
+        }
+      }
+      if (Open)
+        Flush();
+    }
+    llvm::sort(Events, [](const std::pair<SlotIndex, int> &A,
+                          const std::pair<SlotIndex, int> &B) {
+      return A.first < B.first || (A.first == B.first && A.second < B.second);
+    });
+    int Cur = 0, Max = 0;
+    for (auto &[Idx, Delta] : Events) {
+      Cur += Delta;
+      Max = std::max(Max, Cur);
+    }
+    return Max;
+  }
+
+  // Natural (no-hint) MSB-group assignment for the self-benefit baseline: a
+  // linear scan placing each vreg in the lowest free column run and freeing
+  // columns as live ranges end -- an approximation of what the allocator does
+  // unhinted.
+  DenseMap<unsigned, int> computeNaiveMSB(ArrayRef<Register> Regs,
+                                          unsigned EffMSBGroups) const {
+    DenseMap<unsigned, int> MSB;
+    const unsigned Cols = EffMSBGroups * MSBGroupSize;
+    SmallVector<Register, 0> Order(Regs.begin(), Regs.end());
+    llvm::stable_sort(Order, [&](Register A, Register B) {
+      return LIS->getInterval(A).beginIndex() <
+             LIS->getInterval(B).beginIndex();
+    });
+    SmallVector<bool, 0> Free(Cols, true);
+    // Active allocations: (endIndex, startCol, width) to reclaim columns.
+    SmallVector<std::tuple<SlotIndex, unsigned, unsigned>, 0> Active;
+    for (Register R : Order) {
+      SlotIndex Begin = LIS->getInterval(R).beginIndex();
+      // Reclaim columns of ranges that ended before this def.
+      for (unsigned I = 0; I < Active.size();) {
+        if (std::get<0>(Active[I]) <= Begin) {
+          unsigned StartCol = std::get<1>(Active[I]);
+          unsigned RunWidth = std::get<2>(Active[I]);
+          for (unsigned Col = StartCol; Col < StartCol + RunWidth; ++Col)
+            Free[Col] = true;
+          Active[I] = Active.back();
+          Active.pop_back();
+        } else
+          ++I;
+      }
+      unsigned Width = dwords(R);
+      // Lowest free run of Width columns.
+      int Start = -1;
+      for (unsigned Col = 0, Run = 0; Col < Cols; ++Col) {
+        Run = Free[Col] ? Run + 1 : 0;
+        if (Run == Width) {
+          Start = static_cast<int>(Col + 1 - Width);
+          break;
+        }
+      }
+      int Group;
+      if (Start < 0) {
+        // No contiguous run fits: reserve Width columns in the least-occupied
+        // MSB group so this vreg's footprint stays visible (otherwise the
+        // baseline looks artificially uncongested and skews the self-benefit
+        // comparison).
+        unsigned BestGroup = 0, BestFreeCount = 0;
+        for (unsigned Cand = 0; Cand < EffMSBGroups; ++Cand) {
+          unsigned FreeCount = 0;
+          for (unsigned Col = Cand * MSBGroupSize;
+               Col < (Cand + 1) * MSBGroupSize; ++Col)
+            FreeCount += Free[Col];
+          if (FreeCount >= BestFreeCount) {
+            BestFreeCount = FreeCount;
+            BestGroup = Cand;
+          }
+        }
+        for (unsigned Col = BestGroup * MSBGroupSize, Reserved = 0;
+             Col < (BestGroup + 1) * MSBGroupSize && Reserved < Width; ++Col)
+          if (Free[Col]) {
+            Free[Col] = false;
+            ++Reserved;
+          }
+        Group = static_cast<int>(BestGroup);
+      } else {
+        for (unsigned Col = Start; Col < Start + Width; ++Col)
+          Free[Col] = false;
+        Active.emplace_back(LIS->getInterval(R).endIndex(),
+                            static_cast<unsigned>(Start), Width);
+        Group = Start / static_cast<int>(MSBGroupSize);
+      }
+      MSB[R.virtRegIndex()] = Group;
+    }
+    return MSB;
+  }
+
+  // Predicted freq-weighted s_set_vgpr_msb count for a vreg->MSB-group map,
+  // simulated like AMDGPULowerVGPREncoding: walk the stream with sticky
+  // per-slot state (reset per block), charging blockFreq per instruction that
+  // changes a slot's group.
+  uint64_t simSwitchWeight(ArrayRef<MachineBasicBlock *> Blocks,
+                           function_ref<int(Register)> MsbOf,
+                           bool LoopOnly = false) const {
+    uint64_t Sw = 0;
+    for (MachineBasicBlock *MBBp : Blocks) {
+      MachineBasicBlock &MBB = *MBBp;
+      // Realizability/relevance: only in-loop switches recur every iteration
+      // and dominate runtime cost; prologue/epilogue switches fire once.
+      // Scoring the gate on loop blocks only keeps the plan from trading a loop
+      // win for one-time out-of-loop churn (which the whole-function total
+      // misranks).
+      if (LoopOnly && (!MLI || MLI->getLoopDepth(&MBB) == 0))
+        continue;
+      uint64_t Freq = blockFreq(MBB);
+      // Mode is reset to group 0 at a block header (and again at a call /
+      // terminator / VGPR inline asm), matching AMDGPULowerVGPREncoding.
+      int Last[4] = {0, 0, 0, 0};
+      for (MachineInstr &MI : MBB) {
+        if (MI.isMetaInstruction())
+          continue;
+        if (MI.isTerminator() || MI.isCall() ||
+            (MI.isInlineAsm() && TII->hasVGPRUses(MI))) {
+          Last[0] = Last[1] = Last[2] = Last[3] = 0;
+          continue;
+        }
+        auto Ops = AMDGPU::getVGPRLoweringOperandTables(MI.getDesc());
+        if (!Ops.first)
+          continue;
+        int Need[4] = {-1, -1, -1, -1};
+        for (unsigned S = 0; S < 4; ++S) {
+          const MachineOperand *MO = TII->getNamedOperand(MI, Ops.first[S]);
+          if ((!MO || !MO->isReg() || !MO->getReg()) && Ops.second)
+            MO = TII->getNamedOperand(MI, Ops.second[S]);
+          if (!MO || !MO->isReg() || !MO->getReg())
+            continue;
+          Register R = MO->getReg();
+          if (isVGPRVirtReg(R))
+            Need[S] = std::max(0, MsbOf(R));
+          else if (R.isPhysical() && TRI->isVGPR(*MRI, R))
+            Need[S] = static_cast<int>(TRI->getHWRegIndex(R) >> 8);
+        }
+        bool Changed = false;
+        for (unsigned S = 0; S < 4; ++S)
+          if (Need[S] >= 0 && Last[S] != Need[S])
+            Changed = true;
+        if (Changed)
+          Sw += Freq;
+        for (unsigned S = 0; S < 4; ++S)
+          if (Need[S] >= 0)
+            Last[S] = Need[S];
+      }
+    }
+    return Sw;
+  }
+};
+
+class AMDGPUVGPRMSBAffinityLegacy : public MachineFunctionPass {
+public:
+  static char ID;
+
+  AMDGPUVGPRMSBAffinityLegacy() : MachineFunctionPass(ID) {}
+
+  bool runOnMachineFunction(MachineFunction &MF) override {
+    auto *LISW = getAnalysisIfAvailable<LiveIntervalsWrapperPass>();
+    auto *MLIW = getAnalysisIfAvailable<MachineLoopInfoWrapperPass>();
+    return AMDGPUVGPRMSBAffinity().run(MF, LISW ? &LISW->getLIS() : nullptr,
+                                       MLIW ? &MLIW->getLI() : nullptr);
+  }
+
+  StringRef getPassName() const override { return "AMDGPU VGPR MSB Affinity"; }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<LiveIntervalsWrapperPass>();
+    AU.addRequired<MachineLoopInfoWrapperPass>();
+    AU.setPreservesAll();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+};
+
+} // namespace
+
+bool AMDGPUVGPRMSBAffinity::run(MachineFunction &MF, LiveIntervals *LISIn,
+                                MachineLoopInfo *MLIIn) {
+  if (!EnableVGPRMSBAffinity)
+    return false;
+
+  const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
+  if (!ST.has1024AddressableVGPRs() || !LISIn)
+    return false;
+
+  // Only steer compute kernels; graphics shaders are out of scope for the
+  // 1024-VGPR / s_set_vgpr_msb MSB grouping this pass targets.
+  if (!AMDGPU::isCompute(MF.getFunction().getCallingConv()))
+    return false;
+
+  TII = ST.getInstrInfo();
+  TRI = ST.getRegisterInfo();
+  MRI = &MF.getRegInfo();
+  LIS = LISIn;
+  MLI = MLIIn;
+  STI = &ST;
+
+  LLVM_DEBUG(dbgs() << "*** AMDGPUVGPRMSBAffinity on " << MF.getName()
+                    << " ***\n");
+
+  // Coalesce vregs that will share one physreg so the footprint counts them
+  // once.
+  buildValueGroups(MF);
+
+  // If the whole function fits one 256-VGPR group, no S_SET_VGPR_MSB is ever
+  // needed and partitioning would only inflate the VGPR count.
+  SmallVector<Register, 0> AllVGPRs;
+  for (unsigned I = 0, E = MRI->getNumVirtRegs(); I != E; ++I) {
+    Register R = Register::index2VirtReg(I);
+    if (!MRI->reg_nodbg_empty(R) && isVGPRVirtReg(R))
+      AllVGPRs.push_back(R);
+  }
+  unsigned GlobalFP = maxSimultaneousDwords(AllVGPRs);
+  LLVM_DEBUG(dbgs() << "  early-check GlobalFP=" << GlobalFP << "\n");
+  if (GlobalFP <= MSBGroupSize) {
+    LLVM_DEBUG(dbgs() << "  -> return: footprint fits one group\n");
+    return false;
+  }
+
+  // Baseline occupancy: min of the VGPR-limited estimate and MFI's non-VGPR
+  // limit.
+  const SIMachineFunctionInfo *MFIOcc = MF.getInfo<SIMachineFunctionInfo>();
+  unsigned VOcc = STI->getOccupancyWithNumVGPRs(
+      GlobalFP, MFIOcc->getDynamicVGPRBlockSize());
+  unsigned BaseOcc = std::min(VOcc, MFIOcc->getOccupancy());
+  LLVM_DEBUG(dbgs() << "  early-check BaseOcc=" << BaseOcc << " (VOcc=" << VOcc
+                    << " MFIOcc=" << MFIOcc->getOccupancy() << ")\n");
+  if (BaseOcc == 0)
----------------
yuyzhang512 wrote:

Done, changed to an assert

https://github.com/llvm/llvm-project/pull/216909


More information about the llvm-commits mailing list