[llvm] 657adcb - [NFC][regalloc] Move ExtraRegInfo and related to LiveRangeStageManager
Mircea Trofin via llvm-commits
llvm-commits at lists.llvm.org
Mon Dec 13 10:31:45 PST 2021
Author: Mircea Trofin
Date: 2021-12-13T10:10:57-08:00
New Revision: 657adcb0779d2d119737f2ca557a87456024a5ac
URL: https://github.com/llvm/llvm-project/commit/657adcb0779d2d119737f2ca557a87456024a5ac
DIFF: https://github.com/llvm/llvm-project/commit/657adcb0779d2d119737f2ca557a87456024a5ac.diff
LOG: [NFC][regalloc] Move ExtraRegInfo and related to LiveRangeStageManager
This would allow sharing the LiveRangeStageManager between different
RegAllocEvictionAdvisors. One scenario is for ML training, where we want
to capture what the default advisor would do, for bootstrapping (speeds
up training).
Differential Revision: https://reviews.llvm.org/D114831
Added:
Modified:
llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
llvm/lib/CodeGen/RegAllocGreedy.cpp
Removed:
################################################################################
diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
index 85fd3207888b..0e3c24a50cc0 100644
--- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
+++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
@@ -85,6 +85,85 @@ struct EvictionCost {
std::tie(O.BrokenHints, O.MaxWeight);
}
};
+
+/// Track allocation stage and eviction loop prevention during allocation.
+// TODO(mtrofin): Consider exposing RAGreedy in a header instead, and folding
+// this back into it.
+class ExtraRegInfo final {
+ // RegInfo - Keep additional information about each live range.
+ struct RegInfo {
+ LiveRangeStage Stage = RS_New;
+
+ // Cascade - Eviction loop prevention. See
+ // canEvictInterferenceBasedOnCost().
+ unsigned Cascade = 0;
+
+ RegInfo() = default;
+ };
+
+ IndexedMap<RegInfo, VirtReg2IndexFunctor> Info;
+ unsigned NextCascade = 1;
+
+public:
+ ExtraRegInfo() = default;
+ ExtraRegInfo(const ExtraRegInfo &) = delete;
+
+ LiveRangeStage getStage(Register Reg) const { return Info[Reg].Stage; }
+
+ LiveRangeStage getStage(const LiveInterval &VirtReg) const {
+ return getStage(VirtReg.reg());
+ }
+
+ void setStage(Register Reg, LiveRangeStage Stage) {
+ Info.grow(Reg.id());
+ Info[Reg].Stage = Stage;
+ }
+
+ void setStage(const LiveInterval &VirtReg, LiveRangeStage Stage) {
+ setStage(VirtReg.reg(), Stage);
+ }
+
+ /// Return the current stage of the register, if present, otherwise initialize
+ /// it and return that.
+ LiveRangeStage getOrInitStage(Register Reg) {
+ Info.grow(Reg.id());
+ return getStage(Reg);
+ }
+
+ unsigned getCascade(Register Reg) const { return Info[Reg].Cascade; }
+
+ void setCascade(Register Reg, unsigned Cascade) {
+ Info.grow(Reg.id());
+ Info[Reg].Cascade = Cascade;
+ }
+
+ unsigned getOrAssignNewCascade(Register Reg) {
+ unsigned Cascade = getCascade(Reg);
+ if (!Cascade) {
+ Cascade = NextCascade++;
+ setCascade(Reg, Cascade);
+ }
+ return Cascade;
+ }
+
+ unsigned getCascadeOrCurrentNext(Register Reg) const {
+ unsigned Cascade = getCascade(Reg);
+ if (!Cascade)
+ Cascade = NextCascade;
+ return Cascade;
+ }
+
+ template <typename Iterator>
+ void setStage(Iterator Begin, Iterator End, LiveRangeStage NewStage) {
+ for (; Begin != End; ++Begin) {
+ Register Reg = *Begin;
+ Info.grow(Reg.id());
+ if (Info[Reg].Stage == RS_New)
+ Info[Reg].Stage = NewStage;
+ }
+ }
+ void LRE_DidCloneVirtReg(Register New, Register Old);
+};
} // namespace llvm
#endif // LLVM_CODEGEN_REGALLOCEVICTIONADVISOR_H
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp
index 9eab8746982a..796139c932c2 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.cpp
+++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp
@@ -172,8 +172,8 @@ class RAGreedy : public MachineFunctionPass,
// state
std::unique_ptr<Spiller> SpillerInstance;
PQueue Queue;
- unsigned NextCascade;
std::unique_ptr<VirtRegAuxInfo> VRAI;
+ Optional<ExtraRegInfo> ExtraInfo;
// Enum CutOffStage to keep a track whether the register allocation failed
// because of the cutoffs encountered in last chance recoloring.
@@ -195,76 +195,6 @@ class RAGreedy : public MachineFunctionPass,
static const char *const StageName[];
#endif
- // RegInfo - Keep additional information about each live range.
- struct RegInfo {
- LiveRangeStage Stage = RS_New;
-
- // Cascade - Eviction loop prevention. See
- // canEvictInterferenceBasedOnCost().
- unsigned Cascade = 0;
-
- RegInfo() = default;
- };
-
- IndexedMap<RegInfo, VirtReg2IndexFunctor> ExtraRegInfo;
-
- LiveRangeStage getStage(Register Reg) const {
- return ExtraRegInfo[Reg].Stage;
- }
-
- LiveRangeStage getStage(const LiveInterval &VirtReg) const {
- return getStage(VirtReg.reg());
- }
-
- void setStage(Register Reg, LiveRangeStage Stage) {
- ExtraRegInfo.resize(MRI->getNumVirtRegs());
- ExtraRegInfo[Reg].Stage = Stage;
- }
-
- void setStage(const LiveInterval &VirtReg, LiveRangeStage Stage) {
- setStage(VirtReg.reg(), Stage);
- }
-
- /// Return the current stage of the register, if present, otherwise initialize
- /// it and return that.
- LiveRangeStage getOrInitStage(Register Reg) {
- ExtraRegInfo.grow(Reg);
- return getStage(Reg);
- }
-
- unsigned getCascade(Register Reg) const { return ExtraRegInfo[Reg].Cascade; }
-
- void setCascade(Register Reg, unsigned Cascade) {
- ExtraRegInfo.resize(MRI->getNumVirtRegs());
- ExtraRegInfo[Reg].Cascade = Cascade;
- }
-
- unsigned getOrAssignNewCascade(Register Reg) {
- unsigned Cascade = getCascade(Reg);
- if (!Cascade) {
- Cascade = NextCascade++;
- setCascade(Reg, Cascade);
- }
- return Cascade;
- }
-
- unsigned getCascadeOrCurrentNext(Register Reg) const {
- unsigned Cascade = getCascade(Reg);
- if (!Cascade)
- Cascade = NextCascade;
- return Cascade;
- }
-
- template<typename Iterator>
- void setStage(Iterator Begin, Iterator End, LiveRangeStage NewStage) {
- ExtraRegInfo.resize(MRI->getNumVirtRegs());
- for (;Begin != End; ++Begin) {
- Register Reg = *Begin;
- if (ExtraRegInfo[Reg].Stage == RS_New)
- ExtraRegInfo[Reg].Stage = NewStage;
- }
- }
-
/// EvictionTrack - Keeps track of past evictions in order to optimize region
/// split decision.
class EvictionTrack {
@@ -696,22 +626,25 @@ void RAGreedy::LRE_WillShrinkVirtReg(Register VirtReg) {
}
void RAGreedy::LRE_DidCloneVirtReg(Register New, Register Old) {
+ ExtraInfo->LRE_DidCloneVirtReg(New, Old);
+}
+
+void ExtraRegInfo::LRE_DidCloneVirtReg(Register New, Register Old) {
// Cloning a register we haven't even heard about yet? Just ignore it.
- if (!ExtraRegInfo.inBounds(Old))
+ if (!Info.inBounds(Old))
return;
// LRE may clone a virtual register because dead code elimination causes it to
// be split into connected components. The new components are much smaller
// than the original, so they should get a new chance at being assigned.
// same stage as the parent.
- ExtraRegInfo[Old].Stage = RS_Assign;
- ExtraRegInfo.grow(New);
- ExtraRegInfo[New] = ExtraRegInfo[Old];
+ Info[Old].Stage = RS_Assign;
+ Info.grow(New.id());
+ Info[New] = Info[Old];
}
void RAGreedy::releaseMemory() {
SpillerInstance.reset();
- ExtraRegInfo.clear();
GlobalCand.clear();
}
@@ -725,10 +658,10 @@ void RAGreedy::enqueue(PQueue &CurQueue, LiveInterval *LI) {
assert(Reg.isVirtual() && "Can only enqueue virtual registers");
unsigned Prio;
- auto Stage = getOrInitStage(Reg);
+ auto Stage = ExtraInfo->getOrInitStage(Reg);
if (Stage == RS_New) {
Stage = RS_Assign;
- setStage(Reg, Stage);
+ ExtraInfo->setStage(Reg, Stage);
}
if (Stage == RS_Split) {
// Unsplit ranges that couldn't be allocated immediately are deferred until
@@ -891,7 +824,7 @@ Register RAGreedy::canReassign(LiveInterval &VirtReg, Register PrevReg) const {
/// @param BreaksHint True when B is already assigned to its preferred register.
bool RAGreedy::shouldEvict(LiveInterval &A, bool IsHint,
LiveInterval &B, bool BreaksHint) const {
- bool CanSplit = getStage(B) < RS_Spill;
+ bool CanSplit = ExtraInfo->getStage(B) < RS_Spill;
// Be fairly aggressive about following hints as long as the evictee can be
// split.
@@ -941,9 +874,7 @@ bool RAGreedy::canEvictInterferenceBasedOnCost(
//
// This works out so a register without a cascade number is allowed to evict
// anything, and it can be evicted by anything.
- unsigned Cascade = ExtraRegInfo[VirtReg.reg()].Cascade;
- if (!Cascade)
- Cascade = NextCascade;
+ unsigned Cascade = ExtraInfo->getCascadeOrCurrentNext(VirtReg.reg());
EvictionCost Cost;
for (MCRegUnitIterator Units(PhysReg, TRI); Units.isValid(); ++Units) {
@@ -965,7 +896,7 @@ bool RAGreedy::canEvictInterferenceBasedOnCost(
return false;
// Never evict spill products. They cannot split or spill.
- if (getStage(*Intf) == RS_Done)
+ if (ExtraInfo->getStage(*Intf) == RS_Done)
return false;
// Once a live range becomes small enough, it is urgent that we find a
// register for it. This is indicated by an infinite spill weight. These
@@ -980,7 +911,7 @@ bool RAGreedy::canEvictInterferenceBasedOnCost(
RegClassInfo.getNumAllocatableRegs(
MRI->getRegClass(Intf->reg())));
// Only evict older cascades or live ranges without a cascade.
- unsigned IntfCascade = ExtraRegInfo[Intf->reg()].Cascade;
+ unsigned IntfCascade = ExtraInfo->getCascade(Intf->reg());
if (Cascade <= IntfCascade) {
if (!Urgent)
return false;
@@ -1043,7 +974,7 @@ bool RAGreedy::canEvictInterferenceInRange(const LiveInterval &VirtReg,
if (!Register::isVirtualRegister(Intf->reg()))
return false;
// Never evict spill products. They cannot split or spill.
- if (getStage(*Intf) == RS_Done)
+ if (ExtraInfo->getStage(*Intf) == RS_Done)
return false;
// Would this break a satisfied hint?
@@ -1106,7 +1037,7 @@ void RAGreedy::evictInterference(LiveInterval &VirtReg, MCRegister PhysReg,
// Make sure that VirtReg has a cascade number, and assign that cascade
// number to every evicted register. These live ranges than then only be
// evicted by a newer cascade, preventing infinite loops.
- unsigned Cascade = getOrAssignNewCascade(VirtReg.reg());
+ unsigned Cascade = ExtraInfo->getOrAssignNewCascade(VirtReg.reg());
LLVM_DEBUG(dbgs() << "evicting " << printReg(PhysReg, TRI)
<< " interference: Cascade " << Cascade << '\n');
@@ -1132,10 +1063,10 @@ void RAGreedy::evictInterference(LiveInterval &VirtReg, MCRegister PhysReg,
LastEvicted.addEviction(PhysReg, VirtReg.reg(), Intf->reg());
Matrix->unassign(*Intf);
- assert((getCascade(Intf->reg()) < Cascade ||
+ assert((ExtraInfo->getCascade(Intf->reg()) < Cascade ||
VirtReg.isSpillable() < Intf->isSpillable()) &&
"Cannot decrease cascade number, illegal eviction");
- setCascade(Intf->reg(), Cascade);
+ ExtraInfo->setCascade(Intf->reg(), Cascade);
++NumEvicted;
NewVRegs.push_back(Intf->reg());
}
@@ -1817,13 +1748,13 @@ void RAGreedy::splitAroundRegion(LiveRangeEdit &LREdit,
const LiveInterval &Reg = LIS->getInterval(LREdit.get(I));
// Ignore old intervals from DCE.
- if (getOrInitStage(Reg.reg()) != RS_New)
+ if (ExtraInfo->getOrInitStage(Reg.reg()) != RS_New)
continue;
// Remainder interval. Don't try splitting again, spill if it doesn't
// allocate.
if (IntvMap[I] == 0) {
- setStage(Reg, RS_Spill);
+ ExtraInfo->setStage(Reg, RS_Spill);
continue;
}
@@ -1834,7 +1765,7 @@ void RAGreedy::splitAroundRegion(LiveRangeEdit &LREdit,
LLVM_DEBUG(dbgs() << "Main interval covers the same " << OrigBlocks
<< " blocks as original.\n");
// Don't allow repeated splitting as a safe guard against looping.
- setStage(Reg, RS_Split2);
+ ExtraInfo->setStage(Reg, RS_Split2);
}
continue;
}
@@ -2065,8 +1996,8 @@ unsigned RAGreedy::tryBlockSplit(LiveInterval &VirtReg, AllocationOrder &Order,
// goes straight to spilling, the new local ranges get to stay RS_New.
for (unsigned I = 0, E = LREdit.size(); I != E; ++I) {
const LiveInterval &LI = LIS->getInterval(LREdit.get(I));
- if (getOrInitStage(LI.reg()) == RS_New && IntvMap[I] == 0)
- setStage(LI, RS_Spill);
+ if (ExtraInfo->getOrInitStage(LI.reg()) == RS_New && IntvMap[I] == 0)
+ ExtraInfo->setStage(LI, RS_Spill);
}
if (VerifyEnabled)
@@ -2152,7 +2083,7 @@ RAGreedy::tryInstructionSplit(LiveInterval &VirtReg, AllocationOrder &Order,
SE->finish(&IntvMap);
DebugVars->splitRegister(VirtReg.reg(), LREdit.regs(), *LIS);
// Assign all new registers to RS_Spill. This was the last chance.
- setStage(LREdit.begin(), LREdit.end(), RS_Spill);
+ ExtraInfo->setStage(LREdit.begin(), LREdit.end(), RS_Spill);
return 0;
}
@@ -2320,7 +2251,7 @@ unsigned RAGreedy::tryLocalSplit(LiveInterval &VirtReg, AllocationOrder &Order,
// These rules allow a 3 -> 2+3 split once, which we need. They also prevent
// excessive splitting and infinite loops.
//
- bool ProgressRequired = getStage(VirtReg) >= RS_Split2;
+ bool ProgressRequired = ExtraInfo->getStage(VirtReg) >= RS_Split2;
// Best split candidate.
unsigned BestBefore = NumGaps;
@@ -2456,7 +2387,7 @@ unsigned RAGreedy::tryLocalSplit(LiveInterval &VirtReg, AllocationOrder &Order,
assert(!ProgressRequired && "Didn't make progress when it was required.");
for (unsigned I = 0, E = IntvMap.size(); I != E; ++I)
if (IntvMap[I] == 1) {
- setStage(LIS->getInterval(LREdit.get(I)), RS_Split2);
+ ExtraInfo->setStage(LIS->getInterval(LREdit.get(I)), RS_Split2);
LLVM_DEBUG(dbgs() << ' ' << printReg(LREdit.get(I)));
}
LLVM_DEBUG(dbgs() << '\n');
@@ -2477,7 +2408,7 @@ unsigned RAGreedy::trySplit(LiveInterval &VirtReg, AllocationOrder &Order,
SmallVectorImpl<Register> &NewVRegs,
const SmallVirtRegSet &FixedRegisters) {
// Ranges must be Split2 or less.
- if (getStage(VirtReg) >= RS_Spill)
+ if (ExtraInfo->getStage(VirtReg) >= RS_Spill)
return 0;
// Local intervals are handled separately.
@@ -2499,7 +2430,7 @@ unsigned RAGreedy::trySplit(LiveInterval &VirtReg, AllocationOrder &Order,
// First try to split around a region spanning multiple blocks. RS_Split2
// ranges already made dubious progress with region splitting, so they go
// straight to single block splitting.
- if (getStage(VirtReg) < RS_Split2) {
+ if (ExtraInfo->getStage(VirtReg) < RS_Split2) {
MCRegister PhysReg = tryRegionSplit(VirtReg, Order, NewVRegs);
if (PhysReg || !NewVRegs.empty())
return PhysReg;
@@ -2551,7 +2482,7 @@ bool RAGreedy::mayRecolorAllInterferences(
// it would not be recolorable as it is in the same state as VirtReg.
// However, if VirtReg has tied defs and Intf doesn't, then
// there is still a point in examining if it can be recolorable.
- if (((getStage(*Intf) == RS_Done &&
+ if (((ExtraInfo->getStage(*Intf) == RS_Done &&
MRI->getRegClass(Intf->reg()) == CurRC) &&
!(hasTiedDef(MRI, VirtReg.reg()) &&
!hasTiedDef(MRI, Intf->reg()))) ||
@@ -2615,7 +2546,7 @@ unsigned RAGreedy::tryLastChanceRecoloring(LiveInterval &VirtReg,
LLVM_DEBUG(dbgs() << "Try last chance recoloring for " << VirtReg << '\n');
// Ranges must be Done.
- assert((getStage(VirtReg) >= RS_Done || !VirtReg.isSpillable()) &&
+ assert((ExtraInfo->getStage(VirtReg) >= RS_Done || !VirtReg.isSpillable()) &&
"Last chance recoloring should really be last chance");
// Set the max depth to LastChanceRecoloringMaxDepth.
// We may want to reconsider that if we end up with a too large search space
@@ -2806,7 +2737,7 @@ MCRegister
RAGreedy::tryAssignCSRFirstTime(LiveInterval &VirtReg, AllocationOrder &Order,
MCRegister PhysReg, uint8_t &CostPerUseLimit,
SmallVectorImpl<Register> &NewVRegs) {
- if (getStage(VirtReg) == RS_Spill && VirtReg.isSpillable()) {
+ if (ExtraInfo->getStage(VirtReg) == RS_Spill && VirtReg.isSpillable()) {
// We choose spill over using the CSR for the first time if the spill cost
// is lower than CSRCost.
SA->analyze(&VirtReg);
@@ -2818,7 +2749,7 @@ RAGreedy::tryAssignCSRFirstTime(LiveInterval &VirtReg, AllocationOrder &Order,
CostPerUseLimit = 1;
return 0;
}
- if (getStage(VirtReg) < RS_Split) {
+ if (ExtraInfo->getStage(VirtReg) < RS_Split) {
// We choose pre-splitting over using the CSR for the first time if
// the cost of splitting is lower than CSRCost.
SA->analyze(&VirtReg);
@@ -3063,9 +2994,9 @@ MCRegister RAGreedy::selectOrSplitImpl(LiveInterval &VirtReg,
return PhysReg;
}
- LiveRangeStage Stage = getStage(VirtReg);
+ LiveRangeStage Stage = ExtraInfo->getStage(VirtReg);
LLVM_DEBUG(dbgs() << StageName[Stage] << " Cascade "
- << getCascade(VirtReg.reg()) << '\n');
+ << ExtraInfo->getCascade(VirtReg.reg()) << '\n');
// Try to evict a less worthy live range, but only for ranges from the primary
// queue. The RS_Split ranges already failed to do this, and they should not
@@ -3094,7 +3025,7 @@ MCRegister RAGreedy::selectOrSplitImpl(LiveInterval &VirtReg,
// Wait until the second time, when all smaller ranges have been allocated.
// This gives a better picture of the interference to split around.
if (Stage < RS_Split) {
- setStage(VirtReg, RS_Split);
+ ExtraInfo->setStage(VirtReg, RS_Split);
LLVM_DEBUG(dbgs() << "wait for second round\n");
NewVRegs.push_back(VirtReg.reg());
return 0;
@@ -3120,12 +3051,12 @@ MCRegister RAGreedy::selectOrSplitImpl(LiveInterval &VirtReg,
// Finally spill VirtReg itself.
if ((EnableDeferredSpilling ||
TRI->shouldUseDeferredSpillingForVirtReg(*MF, VirtReg)) &&
- getStage(VirtReg) < RS_Memory) {
+ ExtraInfo->getStage(VirtReg) < RS_Memory) {
// TODO: This is experimental and in particular, we do not model
// the live range splitting done by spilling correctly.
// We would need a deep integration with the spiller to do the
// right thing here. Anyway, that is still good for early testing.
- setStage(VirtReg, RS_Memory);
+ ExtraInfo->setStage(VirtReg, RS_Memory);
LLVM_DEBUG(dbgs() << "Do as if this register is in memory\n");
NewVRegs.push_back(VirtReg.reg());
} else {
@@ -3133,7 +3064,7 @@ MCRegister RAGreedy::selectOrSplitImpl(LiveInterval &VirtReg,
TimerGroupDescription, TimePassesIsEnabled);
LiveRangeEdit LRE(&VirtReg, NewVRegs, *MF, *LIS, VRM, this, &DeadRemats);
spiller().spill(LRE);
- setStage(NewVRegs.begin(), NewVRegs.end(), RS_Done);
+ ExtraInfo->setStage(NewVRegs.begin(), NewVRegs.end(), RS_Done);
// Tell LiveDebugVariables about the new ranges. Ranges not being covered by
// the new regs are kept in LDV (still mapping to the old register), until
@@ -3354,8 +3285,7 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
SA.reset(new SplitAnalysis(*VRM, *LIS, *Loops));
SE.reset(new SplitEditor(*SA, *AA, *LIS, *VRM, *DomTree, *MBFI, *VRAI));
- ExtraRegInfo.clear();
- NextCascade = 1;
+ ExtraInfo.emplace();
IntfCache.init(MF, Matrix->getLiveUnions(), Indexes, LIS, TRI);
GlobalCand.resize(32); // This will grow as needed.
SetOfBrokenHints.clear();
More information about the llvm-commits
mailing list