[llvm] 09f8509 - [AMDGPU] Fix speculative register pressure queries (#208574)

[via llvm-commits]

Author: ilia-cher
Date: 2026-07-16T09:45:55+02:00
New Revision: 09f8509bec1fec4636c8e592dcd55a80f39f2667

URL: https://github.com/llvm/llvm-project/commit/09f8509bec1fec4636c8e592dcd55a80f39f2667
DIFF: https://github.com/llvm/llvm-project/commit/09f8509bec1fec4636c8e592dcd55a80f39f2667.diff

LOG: [AMDGPU] Fix speculative register pressure queries (#208574)

There are two issues with the way we currently speculate register
pressure:
1. GCNDownwardRPTracker::advance(with UseInternalIterator=false), which
is called by the scheduler in schedNode, resets LastTrackedMI, so the
tracker (bumpDownwardPressure) doesn't know where the last scheduled
instruction is and falls back to the beginning of the basic block. As a
result, when we estimate RP impact for a given MI, we tend to find uses
that are often already scheduled and should be skipped.
2. When looking for the remaining uses between LastTrackedMI and the
candidate MI we should skip already scheduled instruction.

Added: 
    

Modified: 
    llvm/lib/Target/AMDGPU/GCNRegPressure.cpp
    llvm/unittests/Target/AMDGPU/GCNRegPressureTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AMDGPU/GCNRegPressure.cpp b/llvm/lib/Target/AMDGPU/GCNRegPressure.cpp
index 42a430da0acb8..f4f03a7f1430b 100644
--- a/llvm/lib/Target/AMDGPU/GCNRegPressure.cpp
+++ b/llvm/lib/Target/AMDGPU/GCNRegPressure.cpp
@@ -767,8 +767,14 @@ bool GCNDownwardRPTracker::advance(MachineInstr *MI, bool UseInternalIterator) {
   advanceBeforeNext(MI, UseInternalIterator);
   advanceToNext(MI, UseInternalIterator);
   if (!UseInternalIterator) {
+    const MachineInstr *SavedLastTrackedMI = LastTrackedMI;
     // We must remove any dead def lanes from the current RP
     advanceBeforeNext(MI, true);
+    // Restore LastTrackedMI set by advanceToNext, otherwise
+    // speculative queries (bumpDownwardPressure) don't
+    // know the last scheduled instruction and fail to
+    // correctly estimate pressure change.
+    LastTrackedMI = SavedLastTrackedMI;
   }
   return true;
 }
@@ -821,6 +827,18 @@ GCNDownwardRPTracker::bumpDownwardPressure(const MachineInstr *MI,
   SlotIndex SlotIdx;
   SlotIdx = LIS.getInstructionIndex(*MI).getRegSlot();
 
+  SlotIndex CurrIdx;
+  const MachineBasicBlock *MBB = MI->getParent();
+  MachineBasicBlock::const_iterator StartPos =
+      LastTrackedMI ? std::next(LastTrackedMI->getIterator()) : MBB->begin();
+  MachineBasicBlock::const_iterator IdxPos =
+      skipDebugInstructionsForward(StartPos, MBB->end());
+  if (IdxPos == MBB->end()) {
+    CurrIdx = LIS.getMBBEndIdx(MBB);
+  } else {
+    CurrIdx = LIS.getInstructionIndex(*IdxPos).getRegSlot();
+  }
+
   // Account for register pressure similar to RegPressureTracker::recede().
   RegisterOperands RegOpers;
   RegOpers.collect(*MI, *TRI, *MRI, true, /*IgnoreDead=*/false);
@@ -839,16 +857,6 @@ GCNDownwardRPTracker::bumpDownwardPressure(const MachineInstr *MI,
     // last uses for the current position.
     // FIXME: allow the caller to pass in the list of vreg uses that remain
     // to be bottom-scheduled to avoid searching uses at each query.
-    SlotIndex CurrIdx;
-    const MachineBasicBlock *MBB = MI->getParent();
-    MachineBasicBlock::const_iterator IdxPos = skipDebugInstructionsForward(
-        LastTrackedMI ? LastTrackedMI : MBB->begin(), MBB->end());
-    if (IdxPos == MBB->end()) {
-      CurrIdx = LIS.getMBBEndIdx(MBB);
-    } else {
-      CurrIdx = LIS.getInstructionIndex(*IdxPos).getRegSlot();
-    }
-
     LastUseMask =
         findUseBetween(Reg, LastUseMask, CurrIdx, SlotIdx, *MRI, TRI, &LIS);
     if (LastUseMask.none())

diff  --git a/llvm/unittests/Target/AMDGPU/GCNRegPressureTest.cpp b/llvm/unittests/Target/AMDGPU/GCNRegPressureTest.cpp
index c0acf9e7dde4f..3fe7b6f261015 100644
--- a/llvm/unittests/Target/AMDGPU/GCNRegPressureTest.cpp
+++ b/llvm/unittests/Target/AMDGPU/GCNRegPressureTest.cpp
@@ -8,6 +8,8 @@
 
 #include "GCNRegPressure.h"
 #include "AMDGPUUnitTests.h"
+#include "GCNSubtarget.h"
+#include "SIRegisterInfo.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/MIRParser/MIRParser.h"
 #include "llvm/CodeGen/MachineFunctionAnalysis.h"
@@ -145,3 +147,56 @@ body:             |
   EXPECT_EQ(RPTracker.moveMaxPressure().getVGPRNum(false), 1U);
   EXPECT_EQ(RPTrackerNoLiveIns.moveMaxPressure().getVGPRNum(false), 1U);
 }
+
+// Tests the correct handling of multiple uses of the same virtual register
+// in bumpDownwardPressure (speculative estimate of register pressure).
+TEST_F(GCNRegPressureTest, BumpDownwardPressureLastUseAfterCommit) {
+  StringRef MIR = R"(
+name:            BumpDownwardPressureLastUseAfterCommit
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    %0:vgpr_32 = IMPLICIT_DEF
+    %1:vreg_256_align2 = IMPLICIT_DEF
+    S_NOP 0, implicit %1
+    S_NOP 0, implicit %1
+    S_NOP 0, implicit %0
+    S_ENDPGM 0
+...
+)";
+  ASSERT_TRUE(parseMIR(MIR));
+  MachineFunction &MF = getMF("BumpDownwardPressureLastUseAfterCommit");
+  const LiveIntervals &LIS = MFAM.getResult<LiveIntervalsAnalysis>(MF);
+  const MachineRegisterInfo &MRI = MF.getRegInfo();
+  const SIRegisterInfo *TRI = MF.getSubtarget<GCNSubtarget>().getRegisterInfo();
+
+  MachineBasicBlock &MBB = *MF.getBlockNumbered(0);
+
+  SmallVector<MachineInstr *, 8> Instrs;
+  for (MachineInstr &MI : MBB)
+    Instrs.push_back(&MI);
+  // 0: def %0, 1: def %1, 2: U1 (use %1), 3: U2 (last use %1),
+  // 4: use %0, 5: S_ENDPGM
+  MachineInstr *DefV0 = Instrs[0];
+  MachineInstr *DefV1 = Instrs[1];
+  MachineInstr *U1 = Instrs[2];
+  MachineInstr *U2 = Instrs[3];
+
+  GCNDownwardRPTracker RPTracker(LIS);
+  GCNRPTracker::LiveRegSet Empty;
+  RPTracker.reset(MRI, Empty);
+
+  // Commit the defs and the first use of %1 via the externally-managed
+  // iterator (same as while scheduling).
+  RPTracker.advance(DefV0, /*UseInternalIterator=*/false);
+  RPTracker.advance(DefV1, /*UseInternalIterator=*/false);
+  RPTracker.advance(U1, /*UseInternalIterator=*/false);
+
+  // After committing U1, both %0 (1 VGPR) and %1 (vreg_256 = 8 VGPRs) are live.
+  EXPECT_EQ(RPTracker.getPressure().getArchVGPRNum(), 9U);
+
+  // Speculate the last use of %1. %1 must die here, dropping its 8 VGPRs and
+  // leaving only %0 live.
+  GCNRegPressure P = RPTracker.bumpDownwardPressure(U2, TRI);
+  EXPECT_EQ(P.getArchVGPRNum(), 1U);
+}