[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

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