diff --git a/llvm/include/llvm/CodeGen/CalcSpillWeights.h b/llvm/include/llvm/CodeGen/CalcSpillWeights.h --- a/llvm/include/llvm/CodeGen/CalcSpillWeights.h +++ b/llvm/include/llvm/CodeGen/CalcSpillWeights.h @@ -63,7 +63,7 @@ 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. diff --git a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp --- a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp @@ -233,6 +233,27 @@ 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 &getCachedFeatures() { + return CachedFeatures; + } + +private: + void calculateSpillWeightAndHint(LiveInterval &LI) override { + VirtRegAuxInfo::calculateSpillWeightAndHint(LI); + CachedFeatures.erase(LI.reg().id()); + } + DenseMap CachedFeatures; +}; + using CandidateRegList = std::array, NumberOfInterferences>; using FeaturesListNormalizer = std::array; @@ -290,7 +311,7 @@ FixedRegisters); } - const LIFeatureComponents + const LIFeatureComponents & getLIFeatureComponents(const LiveInterval &LI) const; // Hold on to a default advisor for: @@ -306,6 +327,7 @@ // This could be static and shared, but its initialization is non-trivial. std::bitset DoNotNormalize; const float InitialQSize; + CachingVRAI *CVRAI = nullptr; }; // =================================== @@ -503,6 +525,9 @@ : RegAllocEvictionAdvisor(MF, RA), DefaultAdvisor(MF, RA), Runner(std::move(Runner)), MBFI(MBFI), Loops(Loops), InitialQSize(MLEvictAdvisor::getInitialQueueSize(MF)) { + auto V = std::make_unique(MF, *LIS, *VRM, Loops, MBFI); + CVRAI = V.get(); + VRAI = std::move(V); assert(this->Runner); DoNotNormalize.set(FeatureIDs::mask); DoNotNormalize.set(FeatureIDs::is_free); @@ -686,9 +711,15 @@ 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 Visited; const TargetRegisterInfo &TRI = *MF.getSubtarget().getRegisterInfo(); @@ -769,7 +800,7 @@ 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; diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h --- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h +++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h @@ -90,6 +90,7 @@ /// 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; @@ -114,6 +115,13 @@ /// 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); @@ -144,6 +152,9 @@ /// obtained from the TargetSubtargetInfo. const bool EnableLocalReassign; + /// Must be instantiated by the ctor of each derived class. + std::unique_ptr VRAI; + private: unsigned NextCascade = 1; }; @@ -200,8 +211,7 @@ // 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 &, diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp --- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp +++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp @@ -12,6 +12,7 @@ #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" @@ -122,3 +123,10 @@ EnableLocalReassign(EnableLocalReassignment || MF.getSubtarget().enableRALocalReassignment( MF.getTarget().getOptLevel())) {} + +DefaultEvictionAdvisor::DefaultEvictionAdvisor(MachineFunction &MF, + const RAGreedy &RA) + : RegAllocEvictionAdvisor(MF, RA) { + VRAI = std::make_unique(MF, *LIS, *VRM, *RA.getLoopInfo(), + *RA.getMBFI()); +} diff --git a/llvm/lib/CodeGen/RegAllocGreedy.h b/llvm/lib/CodeGen/RegAllocGreedy.h --- a/llvm/lib/CodeGen/RegAllocGreedy.h +++ b/llvm/lib/CodeGen/RegAllocGreedy.h @@ -157,6 +157,8 @@ 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: @@ -186,7 +188,7 @@ // state std::unique_ptr SpillerInstance; PQueue Queue; - std::unique_ptr VRAI; + VirtRegAuxInfo *VRAI = nullptr; Optional ExtraInfo; std::unique_ptr EvictAdvisor; diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp --- a/llvm/lib/CodeGen/RegAllocGreedy.cpp +++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp @@ -2926,7 +2926,7 @@ EvictAdvisor = getAnalysis().getAdvisor(*MF, *this); - VRAI = std::make_unique(*MF, *LIS, *VRM, *Loops, *MBFI); + VRAI = EvictAdvisor->getVirtRegAuxInfo(); SpillerInstance.reset(createInlineSpiller(*this, *MF, *VRM, *VRAI)); VRAI->calculateSpillWeightsAndHints();