From 5947599705a531b46ff800e46ee64e96c9119450 Mon Sep 17 00:00:00 2001 From: Aiden Grossman Date: Sun, 26 Jan 2025 03:09:21 +0000 Subject: [PATCH 1/5] [MLGO] Count LR Evictions Rather than Relying on Cascade This patch adjusts the mlregalloc-max-cascade flag (renaming it to mlregalloc-max-eviction-count) to actually count evictions rather than just looking at the cascade number. The cascade number is not very representative of how many times a LR has been evicted, which can lead to some problems in certain cases, where we might end up with many eviction problems where we have now masked off all the interferences and are forced to evict the candidate. This is probably what I should've done in the first place. No test case as this only shows up in quite large functions post ThinLTO and it would be hard to construct something that would serve as a nice regression test without being super brittle. I've tested this on the pathological cases that we have come across so far and it works. Fixes #122829 --- llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 29 ++++++++++++++++----- 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp index 9c6487b40d606..e7cdcf9f0d3c0 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,8 @@ class MLEvictAdvisor : public RegAllocEvictionAdvisor { using RegID = unsigned; mutable DenseMap CachedFeatures; + + mutable std::unordered_map VirtRegEvictionCounts; }; #define _DECL_FEATURES(type, name, shape, _) \ @@ -657,7 +659,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 (VirtRegEvictionCounts[Intf->reg().id()] > MaxEvictionCount && !Urgent) return false; // Only evict older cascades or live ranges without a cascade. @@ -803,6 +805,21 @@ MCRegister MLEvictAdvisor::tryFindEvictionCandidate( } assert(CandidatePos < ValidPosLimit); (void)ValidPosLimit; + + // Update information about how many times the virtual registers being + // evicted have been evicted. + if (CandidatePos == CandidateVirtRegPos) { + VirtRegEvictionCounts[VirtReg.reg()] += 1; + } 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)) { + VirtRegEvictionCounts[Intf->reg()] += 1; + } + } + } + return Regs[CandidatePos].first; } From dc4a7c962f575c3ee310d3edb0ad20d0ddcea896 Mon Sep 17 00:00:00 2001 From: Aiden Grossman Date: Mon, 27 Jan 2025 22:37:11 +0000 Subject: [PATCH 2/5] address feedback --- llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 39 ++++++++++++++------- 1 file changed, 26 insertions(+), 13 deletions(-) diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp index e7cdcf9f0d3c0..340a73da71531 100644 --- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp @@ -366,6 +366,9 @@ class MLEvictAdvisor : public RegAllocEvictionAdvisor { mutable DenseMap CachedFeatures; mutable std::unordered_map VirtRegEvictionCounts; + void incrementEvictionCount(const CandidateRegList &Regs, + const LiveInterval &VirtReg, + size_t CandidatePosition) const; }; #define _DECL_FEATURES(type, name, shape, _) \ @@ -658,7 +661,10 @@ bool MLEvictAdvisor::loadInterferenceFeatures( // large amount of compile time being spent in regalloc. If we hit the // 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. + // eviction if the candidate is not spillable. 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). if (VirtRegEvictionCounts[Intf->reg().id()] > MaxEvictionCount && !Urgent) return false; @@ -680,6 +686,22 @@ bool MLEvictAdvisor::loadInterferenceFeatures( return true; } +void MLEvictAdvisor::incrementEvictionCount(const CandidateRegList &Regs, + const LiveInterval &VirtReg, + size_t CandidatePosition) const { + if (CandidatePosition == CandidateVirtRegPos) { + VirtRegEvictionCounts[VirtReg.reg()] += 1; + } else { + for (MCRegUnit Unit : TRI->regunits(Regs[CandidatePosition].first)) { + LiveIntervalUnion::Query &Q = Matrix->query(VirtReg, Unit); + const auto &IFIntervals = Q.interferingVRegs(EvictInterferenceCutoff); + for (const LiveInterval *Intf : reverse(IFIntervals)) { + VirtRegEvictionCounts[Intf->reg()] += 1; + } + } + } +} + MCRegister MLEvictAdvisor::tryFindEvictionCandidate( const LiveInterval &VirtReg, const AllocationOrder &Order, uint8_t CostPerUseLimit, const SmallVirtRegSet &FixedRegisters) const { @@ -807,18 +829,9 @@ MCRegister MLEvictAdvisor::tryFindEvictionCandidate( (void)ValidPosLimit; // Update information about how many times the virtual registers being - // evicted have been evicted. - if (CandidatePos == CandidateVirtRegPos) { - VirtRegEvictionCounts[VirtReg.reg()] += 1; - } 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)) { - VirtRegEvictionCounts[Intf->reg()] += 1; - } - } - } + // evicted have been evicted so that we can prevent the model from evicting + // the same ranges continually and eating compile time. + incrementEvictionCount(Regs, VirtReg, CandidatePos); return Regs[CandidatePos].first; } From ff42e631e49efea69d553849763aaab1cc94eb45 Mon Sep 17 00:00:00 2001 From: Aiden Grossman Date: Mon, 27 Jan 2025 23:14:42 +0000 Subject: [PATCH 3/5] Address feedback --- llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 53 +++++++++++---------- 1 file changed, 28 insertions(+), 25 deletions(-) diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp index 340a73da71531..a689abb7acd6c 100644 --- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp @@ -366,9 +366,21 @@ class MLEvictAdvisor : public RegAllocEvictionAdvisor { mutable DenseMap CachedFeatures; mutable std::unordered_map VirtRegEvictionCounts; - void incrementEvictionCount(const CandidateRegList &Regs, - const LiveInterval &VirtReg, - size_t CandidatePosition) const; + + 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, _) \ @@ -661,11 +673,8 @@ bool MLEvictAdvisor::loadInterferenceFeatures( // large amount of compile time being spent in regalloc. If we hit the // 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 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). - if (VirtRegEvictionCounts[Intf->reg().id()] > MaxEvictionCount && !Urgent) + // eviction if the candidate is not spillable. + if (getEvictionCount(Intf->reg()) > MaxEvictionCount && !Urgent) return false; // Only evict older cascades or live ranges without a cascade. @@ -686,22 +695,6 @@ bool MLEvictAdvisor::loadInterferenceFeatures( return true; } -void MLEvictAdvisor::incrementEvictionCount(const CandidateRegList &Regs, - const LiveInterval &VirtReg, - size_t CandidatePosition) const { - if (CandidatePosition == CandidateVirtRegPos) { - VirtRegEvictionCounts[VirtReg.reg()] += 1; - } else { - for (MCRegUnit Unit : TRI->regunits(Regs[CandidatePosition].first)) { - LiveIntervalUnion::Query &Q = Matrix->query(VirtReg, Unit); - const auto &IFIntervals = Q.interferingVRegs(EvictInterferenceCutoff); - for (const LiveInterval *Intf : reverse(IFIntervals)) { - VirtRegEvictionCounts[Intf->reg()] += 1; - } - } - } -} - MCRegister MLEvictAdvisor::tryFindEvictionCandidate( const LiveInterval &VirtReg, const AllocationOrder &Order, uint8_t CostPerUseLimit, const SmallVirtRegSet &FixedRegisters) const { @@ -831,7 +824,17 @@ MCRegister MLEvictAdvisor::tryFindEvictionCandidate( // 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. - incrementEvictionCount(Regs, VirtReg, CandidatePos); + 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; } From 462fcfcc89f25301ae814b368fd3d4f9ce2a6657 Mon Sep 17 00:00:00 2001 From: Aiden Grossman Date: Mon, 27 Jan 2025 23:18:16 +0000 Subject: [PATCH 4/5] formatting --- llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp index a689abb7acd6c..c79127d6b8c9b 100644 --- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp @@ -673,7 +673,7 @@ bool MLEvictAdvisor::loadInterferenceFeatures( // large amount of compile time being spent in regalloc. If we hit the // 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. + // eviction if the candidate is not spillable. if (getEvictionCount(Intf->reg()) > MaxEvictionCount && !Urgent) return false; From 48bc71bd94561e8e14c28072bc232df75526349f Mon Sep 17 00:00:00 2001 From: Aiden Grossman Date: Mon, 27 Jan 2025 23:20:07 +0000 Subject: [PATCH 5/5] feedback --- llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp index c79127d6b8c9b..51ef625335ea9 100644 --- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp @@ -376,9 +376,8 @@ class MLEvictAdvisor : public RegAllocEvictionAdvisor { unsigned getEvictionCount(Register Reg) const { auto EvictionCountIt = VirtRegEvictionCounts.find(Reg.id()); - if (EvictionCountIt != VirtRegEvictionCounts.end()) { + if (EvictionCountIt != VirtRegEvictionCounts.end()) return EvictionCountIt->second; - } return 0; } };