[NFC][regalloc] Move ExtraRegInfo and related to LiveRangeStageManager

[mtrofin]

Summary

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).

RFC: https://lists.llvm.org/pipermail/llvm-dev/2021-November/153639.html

[mtrofin]

Gentle reminder - thanks!

[MatzeB]

Is there any way to see a complete stack to make it easier to comment on the approach as a whole?

[mtrofin]

In D114831#3167586, @MatzeB wrote:

Is there any way to see a complete stack to make it easier to comment on the approach as a whole?

I sent a while back D113525 as an overview for the effort of factoring eviction as an analysis. This particular patch slightly deviates from that patch: rather than mashing the live range state tracking in the analysis, for the reason outlined in the review message, we factor it as its own class.

[MatzeB]

I cannot shake the feeling that we should put the interfaces into a RegAllocGreedy.h, the interface doesn't feel generic enough for anything else (and FWIW I am not advocating to make it more generic, just set the right expectations that this is meant to be used in combination with RegAllocGreedy). And with expectations managed maybe it's good enough to just hand out a reference to ExtraRegInfo to some friend implementations of the eviction interface instead of introducing more abstractions and "Manager" classes. How do others working on the regallocs feel?

llvm/lib/CodeGen/RegAllocEvictionAdvisor.h

89

Add a comment to document the role of this class and try to find a better name (ExtraRegInfos to match the ExtraRegInfo map this is appears to be wrapping?)

102–107

Can we keep the MRI reference out of the interface?

119–137

Wondering if we can go without an MRI reference

llvm/lib/CodeGen/RegAllocGreedy.cpp

176

Is it possible to embed this by value to avoid the extra indirections?

[mtrofin]

In D114831#3167768, @MatzeB wrote:

I cannot shake the feeling that we should put the interfaces into a RegAllocGreedy.h, the interface doesn't feel generic enough for anything else (and FWIW I am not advocating to make it more generic, just set the right expectations that this is meant to be used in combination with RegAllocGreedy). And with expectations managed maybe it's good enough to just hand out a reference to ExtraRegInfo to some friend implementations of the eviction interface instead of introducing more abstractions and "Manager" classes. How do others working on the regallocs feel?

We'll want to use ExtraRegInfo & associated APIs in the ML-based implementation, hence the factoring.

[mtrofin]

llvm/lib/CodeGen/RegAllocEvictionAdvisor.h

89

Hmm, I thought both the Stage and the Cascade relate to tracking the transitions of the virtual register through the allocation; but I can see how, if we move this to a RegAllocGreedy.h, and we want to track more state later, ExtraRegInfo may work better. How about we use that and rename the vector to just Info (it's private anyway)?

102–107

The goal was to avoid needing to pass the exact same thing (the MRI->getNumVirtRegs()) at the points where we need to resize.

So I learn: what's the downside of MRI here?

Would passing the VirtRegMap be a reasonable alternative?

119–137

ack (consolidating this in the comment above)

llvm/lib/CodeGen/RegAllocGreedy.cpp

176

Oh, is this the reason for avoiding nondefault ctor? Let me fire up a https://llvm-compile-time-tracker.com/ experiment, if overhead is a concern.

[mtrofin]

llvm/lib/CodeGen/RegAllocGreedy.cpp

176

Here it is - arguably, the indirection has no impact.

[mtrofin]

llvm/lib/CodeGen/RegAllocGreedy.cpp

176

Actually - I think we can avoid the heap allocation, and not require default ctor, too, if we used Optional. I'll try that.

[mtrofin]

incorporated feedback:

• using ExtraRegInfo for type name
• dropped heap allocation in favor of Optional

[MatzeB]

We'll want to use ExtraRegInfo & associated APIs in the ML-based implementation, hence the factoring.

I understand that. I am hinting at possible designs on how much or how little we abstract:

1. Treat "eviction" as a strategy that is part of RegAllocGreedy and specifically embrace the fact that things only work in combination with RegAllocGreedy:

// RegAllocGreedy.h
class RegAllocGreedy;

// This is interface only works for RAGreedy, hence it is in RegAllocGreedy.h. Benefit: we can directly
// re-use datastructures and call methods on RAGreedy.
class EvictionStrategy {
  RegAllocGreedy *RAGreedy;
  EvictionStrategy(RegAllocGreedy* RAGreedy) : RAGreedy(RAGreedy) {}

  void doStuff() {
    // ...
    RAGreedy->setStage(...);
    // ...
  }
};

class RegAllocGreedy {

  void setStage(...);
  friend class EvictionStrategy;
  // ...
  EvictionStrategy *Eviction;
};

This will require fewer abstraction to be built up, with the downside of having a higher entanglement between the two.

2. The other end of the spectrum is pushing for more abstraction modeling eviction independently of the rest of the register allocator. This usually comes with the goal of having a cleaner separation of concerns and better re-usability of the algorithm (for other regallocs or compilers?).

Now my point is that the current patches appear to aiming more at option 2) while at the same time the interfaces are clearly heavily inspired/only usable for RegAllocGreedy. So I was posing the question if we should rather acknowledge that reality and aim more for a 1) design instead...

