1 files changed, 85 insertions, 50 deletions

diff --git a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
index a74c57690640..33782c755eb0 100644
--- a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
@@ -220,6 +220,19 @@ void resetInputs(MLModelRunner &Runner) {
 #undef _RESET
 }
 
+// Per-live interval components that get aggregated into the feature values that
+// will be passed to the evaluator.
+struct LIFeatureComponents {
+  double R = 0;
+  double W = 0;
+  double RW = 0;
+  double IndVarUpdates = 0;
+  double HintWeights = 0.0;
+  int64_t NrDefsAndUses = 0;
+  float HottestBlockFreq = 0.0;
+  bool IsRemat = false;
+};
+
 using CandidateRegList =
     std::array<std::pair<MCRegister, bool>, NumberOfInterferences>;
 using FeaturesListNormalizer = std::array<float, FeatureIDs::FeatureCount>;
@@ -227,8 +240,8 @@ using FeaturesListNormalizer = std::array<float, FeatureIDs::FeatureCount>;
 /// The ML evictor (commonalities between release and development mode)
 class MLEvictAdvisor : public RegAllocEvictionAdvisor {
 public:
-  MLEvictAdvisor(const MachineFunction &MF, const RAGreedy &RA,
-                 MLModelRunner *Runner, const MachineBlockFrequencyInfo &MBFI,
+  MLEvictAdvisor(MachineFunction &MF, const RAGreedy &RA, MLModelRunner *Runner,
+                 const MachineBlockFrequencyInfo &MBFI,
                  const MachineLoopInfo &Loops);
 
 protected:
@@ -277,6 +290,9 @@ private:
                                                         FixedRegisters);
   }
 
+  const LIFeatureComponents
+  getLIFeatureComponents(const LiveInterval &LI) const;
+
   // Hold on to a default advisor for:
   // 1) the implementation of canEvictHintInterference, because we didn't learn
   // that nuance yet;
