[llvm] [AMDGPU] Add VGPR MSB-affinity pass to reduce s_set_vgpr_msb on gfx1250 (PR #216909)

[Yuyang Zhang via llvm-commits]

================
@@ -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) {
----------------
yuyzhang512 wrote:

Done — clusterByWeight, buildValueGroups, maxSimultaneousDwords, computeNaiveMSB and simSwitchWeight are now defined out-of-line

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