diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp index 9c6487b40d606..51ef625335ea9 100644 --- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp @@ -63,11 +63,11 @@ static cl::opt InteractiveChannelBaseName( "outgoing name should be " ".out")); -static cl::opt - MaxCascade("mlregalloc-max-cascade", cl::Hidden, - cl::desc("The maximum number of times a live range can be " - "evicted before preventing it from being evicted"), - cl::init(20)); +static cl::opt MaxEvictionCount( + "mlregalloc-max-eviction-count", cl::Hidden, + cl::desc("The maximum number of times a live range can be " + "evicted before preventing it from being evicted"), + cl::init(100)); // Options that only make sense in development mode #ifdef LLVM_HAVE_TFLITE @@ -364,6 +364,22 @@ class MLEvictAdvisor : public RegAllocEvictionAdvisor { using RegID = unsigned; mutable DenseMap CachedFeatures; + + mutable std::unordered_map VirtRegEvictionCounts; + + void onEviction(Register RegBeingEvicted) const { + // If we cannot find the virtual register in the map, we just assume it has + // not been evicted before and thus has a value of zero (which is what the + // subscript operator returns by default). + ++VirtRegEvictionCounts[RegBeingEvicted.id()]; + } + + unsigned getEvictionCount(Register Reg) const { + auto EvictionCountIt = VirtRegEvictionCounts.find(Reg.id()); + if (EvictionCountIt != VirtRegEvictionCounts.end()) + return EvictionCountIt->second; + return 0; + } }; #define _DECL_FEATURES(type, name, shape, _) \ @@ -657,7 +673,7 @@ bool MLEvictAdvisor::loadInterferenceFeatures( // threshold, prevent the range from being evicted. We still let the // range through if it is urgent as we are required to produce an // eviction if the candidate is not spillable. - if (IntfCascade >= MaxCascade && !Urgent) + if (getEvictionCount(Intf->reg()) > MaxEvictionCount && !Urgent) return false; // Only evict older cascades or live ranges without a cascade. @@ -803,6 +819,22 @@ MCRegister MLEvictAdvisor::tryFindEvictionCandidate( } assert(CandidatePos < ValidPosLimit); (void)ValidPosLimit; + + // Update information about how many times the virtual registers being + // evicted have been evicted so that we can prevent the model from evicting + // the same ranges continually and eating compile time. + if (CandidatePos == CandidateVirtRegPos) { + onEviction(VirtReg.reg()); + } else { + for (MCRegUnit Unit : TRI->regunits(Regs[CandidatePos].first)) { + LiveIntervalUnion::Query &Q = Matrix->query(VirtReg, Unit); + const auto &IFIntervals = Q.interferingVRegs(EvictInterferenceCutoff); + for (const LiveInterval *Intf : reverse(IFIntervals)) { + onEviction(Intf->reg()); + } + } + } + return Regs[CandidatePos].first; }