[llvm] 8a24e0c - [nfc][mlgo] Lazily compute the regalloc reward

[Mircea Trofin via llvm-commits]

Author: Mircea Trofin
Date: 2022-09-26T15:34:29-07:00
New Revision: 8a24e0cb5aeecef63fb77830e437bcc116985af1

URL: https://github.com/llvm/llvm-project/commit/8a24e0cb5aeecef63fb77830e437bcc116985af1
DIFF: https://github.com/llvm/llvm-project/commit/8a24e0cb5aeecef63fb77830e437bcc116985af1.diff

LOG: [nfc][mlgo] Lazily compute the regalloc reward

Differential Revision: https://reviews.llvm.org/D134664

Added: 
    

Modified: 
    llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
    llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
    llvm/lib/CodeGen/RegAllocPriorityAdvisor.h

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
index 86c28ab0d309c..e48aac4136d8e 100644
--- a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
@@ -461,9 +461,10 @@ class DevelopmentModeEvictionAdvisorAnalysis final
     return I->second.get();
   }
 
-  void logRewardIfNeeded(const MachineFunction &MF, float Reward) override {
+  void logRewardIfNeeded(const MachineFunction &MF,
+                         llvm::function_ref<float()> GetReward) override {
     if (auto *Log = this->getLogger(MF))
-      Log->logFloatFinalReward(Reward);
+      Log->logFloatFinalReward(GetReward());
   }
 
 private:
@@ -1067,13 +1068,19 @@ int64_t DevelopmentModeEvictAdvisor::tryFindEvictionCandidatePosition(
 }
 
 bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
-  float Reward = static_cast<float>(
-      calculateRegAllocScore(MF, getAnalysis<MachineBlockFrequencyInfo>())
-          .getScore());
-
-  getAnalysis<RegAllocEvictionAdvisorAnalysis>().logRewardIfNeeded(MF, Reward);
-  getAnalysis<RegAllocPriorityAdvisorAnalysis>().logRewardIfNeeded(MF, Reward);
-
+  Optional<float> CachedReward;
+  auto GetReward = [&]() {
+    if (!CachedReward)
+      CachedReward = static_cast<float>(
+          calculateRegAllocScore(MF, getAnalysis<MachineBlockFrequencyInfo>())
+              .getScore());
+    return *CachedReward;
+  };
+
+  getAnalysis<RegAllocEvictionAdvisorAnalysis>().logRewardIfNeeded(MF,
+                                                                   GetReward);
+  getAnalysis<RegAllocPriorityAdvisorAnalysis>().logRewardIfNeeded(MF,
+                                                                   GetReward);
   return false;
 }
 #endif // #ifdef LLVM_HAVE_TF_API

diff  --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
index c7f48bc9ea9d7..a3936eacf560c 100644
--- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
+++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
@@ -177,7 +177,8 @@ class RegAllocEvictionAdvisorAnalysis : public ImmutablePass {
   virtual std::unique_ptr<RegAllocEvictionAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
   AdvisorMode getAdvisorMode() const { return Mode; }
-  virtual void logRewardIfNeeded(const MachineFunction &MF, float Reward){};
+  virtual void logRewardIfNeeded(const MachineFunction &MF,
+                                 llvm::function_ref<float()> GetReward){};
 
 protected:
   // This analysis preserves everything, and subclasses may have additional

diff  --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
index 520d8d836fe47..1e9fa967214cc 100644
--- a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
+++ b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
@@ -68,7 +68,8 @@ class RegAllocPriorityAdvisorAnalysis : public ImmutablePass {
   virtual std::unique_ptr<RegAllocPriorityAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
   AdvisorMode getAdvisorMode() const { return Mode; }
-  virtual void logRewardIfNeeded(const MachineFunction &MF, float Reward){};
+  virtual void logRewardIfNeeded(const MachineFunction &MF,
+                                 llvm::function_ref<float()> GetReward){};
 
 protected:
   // This analysis preserves everything, and subclasses may have additional