[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

250–253

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

250–253

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/include/llvm/CodeGen/CalcSpillWeights.h

public:
     VirtRegAuxInfo(MachineFunction &MF, LiveIntervals &LIS,
                    const VirtRegMap &VRM, const MachineLoopInfo &Loops,
                    const MachineBlockFrequencyInfo &MBFI)
         : MF(MF), LIS(LIS), VRM(VRM), Loops(Loops), MBFI(MBFI) {}
 
     virtual ~VirtRegAuxInfo() = default;
 
     /// (re)compute li's spill weight and allocation hint.
-    void calculateSpillWeightAndHint(LiveInterval &LI);
+    virtual void calculateSpillWeightAndHint(LiveInterval &LI);
 
     /// Compute future expected spill weight of a split artifact of LI
     /// that will span between start and end slot indexes.
     /// \param LI     The live interval to be split.
     /// \param Start  The expected beginning of the split artifact. Instructions
     ///               before start will not affect the weight.
     /// \param End    The expected end of the split artifact. Instructions
     ///               after end will not affect the weight.

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

struct LIFeatureComponents {
   double RW = 0;
   double IndVarUpdates = 0;
   double HintWeights = 0.0;
   int64_t NrDefsAndUses = 0;
   float HottestBlockFreq = 0.0;
   bool IsRemat = false;
 };
 
+class CachingVRAI : public VirtRegAuxInfo {
+public:
+  CachingVRAI(MachineFunction &MF, LiveIntervals &LIS, const VirtRegMap &VRM,
+              const MachineLoopInfo &Loops,
+              const MachineBlockFrequencyInfo &MBFI)
+      : VirtRegAuxInfo(MF, LIS, VRM, Loops, MBFI) {}
+
+  using RegID = unsigned;
+
+  DenseMap<RegID, LIFeatureComponents> &getCachedFeatures() {
+    return CachedFeatures;
+  }
+
+private:
+  void calculateSpillWeightAndHint(LiveInterval &LI) override {
+    VirtRegAuxInfo::calculateSpillWeightAndHint(LI);
+    CachedFeatures.erase(LI.reg().id());
+  }
+  DenseMap<RegID, LIFeatureComponents> CachedFeatures;
+};
+
 using CandidateRegList =
     std::array<std::pair<MCRegister, bool>, NumberOfInterferences>;
 using FeaturesListNormalizer = std::array<float, FeatureIDs::FeatureCount>;
 
 /// The ML evictor (commonalities between release and development mode)
 class MLEvictAdvisor : public RegAllocEvictionAdvisor {
 public:
   MLEvictAdvisor(MachineFunction &MF, const RAGreedy &RA, MLModelRunner *Runner,
private:
   // Point-in-time: we didn't learn this, so we always delegate to the default.
   bool canEvictHintInterference(
       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;
+  CachingVRAI *CVRAI = nullptr;
 };
 
 // ===================================
 // Release (AOT) - specifics
 // ===================================
 #if defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL)
 const std::array<std::string, FeatureIDs::FeatureCount> FeatureNames{
 #define _GETNAME(_, NAME, __, ___) #NAME,

 
 MLEvictAdvisor::MLEvictAdvisor(MachineFunction &MF, const RAGreedy &RA,
                                MLModelRunner *Runner,
                                const MachineBlockFrequencyInfo &MBFI,
                                const MachineLoopInfo &Loops)
     : RegAllocEvictionAdvisor(MF, RA), DefaultAdvisor(MF, RA),
       Runner(std::move(Runner)), MBFI(MBFI), Loops(Loops),
       InitialQSize(MLEvictAdvisor::getInitialQueueSize(MF)) {
+  auto V = std::make_unique<CachingVRAI>(MF, *LIS, *VRM, Loops, MBFI);
+  CVRAI = V.get();
+  VRAI = std::move(V);
   assert(this->Runner);
   DoNotNormalize.set(FeatureIDs::mask);
   DoNotNormalize.set(FeatureIDs::is_free);
   DoNotNormalize.set(FeatureIDs::is_hint);
   DoNotNormalize.set(FeatureIDs::is_local);
   DoNotNormalize.set(FeatureIDs::min_stage);
   DoNotNormalize.set(FeatureIDs::max_stage);
   DoNotNormalize.set(FeatureIDs::progress);
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;
+  CachingVRAI::RegID ID = LI.reg().id();
+  LIFeatureComponents Empty;
+  auto I = CVRAI->getCachedFeatures().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;
 

llvm/lib/CodeGen/RegAllocEvictionAdvisor.h

bool operator<(const EvictionCost &O) const {
     return std::tie(BrokenHints, MaxWeight) <
            std::tie(O.BrokenHints, O.MaxWeight);
   }
 };
 
 /// Interface to the eviction advisor, which is responsible for making a
 /// decision as to which live ranges should be evicted (if any).
 class RAGreedy;
+class VirtRegAuxInfo;
 class RegAllocEvictionAdvisor {
 public:
   RegAllocEvictionAdvisor(const RegAllocEvictionAdvisor &) = delete;
   RegAllocEvictionAdvisor(RegAllocEvictionAdvisor &&) = delete;
   virtual ~RegAllocEvictionAdvisor() = default;
 
   /// Find a physical register that can be freed by evicting the FixedRegisters,
   /// or return NoRegister. The eviction decision is assumed to be correct (i.e.
   /// no fixed live ranges are evicted) and profitable.
   virtual MCRegister
   tryFindEvictionCandidate(LiveInterval &VirtReg, const AllocationOrder &Order,
                            uint8_t CostPerUseLimit,
                            const SmallVirtRegSet &FixedRegisters) const = 0;
 
   /// Find out if we can evict the live ranges occupying the given PhysReg,
   /// which is a hint (preferred register) for VirtReg.
   virtual bool
   canEvictHintInterference(LiveInterval &VirtReg, MCRegister PhysReg,
                            const SmallVirtRegSet &FixedRegisters) const = 0;
 
   /// Returns true if the given \p PhysReg is a callee saved register and has
   /// not been used for allocation yet.
   bool isUnusedCalleeSavedReg(MCRegister PhysReg) const;
 
+  /// Get the live range aux info object specific to this advisor.
+  VirtRegAuxInfo *getVirtRegAuxInfo() {
+    assert(VRAI && "Each RegAllocEvictionAdvisor specialization should have "
+                   "constructed a VRAI");
+    return VRAI.get();
+  }
+
 protected:
   RegAllocEvictionAdvisor(MachineFunction &MF, const RAGreedy &RA);
 
   Register canReassign(LiveInterval &VirtReg, Register PrevReg) const;
 
   // Get the upper limit of elements in the given Order we need to analize.
   // TODO: is this heuristic,  we could consider learning it.
   Optional<unsigned> getOrderLimit(const LiveInterval &VirtReg,
protected:
   const TargetRegisterInfo *const TRI;
   const RegisterClassInfo &RegClassInfo;
   const ArrayRef<uint8_t> RegCosts;
 
   /// Run or not the local reassignment heuristic. This information is
   /// obtained from the TargetSubtargetInfo.
   const bool EnableLocalReassign;
 
+  /// Must be instantiated by the ctor of each derived class.
+  std::unique_ptr<VirtRegAuxInfo> VRAI;
+
 private:
   unsigned NextCascade = 1;
 };
 
 /// ImmutableAnalysis abstraction for fetching the Eviction Advisor. We model it
 /// as an analysis to decouple the user from the implementation insofar as
 /// dependencies on other analyses goes. The motivation for it being an
 /// immutable pass is twofold:

 RegAllocEvictionAdvisorAnalysis *createReleaseModeAdvisor();
 
 RegAllocEvictionAdvisorAnalysis *createDevelopmentModeAdvisor();
 
 // TODO: move to RegAllocEvictionAdvisor.cpp when we move implementation
 // out of RegAllocGreedy.cpp
 class DefaultEvictionAdvisor : public RegAllocEvictionAdvisor {
 public:
-  DefaultEvictionAdvisor(MachineFunction &MF, const RAGreedy &RA)
-      : RegAllocEvictionAdvisor(MF, RA) {}
+  DefaultEvictionAdvisor(MachineFunction &MF, const RAGreedy &RA);
 
 private:
   MCRegister tryFindEvictionCandidate(LiveInterval &, const AllocationOrder &,
                                       uint8_t,
                                       const SmallVirtRegSet &) const override;
   bool canEvictHintInterference(LiveInterval &, MCRegister,
                                 const SmallVirtRegSet &) const override;
   bool canEvictInterferenceBasedOnCost(LiveInterval &, MCRegister, bool,
                                        EvictionCost &,
                                        const SmallVirtRegSet &) const;
   bool shouldEvict(LiveInterval &A, bool, LiveInterval &B, bool) const;
 };
 } // namespace llvm
 
 #endif // LLVM_CODEGEN_REGALLOCEVICTIONADVISOR_H

llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp

 //===- RegAllocEvictionAdvisor.cpp - eviction advisor ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Implementation of the default eviction advisor and of the Analysis pass.
 //
 //===----------------------------------------------------------------------===//
 
 #include "RegAllocEvictionAdvisor.h"
 #include "RegAllocGreedy.h"
+#include "llvm/CodeGen/CalcSpillWeights.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
 #include "llvm/CodeGen/VirtRegMap.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/PassRegistry.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ErrorHandling.h"
RegAllocEvictionAdvisor::RegAllocEvictionAdvisor(MachineFunction &MF,
                                                  const RAGreedy &RA)
     : MF(MF), RA(RA), Matrix(RA.getInterferenceMatrix()),
       LIS(RA.getLiveIntervals()), VRM(RA.getVirtRegMap()),
       MRI(&VRM->getRegInfo()), TRI(MF.getSubtarget().getRegisterInfo()),
       RegClassInfo(RA.getRegClassInfo()), RegCosts(TRI->getRegisterCosts(MF)),
       EnableLocalReassign(EnableLocalReassignment ||
                           MF.getSubtarget().enableRALocalReassignment(
                               MF.getTarget().getOptLevel())) {}
+
+DefaultEvictionAdvisor::DefaultEvictionAdvisor(MachineFunction &MF,
+                                               const RAGreedy &RA)
+    : RegAllocEvictionAdvisor(MF, RA) {
+  VRAI = std::make_unique<VirtRegAuxInfo>(MF, *LIS, *VRM, *RA.getLoopInfo(),
+                                          *RA.getMBFI());
+}

llvm/lib/CodeGen/RegAllocGreedy.h

public:
   };
 
   LiveRegMatrix *getInterferenceMatrix() const { return Matrix; }
   LiveIntervals *getLiveIntervals() const { return LIS; }
   VirtRegMap *getVirtRegMap() const { return VRM; }
   const RegisterClassInfo &getRegClassInfo() const { return RegClassInfo; }
   const ExtraRegInfo &getExtraInfo() const { return *ExtraInfo; }
   size_t getQueueSize() const { return Queue.size(); }
+  MachineLoopInfo *getLoopInfo() const { return Loops; }
+  MachineBlockFrequencyInfo *getMBFI() const { return MBFI; }
   // end (interface to eviction advisers)
 
 private:
   // Convenient shortcuts.
   using PQueue = std::priority_queue<std::pair<unsigned, unsigned>>;
   using SmallLISet = SmallPtrSet<LiveInterval *, 4>;
 
   // context
private:
   EdgeBundles *Bundles;
   SpillPlacement *SpillPlacer;
   LiveDebugVariables *DebugVars;
   AliasAnalysis *AA;
 
   // state
   std::unique_ptr<Spiller> SpillerInstance;
   PQueue Queue;
-  std::unique_ptr<VirtRegAuxInfo> VRAI;
+  VirtRegAuxInfo *VRAI = nullptr;
   Optional<ExtraRegInfo> ExtraInfo;
   std::unique_ptr<RegAllocEvictionAdvisor> EvictAdvisor;
 
   // 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

llvm/lib/CodeGen/RegAllocGreedy.cpp

bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
   initializeCSRCost();
 
   RegCosts = TRI->getRegisterCosts(*MF);
 
   ExtraInfo.emplace();
   EvictAdvisor =
       getAnalysis<RegAllocEvictionAdvisorAnalysis>().getAdvisor(*MF, *this);
 
-  VRAI = std::make_unique<VirtRegAuxInfo>(*MF, *LIS, *VRM, *Loops, *MBFI);
+  VRAI = EvictAdvisor->getVirtRegAuxInfo();
   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));