While I personally think 1) is the better choice in this context, I wouldn't be surprised if different people have different opinions here, so I was fishing for opinions...

[MatzeB]

llvm/lib/CodeGen/RegAllocGreedy.cpp

176

Yes I highly doubt you can actually measure any difference in this particular case... This was more a matter of principle and setting examples for other passes ;-)

[MatzeB]

llvm/lib/CodeGen/RegAllocEvictionAdvisor.h

102–107

So I learn: what's the downside of MRI here?

No real downside.

I would just say that reducing dependencies/entanglement with other classes (switching to VirtRegMap is as bad in my view). I just think it's a worthwhile general goal when creating new abstractions. And this seemed to be fixable without any real efforts/or lost convenience.

The goal was to avoid needing to pass the exact same thing (the MRI->getNumVirtRegs()) at the points where we need to resize.

It just seemed to me like we can simply get rid of the need for getNumVirtRegs() in the setStage/setCascade functions by switching to the grow() interface. I hope you can see the code suggestions I made in phabricator that show that (I think phab sometimes does not reproduce them properly in the E-Mails, so just making sure we're not talking past each other because you don't see the suggestions)

Anyway these here are nitpicks/suggestions, no harm done if we cannot implement them after all.

[mtrofin]

In D114831#3168361, @MatzeB wrote:

We'll want to use ExtraRegInfo & associated APIs in the ML-based implementation, hence the factoring.

I understand that. I am hinting at possible designs on how much or how little we abstract:

1. Treat "eviction" as a strategy that is part of RegAllocGreedy and specifically embrace the fact that things only work in combination with RegAllocGreedy:

// RegAllocGreedy.h
class RegAllocGreedy;

// This is interface only works for RAGreedy, hence it is in RegAllocGreedy.h. Benefit: we can directly
// re-use datastructures and call methods on RAGreedy.
class EvictionStrategy {
  RegAllocGreedy *RAGreedy;
  EvictionStrategy(RegAllocGreedy* RAGreedy) : RAGreedy(RAGreedy) {}

  void doStuff() {
    // ...
    RAGreedy->setStage(...);
    // ...
  }
};

class RegAllocGreedy {

  void setStage(...);
  friend class EvictionStrategy;
  // ...
  EvictionStrategy *Eviction;
};

This will require fewer abstraction to be built up, with the downside of having a higher entanglement between the two.

2. The other end of the spectrum is pushing for more abstraction modeling eviction independently of the rest of the register allocator. This usually comes with the goal of having a cleaner separation of concerns and better re-usability of the algorithm (for other regallocs or compilers?).

Now my point is that the current patches appear to aiming more at option 2) while at the same time the interfaces are clearly heavily inspired/only usable for RegAllocGreedy. So I was posing the question if we should rather acknowledge that reality and aim more for a 1) design instead...

While I personally think 1) is the better choice in this context, I wouldn't be surprised if different people have different opinions here, so I was fishing for opinions...

Option 1 would mean extracting RAGreedy in that header, and that looks like a larger factoring in a .h, in comparison to the ExtraRegInfo stuff.

[mtrofin]

llvm/lib/CodeGen/RegAllocEvictionAdvisor.h

102–107

I like the .grow() idea, my concern was wether we'd see any compile time impact. Let me try it and see.

[MatzeB]

Option 1 would mean extracting RAGreedy in that header, and that looks like a larger factoring in a .h, in comparison to the ExtraRegInfo stuff.

Yes, but it's just moving the declaration around without changing things otherwise. Which seems like the "smaller" factoring to me. (at the price of higher entanglement as I stated)

