[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);
----------------
yuyzhang512 wrote:

Done

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