[nfc][mlgo] Lazily compute the regalloc reward

[jacobhegna]

so the point of this is just to avoid recomputing the reward for the two different regalloc models? the reward will always be computed the same way for both models?

[mtrofin]

In D134664#3816089, @jacobhegna wrote:

so the point of this is just to avoid recomputing the reward for the two different regalloc models? the reward will always be computed the same way for both models?

the point is to avoid computing if noone needs it - which is always the case except when training.

[jacobhegna]

In D134664#3816107, @mtrofin wrote:

In D134664#3816089, @jacobhegna wrote:

so the point of this is just to avoid recomputing the reward for the two different regalloc models? the reward will always be computed the same way for both models?

the point is to avoid computing if noone needs it - which is always the case except when training.

ah ah ok, makes sense

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

public:
   /// one. This is used to inject the score by the RegAllocScoring pass.
   Logger *getLogger(const MachineFunction &MF) const {
     auto I = LogMap.find(MF.getName());
     if (I == LogMap.end())
       return nullptr;
     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:
   std::vector<TensorSpec> InputFeatures;
   std::vector<TensorSpec> TrainingInputFeatures;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<MachineBlockFrequencyInfo>();
for (size_t I = 1; I < MUTR->outputLoggedFeatureSpecs().size();
           reinterpret_cast<const char *>(
               MUTR->lastEvaluationResult()->getUntypedTensorValue(I)));
   // The output is right after the features and the extra outputs
   Log->logInt64Value(CurrentFeature, &Ret);
   return Ret;
 }
 
 bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
-  float Reward = static_cast<float>(
+  Optional<float> CachedReward;
+  auto GetReward = [&]() {
+    if (!CachedReward)
+      CachedReward = static_cast<float>(
           calculateRegAllocScore(MF, getAnalysis<MachineBlockFrequencyInfo>())
               .getScore());
+    return *CachedReward;
+  };
 
-  getAnalysis<RegAllocEvictionAdvisorAnalysis>().logRewardIfNeeded(MF, Reward);
-  getAnalysis<RegAllocPriorityAdvisorAnalysis>().logRewardIfNeeded(MF, Reward);
-
+  getAnalysis<RegAllocEvictionAdvisorAnalysis>().logRewardIfNeeded(MF,
+                                                                   GetReward);
+  getAnalysis<RegAllocPriorityAdvisorAnalysis>().logRewardIfNeeded(MF,
+                                                                   GetReward);
   return false;
 }
 #endif // #ifdef LLVM_HAVE_TF_API
 
 RegAllocEvictionAdvisorAnalysis *llvm::createReleaseModeAdvisor() {
   return new ReleaseModeEvictionAdvisorAnalysis();
 }
 
 // In all cases except development mode, we don't need scoring.
 #if !defined(LLVM_HAVE_TF_API)
 bool RegAllocScoring::runOnMachineFunction(MachineFunction &) { return false; }
 #endif

llvm/lib/CodeGen/RegAllocEvictionAdvisor.h

public:
   RegAllocEvictionAdvisorAnalysis(AdvisorMode Mode)
       : ImmutablePass(ID), Mode(Mode){};
   static char ID;
 
   /// Get an advisor for the given context (i.e. machine function, etc)
   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
   // requirements.
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
   }
 

llvm/lib/CodeGen/RegAllocPriorityAdvisor.h

public:
   RegAllocPriorityAdvisorAnalysis(AdvisorMode Mode)
       : ImmutablePass(ID), Mode(Mode){};
   static char ID;
 
   /// Get an advisor for the given context (i.e. machine function, etc)
   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
   // requirements.
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
   }