@@ -319,7 +335,7 @@ private:
   }
 
   std::unique_ptr<RegAllocEvictionAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+  getAdvisor(MachineFunction &MF, const RAGreedy &RA) override {
     if (!Runner)
       Runner = std::make_unique<ReleaseModeModelRunner<RegallocEvictModel>>(
           MF.getFunction().getContext(), FeatureNames, DecisionName);
@@ -364,7 +380,7 @@ static const std::vector<TensorSpec> TrainingInputFeatures{
 
 class DevelopmentModeEvictAdvisor : public MLEvictAdvisor {
 public:
-  DevelopmentModeEvictAdvisor(const MachineFunction &MF, const RAGreedy &RA,
+  DevelopmentModeEvictAdvisor(MachineFunction &MF, const RAGreedy &RA,
                               MLModelRunner *Runner,
                               const MachineBlockFrequencyInfo &MBFI,
                               const MachineLoopInfo &Loops, Logger *Log)
@@ -420,7 +436,7 @@ private:
   }
 
   std::unique_ptr<RegAllocEvictionAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+  getAdvisor(MachineFunction &MF, const RAGreedy &RA) override {
     LLVMContext &Ctx = MF.getFunction().getContext();
     if (ModelUnderTraining.empty() && TrainingLog.empty()) {
       Ctx.emitError("Regalloc development mode should be requested with at "
@@ -480,7 +496,7 @@ float MLEvictAdvisor::getInitialQueueSize(const MachineFunction &MF) {
   return Ret;
 }
 
-MLEvictAdvisor::MLEvictAdvisor(const MachineFunction &MF, const RAGreedy &RA,
+MLEvictAdvisor::MLEvictAdvisor(MachineFunction &MF, const RAGreedy &RA,
                                MLModelRunner *Runner,
                                const MachineBlockFrequencyInfo &MBFI,
                                const MachineLoopInfo &Loops)
@@ -615,16 +631,15 @@ MCRegister MLEvictAdvisor::tryFindEvictionCandidate(
   for (auto I = Order.begin(), E = Order.getOrderLimitEnd(OrderLimit); I != E;
        ++I, ++Pos) {
     MCRegister PhysReg = *I;
-    Regs[Pos] = std::make_pair(PhysReg, true);
+    assert(!Regs[Pos].second);
     assert(PhysReg);
     if (!canAllocatePhysReg(CostPerUseLimit, PhysReg)) {
-      Regs[Pos].second = false;
       continue;
     }
     if (loadInterferenceFeatures(VirtReg, PhysReg, I.isHint(), FixedRegisters,
                                  Largest, Pos)) {
       ++Available;
-      Regs[Pos].second = true;
+      Regs[Pos] = std::make_pair(PhysReg, true);
     }
   }
   if (Available == 0) {
@@ -632,6 +647,7 @@ MCRegister MLEvictAdvisor::tryFindEvictionCandidate(
     assert(!MustFindEviction);
     return MCRegister::NoRegister;
   }
+  const size_t ValidPosLimit = Pos;
   // If we must find eviction, the candidate should be masked out of the
   // decision making process.
   Regs[CandidateVirtRegPos].second = !MustFindEviction;
@@ -665,9 +681,55 @@ MCRegister MLEvictAdvisor::tryFindEvictionCandidate(
     assert(!MustFindEviction);
     return MCRegister::NoRegister;
   }
+  assert(CandidatePos < ValidPosLimit);
+  (void)ValidPosLimit;
   return Regs[CandidatePos].first;
 }
 
+const LIFeatureComponents
+MLEvictAdvisor::getLIFeatureComponents(const LiveInterval &LI) const {
+  LIFeatureComponents Ret;
+  SmallPtrSet<MachineInstr *, 8> Visited;
+  const TargetRegisterInfo &TRI = *MF.getSubtarget().getRegisterInfo();
+
+  for (MachineRegisterInfo::reg_instr_nodbg_iterator
+           I = MRI->reg_instr_nodbg_begin(LI.reg()),
+           E = MRI->reg_instr_nodbg_end();
+       I != E;) {
+    MachineInstr *MI = &*(I++);
+
+    ++Ret.NrDefsAndUses;
+    if (!Visited.insert(MI).second)
+      continue;
+
+    if (MI->isIdentityCopy() || MI->isImplicitDef())
+      continue;
+
+    bool Reads, Writes;
+    std::tie(Reads, Writes) = MI->readsWritesVirtualRegister(LI.reg());
+
+    float Freq = MBFI.getBlockFreqRelativeToEntryBlock(MI->getParent());
+    Ret.HottestBlockFreq = std::max(Freq, Ret.HottestBlockFreq);
+
+    Ret.R += (Reads && !Writes) * Freq;
+    Ret.W += (!Reads && Writes) * Freq;
+    Ret.RW += (Reads && Writes) * Freq;
+
+    auto *MBB = MI->getParent();
+    auto *Loop = Loops.getLoopFor(MBB);
+    bool IsExiting = Loop ? Loop->isLoopExiting(MBB) : false;
+
+    if (Writes && IsExiting && LIS->isLiveOutOfMBB(LI, MBB))
+      Ret.IndVarUpdates += Freq;
+
+    if (MI->isCopy() && VirtRegAuxInfo::copyHint(MI, LI.reg(), TRI, *MRI))
+      Ret.HintWeights += Freq;
+  }
+  Ret.IsRemat = VirtRegAuxInfo::isRematerializable(
+      LI, *LIS, *VRM, *MF.getSubtarget().getInstrInfo());
+  return Ret;
+}
+
 // Overall, this currently mimics what we do for weight calculation, but instead
 // of accummulating the various features, we keep them separate.
 void MLEvictAdvisor::extractFeatures(
@@ -676,11 +738,11 @@ void MLEvictAdvisor::extractFeatures(
     int64_t IsHint, int64_t LocalIntfsCount, float NrUrgent) const {
   int64_t NrDefsAndUses = 0;
   int64_t NrBrokenHints = 0;
-  float R = 0;
-  float W = 0;
-  float RW = 0;
-  float IndVarUpdates = 0;
-  float HintWeights = 0.0;
+  double R = 0.0;
+  double W = 0.0;
+  double RW = 0.0;
+  double IndVarUpdates = 0.0;
+  double HintWeights = 0.0;
   float StartBBFreq = 0.0;
   float EndBBFreq = 0.0;
   float HottestBlockFreq = 0.0;
@@ -707,46 +769,19 @@ void MLEvictAdvisor::extractFeatures(
 
     if (LI.endIndex() > EndSI)
       EndSI = LI.endIndex();
-
-    SmallPtrSet<MachineInstr *, 8> Visited;
-    const TargetRegisterInfo &TRI = *MF.getSubtarget().getRegisterInfo();
+    const LIFeatureComponents LIFC = getLIFeatureComponents(LI);
     NrBrokenHints += VRM->hasPreferredPhys(LI.reg());
 
-    for (MachineRegisterInfo::reg_instr_nodbg_iterator
-             I = MRI->reg_instr_nodbg_begin(LI.reg()),
-             E = MRI->reg_instr_nodbg_end();
-         I != E;) {
-      MachineInstr *MI = &*(I++);
+    NrDefsAndUses += LIFC.NrDefsAndUses;
+    HottestBlockFreq = std::max(HottestBlockFreq, LIFC.HottestBlockFreq);
+    R += LIFC.R;
+    W += LIFC.W;
+    RW += LIFC.RW;
 
-      ++NrDefsAndUses;
-      if (!Visited.insert(MI).second)
-        continue;
+    IndVarUpdates += LIFC.IndVarUpdates;
 
-      if (MI->isIdentityCopy() || MI->isImplicitDef())
-        continue;
-
-      bool Reads, Writes;
-      std::tie(Reads, Writes) = MI->readsWritesVirtualRegister(LI.reg());
-
-      float Freq = MBFI.getBlockFreqRelativeToEntryBlock(MI->getParent());
-      if (Freq > HottestBlockFreq)
-        HottestBlockFreq = Freq;
-      R += (Reads && !Writes) * Freq;
-      W += (!Reads && Writes) * Freq;
-      RW += (Reads && Writes) * Freq;
-
-      auto *MBB = MI->getParent();
-      auto *Loop = Loops.getLoopFor(MBB);
-      bool IsExiting = Loop ? Loop->isLoopExiting(MBB) : false;
-
-      if (Writes && IsExiting && LIS->isLiveOutOfMBB(LI, MBB))
-        IndVarUpdates += Freq;
-
-      if (MI->isCopy() && VirtRegAuxInfo::copyHint(MI, LI.reg(), TRI, *MRI))
-        HintWeights += Freq;
-    }
-    NrRematerializable += VirtRegAuxInfo::isRematerializable(
-        LI, *LIS, *VRM, *MF.getSubtarget().getInstrInfo());
+    HintWeights += LIFC.HintWeights;
+    NrRematerializable += LIFC.IsRemat;
   }
   size_t Size = 0;
   if (!Intervals.empty()) {