[llvm] ad3c19a - [AMDGPU][CodeGen] Incrementally update reserved regs for SIPreAllocateWWMRegs pass in RegisterClassInfo (#212201)
via llvm-commits
llvm-commits at lists.llvm.org
Fri Aug 28 08:36:02 PDT 2026
Author: Nikhil Kotikalapudi
Date: 2026-08-28T17:35:57+02:00
New Revision: ad3c19a8984263286bc930c4d7a127cced8200ec
URL: https://github.com/llvm/llvm-project/commit/ad3c19a8984263286bc930c4d7a127cced8200ec
DIFF: https://github.com/llvm/llvm-project/commit/ad3c19a8984263286bc930c4d7a127cced8200ec.diff
LOG: [AMDGPU][CodeGen] Incrementally update reserved regs for SIPreAllocateWWMRegs pass in RegisterClassInfo (#212201)
Based on this TODO comment in SIPreAllocateWWMRegs.cpp:
76 // TODO: Update RCI with the additional reserved registers the pass sets.
77 AU.addRequired<MachineRegisterClassInfoWrapperPass>();
Creates function updateReservedRegs in RegisterClassInfo.cpp which takes
the BitVector containing reserved register information without requiring
any unnecessary recomputations of the entire RegisterClassInfo object.
It:
- Compares the new reserved-register set with the cached set.
- Removes newly reserved registers from cached allocation orders.
- Updates register counts, costs, subclass information, and pressure
limits.
- Invalidates entries when incremental updating is unsafe; they are
recomputed only when next requested.
- Allows `SIPreAllocateWWMRegs` to preserve `RegisterClassInfo` after
reserving WWM registers.
Also drops the unconditional compute(RC) in computePSetLimit, which was
redundant since getNumAllocatableRegs already recomputes stale entries.
This is within the scope of my PR because it would unnecessarily rebuild
an already-valid entry, rendering the incremental update pointless.
---------
Co-authored-by: Matt Arsenault <arsenm2 at gmail.com>
Added:
llvm/test/CodeGen/AMDGPU/si-pre-allocate-wwm-regs-preserve-rci.mir
llvm/unittests/Target/AMDGPU/RCIUpdateReservedRegsTest.cpp
Modified:
llvm/include/llvm/CodeGen/RegisterClassInfo.h
llvm/lib/CodeGen/RegisterClassInfo.cpp
llvm/lib/Target/AMDGPU/SIPreAllocateWWMRegs.cpp
llvm/unittests/Target/AMDGPU/CMakeLists.txt
Removed:
llvm/test/CodeGen/AMDGPU/si-pre-allocate-wwm-regs-invalidate-rci.mir
################################################################################
diff --git a/llvm/include/llvm/CodeGen/RegisterClassInfo.h b/llvm/include/llvm/CodeGen/RegisterClassInfo.h
index 256277832db24..ab286d8da1e07 100644
--- a/llvm/include/llvm/CodeGen/RegisterClassInfo.h
+++ b/llvm/include/llvm/CodeGen/RegisterClassInfo.h
@@ -98,6 +98,15 @@ class RegisterClassInfo {
LLVM_ABI void runOnMachineFunction(const MachineFunction &MF,
bool Rev = false);
+ /// Update cached register class information using ReservedInput, MRI's
+ /// current reserved-register set. Cached orders are compacted when
+ /// only registers are added and `-stress-regalloc` is disabled; otherwise,
+ /// they are invalidated and recomputed on demand.
+ ///
+ /// RegisterClassInfo must be initialized. Target information, callee-saved
+ /// registers, register costs, and allocation orders must remain unchanged.
+ LLVM_ABI void updateReservedRegs(const BitVector &ReservedInput);
+
LLVM_ABI bool invalidate(MachineFunction &, const PreservedAnalyses &PA,
MachineFunctionAnalysisManager::Invalidator &) {
auto PAC = PA.getChecker<MachineRegisterClassAnalysis>();
diff --git a/llvm/lib/CodeGen/RegisterClassInfo.cpp b/llvm/lib/CodeGen/RegisterClassInfo.cpp
index da25c9b8c29c5..c3872e155098f 100644
--- a/llvm/lib/CodeGen/RegisterClassInfo.cpp
+++ b/llvm/lib/CodeGen/RegisterClassInfo.cpp
@@ -120,6 +120,67 @@ void RegisterClassInfo::runOnMachineFunction(const MachineFunction &mf,
}
}
+void RegisterClassInfo::updateReservedRegs(const BitVector &ReservedInput) {
+ assert(MF && TRI && RegClass &&
+ "RegisterClassInfo must be initialized before updating reserved regs");
+ assert(ReservedInput.size() == Reserved.size() &&
+ "Reserved register bit vectors must have the same size");
+ if (ReservedInput == Reserved)
+ return;
+
+ // Cached orders cannot regain unreserved registers; recompute them lazily.
+ bool OnlyNewReservations = Reserved.subsetOf(ReservedInput);
+
+ Reserved = ReservedInput;
+
+ // Pressure limits depend on the number of allocatable registers.
+ std::fill_n(PSetLimits.get(), TRI->getNumRegPressureSets(), 0);
+
+ // NumRegs may hide entries beyond the stress limit, so those orders cannot
+ // safely be compacted using only their visible prefix.
+ if (!OnlyNewReservations || StressRA) {
+ ++Tag;
+ return;
+ }
+
+ for (const TargetRegisterClass &RC : TRI->regclasses()) {
+ RCInfo &Info = RegClass[RC.getID()];
+
+ // Skip stale class information.
+ if (Info.Tag != Tag)
+ continue;
+
+ unsigned NewNumRegs = 0;
+ uint8_t MinCost = uint8_t(~0u);
+ uint8_t LastCost = uint8_t(~0u);
+ unsigned LastCostChange = 0;
+
+ for (unsigned I = 0; I != Info.NumRegs; ++I) {
+ MCPhysReg PhysReg = Info.Order[I];
+ if (Reserved.test(PhysReg))
+ continue;
+
+ uint8_t Cost = RegCosts[PhysReg];
+ MinCost = std::min(MinCost, Cost);
+ if (Cost != LastCost)
+ LastCostChange = NewNumRegs;
+
+ Info.Order[NewNumRegs++] = PhysReg;
+ LastCost = Cost;
+ }
+
+ Info.NumRegs = NewNumRegs;
+ Info.MinCost = MinCost;
+ Info.LastCostChange = LastCostChange;
+
+ Info.ProperSubClass = false;
+ if (const TargetRegisterClass *Super =
+ TRI->getLargestLegalSuperClass(&RC, *MF))
+ if (Super != &RC && getNumAllocatableRegs(Super) > Info.NumRegs)
+ Info.ProperSubClass = true;
+ }
+}
+
/// compute - Compute the preferred allocation order for RC with reserved
/// registers filtered out. Volatile registers come first followed by CSR
/// aliases ordered according to the CSR order specified by the target.
@@ -203,7 +264,6 @@ void RegisterClassInfo::compute(const TargetRegisterClass *RC) const {
unsigned RegisterClassInfo::computePSetLimit(unsigned Idx) const {
const TargetRegisterClass *RC = TRI->getLargestRegClassForRegPressureSet(Idx);
assert(RC && "Failed to find register class");
- compute(RC);
unsigned NAllocatableRegs = getNumAllocatableRegs(RC);
unsigned RegPressureSetLimit = TRI->getRegPressureSetLimit(*MF, Idx);
// If all the regs are reserved, return raw RegPressureSetLimit.
diff --git a/llvm/lib/Target/AMDGPU/SIPreAllocateWWMRegs.cpp b/llvm/lib/Target/AMDGPU/SIPreAllocateWWMRegs.cpp
index 0d9ef77b70562..3a749d040f8e6 100644
--- a/llvm/lib/Target/AMDGPU/SIPreAllocateWWMRegs.cpp
+++ b/llvm/lib/Target/AMDGPU/SIPreAllocateWWMRegs.cpp
@@ -17,13 +17,11 @@
#include "MCTargetDesc/AMDGPUMCTargetDesc.h"
#include "SIMachineFunctionInfo.h"
#include "llvm/ADT/PostOrderIterator.h"
-#include "llvm/CodeGen/LiveDebugVariables.h"
#include "llvm/CodeGen/LiveIntervals.h"
#include "llvm/CodeGen/LiveRegMatrix.h"
#include "llvm/CodeGen/MachineFrameInfo.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/RegisterClassInfo.h"
-#include "llvm/CodeGen/SlotIndexes.h"
#include "llvm/CodeGen/VirtRegMap.h"
#include "llvm/InitializePasses.h"
@@ -50,7 +48,7 @@ class SIPreAllocateWWMRegs {
LiveIntervals *LIS;
LiveRegMatrix *Matrix;
VirtRegMap *VRM;
- const RegisterClassInfo &RegClassInfo;
+ RegisterClassInfo &RCI;
std::vector<unsigned> RegsToRewrite;
#ifndef NDEBUG
@@ -61,8 +59,8 @@ class SIPreAllocateWWMRegs {
public:
SIPreAllocateWWMRegs(LiveIntervals *LIS, LiveRegMatrix *Matrix,
- VirtRegMap *VRM, const RegisterClassInfo &RCI)
- : LIS(LIS), Matrix(Matrix), VRM(VRM), RegClassInfo(RCI) {}
+ VirtRegMap *VRM, RegisterClassInfo &RCI)
+ : LIS(LIS), Matrix(Matrix), VRM(VRM), RCI(RCI) {}
bool run(MachineFunction &MF);
};
@@ -78,14 +76,8 @@ class SIPreAllocateWWMRegsLegacy : public MachineFunctionPass {
AU.addRequired<LiveIntervalsWrapperPass>();
AU.addRequired<VirtRegMapWrapperLegacy>();
AU.addRequired<LiveRegMatrixWrapperLegacy>();
- // TODO: Update RCI with the additional reserved registers the pass sets.
AU.addRequired<MachineRegisterClassInfoWrapperPass>();
- AU.setPreservesCFG();
- AU.addPreserved<LiveIntervalsWrapperPass>();
- AU.addPreserved<SlotIndexesWrapperPass>();
- AU.addPreserved<VirtRegMapWrapperLegacy>();
- AU.addPreserved<LiveRegMatrixWrapperLegacy>();
- AU.addPreserved<LiveDebugVariablesWrapperLegacy>();
+ AU.setPreservesAll();
MachineFunctionPass::getAnalysisUsage(AU);
}
};
@@ -122,7 +114,7 @@ bool SIPreAllocateWWMRegs::processDef(MachineOperand &MO) {
LiveInterval &LI = LIS->getInterval(Reg);
- for (MCRegister PhysReg : RegClassInfo.getOrder(MRI->getRegClass(Reg))) {
+ for (MCRegister PhysReg : RCI.getOrder(MRI->getRegClass(Reg))) {
if (!MRI->isPhysRegUsed(PhysReg, /*SkipRegMaskTest=*/true) &&
Matrix->checkInterference(LI, PhysReg) == LiveRegMatrix::IK_Free) {
Matrix->assign(LI, PhysReg);
@@ -180,8 +172,10 @@ void SIPreAllocateWWMRegs::rewriteRegs(MachineFunction &MF) {
RegsToRewrite.clear();
- // Update the set of reserved registers to include WWM ones.
+ // Update the set of reserved registers to include WWM ones
+ // without unnecessarily invalidating RegClassInfo.
MRI->freezeReservedRegs();
+ RCI.updateReservedRegs(MRI->getReservedRegs());
}
#ifndef NDEBUG
@@ -213,7 +207,7 @@ bool SIPreAllocateWWMRegsLegacy::runOnMachineFunction(MachineFunction &MF) {
auto *LIS = &getAnalysis<LiveIntervalsWrapperPass>().getLIS();
auto *Matrix = &getAnalysis<LiveRegMatrixWrapperLegacy>().getLRM();
auto *VRM = &getAnalysis<VirtRegMapWrapperLegacy>().getVRM();
- const auto &RCI = getAnalysis<MachineRegisterClassInfoWrapperPass>().getRCI();
+ auto &RCI = getAnalysis<MachineRegisterClassInfoWrapperPass>().getRCI();
return SIPreAllocateWWMRegs(LIS, Matrix, VRM, RCI).run(MF);
}
@@ -283,12 +277,7 @@ SIPreAllocateWWMRegsPass::run(MachineFunction &MF,
auto *LIS = &MFAM.getResult<LiveIntervalsAnalysis>(MF);
auto *Matrix = &MFAM.getResult<LiveRegMatrixAnalysis>(MF);
auto *VRM = &MFAM.getResult<VirtRegMapAnalysis>(MF);
- const auto &RCI = MFAM.getResult<MachineRegisterClassAnalysis>(MF);
+ auto &RCI = MFAM.getResult<MachineRegisterClassAnalysis>(MF);
SIPreAllocateWWMRegs(LIS, Matrix, VRM, RCI).run(MF);
- // The pass reserves WWM registers, invalidating RegisterClassInfo's
- // allocation order, so it cannot be preserved (see the legacy
- // getAnalysisUsage above).
- PreservedAnalyses PA = PreservedAnalyses::all();
- PA.abandon<MachineRegisterClassAnalysis>();
- return PA;
+ return PreservedAnalyses::all();
}
diff --git a/llvm/test/CodeGen/AMDGPU/si-pre-allocate-wwm-regs-invalidate-rci.mir b/llvm/test/CodeGen/AMDGPU/si-pre-allocate-wwm-regs-invalidate-rci.mir
deleted file mode 100644
index 6571294bac741..0000000000000
--- a/llvm/test/CodeGen/AMDGPU/si-pre-allocate-wwm-regs-invalidate-rci.mir
+++ /dev/null
@@ -1,32 +0,0 @@
-# RUN: llc -mtriple=amdgpu7.00-amd-amdhsa -passes="require<machine-register-class-info>,si-pre-allocate-wwm-regs,require<machine-register-class-info>" -debug-pass-manager -filetype=null %s 2>&1 | FileCheck %s
-
-# INFO: Test that MachineRegisterClassInfo is not preserved in WWM preallocation
-
-# CHECK: Running analysis: MachineRegisterClassAnalysis on test_wwm_reserved
-# CHECK: Running pass: SIPreAllocateWWMRegsPass on test_wwm_reserved
-# CHECK: Invalidating analysis: MachineRegisterClassAnalysis on test_wwm_reserved
-# CHECK: Running analysis: MachineRegisterClassAnalysis on test_wwm_reserved
-
----
-name: test_wwm_reserved
-tracksRegLiveness: true
-frameInfo:
- maxAlignment: 4
-stack:
- - { id: 0, type: spill-slot, size: 4, alignment: 4, stack-id: sgpr-spill }
-machineFunctionInfo:
- isEntryFunction: false
- scratchRSrcReg: '$sgpr0_sgpr1_sgpr2_sgpr3'
- stackPtrOffsetReg: '$sgpr32'
- frameOffsetReg: '$sgpr33'
- hasSpilledSGPRs: true
-body: |
- bb.0:
- liveins: $sgpr4, $vgpr2_vgpr3
- SI_SPILL_S32_SAVE killed $sgpr4, %stack.0, implicit $exec, implicit $sgpr0_sgpr1_sgpr2_sgpr3, implicit $sgpr32
- S_NOP 0
- renamable $sgpr4 = SI_SPILL_S32_RESTORE %stack.0, implicit $exec, implicit $sgpr0_sgpr1_sgpr2_sgpr3, implicit $sgpr32
- %0:vgpr_32 = V_MOV_B32_e32 20, implicit $exec
- GLOBAL_STORE_DWORD $vgpr2_vgpr3, %0:vgpr_32, 0, 0, implicit $exec
- SI_RETURN
-...
diff --git a/llvm/test/CodeGen/AMDGPU/si-pre-allocate-wwm-regs-preserve-rci.mir b/llvm/test/CodeGen/AMDGPU/si-pre-allocate-wwm-regs-preserve-rci.mir
new file mode 100644
index 0000000000000..4c990d3ec68ad
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/si-pre-allocate-wwm-regs-preserve-rci.mir
@@ -0,0 +1,34 @@
+# RUN: llc -mtriple=amdgpu9.0a-amd-amdhsa \
+# RUN: -passes="require<machine-register-class-info>,si-pre-allocate-wwm-regs,require<machine-register-class-info>" \
+# RUN: -debug-pass-manager -filetype=null %s 2>&1 | FileCheck %s
+# RUN: llc -mtriple=amdgpu9.0a-amd-amdhsa \
+# RUN: -passes=si-pre-allocate-wwm-regs -o - %s | FileCheck %s --check-prefix=MIR
+
+# Verify that SIPreAllocateWWMRegsPass updates the
+# MachineRegisterClassAnalysis result in place. The second require must reuse
+# the result computed by the first.
+
+# CHECK: Running pass: RequireAnalysisPass<llvm::MachineRegisterClassAnalysis, llvm::MachineFunction{{.*}}> on test_wwm_reserved
+# CHECK-NEXT: Running analysis: MachineRegisterClassAnalysis on test_wwm_reserved
+# CHECK-NEXT: Running pass: SIPreAllocateWWMRegsPass on test_wwm_reserved
+# CHECK: Running pass: RequireAnalysisPass<llvm::MachineRegisterClassAnalysis, llvm::MachineFunction{{.*}}> on test_wwm_reserved
+# CHECK-NEXT: Running pass: PrintMIRPreparePass on [module]
+
+# MIR: wwmReservedRegs:
+# MIR-NEXT: - '$vgpr0'
+
+---
+name: test_wwm_reserved
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $sgpr1
+ %0:vgpr_32 = IMPLICIT_DEF
+ renamable $sgpr4_sgpr5 = ENTER_STRICT_WWM -1, implicit-def $exec, implicit-def $scc, implicit $exec
+ %1:vgpr_32 = V_MOV_B32_e32 0, implicit $exec
+ %2:vgpr_32 = V_MOV_B32_dpp %1, %0, 323, 12, 15, 0, implicit $exec
+ $exec = EXIT_STRICT_WWM killed renamable $sgpr4_sgpr5
+ %3:vgpr_32 = COPY %0
+ %4:vreg_64 = IMPLICIT_DEF
+ S_NOP 0, implicit %4
+...
diff --git a/llvm/unittests/Target/AMDGPU/CMakeLists.txt b/llvm/unittests/Target/AMDGPU/CMakeLists.txt
index 39cced662567d..a2d0a0797b1fe 100644
--- a/llvm/unittests/Target/AMDGPU/CMakeLists.txt
+++ b/llvm/unittests/Target/AMDGPU/CMakeLists.txt
@@ -32,5 +32,6 @@ add_llvm_target_unittest(AMDGPUTests
InstSizes.cpp
LiveRegUnits.cpp
PALMetadata.cpp
+ RCIUpdateReservedRegsTest.cpp
UniformityAnalysisTest.cpp
)
diff --git a/llvm/unittests/Target/AMDGPU/RCIUpdateReservedRegsTest.cpp b/llvm/unittests/Target/AMDGPU/RCIUpdateReservedRegsTest.cpp
new file mode 100644
index 0000000000000..888f28c14ec2d
--- /dev/null
+++ b/llvm/unittests/Target/AMDGPU/RCIUpdateReservedRegsTest.cpp
@@ -0,0 +1,115 @@
+//===- RegisterClassInfoTest.cpp - RegisterClassInfo tests ----------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPUUnitTests.h"
+#include "MCTargetDesc/AMDGPUMCTargetDesc.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/RegisterClassInfo.h"
+#include "llvm/CodeGen/TargetRegisterInfo.h"
+#include "llvm/MC/MCRegisterInfo.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+class RegisterClassInfoTest : public AMDGPUCodeGenTestBase {
+public:
+ /// Use a GCN triple so VGPR registers are available to the test.
+ void SetUp() override { setUpImpl("amdgpu9.00-amd-amdhsa", "", /*FS=*/""); }
+};
+
+/// Force every RegisterClassInfo cache entry to be populated so
+/// updateReservedRegs exercises incremental compaction instead of lazy
+/// recomputation on the next query.
+static void materializeAll(RegisterClassInfo &RCI,
+ const TargetRegisterInfo &TRI) {
+ for (const TargetRegisterClass &RC : TRI.regclasses()) {
+ (void)RCI.getOrder(&RC);
+ (void)RCI.getNumAllocatableRegs(&RC);
+ (void)RCI.isProperSubClass(&RC);
+ (void)RCI.getMinCost(&RC);
+ (void)RCI.getLastCostChange(&RC);
+ }
+
+ for (unsigned I = 0; I != TRI.getNumRegPressureSets(); ++I)
+ (void)RCI.getRegPressureSetLimit(I);
+}
+
+/// Compare every cached RegisterClassInfo field against a freshly built object.
+static void expectEqual(RegisterClassInfo &Incremental,
+ RegisterClassInfo &Recomputed,
+ const TargetRegisterInfo &TRI) {
+ for (const TargetRegisterClass &RC : TRI.regclasses()) {
+ SCOPED_TRACE(TRI.getRegClassName(&RC));
+ EXPECT_EQ(Incremental.getOrder(&RC), Recomputed.getOrder(&RC));
+ EXPECT_EQ(Incremental.getNumAllocatableRegs(&RC),
+ Recomputed.getNumAllocatableRegs(&RC));
+ EXPECT_EQ(Incremental.isProperSubClass(&RC),
+ Recomputed.isProperSubClass(&RC));
+ EXPECT_EQ(Incremental.getMinCost(&RC), Recomputed.getMinCost(&RC));
+ EXPECT_EQ(Incremental.getLastCostChange(&RC),
+ Recomputed.getLastCostChange(&RC));
+ }
+
+ for (unsigned I = 0; I != TRI.getNumRegPressureSets(); ++I) {
+ SCOPED_TRACE(I);
+ EXPECT_EQ(Incremental.getRegPressureSetLimit(I),
+ Recomputed.getRegPressureSetLimit(I));
+ }
+}
+
+/// Verify that updateReservedRegs matches rebuilding RegisterClassInfo from
+/// scratch.
+TEST_F(RegisterClassInfoTest, IncrementalUpdateMatchesRecompute) {
+ ASSERT_TRUE(parseMIR(R"MIR(
+---
+name: func
+tracksRegLiveness: true
+machineFunctionInfo:
+ isEntryFunction: true
+body: |
+ bb.0:
+ S_ENDPGM 0
+...
+)MIR"));
+
+ MachineFunction &MF = getMF("func");
+ MachineRegisterInfo &MRI = MF.getRegInfo();
+ const TargetRegisterInfo &TRI = *MF.getSubtarget().getRegisterInfo();
+ MRI.freezeReservedRegs();
+
+ RegisterClassInfo Incremental;
+ Incremental.runOnMachineFunction(MF);
+ materializeAll(Incremental, TRI);
+
+ MCRegister VGPR0 = AMDGPU::VGPR0;
+ ASSERT_FALSE(MRI.isReserved(VGPR0));
+ ASSERT_TRUE(
+ llvm::any_of(TRI.regclasses(), [&](const TargetRegisterClass &RC) {
+ return llvm::is_contained(Incremental.getOrder(&RC),
+ static_cast<MCPhysReg>(VGPR0.id()));
+ }));
+
+ MRI.reserveReg(VGPR0, &TRI);
+ for (MCRegAliasIterator Alias(VGPR0, &TRI, /*IncludeSubRegs=*/true);
+ Alias.isValid(); ++Alias)
+ EXPECT_TRUE(MRI.isReserved(*Alias)) << TRI.getName(*Alias);
+
+ Incremental.updateReservedRegs(MRI.getReservedRegs());
+
+ // Construct an independent baseline from the updated reserved-register set.
+ RegisterClassInfo Recomputed;
+ Recomputed.runOnMachineFunction(MF);
+
+ expectEqual(Incremental, Recomputed, TRI);
+}
+
+} // namespace
More information about the llvm-commits
mailing list