[nfc][mlgo][regalloc] Cache live interval feature components

[mtrofin]

Summary

Lazily cache the feature components of a LiveInterval.

[mtrofin]

Gentle reminder - I don't think this change is too controversial, but lmk if advisors owning VRAIs seems too gauche.

[mtrofin]

rebased

[MatzeB]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

252–255

Isn't this only ever computed once per live interval? I don't see any code computing this a 2nd time making it necessary to invalidate a cache...

The only user of this function is VirtRegAuxInfo::calculateSpillWeightsAndHints computing the initial spill weights when initializing the regalloc; and LiveRangeEdit::calculateRegClassAndHint which calls this function only for the newly created live intervals.

Why would you need to invalidate a cache then?

[mtrofin]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

252–255

Oh, true - so much simpler.

(separately, I wonder if I can remove setWeight and just ctor the LI with it pre-calculated)

[mtrofin]

simplification

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

 public:
   MLEvictAdvisor(const MachineFunction &MF, const RAGreedy &RA,
                  MLModelRunner *Runner, const MachineBlockFrequencyInfo &MBFI,
                  const MachineLoopInfo &Loops);
 
 protected:
   const RegAllocEvictionAdvisor &getDefaultAdvisor() const {
     return static_cast<const RegAllocEvictionAdvisor &>(DefaultAdvisor);
   }
 
   // The assumption is that if the Runner could not be constructed, we emit-ed
   // error, and we shouldn't be asking for it here.
   const MLModelRunner &getRunner() const { return *Runner; }
 
   /// This just calls Evaluate on the Runner, but in the development mode case,
   /// if we're just capturing the log of the default advisor, it needs to call
   /// the latter instead, so we need to pass all the necessary parameters for
   /// it. In the development case, it will also log.
   virtual int64_t
   tryFindEvictionCandidatePosition(const LiveInterval &VirtReg,
private:
   // Point-in-time: we didn't learn this, so we always delegate to the default.
   bool canEvictHintInterference(
       const LiveInterval &VirtReg, MCRegister PhysReg,
       const SmallVirtRegSet &FixedRegisters) const override {
     return getDefaultAdvisor().canEvictHintInterference(VirtReg, PhysReg,
                                                         FixedRegisters);
   }
 
-  const LIFeatureComponents
+  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;
   // 2) for bootstrapping (logging) in the development mode case.
   const DefaultEvictionAdvisor DefaultAdvisor;
   MLModelRunner *const Runner;
   const MachineBlockFrequencyInfo &MBFI;
   const MachineLoopInfo &Loops;
 
   // Indices of those features we don't want to normalize.
   // This could be static and shared, but its initialization is non-trivial.
   std::bitset<FeatureIDs::FeatureCount> DoNotNormalize;
   const float InitialQSize;
+
+  using RegID = unsigned;
+  mutable DenseMap<RegID, LIFeatureComponents> CachedFeatures;
 };
 
 // ===================================
 // Release (AOT) - specifics
 // ===================================
 #if defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL)
 const std::array<std::string, FeatureIDs::FeatureCount> FeatureNames{
 #define _GETNAME(_, NAME, __, ___) #NAME,
if (CandidatePos == CandidateVirtRegPos) {
     assert(!MustFindEviction);
     return MCRegister::NoRegister;
   }
   assert(CandidatePos < ValidPosLimit);
   (void)ValidPosLimit;
   return Regs[CandidatePos].first;
 }
 
-const LIFeatureComponents
+const LIFeatureComponents &
 MLEvictAdvisor::getLIFeatureComponents(const LiveInterval &LI) const {
-  LIFeatureComponents Ret;
+  RegID ID = LI.reg().id();
+  LIFeatureComponents Empty;
+  auto I = CachedFeatures.insert(std::make_pair(ID, Empty));
+  LIFeatureComponents &Ret = I.first->getSecond();
+  if (!I.second)
+    return 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++);
for (const auto *L : Intervals) {
 
     TotalWeight = std::max(TotalWeight, LI.weight());
 
     if (LI.beginIndex() < StartSI)
       StartSI = LI.beginIndex();
 
     if (LI.endIndex() > EndSI)
       EndSI = LI.endIndex();
-    const LIFeatureComponents LIFC = getLIFeatureComponents(LI);
+    const LIFeatureComponents &LIFC = getLIFeatureComponents(LI);
     NrBrokenHints += VRM->hasPreferredPhys(LI.reg());
 
     NrDefsAndUses += LIFC.NrDefsAndUses;
     HottestBlockFreq = std::max(HottestBlockFreq, LIFC.HottestBlockFreq);
     R += LIFC.R;
     W += LIFC.W;
     RW += LIFC.RW;