[mtrofin]

In D114831#3168425, @MatzeB wrote:

Option 1 would mean extracting RAGreedy in that header, and that looks like a larger factoring in a .h, in comparison to the ExtraRegInfo stuff.

Yes, but it's just moving the declaration around without changing things otherwise. Which seems like the "smaller" factoring to me. (at the price of higher entanglement as I stated)

Ack.

I wonder if we can decouple this design problem from this patch. Should there end up being a strong preference for going to a RegAllocGreedy.h, that should be easy to do at any later point, wdyt?

[mtrofin]

Instead of resizing, just rely on grow().

This doesn't seem to impact compiletime:
https://llvm-compile-time-tracker.com/compare.php?from=e0b259f22c003ffe909693c6ab0d508d1814434d&to=447c8d6c740e186aa2098579eb532f74760e371b&stat=instructions

[mtrofin]

llvm/lib/CodeGen/RegAllocEvictionAdvisor.h

102–107

Yup, no measurable effect when removing resize and just grow-ing.

[mtrofin]

missed feedback about adding a comment

[mtrofin]

Gentle reminder - could we defer deciding on a RegAllocGreedy.h?

Thanks!

[mtrofin]

Gentle reminder - @MatzeB, @qcolombet, please let me know how to move this forward - thanks!

[wenlei]

In D114831#3168435, @mtrofin wrote:

In D114831#3168425, @MatzeB wrote:

Option 1 would mean extracting RAGreedy in that header, and that looks like a larger factoring in a .h, in comparison to the ExtraRegInfo stuff.

Yes, but it's just moving the declaration around without changing things otherwise. Which seems like the "smaller" factoring to me. (at the price of higher entanglement as I stated)

Ack.

I wonder if we can decouple this design problem from this patch. Should there end up being a strong preference for going to a RegAllocGreedy.h, that should be easy to do at any later point, wdyt?

I agree with Matthias that the advisor now is really tied to Greedy and #1 makes sense. Handle this in a later patch is also reasonable, but leaving it in current state for too long probably isn't ideal. If you plan to handle it later, perhaps stick in a TODO comment?

Other than that, the refactoring seems straightforward.

[mtrofin]

added todo to expose RAGreedy in a header instead

[mtrofin]

In D114831#3188279, @wenlei wrote:

In D114831#3168435, @mtrofin wrote:

In D114831#3168425, @MatzeB wrote:

Option 1 would mean extracting RAGreedy in that header, and that looks like a larger factoring in a .h, in comparison to the ExtraRegInfo stuff.

Yes, but it's just moving the declaration around without changing things otherwise. Which seems like the "smaller" factoring to me. (at the price of higher entanglement as I stated)

Ack.

I wonder if we can decouple this design problem from this patch. Should there end up being a strong preference for going to a RegAllocGreedy.h, that should be easy to do at any later point, wdyt?

I agree with Matthias that the advisor now is really tied to Greedy and #1 makes sense. Handle this in a later patch is also reasonable, but leaving it in current state for too long probably isn't ideal. If you plan to handle it later, perhaps stick in a TODO comment?

Other than that, the refactoring seems straightforward.

Done - I'll address it once everything is settled (i.e. all the components landed, incl. the ML advisor) - it will probably help with any nuances of the refactoring, too.

[wenlei]

