[PATCH] D134232: [NFC][MLGO] Introduce logRewardIfNeeded method

[Eric Wang via Phabricator via llvm-commits]

Flpha0830 created this revision.
Herald added subscribers: mtrofin, hiraditya, qcolombet, MatzeB.
Herald added a project: All.
Flpha0830 requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

This patch introduces a logRewardIfNeeded method to reuse regallocscoring.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D134232

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


Index: llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
===================================================================
--- llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
+++ llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
@@ -68,6 +68,7 @@
   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){};
 
 protected:
   // This analysis preserves everything, and subclasses may have additional
Index: llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
===================================================================
--- llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
+++ llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
@@ -177,6 +177,7 @@
   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){};
 
 protected:
   // This analysis preserves everything, and subclasses may have additional
Index: llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
===================================================================
--- llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
+++ llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
@@ -90,6 +90,7 @@
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
     AU.addRequired<RegAllocEvictionAdvisorAnalysis>();
+    AU.addRequired<RegAllocPriorityAdvisorAnalysis>();
     AU.addRequired<MachineBlockFrequencyInfo>();
     MachineFunctionPass::getAnalysisUsage(AU);
   }
@@ -420,6 +421,11 @@
     return I->second.get();
   }
 
+  void logRewardIfNeeded(const MachineFunction &MF, float Reward) override {
+    if (auto *Log = this->getLogger(MF))
+      Log->logFloatFinalReward(Reward);
+  }
+
 private:
   std::vector<TensorSpec> InputFeatures;
   std::vector<TensorSpec> TrainingInputFeatures;
@@ -890,12 +896,12 @@
 }
 
 bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
-  if (auto *DevModeAnalysis = dyn_cast<DevelopmentModeEvictionAdvisorAnalysis>(
-          &getAnalysis<RegAllocEvictionAdvisorAnalysis>()))
-    if (auto *Log = DevModeAnalysis->getLogger(MF))
-      Log->logFloatFinalReward(static_cast<float>(
-          calculateRegAllocScore(MF, getAnalysis<MachineBlockFrequencyInfo>())
-              .getScore()));
+  float Reward = static_cast<float>(
+      calculateRegAllocScore(MF, getAnalysis<MachineBlockFrequencyInfo>())
+          .getScore());
+
+  getAnalysis<RegAllocEvictionAdvisorAnalysis>().logRewardIfNeeded(MF, Reward);
+  getAnalysis<RegAllocPriorityAdvisorAnalysis>().logRewardIfNeeded(MF, Reward);
 
   return false;
 }


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D134232.461392.patch
Type: text/x-patch
Size: 2785 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20220919/80be0319/attachment.bin>