[AMDGPU][Scheduler] Add GCNRegPressure-based methods to GCNRPTarget - #182853
Merged
Merged
Conversation
This adds a few methods to `GCNRPTarget` that can estimate/perform RP savings based on `GCNRegPressure` instead of a single `Register`, opening the door to model/incorporate more complex savings made up of multiple registers of potentially different classes. The scheduler's rematerialization stage now uses this new API. Although there are no test changes this is not really NFC since register pressure savings in the rematerialization stage are now computed through `GCNRegPressure` instead of the stage itself. If anything this makes them more consistent with the rest of the RP-tracking infrastructure.
Member
|
@llvm/pr-subscribers-backend-amdgpu Author: Lucas Ramirez (lucas-rami) ChangesThis adds a few methods to Although there are no test changes this is not really NFC since register pressure savings in the rematerialization stage are now computed through Full diff: https://github.com/llvm/llvm-project/pull/182853.diff 4 Files Affected:
diff --git a/llvm/lib/Target/AMDGPU/GCNRegPressure.cpp b/llvm/lib/Target/AMDGPU/GCNRegPressure.cpp
index 89307ef9767b7..500c00dec22cf 100644
--- a/llvm/lib/Target/AMDGPU/GCNRegPressure.cpp
+++ b/llvm/lib/Target/AMDGPU/GCNRegPressure.cpp
@@ -424,10 +424,34 @@ bool GCNRPTarget::isSaveBeneficial(Register Reg) const {
return (UnifiedRF && Excess.VGPR) || Excess.ArchVGPR;
}
-bool GCNRPTarget::satisfied() const {
- if (RP.getSGPRNum() > MaxSGPRs || RP.getVGPRNum(false) > MaxVGPRs)
+unsigned GCNRPTarget::getNumRegsBenefit(const GCNRegPressure &SaveRP) const {
+ RegExcess Excess(MF, RP, *this);
+ unsigned NumRegsSaved = 0;
+ unsigned NumVGPRAboveAddrLimit = 0;
+
+ if (Excess.SGPR)
+ NumRegsSaved += std::min(Excess.SGPR, SaveRP.getSGPRNum());
+ if (Excess.ArchVGPR)
+ NumVGPRAboveAddrLimit += std::min(Excess.ArchVGPR, SaveRP.getArchVGPRNum());
+ if (Excess.AGPR)
+ NumVGPRAboveAddrLimit += std::min(Excess.AGPR, SaveRP.getAGPRNum());
+ NumRegsSaved += NumVGPRAboveAddrLimit;
+
+ if (UnifiedRF && Excess.VGPR) {
+ // Do not double-count VGPRs that are both above the addressable limit in
+ // their respective class and contribute to an overall excess in VGPR.
+ const unsigned VGPRSave = SaveRP.getVGPRNum(true);
+ if (NumVGPRAboveAddrLimit < VGPRSave)
+ NumRegsSaved += std::min(Excess.VGPR, VGPRSave - NumVGPRAboveAddrLimit);
+ }
+
+ return NumRegsSaved;
+}
+
+bool GCNRPTarget::satisfied(const GCNRegPressure &TestRP) const {
+ if (TestRP.getSGPRNum() > MaxSGPRs || TestRP.getVGPRNum(false) > MaxVGPRs)
return false;
- if (UnifiedRF && RP.getVGPRNum(true) > MaxUnifiedVGPRs)
+ if (UnifiedRF && TestRP.getVGPRNum(true) > MaxUnifiedVGPRs)
return false;
return true;
}
diff --git a/llvm/lib/Target/AMDGPU/GCNRegPressure.h b/llvm/lib/Target/AMDGPU/GCNRegPressure.h
index c55796c37f287..80121d6eb7def 100644
--- a/llvm/lib/Target/AMDGPU/GCNRegPressure.h
+++ b/llvm/lib/Target/AMDGPU/GCNRegPressure.h
@@ -255,8 +255,17 @@ class GCNRPTarget {
RP.inc(Reg, Mask, LaneBitmask::getNone(), MRI);
}
+ /// Returns the benefit towards achieving the RP target that saving \p SaveRP
+ /// represents, in total number of registers saved across all classes.
+ unsigned getNumRegsBenefit(const GCNRegPressure &SaveRP) const;
+
+ /// Saves a total pressure of \p SaveRP.
+ void saveRP(const GCNRegPressure &SaveRP) { RP -= SaveRP; }
+
+ /// Whether \p TestRP is at or below the defined pressure target.
+ bool satisfied(const GCNRegPressure &TestRP) const;
/// Whether the current RP is at or below the defined pressure target.
- bool satisfied() const;
+ bool satisfied() const { return satisfied(RP); }
bool hasVectorRegisterExcess() const;
unsigned getMaxSGPRs() const { return MaxSGPRs; }
diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
index b0441955e70b9..62969b89f8a8f 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
@@ -1566,15 +1566,36 @@ bool PreRARematStage::initGCNSchedStage() {
break;
REMAT_DEBUG(dbgs() << "** REMAT " << PrintRematReg(Remat) << '\n';);
+ MachineInstr *RematMI =
+ Candidate.rematerialize(RecomputeRP, RPTargets, DAG);
+ RescheduleRegions |= Remat.Live;
+
// Every rematerialization we do here is likely to move the instruction
// into a higher frequency region, increasing the total sum latency of the
// instruction itself. This is acceptable if we are eliminating a spill in
// the process, but when the goal is increasing occupancy we get nothing
// out of rematerialization if occupancy is not increased in the end; in
// such cases we want to roll back the rematerialization.
- RollbackInfo *Rollback =
- TargetOcc ? &Rollbacks.emplace_back(&Remat) : nullptr;
- rematerialize(Remat, RecomputeRP, Rollback);
+ if (TargetOcc) {
+ RollbackInfo &Rollback = Rollbacks.emplace_back(&Remat);
+ Rollback.RematMI = RematMI;
+ // Make the original MI a debug value so that it does not influence
+ // scheduling and replace all read registers with a sentinel register to
+ // prevent operands to appear in use-lists of other MIs during LIS
+ // updates. Store mappings between operand indices and original
+ // registers for potential rollback.
+ Remat.DefMI->setDesc(DAG.TII->get(TargetOpcode::DBG_VALUE));
+ for (auto [Idx, MO] : enumerate(Remat.DefMI->operands())) {
+ if (MO.isReg() && MO.readsReg()) {
+ Rollback.RegMap.insert({Idx, MO.getReg()});
+ MO.setReg(Register());
+ }
+ }
+ } else {
+ // Just delete the original instruction if it cannot be rolled back.
+ DAG.deleteMI(Remat.DefRegion, Remat.DefMI);
+ }
+
unsetSatisifedRPTargets(Remat.Live);
}
@@ -2885,21 +2906,8 @@ PreRARematStage::ScoredRemat::FreqInfo::FreqInfo(
PreRARematStage::ScoredRemat::ScoredRemat(RematReg *Remat, const FreqInfo &Freq,
const GCNScheduleDAGMILive &DAG)
- : Remat(Remat), NumRegs(getNumRegs(DAG)), FreqDiff(getFreqDiff(Freq)) {}
-
-unsigned PreRARematStage::ScoredRemat::getNumRegs(
- const GCNScheduleDAGMILive &DAG) const {
- const TargetRegisterClass &RC = *DAG.MRI.getRegClass(Remat->getReg());
- unsigned RegSize = DAG.TRI->getRegSizeInBits(RC);
- if (unsigned SubIdx = Remat->DefMI->getOperand(0).getSubReg()) {
- // The following may return -1 (i.e., a large unsigned number) on indices
- // that may be used to access subregisters of multiple sizes; in such cases
- // fallback on the size derived from the register class.
- unsigned SubRegSize = DAG.TRI->getSubRegIdxSize(SubIdx);
- if (SubRegSize < RegSize)
- RegSize = SubRegSize;
- }
- return divideCeil(RegSize, 32);
+ : Remat(Remat), FreqDiff(getFreqDiff(Freq)) {
+ RPSave.inc(Remat->getReg(), LaneBitmask::getNone(), Remat->Mask, DAG.MRI);
}
int64_t PreRARematStage::ScoredRemat::getFreqDiff(const FreqInfo &Freq) const {
@@ -2924,12 +2932,21 @@ void PreRARematStage::ScoredRemat::update(const BitVector &TargetRegions,
MaxFreq = 0;
RegionImpact = 0;
for (unsigned I : TargetRegions.set_bits()) {
- if (!Remat->Live[I] || !RPTargets[I].isSaveBeneficial(Remat->getReg()))
+ if (!Remat->Live[I])
continue;
+
+ // The rematerialization must contribute positively in at least one
+ // register class with usage above the RP target for this region to
+ // contribute to the score.
+ const GCNRPTarget &RegionTarget = RPTargets[I];
+ const unsigned NumRegsBenefit = RegionTarget.getNumRegsBenefit(RPSave);
+ if (!NumRegsBenefit)
+ continue;
+
bool UnusedLT = Remat->isUnusedLiveThrough(I);
// Regions in which RP is guaranteed to decrease have more weight.
- RegionImpact += UnusedLT ? 2 : 1;
+ RegionImpact += (UnusedLT ? 2 : 1) * NumRegsBenefit;
if (ReduceSpill) {
uint64_t Freq = FreqInfo.Regions[I];
@@ -2941,41 +2958,22 @@ void PreRARematStage::ScoredRemat::update(const BitVector &TargetRegions,
MaxFreq = std::max(MaxFreq, Freq);
}
}
- RegionImpact *= NumRegs;
}
-void PreRARematStage::rematerialize(const RematReg &Remat,
- BitVector &RecomputeRP,
- RollbackInfo *Rollback) {
- const SIInstrInfo *TII = MF.getSubtarget<GCNSubtarget>().getInstrInfo();
- MachineInstr &DefMI = *Remat.DefMI;
+MachineInstr *PreRARematStage::ScoredRemat::rematerialize(
+ BitVector &RecomputeRP, SmallVectorImpl<GCNRPTarget> &RPTargets,
+ GCNScheduleDAGMILive &DAG) const {
+ const SIInstrInfo *TII = DAG.MF.getSubtarget<GCNSubtarget>().getInstrInfo();
+ MachineInstr &DefMI = *Remat->DefMI;
Register Reg = DefMI.getOperand(0).getReg();
Register NewReg = DAG.MRI.cloneVirtualRegister(Reg);
// Rematerialize the register in the region where it is used.
- MachineBasicBlock::iterator InsertPos = Remat.UseMI;
+ MachineBasicBlock::iterator InsertPos = Remat->UseMI;
TII->reMaterialize(*InsertPos->getParent(), InsertPos, NewReg, 0, DefMI);
MachineInstr *RematMI = &*std::prev(InsertPos);
- Remat.UseMI->substituteRegister(Reg, NewReg, 0, *DAG.TRI);
- Remat.insertMI(Remat.UseRegion, RematMI, DAG);
- if (Rollback) {
- Rollback->RematMI = RematMI;
- // Make the original MI a debug value so that it does not influence
- // scheduling and replace all read registers with a sentinel register to
- // prevent operands to appear in use-lists of other MIs during LIS
- // updates. Store mappings between operand indices and original registers
- // for potential rollback.
- DefMI.setDesc(TII->get(TargetOpcode::DBG_VALUE));
- for (auto [Idx, MO] : enumerate(Remat.DefMI->operands())) {
- if (MO.isReg() && MO.readsReg()) {
- Rollback->RegMap.insert({Idx, MO.getReg()});
- MO.setReg(Register());
- }
- }
- } else {
- // Just delete the original instruction if it cannot be rolled back.
- DAG.deleteMI(Remat.DefRegion, &DefMI);
- }
+ Remat->UseMI->substituteRegister(Reg, NewReg, 0, *DAG.TRI);
+ Remat->insertMI(Remat->UseRegion, RematMI, DAG);
#ifdef EXPENSIVE_CHECKS
// All uses are known to be available / live at the remat point. Thus,
@@ -3010,15 +3008,15 @@ void PreRARematStage::rematerialize(const RematReg &Remat,
// and adjust RP targets. The save is guaranteed in regions in which the
// register is live-through and unused but optimistic in all other regions
// where the register is live.
- for (unsigned I : Remat.Live.set_bits()) {
- RPTargets[I].saveReg(Reg, Remat.Mask, DAG.MRI);
+ for (unsigned I : Remat->Live.set_bits()) {
+ RPTargets[I].saveRP(RPSave);
DAG.LiveIns[I].erase(Reg);
DAG.RegionLiveOuts.getLiveRegsForRegionIdx(I).erase(Reg);
- if (!Remat.isUnusedLiveThrough(I))
+ if (!Remat->isUnusedLiveThrough(I))
RecomputeRP.set(I);
}
- RescheduleRegions |= Remat.Live;
+ return RematMI;
}
void PreRARematStage::commitRematerializations() const {
diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
index 6b6a40365b52d..eb283c6b37805 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
@@ -596,6 +596,14 @@ class PreRARematStage : public GCNSchedStage {
ScoredRemat(RematReg *Remat, const FreqInfo &Freq,
const GCNScheduleDAGMILive &DAG);
+ /// Rematerializes the candidate and returns the new MI. This removes the
+ /// rematerialized register from live-in/out lists in the \p DAG and updates
+ /// \p RPTargets in all affected regions. Regions in which RP savings are
+ /// not guaranteed are set in \p RecomputeRP.
+ MachineInstr *rematerialize(BitVector &RecomputeRP,
+ SmallVectorImpl<GCNRPTarget> &RPTargets,
+ GCNScheduleDAGMILive &DAG) const;
+
/// Updates the rematerialization's score w.r.t. the current \p RPTargets.
/// \p RegionFreq indicates the frequency of each region
void update(const BitVector &TargetRegions, ArrayRef<GCNRPTarget> RPTargets,
@@ -631,8 +639,9 @@ class PreRARematStage : public GCNSchedStage {
#endif
private:
- /// Number of 32-bit registers this rematerialization covers.
- unsigned NumRegs;
+ /// Expected register pressure decrease induced by rematerializing this
+ /// candidate.
+ GCNRegPressure RPSave;
// The three members below are the scoring components, top to bottom from
// most important to least important when comparing candidates.
@@ -649,8 +658,6 @@ class PreRARematStage : public GCNSchedStage {
/// scaled by the size of the register being rematerialized.
unsigned RegionImpact;
- unsigned getNumRegs(const GCNScheduleDAGMILive &DAG) const;
-
int64_t getFreqDiff(const FreqInfo &Freq) const;
};
@@ -729,15 +736,6 @@ class PreRARematStage : public GCNSchedStage {
/// rematerializable register was found.
bool collectRematRegs(const DenseMap<MachineInstr *, unsigned> &MIRegion);
- /// Rematerializes \p Remat. This removes the rematerialized register from
- /// live-in/out lists in the DAG and updates RP targets in all affected
- /// regions, which are also marked in \ref RescheduleRegions. Regions in which
- /// RP savings are not guaranteed are set in \p RecomputeRP. When \p Rollback
- /// is non-null, fills it with required information to be able to rollback the
- /// rematerialization post-rescheduling.
- void rematerialize(const RematReg &Remat, BitVector &RecomputeRP,
- RollbackInfo *Rollback);
-
/// Deletes all rematerialized MIs from the MIR when they were kept around for
/// potential rollback.
void commitRematerializations() const;
|
qcolombet
approved these changes
Feb 27, 2026
qcolombet
left a comment
qcolombet
left a comment
Contributor
There was a problem hiding this comment.
You're the expert here and the refactoring looks fine at the high level, so approving.
There are a few things that I highlighted that made me pause.
| // Just delete the original instruction if it cannot be rolled back. | ||
| DAG.deleteMI(Remat.DefRegion, Remat.DefMI); | ||
| } | ||
|
|
||
| unsetSatisifedRPTargets(Remat.Live); |
Contributor
There was a problem hiding this comment.
While here, maybe we can fix the typo in the here :)
Satisified => Satisfied.
Contributor
Author
There was a problem hiding this comment.
Missed that, thanks for the catch :)
| NumVGPRAboveAddrLimit += std::min(Excess.AGPR, SaveRP.getAGPRNum()); | ||
| NumRegsSaved += NumVGPRAboveAddrLimit; | ||
|
|
||
| if (UnifiedRF && Excess.VGPR) { |
Contributor
There was a problem hiding this comment.
I think the special case for UnifiedRF would benefit a quick reminder of why we're doing this.
I feel the code work because you know the implementation of getVGPRNum(true), but for me who's not familiar with all that, this just looks like black magic.
Contributor
Author
There was a problem hiding this comment.
I tried to clarify a little bit the intent there. I also swapped SaveRPgetVGPRNum(true) for SaveRP.getArchVGPRNum() + SaveRP.getAGPRNum(), which in practice is almost identical. The former just introduces an allocation granule in the calculation which is relevant for actual occupancy estimates (i.e., the hardware can only allocate ArchVGPR in increments of this granule) but not in this saving calculation.
| unsigned getNumRegsBenefit(const GCNRegPressure &SaveRP) const; | ||
|
|
||
| /// Saves a total pressure of \p SaveRP. | ||
| void saveRP(const GCNRegPressure &SaveRP) { RP -= SaveRP; } |
Contributor
There was a problem hiding this comment.
Should we assert that RP is >= than SaveRP?
jmmartinez
reviewed
Feb 27, 2026
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/187/builds/17282 Here is the relevant piece of the build log for the reference |
lucas-rami
added a commit
that referenced
this pull request
Feb 27, 2026
Bug introduced by #182853 (`Remat` is now a pointer).
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/157/builds/44873 Here is the relevant piece of the build log for the reference |
sujianIBM
pushed a commit
to sujianIBM/llvm-project
that referenced
this pull request
Mar 5, 2026
…t` (llvm#182853) This adds a few methods to `GCNRPTarget` that can estimate/perform RP savings based on `GCNRegPressure` instead of a single `Register`, opening the door to model/incorporate more complex savings made up of multiple registers of potentially different classes. The scheduler's rematerialization stage now uses this new API. Although there are no test changes this is not really NFC since register pressure savings in the rematerialization stage are now computed through `GCNRegPressure` instead of the stage itself. If anything this makes them more consistent with the rest of the RP-tracking infrastructure.
sujianIBM
pushed a commit
to sujianIBM/llvm-project
that referenced
this pull request
Mar 5, 2026
…3745) Bug introduced by llvm#182853 (`Remat` is now a pointer).
This was referenced Apr 15, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This adds a few methods to
GCNRPTargetthat can estimate/perform RP savings based onGCNRegPressureinstead of a singleRegister, opening the door to model/incorporate more complex savings made up of multiple registers of potentially different classes. The scheduler's rematerialization stage now uses this new API.Although there are no test changes this is not really NFC since register pressure savings in the rematerialization stage are now computed through
GCNRegPressureinstead of the stage itself. If anything this makes them more consistent with the rest of the RP-tracking infrastructure.