lgtm, thanks. (btw Matthias is on PTO for the holidays and won't be back until Jan)

llvm/lib/CodeGen/RegAllocEvictionAdvisor.h

struct EvictionCost {
 
   void setBrokenHints(unsigned NHints) { BrokenHints = NHints; }
 
   bool operator<(const EvictionCost &O) const {
     return std::tie(BrokenHints, MaxWeight) <
            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

llvm/lib/CodeGen/RegAllocGreedy.cpp

class RAGreedy : public MachineFunctionPass,
   EdgeBundles *Bundles;
   SpillPlacement *SpillPlacer;
   LiveDebugVariables *DebugVars;
   AliasAnalysis *AA;
 
   // 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.
   // Note: This is used as bitmask. New value should be next power of 2.
   enum CutOffStage {
     // No cutoffs encountered
     CO_None = 0,
 
     // lcr-max-depth cutoff encountered
     CO_Depth = 1,
 
     // lcr-max-interf cutoff encountered
     CO_Interf = 2
   };
 
   uint8_t CutOffInfo;
 
 #ifndef NDEBUG
   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 {
 
   public:
     using EvictorInfo =
         std::pair<Register /* evictor */, MCRegister /* physreg */>;
     using EvicteeInfo = llvm::DenseMap<Register /* evictee */, EvictorInfo>;
void RAGreedy::LRE_WillShrinkVirtReg(Register VirtReg) {
 
   // Register is assigned, put it back on the queue for reassignment.
   LiveInterval &LI = LIS->getInterval(VirtReg);
   Matrix->unassign(LI);
   RegAllocBase::enqueue(&LI);
 }
 
 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();
 }
 
 void RAGreedy::enqueueImpl(LiveInterval *LI) { enqueue(Queue, LI); }
 
 void RAGreedy::enqueue(PQueue &CurQueue, LiveInterval *LI) {
   // Prioritize live ranges by size, assigning larger ranges first.
   // The queue holds (size, reg) pairs.
   const unsigned Size = LI->getSize();
   const Register Reg = LI->reg();
   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
     // everything else has been allocated.
     Prio = Size;
   } else if (Stage == RS_Memory) {
     // Memory operand should be considered last.
     // Change the priority such that Memory operand are assigned in

 ///
 /// @param A          The live range to be assigned.
 /// @param IsHint     True when A is about to be assigned to its preferred
 ///                   register.
 /// @param B          The live range to be evicted.
 /// @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.
   if (CanSplit && IsHint && !BreaksHint)
     return true;
 
   if (A.weight() > B.weight()) {
     LLVM_DEBUG(dbgs() << "should evict: " << B << " w= " << B.weight() << '\n');
bool RAGreedy::canEvictInterferenceBasedOnCost(
 
   // Find VirtReg's cascade number. This will be unassigned if VirtReg was never
   // involved in an eviction before. If a cascade number was assigned, deny
   // evicting anything with the same or a newer cascade number. This prevents
   // infinite eviction loops.
   //
   // 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) {
     LiveIntervalUnion::Query &Q = Matrix->query(VirtReg, *Units);
     // If there is 10 or more interferences, chances are one is heavier.
     const auto &Interferences = Q.interferingVRegs(10);
     if (Interferences.size() >= 10)
       return false;
 
     // Check if any interfering live range is heavier than MaxWeight.
     for (LiveInterval *Intf : reverse(Interferences)) {
       assert(Register::isVirtualRegister(Intf->reg()) &&
              "Only expecting virtual register interference from query");
 
       // Do not allow eviction of a virtual register if we are in the middle
       // of last-chance recoloring and this virtual register is one that we
       // have scavenged a physical register for.
       if (FixedRegisters.count(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;
       // 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
       // urgent live ranges get to evict almost anything.
       //
       // Also allow urgent evictions of unspillable ranges from a strictly
       // larger allocation order.
       bool Urgent =
           !VirtReg.isSpillable() &&
           (Intf->isSpillable() ||
            RegClassInfo.getNumAllocatableRegs(MRI->getRegClass(VirtReg.reg())) <
                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;
         // We permit breaking cascades for urgent evictions. It should be the
         // last resort, though, so make it really expensive.
         Cost.BrokenHints += 10;
       }
       // Would this break a satisfied hint?
for (const LiveInterval *Intf : reverse(Q.interferingVRegs())) {
       // Check if interference overlast the segment in interest.
       if (!Intf->overlaps(Start, End))
         continue;
 
       // Cannot evict non virtual reg interference.
       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?
       bool BreaksHint = VRM->hasPreferredPhys(Intf->reg());
       // Update eviction cost.
       Cost.BrokenHints += BreaksHint;
       Cost.MaxWeight = std::max(Cost.MaxWeight, Intf->weight());
       // Abort if this would be too expensive.

 /// evictInterference - Evict any interferring registers that prevent VirtReg
 /// from being assigned to Physreg. This assumes that canEvictInterference
 /// returned true.
 void RAGreedy::evictInterference(LiveInterval &VirtReg, MCRegister PhysReg,
                                  SmallVectorImpl<Register> &NewVRegs) {
   // 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');
 
   // Collect all interfering virtregs first.
   SmallVector<LiveInterval*, 8> Intfs;
   for (MCRegUnitIterator Units(PhysReg, TRI); Units.isValid(); ++Units) {
     LiveIntervalUnion::Query &Q = Matrix->query(VirtReg, *Units);
void RAGreedy::evictInterference(LiveInterval &VirtReg, MCRegister PhysReg,
   for (LiveInterval *Intf : Intfs) {
     // The same VirtReg may be present in multiple RegUnits. Skip duplicates.
     if (!VRM->hasPhys(Intf->reg()))
       continue;
 
     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());
   }
 }
 
 /// Returns true if the given \p PhysReg is a callee saved register and has not
 /// been used for allocation yet.
 bool RAGreedy::isUnusedCalleeSavedReg(MCRegister PhysReg) const {
void RAGreedy::splitAroundRegion(LiveRangeEdit &LREdit,
   // - Remainder intervals should not be split again.
   // - Candidate intervals can be assigned to Cand.PhysReg.
   // - Block-local splits are candidates for local splitting.
   // - DCE leftovers should go back on the queue.
   for (unsigned I = 0, E = LREdit.size(); I != E; ++I) {
     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;
     }
 
     // Global intervals. Allow repeated splitting as long as the number of live
     // blocks is strictly decreasing.
     if (IntvMap[I] < NumGlobalIntvs) {
       if (SA->countLiveBlocks(&Reg) >= OrigBlocks) {
         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;
     }
 
     // Other intervals are treated as new. This includes local intervals created
     // for blocks with multiple uses, and anything created by DCE.
   }
 
unsigned RAGreedy::tryBlockSplit(LiveInterval &VirtReg, AllocationOrder &Order,
 
   // Tell LiveDebugVariables about the new ranges.
   DebugVars->splitRegister(Reg, LREdit.regs(), *LIS);
 
   // Sort out the new intervals created by splitting. The remainder interval
   // 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)
     MF->verify(this, "After splitting live range around basic blocks");
   return 0;
 }
 
 //===----------------------------------------------------------------------===//
if (LREdit.empty()) {
     LLVM_DEBUG(dbgs() << "All uses were copies.\n");
     return 0;
   }
 
   SmallVector<unsigned, 8> IntvMap;
   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;
 }
 
 //===----------------------------------------------------------------------===//
 //                             Local Splitting
 //===----------------------------------------------------------------------===//
 
 /// calcGapWeights - Compute the maximum spill weight that needs to be evicted
unsigned RAGreedy::tryLocalSplit(LiveInterval &VirtReg, AllocationOrder &Order,
   // 2. Require progress be made for ranges with getStage() == RS_Split2. All
   //    the new ranges must have fewer instructions than before the split.
   // 3. New ranges with the same number of instructions are marked RS_Split2,
   //    smaller ranges are marked RS_New.
   //
   // 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;
   unsigned BestAfter = 0;
   float BestDiff = 0;
 
   const float blockFreq =
     SpillPlacer->getBlockFrequency(BI.MBB->getNumber()).getFrequency() *
unsigned RAGreedy::tryLocalSplit(LiveInterval &VirtReg, AllocationOrder &Order,
   bool LiveBefore = BestBefore != 0 || BI.LiveIn;
   bool LiveAfter = BestAfter != NumGaps || BI.LiveOut;
   unsigned NewGaps = LiveBefore + BestAfter - BestBefore + LiveAfter;
   if (NewGaps >= NumGaps) {
     LLVM_DEBUG(dbgs() << "Tagging non-progress ranges:");
     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');
   }
   ++NumLocalSplits;
 
   return 0;
 }
 
 //===----------------------------------------------------------------------===//
 //                          Live Range Splitting
 //===----------------------------------------------------------------------===//
 
 /// trySplit - Try to split VirtReg or one of its interferences, making it
 /// assignable.
 /// @return Physreg when VirtReg may be assigned and/or new NewVRegs.
 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.
   if (LIS->intervalIsInOneMBB(VirtReg)) {
     NamedRegionTimer T("local_split", "Local Splitting", TimerGroupName,
                        TimerGroupDescription, TimePassesIsEnabled);
     SA->analyze(&VirtReg);
     Register PhysReg = tryLocalSplit(VirtReg, Order, NewVRegs);
     if (PhysReg || !NewVRegs.empty())
       return PhysReg;
     return tryInstructionSplit(VirtReg, Order, NewVRegs);
   }
 
   NamedRegionTimer T("global_split", "Global Splitting", TimerGroupName,
                      TimerGroupDescription, TimePassesIsEnabled);
 
   SA->analyze(&VirtReg);
 
   // 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;
   }
 
   // Then isolate blocks.
   return tryBlockSplit(VirtReg, Order, NewVRegs);
 }
if (Q.interferingVRegs(LastChanceRecoloringMaxInterference).size() >=
       CutOffInfo |= CO_Interf;
       return false;
     }
     for (LiveInterval *Intf : reverse(Q.interferingVRegs())) {
       // If Intf is done and sit on the same register class as VirtReg,
       // 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()))) ||
           FixedRegisters.count(Intf->reg())) {
         LLVM_DEBUG(
             dbgs() << "Early abort: the interference is not recolorable.\n");
         return false;
       }
unsigned RAGreedy::tryLastChanceRecoloring(LiveInterval &VirtReg,
                                            SmallVectorImpl<Register> &NewVRegs,
                                            SmallVirtRegSet &FixedRegisters,
                                            unsigned Depth) {
   if (!TRI->shouldUseLastChanceRecoloringForVirtReg(*MF, VirtReg))
     return ~0u;
 
   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
   // for target with hundreds of registers.
   // Indeed, in that case we may want to cut the search space earlier.
   if (Depth >= LastChanceRecoloringMaxDepth && !ExhaustiveSearch) {
     LLVM_DEBUG(dbgs() << "Abort because max depth has been reached.\n");
     CutOffInfo |= CO_Depth;

 /// range can have lower cost than using the CSR for the first time;
 /// Spilling a live range in the cold path can have lower cost than using
 /// the CSR for the first time. Returns the physical register if we decide
 /// to use the CSR; otherwise return 0.
 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);
     if (calcSpillCost() >= CSRCost)
       return PhysReg;
 
     // We are going to spill, set CostPerUseLimit to 1 to make sure that
     // we will not use a callee-saved register in tryEvict.
     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);
     unsigned NumCands = 0;
     BlockFrequency BestCost = CSRCost; // Don't modify CSRCost.
     unsigned BestCand = calculateRegionSplitCost(VirtReg, Order, BestCost,
                                                  NumCands, true /*IgnoreCSR*/);
     if (BestCand == NoCand)
if (CSRCost.getFrequency() && isUnusedCalleeSavedReg(PhysReg) &&
       if (CSRReg || !NewVRegs.empty())
         // Return now if we decide to use a CSR or create new vregs due to
         // pre-splitting.
         return CSRReg;
     } else
       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
   // get a second chance until they have been split.
   if (Stage != RS_Split)
     if (Register PhysReg =
             tryEvict(VirtReg, Order, NewVRegs, CostPerUseLimit,
                      FixedRegisters)) {
if (Stage != RS_Split)
     }
 
   assert((NewVRegs.empty() || Depth) && "Cannot append to existing NewVRegs");
 
   // The first time we see a live range, don't try to split or spill.
   // 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;
   }
 
   if (Stage < RS_Spill) {
     // Try splitting VirtReg or interferences.
     unsigned NewVRegSizeBefore = NewVRegs.size();
MCRegister RAGreedy::selectOrSplitImpl(LiveInterval &VirtReg,
   // invalid inline assembly. The base class will report it.
   if (Stage >= RS_Done || !VirtReg.isSpillable())
     return tryLastChanceRecoloring(VirtReg, Order, NewVRegs, FixedRegisters,
                                    Depth);
 
   // 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 {
     NamedRegionTimer T("spill", "Spiller", TimerGroupName,
                        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
     // we rewrite spilled locations in LDV at a later stage.
     DebugVars->splitRegister(VirtReg.reg(), LRE.regs(), *LIS);
 
     if (VerifyEnabled)
       MF->verify(this, "After spilling");
bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
   SpillerInstance.reset(createInlineSpiller(*this, *MF, *VRM, *VRAI));
 
   VRAI->calculateSpillWeightsAndHints();
 
   LLVM_DEBUG(LIS->dump());
 
   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();
   LastEvicted.clear();
 
   allocatePhysRegs();
   tryHintsRecoloring();
 
   if (VerifyEnabled)
     MF->verify(this, "Before post optimization");
   postOptimization();
   reportStats();
 
   releaseMemory();
   return true;
 }