[llvm-branch-commits] [llvm] [CodeGen][NewPM] Port RegAllocPriorityAdvisor analysis to NPM (PR #118462)

Akshat Oke via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Sat Feb 1 10:24:29 PST 2025


https://github.com/optimisan updated https://github.com/llvm/llvm-project/pull/118462

>From a5c1a65dd4e7b3d77a7cc74febe17437144db46e Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Tue, 3 Dec 2024 10:12:36 +0000
Subject: [PATCH 1/6] [CodeGen][NewPM] Port RegAllocPriorityAdvisor analysis to
 NPM

---
 .../llvm}/CodeGen/RegAllocPriorityAdvisor.h   |  78 +++++++-
 llvm/include/llvm/InitializePasses.h          |   2 +-
 .../llvm/Passes/MachinePassRegistry.def       |   1 +
 llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp   |   6 +-
 .../lib/CodeGen/MLRegAllocPriorityAdvisor.cpp | 184 +++++++++++-------
 llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp  |   2 +-
 llvm/lib/CodeGen/RegAllocGreedy.cpp           |   9 +-
 llvm/lib/CodeGen/RegAllocGreedy.h             |   2 +-
 llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp  | 155 +++++++++++----
 llvm/lib/Passes/PassBuilder.cpp               |   1 +
 10 files changed, 320 insertions(+), 120 deletions(-)
 rename llvm/{lib => include/llvm}/CodeGen/RegAllocPriorityAdvisor.h (57%)

diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
similarity index 57%
rename from llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
rename to llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
index 0758743c2b1403e..a53739fdc3fc40b 100644
--- a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
+++ b/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
@@ -9,8 +9,10 @@
 #ifndef LLVM_CODEGEN_REGALLOCPRIORITYADVISOR_H
 #define LLVM_CODEGEN_REGALLOCPRIORITYADVISOR_H
 
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
 #include "llvm/CodeGen/SlotIndexes.h"
+#include "llvm/IR/PassManager.h"
 #include "llvm/Pass.h"
 
 namespace llvm {
@@ -68,12 +70,72 @@ class DummyPriorityAdvisor : public RegAllocPriorityAdvisor {
   unsigned getPriority(const LiveInterval &LI) const override;
 };
 
-class RegAllocPriorityAdvisorAnalysis : public ImmutablePass {
+/// Common provider for getting the priority advisor and logging rewards.
+/// Legacy analysis forwards all calls to this provider.
+/// New analysis serves the provider as the analysis result.
+/// Expensive setup is done in the constructor, so that the advisor can be
+/// created quickly for every machine function.
+/// TODO: Remove once legacy PM support is dropped.
+class RegAllocPriorityAdvisorProvider {
 public:
   enum class AdvisorMode : int { Default, Release, Development, Dummy };
 
-  RegAllocPriorityAdvisorAnalysis(AdvisorMode Mode)
-      : ImmutablePass(ID), Mode(Mode){};
+  RegAllocPriorityAdvisorProvider(AdvisorMode Mode) : Mode(Mode) {}
+
+  virtual ~RegAllocPriorityAdvisorProvider() = default;
+
+  virtual void logRewardIfNeeded(const MachineFunction &MF,
+                                 llvm::function_ref<float()> GetReward) {};
+
+  virtual std::unique_ptr<RegAllocPriorityAdvisor>
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
+
+  void setAnalyses(SlotIndexes *SI) { this->SI = SI; }
+
+  AdvisorMode getAdvisorMode() const { return Mode; }
+
+protected:
+  SlotIndexes *SI;
+
+private:
+  const AdvisorMode Mode;
+};
+
+RegAllocPriorityAdvisorProvider *createReleaseModePriorityAdvisorProvider();
+
+RegAllocPriorityAdvisorProvider *
+createDevelopmentModePriorityAdvisorProvider(LLVMContext &Ctx);
+
+class RegAllocPriorityAdvisorAnalysis
+    : public AnalysisInfoMixin<RegAllocPriorityAdvisorAnalysis> {
+  static AnalysisKey Key;
+  friend AnalysisInfoMixin<RegAllocPriorityAdvisorAnalysis>;
+
+public:
+  struct Result {
+    // Owned by this analysis.
+    RegAllocPriorityAdvisorProvider *Provider;
+
+    bool invalidate(MachineFunction &MF, const PreservedAnalyses &PA,
+                    MachineFunctionAnalysisManager::Invalidator &Inv) {
+      auto PAC = PA.getChecker<RegAllocPriorityAdvisorAnalysis>();
+      return !PAC.preservedWhenStateless() ||
+             Inv.invalidate<SlotIndexesAnalysis>(MF, PA);
+    }
+  };
+
+  Result run(MachineFunction &MF, MachineFunctionAnalysisManager &MFAM);
+
+private:
+  void initializeProvider(LLVMContext &Ctx);
+  std::unique_ptr<RegAllocPriorityAdvisorProvider> Provider;
+};
+
+class RegAllocPriorityAdvisorAnalysisLegacy : public ImmutablePass {
+public:
+  using AdvisorMode = RegAllocPriorityAdvisorProvider::AdvisorMode;
+  RegAllocPriorityAdvisorAnalysisLegacy(AdvisorMode Mode)
+      : ImmutablePass(ID), Mode(Mode) {};
   static char ID;
 
   /// Get an advisor for the given context (i.e. machine function, etc)
@@ -81,7 +143,7 @@ class RegAllocPriorityAdvisorAnalysis : public ImmutablePass {
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
   AdvisorMode getAdvisorMode() const { return Mode; }
   virtual void logRewardIfNeeded(const MachineFunction &MF,
-                                 llvm::function_ref<float()> GetReward){};
+                                 llvm::function_ref<float()> GetReward) {};
 
 protected:
   // This analysis preserves everything, and subclasses may have additional
@@ -97,11 +159,13 @@ class RegAllocPriorityAdvisorAnalysis : public ImmutablePass {
 
 /// Specialization for the API used by the analysis infrastructure to create
 /// an instance of the priority advisor.
-template <> Pass *callDefaultCtor<RegAllocPriorityAdvisorAnalysis>();
+template <> Pass *callDefaultCtor<RegAllocPriorityAdvisorAnalysisLegacy>();
 
-RegAllocPriorityAdvisorAnalysis *createReleaseModePriorityAdvisor();
+RegAllocPriorityAdvisorAnalysisLegacy *
+createReleaseModePriorityAdvisorAnalysis();
 
-RegAllocPriorityAdvisorAnalysis *createDevelopmentModePriorityAdvisor();
+RegAllocPriorityAdvisorAnalysisLegacy *
+createDevelopmentModePriorityAdvisorAnalysis();
 
 } // namespace llvm
 
diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
index 763ebdfab0693d3..1bbfaedc9f992dd 100644
--- a/llvm/include/llvm/InitializePasses.h
+++ b/llvm/include/llvm/InitializePasses.h
@@ -255,7 +255,7 @@ void initializeReachingDefAnalysisPass(PassRegistry &);
 void initializeReassociateLegacyPassPass(PassRegistry &);
 void initializeRegAllocEvictionAdvisorAnalysisLegacyPass(PassRegistry &);
 void initializeRegAllocFastPass(PassRegistry &);
-void initializeRegAllocPriorityAdvisorAnalysisPass(PassRegistry &);
+void initializeRegAllocPriorityAdvisorAnalysisLegacyPass(PassRegistry &);
 void initializeRegAllocScoringPass(PassRegistry &);
 void initializeRegBankSelectPass(PassRegistry &);
 void initializeRegToMemWrapperPassPass(PassRegistry &);
diff --git a/llvm/include/llvm/Passes/MachinePassRegistry.def b/llvm/include/llvm/Passes/MachinePassRegistry.def
index 08ced7e4558a60b..509895e956f6afc 100644
--- a/llvm/include/llvm/Passes/MachinePassRegistry.def
+++ b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -115,6 +115,7 @@ MACHINE_FUNCTION_ANALYSIS("machine-post-dom-tree",
 MACHINE_FUNCTION_ANALYSIS("machine-trace-metrics", MachineTraceMetricsAnalysis())
 MACHINE_FUNCTION_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis(PIC))
 MACHINE_FUNCTION_ANALYSIS("regalloc-evict", RegAllocEvictionAdvisorAnalysis())
+MACHINE_FUNCTION_ANALYSIS("regalloc-priority", RegAllocPriorityAdvisorAnalysis())
 MACHINE_FUNCTION_ANALYSIS("slot-indexes", SlotIndexesAnalysis())
 MACHINE_FUNCTION_ANALYSIS("spill-code-placement", SpillPlacementAnalysis())
 MACHINE_FUNCTION_ANALYSIS("virtregmap", VirtRegMapAnalysis())
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 1a8e11de909e894..8215d07b74dba39 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -116,7 +116,7 @@ class RegAllocScoring : public MachineFunctionPass {
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
     AU.addRequired<RegAllocEvictionAdvisorAnalysisLegacy>();
-    AU.addRequired<RegAllocPriorityAdvisorAnalysis>();
+    AU.addRequired<RegAllocPriorityAdvisorAnalysisLegacy>();
     AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
     MachineFunctionPass::getAnalysisUsage(AU);
   }
@@ -1242,8 +1242,8 @@ bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
 
   getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().logRewardIfNeeded(
       MF, GetReward);
-  getAnalysis<RegAllocPriorityAdvisorAnalysis>().logRewardIfNeeded(MF,
-                                                                   GetReward);
+  getAnalysis<RegAllocPriorityAdvisorAnalysisLegacy>().logRewardIfNeeded(
+      MF, GetReward);
   return false;
 }
 #endif // #ifdef LLVM_HAVE_TFLITE
diff --git a/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
index 9638df81770c1ef..dfc06eaa0bfe694 100644
--- a/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
@@ -12,7 +12,6 @@
 
 #include "AllocationOrder.h"
 #include "RegAllocGreedy.h"
-#include "RegAllocPriorityAdvisor.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/InteractiveModelRunner.h"
 #include "llvm/Analysis/MLModelRunner.h"
@@ -25,6 +24,7 @@
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/Passes.h"
+#include "llvm/CodeGen/RegAllocPriorityAdvisor.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
 #include "llvm/CodeGen/SlotIndexes.h"
 #include "llvm/CodeGen/VirtRegMap.h"
@@ -121,23 +121,11 @@ static const std::vector<TensorSpec> InputFeatures{
 // ===================================
 // Release (AOT) - specifics
 // ===================================
-class ReleaseModePriorityAdvisorAnalysis final
-    : public RegAllocPriorityAdvisorAnalysis {
+class ReleaseModePriorityAdvisorProvider final
+    : public RegAllocPriorityAdvisorProvider {
 public:
-  ReleaseModePriorityAdvisorAnalysis()
-      : RegAllocPriorityAdvisorAnalysis(AdvisorMode::Release) {}
-  // support for isa<> and dyn_cast.
-  static bool classof(const RegAllocPriorityAdvisorAnalysis *R) {
-    return R->getAdvisorMode() == AdvisorMode::Release;
-  }
-
-private:
-  void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.setPreservesAll();
-    AU.addRequired<SlotIndexesWrapperPass>();
-    RegAllocPriorityAdvisorAnalysis::getAnalysisUsage(AU);
-  }
-
+  ReleaseModePriorityAdvisorProvider()
+      : RegAllocPriorityAdvisorProvider(AdvisorMode::Release) {}
   std::unique_ptr<RegAllocPriorityAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
     if (!Runner) {
@@ -150,12 +138,44 @@ class ReleaseModePriorityAdvisorAnalysis final
             InteractiveChannelBaseName + ".out",
             InteractiveChannelBaseName + ".in");
     }
-    return std::make_unique<MLPriorityAdvisor>(
-        MF, RA, &getAnalysis<SlotIndexesWrapperPass>().getSI(), Runner.get());
+    assert(SI && "SlotIndexes result must be set");
+    return std::make_unique<MLPriorityAdvisor>(MF, RA, SI, Runner.get());
   }
+
+private:
   std::unique_ptr<MLModelRunner> Runner;
 };
 
+RegAllocPriorityAdvisorProvider *createReleaseModePriorityAdvisorProvider() {
+  return new ReleaseModePriorityAdvisorProvider();
+}
+
+class ReleaseModePriorityAdvisorAnalysisLegacy final
+    : public RegAllocPriorityAdvisorAnalysisLegacy {
+public:
+  ReleaseModePriorityAdvisorAnalysisLegacy()
+      : RegAllocPriorityAdvisorAnalysisLegacy(AdvisorMode::Release) {}
+  // support for isa<> and dyn_cast.
+  static bool classof(const RegAllocPriorityAdvisorAnalysisLegacy *R) {
+    return R->getAdvisorMode() == AdvisorMode::Release;
+  }
+
+private:
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesAll();
+    AU.addRequired<SlotIndexesWrapperPass>();
+    RegAllocPriorityAdvisorAnalysisLegacy::getAnalysisUsage(AU);
+  }
+
+  std::unique_ptr<RegAllocPriorityAdvisor>
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+    Provider->setAnalyses(&getAnalysis<SlotIndexesWrapperPass>().getSI());
+    return Provider->getAdvisor(MF, RA);
+  }
+
+  std::unique_ptr<ReleaseModePriorityAdvisorProvider> Provider;
+};
+
 // ===================================
 // Development mode-specifics
 // ===================================
@@ -186,46 +206,17 @@ class DevelopmentModePriorityAdvisor : public MLPriorityAdvisor {
   Logger *const Log;
 };
 
-class DevelopmentModePriorityAdvisorAnalysis final
-    : public RegAllocPriorityAdvisorAnalysis {
-public:
-  DevelopmentModePriorityAdvisorAnalysis()
-      : RegAllocPriorityAdvisorAnalysis(AdvisorMode::Development) {}
-  // support for isa<> and dyn_cast.
-  static bool classof(const RegAllocPriorityAdvisorAnalysis *R) {
-    return R->getAdvisorMode() == AdvisorMode::Development;
-  }
-
-  void logRewardIfNeeded(const MachineFunction &MF,
-                         llvm::function_ref<float()> GetReward) override {
-    if (!Log || !Log->hasAnyObservationForContext(MF.getName()))
-      return;
-    // The function pass manager would run all the function passes for a
-    // function, so we assume the last context belongs to this function. If
-    // this invariant ever changes, we can implement at that time switching
-    // contexts. At this point, it'd be an error
-    if (Log->currentContext() != MF.getName()) {
-      MF.getFunction().getContext().emitError(
-          "The training log context shouldn't have had changed.");
-    }
-    if (Log->hasObservationInProgress())
-      Log->logReward<float>(GetReward());
-  }
-
-private:
-  void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.setPreservesAll();
-    AU.addRequired<SlotIndexesWrapperPass>();
-    RegAllocPriorityAdvisorAnalysis::getAnalysisUsage(AU);
-  }
+class DevelopmentModePriorityAdvisorProvider final
+    : public RegAllocPriorityAdvisorProvider {
 
+public:
   // Save all the logs (when requested).
-  bool doInitialization(Module &M) override {
-    LLVMContext &Ctx = M.getContext();
+  DevelopmentModePriorityAdvisorProvider(LLVMContext &Ctx)
+      : RegAllocPriorityAdvisorProvider(AdvisorMode::Development) {
     if (ModelUnderTraining.empty() && TrainingLog.empty()) {
       Ctx.emitError("Regalloc development mode should be requested with at "
                     "least logging enabled and/or a training model");
-      return false;
+      return;
     }
     if (ModelUnderTraining.empty())
       Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
@@ -234,15 +225,15 @@ class DevelopmentModePriorityAdvisorAnalysis final
           Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
     if (!Runner) {
       Ctx.emitError("Regalloc: could not set up the model runner");
-      return false;
+      return;
     }
     if (TrainingLog.empty())
-      return false;
+      return;
     std::error_code EC;
     auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
     if (EC) {
-      M.getContext().emitError(EC.message() + ":" + TrainingLog);
-      return false;
+      Ctx.emitError(EC.message() + ":" + TrainingLog);
+      return;
     }
     std::vector<TensorSpec> LFS = InputFeatures;
     if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
@@ -254,7 +245,22 @@ class DevelopmentModePriorityAdvisorAnalysis final
 
     Log = std::make_unique<Logger>(std::move(OS), LFS, Reward,
                                    /*IncludeReward*/ true);
-    return false;
+  }
+
+  void logRewardIfNeeded(const MachineFunction &MF,
+                         llvm::function_ref<float()> GetReward) override {
+    if (!Log || !Log->hasAnyObservationForContext(MF.getName()))
+      return;
+    // The function pass manager would run all the function passes for a
+    // function, so we assume the last context belongs to this function. If
+    // this invariant ever changes, we can implement at that time switching
+    // contexts. At this point, it'd be an error
+    if (Log->currentContext() != MF.getName()) {
+      MF.getFunction().getContext().emitError(
+          "The training log context shouldn't have had changed.");
+    }
+    if (Log->hasObservationInProgress())
+      Log->logReward<float>(GetReward());
   }
 
   std::unique_ptr<RegAllocPriorityAdvisor>
@@ -264,23 +270,68 @@ class DevelopmentModePriorityAdvisorAnalysis final
     if (Log) {
       Log->switchContext(MF.getName());
     }
-
+    assert(SI && "SlotIndexes result must be set");
     return std::make_unique<DevelopmentModePriorityAdvisor>(
-        MF, RA, &getAnalysis<SlotIndexesWrapperPass>().getSI(), Runner.get(),
-        Log.get());
+        MF, RA, SI, Runner.get(), Log.get());
   }
 
   std::unique_ptr<MLModelRunner> Runner;
   std::unique_ptr<Logger> Log;
 };
+
+RegAllocPriorityAdvisorProvider *
+createDevelopmentModePriorityAdvisorProvider(LLVMContext &Ctx) {
+  return new DevelopmentModePriorityAdvisorProvider(Ctx);
+}
+
+class DevelopmentModePriorityAdvisorAnalysisLegacy final
+    : public RegAllocPriorityAdvisorAnalysisLegacy {
+public:
+  DevelopmentModePriorityAdvisorAnalysisLegacy()
+      : RegAllocPriorityAdvisorAnalysisLegacy(AdvisorMode::Development) {}
+
+  // support for isa<> and dyn_cast.
+  static bool classof(const RegAllocPriorityAdvisorAnalysisLegacy *R) {
+    return R->getAdvisorMode() == AdvisorMode::Development;
+  }
+
+  void logRewardIfNeeded(const MachineFunction &MF,
+                         llvm::function_ref<float()> GetReward) override {
+    Provider->logRewardIfNeeded(MF, GetReward);
+  }
+
+private:
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesAll();
+    AU.addRequired<SlotIndexesWrapperPass>();
+    RegAllocPriorityAdvisorAnalysisLegacy::getAnalysisUsage(AU);
+  }
+
+  // Save all the logs (when requested).
+  bool doInitialization(Module &M) override {
+    Provider = std::make_unique<DevelopmentModePriorityAdvisorProvider>(
+        M.getContext());
+    return false;
+    ;
+  }
+
+  std::unique_ptr<RegAllocPriorityAdvisor>
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+    Provider->setAnalyses(&getAnalysis<SlotIndexesWrapperPass>().getSI());
+    return Provider->getAdvisor(MF, RA);
+  }
+
+  std::unique_ptr<DevelopmentModePriorityAdvisorProvider> Provider;
+};
 #endif //#ifdef LLVM_HAVE_TFLITE
 
 } // namespace llvm
 
-RegAllocPriorityAdvisorAnalysis *llvm::createReleaseModePriorityAdvisor() {
+RegAllocPriorityAdvisorAnalysisLegacy *
+llvm::createReleaseModePriorityAdvisorAnalysis() {
   return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
                  !InteractiveChannelBaseName.empty()
-             ? new ReleaseModePriorityAdvisorAnalysis()
+             ? new ReleaseModePriorityAdvisorAnalysisLegacy()
              : nullptr;
 }
 
@@ -310,8 +361,9 @@ unsigned MLPriorityAdvisor::getPriority(const LiveInterval &LI) const {
 }
 
 #ifdef LLVM_HAVE_TFLITE
-RegAllocPriorityAdvisorAnalysis *llvm::createDevelopmentModePriorityAdvisor() {
-  return new DevelopmentModePriorityAdvisorAnalysis();
+RegAllocPriorityAdvisorAnalysisLegacy *
+llvm::createDevelopmentModePriorityAdvisorAnalysis() {
+  return new DevelopmentModePriorityAdvisorAnalysisLegacy();
 }
 
 unsigned
diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
index d042ca27694311c..c6debbdc7dade1d 100644
--- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
@@ -12,11 +12,11 @@
 #include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
 #include "AllocationOrder.h"
 #include "RegAllocGreedy.h"
-#include "RegAllocPriorityAdvisor.h"
 #include "llvm/CodeGen/LiveRegMatrix.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
+#include "llvm/CodeGen/RegAllocPriorityAdvisor.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
 #include "llvm/CodeGen/VirtRegMap.h"
 #include "llvm/IR/Module.h"
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp
index ae4ec6578c22ae2..de4e9b889ee1517 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.cpp
+++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp
@@ -15,7 +15,6 @@
 #include "AllocationOrder.h"
 #include "InterferenceCache.h"
 #include "RegAllocBase.h"
-#include "RegAllocPriorityAdvisor.h"
 #include "SplitKit.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
@@ -46,6 +45,7 @@
 #include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
+#include "llvm/CodeGen/RegAllocPriorityAdvisor.h"
 #include "llvm/CodeGen/RegAllocRegistry.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
 #include "llvm/CodeGen/SlotIndexes.h"
@@ -165,7 +165,7 @@ INITIALIZE_PASS_DEPENDENCY(EdgeBundlesWrapperLegacy)
 INITIALIZE_PASS_DEPENDENCY(SpillPlacementWrapperLegacy)
 INITIALIZE_PASS_DEPENDENCY(MachineOptimizationRemarkEmitterPass)
 INITIALIZE_PASS_DEPENDENCY(RegAllocEvictionAdvisorAnalysisLegacy)
-INITIALIZE_PASS_DEPENDENCY(RegAllocPriorityAdvisorAnalysis)
+INITIALIZE_PASS_DEPENDENCY(RegAllocPriorityAdvisorAnalysisLegacy)
 INITIALIZE_PASS_END(RAGreedy, "greedy",
                 "Greedy Register Allocator", false, false)
 
@@ -220,7 +220,7 @@ void RAGreedy::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequired<SpillPlacementWrapperLegacy>();
   AU.addRequired<MachineOptimizationRemarkEmitterPass>();
   AU.addRequired<RegAllocEvictionAdvisorAnalysisLegacy>();
-  AU.addRequired<RegAllocPriorityAdvisorAnalysis>();
+  AU.addRequired<RegAllocPriorityAdvisorAnalysisLegacy>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
@@ -2771,7 +2771,8 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
   EvictAdvisor = EvictAdvisorProvider.getAdvisor(*MF, *this, MBFI, Loops);
 
   PriorityAdvisor =
-      getAnalysis<RegAllocPriorityAdvisorAnalysis>().getAdvisor(*MF, *this);
+      getAnalysis<RegAllocPriorityAdvisorAnalysisLegacy>().getAdvisor(*MF,
+                                                                      *this);
 
   VRAI = std::make_unique<VirtRegAuxInfo>(*MF, *LIS, *VRM, *Loops, *MBFI);
   SpillerInstance.reset(
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.h b/llvm/lib/CodeGen/RegAllocGreedy.h
index 001a1e86a969141..89ceefd37795bcd 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.h
+++ b/llvm/lib/CodeGen/RegAllocGreedy.h
@@ -14,7 +14,6 @@
 
 #include "InterferenceCache.h"
 #include "RegAllocBase.h"
-#include "RegAllocPriorityAdvisor.h"
 #include "SplitKit.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
@@ -28,6 +27,7 @@
 #include "llvm/CodeGen/LiveRangeEdit.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/RegAllocPriorityAdvisor.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
 #include "llvm/CodeGen/SpillPlacement.h"
 #include "llvm/CodeGen/Spiller.h"
diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
index 4525b8fc5a383f4..0bb1219333237ff 100644
--- a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
@@ -10,7 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "RegAllocPriorityAdvisor.h"
+#include "llvm/CodeGen/RegAllocPriorityAdvisor.h"
 #include "RegAllocGreedy.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/VirtRegMap.h"
@@ -20,72 +20,113 @@
 
 using namespace llvm;
 
-static cl::opt<RegAllocPriorityAdvisorAnalysis::AdvisorMode> Mode(
+static cl::opt<RegAllocPriorityAdvisorProvider::AdvisorMode> Mode(
     "regalloc-enable-priority-advisor", cl::Hidden,
-    cl::init(RegAllocPriorityAdvisorAnalysis::AdvisorMode::Default),
+    cl::init(RegAllocPriorityAdvisorProvider::AdvisorMode::Default),
     cl::desc("Enable regalloc advisor mode"),
     cl::values(
-        clEnumValN(RegAllocPriorityAdvisorAnalysis::AdvisorMode::Default,
+        clEnumValN(RegAllocPriorityAdvisorProvider::AdvisorMode::Default,
                    "default", "Default"),
-        clEnumValN(RegAllocPriorityAdvisorAnalysis::AdvisorMode::Release,
+        clEnumValN(RegAllocPriorityAdvisorProvider::AdvisorMode::Release,
                    "release", "precompiled"),
-        clEnumValN(RegAllocPriorityAdvisorAnalysis::AdvisorMode::Development,
+        clEnumValN(RegAllocPriorityAdvisorProvider::AdvisorMode::Development,
                    "development", "for training"),
         clEnumValN(
-            RegAllocPriorityAdvisorAnalysis::AdvisorMode::Dummy, "dummy",
+            RegAllocPriorityAdvisorProvider::AdvisorMode::Dummy, "dummy",
             "prioritize low virtual register numbers for test and debug")));
 
-char RegAllocPriorityAdvisorAnalysis::ID = 0;
-INITIALIZE_PASS(RegAllocPriorityAdvisorAnalysis, "regalloc-priority",
+char RegAllocPriorityAdvisorAnalysisLegacy::ID = 0;
+INITIALIZE_PASS(RegAllocPriorityAdvisorAnalysisLegacy, "regalloc-priority",
                 "Regalloc priority policy", false, true)
 
 namespace {
-class DefaultPriorityAdvisorAnalysis final
-    : public RegAllocPriorityAdvisorAnalysis {
+
+class DefaultPriorityAdvisorProvider final
+    : public RegAllocPriorityAdvisorProvider {
+public:
+  DefaultPriorityAdvisorProvider(bool NotAsRequested, LLVMContext &Ctx)
+      : RegAllocPriorityAdvisorProvider(AdvisorMode::Default) {
+    if (NotAsRequested)
+      Ctx.emitError("Requested regalloc priority advisor analysis "
+                    "could be created. Using default");
+  }
+
+  // support for isa<> and dyn_cast.
+  static bool classof(const RegAllocPriorityAdvisorProvider *R) {
+    return R->getAdvisorMode() == AdvisorMode::Default;
+  }
+
+  std::unique_ptr<RegAllocPriorityAdvisor>
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+    assert(SI && "SlotIndexes result must be set");
+    return std::make_unique<DefaultPriorityAdvisor>(MF, RA, SI);
+  }
+};
+
+class DummyPriorityAdvisorProvider final : public RegAllocPriorityAdvisorProvider {
+public:
+  DummyPriorityAdvisorProvider() : RegAllocPriorityAdvisorProvider(AdvisorMode::Dummy) {}
+
+  static bool classof(const RegAllocPriorityAdvisorProvider *R) {
+    return R->getAdvisorMode() == AdvisorMode::Dummy;
+  }
+
+  std::unique_ptr<RegAllocPriorityAdvisor>
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+    assert(SI && "SlotIndexes result must be set");
+    return std::make_unique<DummyPriorityAdvisor>(MF, RA, SI);
+  }
+};
+
+class DefaultPriorityAdvisorAnalysisLegacy final
+    : public RegAllocPriorityAdvisorAnalysisLegacy {
 public:
-  DefaultPriorityAdvisorAnalysis(bool NotAsRequested)
-      : RegAllocPriorityAdvisorAnalysis(AdvisorMode::Default),
+  DefaultPriorityAdvisorAnalysisLegacy(bool NotAsRequested)
+      : RegAllocPriorityAdvisorAnalysisLegacy(AdvisorMode::Default),
         NotAsRequested(NotAsRequested) {}
 
   // support for isa<> and dyn_cast.
-  static bool classof(const RegAllocPriorityAdvisorAnalysis *R) {
+  static bool classof(const RegAllocPriorityAdvisorAnalysisLegacy *R) {
     return R->getAdvisorMode() == AdvisorMode::Default;
   }
 
 private:
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<SlotIndexesWrapperPass>();
-    RegAllocPriorityAdvisorAnalysis::getAnalysisUsage(AU);
+    RegAllocPriorityAdvisorAnalysisLegacy::getAnalysisUsage(AU);
   }
+
   std::unique_ptr<RegAllocPriorityAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
-    return std::make_unique<DefaultPriorityAdvisor>(
-        MF, RA, &getAnalysis<SlotIndexesWrapperPass>().getSI());
+    Provider->setAnalyses(&getAnalysis<SlotIndexesWrapperPass>().getSI());
+    return Provider->getAdvisor(MF, RA);
   }
   bool doInitialization(Module &M) override {
-    if (NotAsRequested)
-      M.getContext().emitError("Requested regalloc priority advisor analysis "
-                               "could be created. Using default");
-    return RegAllocPriorityAdvisorAnalysis::doInitialization(M);
+    Provider.reset(
+        new DefaultPriorityAdvisorProvider(NotAsRequested, M.getContext()));
+    return false;
   }
+
   const bool NotAsRequested;
+  std::unique_ptr<DefaultPriorityAdvisorProvider> Provider;
 };
 
 class DummyPriorityAdvisorAnalysis final
-    : public RegAllocPriorityAdvisorAnalysis {
+    : public RegAllocPriorityAdvisorAnalysisLegacy {
 public:
+  using RegAllocPriorityAdvisorAnalysisLegacy::AdvisorMode;
   DummyPriorityAdvisorAnalysis()
-      : RegAllocPriorityAdvisorAnalysis(AdvisorMode::Dummy) {}
+      : RegAllocPriorityAdvisorAnalysisLegacy(AdvisorMode::Dummy) {}
 
   // support for isa<> and dyn_cast.
-  static bool classof(const RegAllocPriorityAdvisorAnalysis *R) {
+  static bool classof(const RegAllocPriorityAdvisorAnalysisLegacy *R) {
     return R->getAdvisorMode() == AdvisorMode::Dummy;
   }
 
 private:
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<SlotIndexesWrapperPass>();
-    RegAllocPriorityAdvisorAnalysis::getAnalysisUsage(AU);
+    RegAllocPriorityAdvisorAnalysisLegacy::getAnalysisUsage(AU);
   }
 
   std::unique_ptr<RegAllocPriorityAdvisor>
@@ -97,30 +138,70 @@ class DummyPriorityAdvisorAnalysis final
 
 } // namespace
 
-template <> Pass *llvm::callDefaultCtor<RegAllocPriorityAdvisorAnalysis>() {
-  Pass *Ret = nullptr;
+void RegAllocPriorityAdvisorAnalysis::initializeProvider(LLVMContext &Ctx) {
+  if (Provider)
+    return;
+
   switch (Mode) {
-  case RegAllocPriorityAdvisorAnalysis::AdvisorMode::Default:
-    Ret = new DefaultPriorityAdvisorAnalysis(/*NotAsRequested*/ false);
+  case RegAllocPriorityAdvisorProvider::AdvisorMode::Default:
+    Provider.reset(
+        new DefaultPriorityAdvisorProvider(/*NotAsRequested*/ false, Ctx));
     break;
-  case RegAllocPriorityAdvisorAnalysis::AdvisorMode::Dummy:
-    Ret = new DummyPriorityAdvisorAnalysis();
+  case RegAllocPriorityAdvisorProvider::AdvisorMode::Dummy:
+    Provider.reset(new DummyPriorityAdvisorProvider());
     break;
-  case RegAllocPriorityAdvisorAnalysis::AdvisorMode::Development:
+  case RegAllocPriorityAdvisorProvider::AdvisorMode::Development:
 #if defined(LLVM_HAVE_TFLITE)
-    Ret = createDevelopmentModePriorityAdvisor();
+    Provider.reset(createDevelopmentModePriorityAdvisorProvider(Ctx));
 #endif
     break;
-  case RegAllocPriorityAdvisorAnalysis::AdvisorMode::Release:
-    Ret = createReleaseModePriorityAdvisor();
+  case RegAllocPriorityAdvisorProvider::AdvisorMode::Release:
+    Provider.reset(createReleaseModePriorityAdvisorProvider());
+    break;
+  }
+  if (!Provider)
+    Provider.reset(
+        new DefaultPriorityAdvisorProvider(/*NotAsRequested*/ true, Ctx));
+}
+
+AnalysisKey RegAllocPriorityAdvisorAnalysis::Key;
+
+RegAllocPriorityAdvisorAnalysis::Result
+RegAllocPriorityAdvisorAnalysis::run(MachineFunction &MF,
+                                     MachineFunctionAnalysisManager &MFAM) {
+  // Lazily initialize the provider.
+  initializeProvider(MF.getFunction().getContext());
+  // On each run, update the analysis for the provider.
+  Provider->setAnalyses(&MFAM.getResult<SlotIndexesAnalysis>(MF));
+  // The requiring analysis will construct the advisor.
+  return Result{Provider.get()};
+}
+
+template <>
+Pass *llvm::callDefaultCtor<RegAllocPriorityAdvisorAnalysisLegacy>() {
+  Pass *Ret = nullptr;
+  switch (Mode) {
+  case RegAllocPriorityAdvisorProvider::AdvisorMode::Default:
+    Ret = new DefaultPriorityAdvisorAnalysisLegacy(/*NotAsRequested*/ false);
+    break;
+  case RegAllocPriorityAdvisorProvider::AdvisorMode::Development:
+#if defined(LLVM_HAVE_TFLITE)
+    Ret = createDevelopmentModePriorityAdvisorAnalysis();
+#endif
+    break;
+  case RegAllocPriorityAdvisorProvider::AdvisorMode::Release:
+    Ret = createReleaseModePriorityAdvisorAnalysis();
+    break;
+  case RegAllocPriorityAdvisorProvider::AdvisorMode::Dummy:
+    Ret = new DummyPriorityAdvisorAnalysis();
     break;
   }
   if (Ret)
     return Ret;
-  return new DefaultPriorityAdvisorAnalysis(/*NotAsRequested*/ true);
+  return new DefaultPriorityAdvisorAnalysisLegacy(/*NotAsRequested*/ true);
 }
 
-StringRef RegAllocPriorityAdvisorAnalysis::getPassName() const {
+StringRef RegAllocPriorityAdvisorAnalysisLegacy::getPassName() const {
   switch (getAdvisorMode()) {
   case AdvisorMode::Default:
     return "Default Regalloc Priority Advisor";
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index a1eb9b4a9cc9f1c..2e2193575653a29 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -126,6 +126,7 @@
 #include "llvm/CodeGen/PreISelIntrinsicLowering.h"
 #include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
 #include "llvm/CodeGen/RegAllocFast.h"
+#include "llvm/CodeGen/RegAllocPriorityAdvisor.h"
 #include "llvm/CodeGen/RegUsageInfoCollector.h"
 #include "llvm/CodeGen/RegUsageInfoPropagate.h"
 #include "llvm/CodeGen/RegisterUsageInfo.h"

>From f92d72fe7f77632f0a205316cafee6b8ea8087d5 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Wed, 11 Dec 2024 09:03:53 +0000
Subject: [PATCH 2/6] use the provider as the analysis result

Avoid setting the advisor from the legacy wrapper after setting all
other analyses.
---
 .../llvm/CodeGen/RegAllocPriorityAdvisor.h    |  5 +++--
 .../lib/CodeGen/MLRegAllocPriorityAdvisor.cpp | 19 ++++++++++---------
 llvm/lib/CodeGen/RegAllocGreedy.cpp           |  6 +++---
 llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp  |  8 ++++----
 4 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
index a53739fdc3fc40b..600e63110ea915c 100644
--- a/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
+++ b/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
@@ -139,8 +139,7 @@ class RegAllocPriorityAdvisorAnalysisLegacy : public ImmutablePass {
   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;
+  virtual std::unique_ptr<RegAllocPriorityAdvisorProvider> &getProvider() = 0;
   AdvisorMode getAdvisorMode() const { return Mode; }
   virtual void logRewardIfNeeded(const MachineFunction &MF,
                                  llvm::function_ref<float()> GetReward) {};
@@ -152,6 +151,8 @@ class RegAllocPriorityAdvisorAnalysisLegacy : public ImmutablePass {
     AU.setPreservesAll();
   }
 
+  std::unique_ptr<RegAllocPriorityAdvisorProvider> Provider;
+
 private:
   StringRef getPassName() const override;
   const AdvisorMode Mode;
diff --git a/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
index dfc06eaa0bfe694..dd598767f1bbc21 100644
--- a/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
@@ -167,13 +167,17 @@ class ReleaseModePriorityAdvisorAnalysisLegacy final
     RegAllocPriorityAdvisorAnalysisLegacy::getAnalysisUsage(AU);
   }
 
-  std::unique_ptr<RegAllocPriorityAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+  std::unique_ptr<RegAllocPriorityAdvisorProvider> &getProvider() override {
     Provider->setAnalyses(&getAnalysis<SlotIndexesWrapperPass>().getSI());
-    return Provider->getAdvisor(MF, RA);
+    return Provider;
   }
 
-  std::unique_ptr<ReleaseModePriorityAdvisorProvider> Provider;
+  bool doInitialization(Module &M) override {
+    Provider = std::make_unique<ReleaseModePriorityAdvisorProvider>();
+    return false;
+  }
+
+  // std::unique_ptr<ReleaseModePriorityAdvisorProvider> Provider;
 };
 
 // ===================================
@@ -315,13 +319,10 @@ class DevelopmentModePriorityAdvisorAnalysisLegacy final
     ;
   }
 
-  std::unique_ptr<RegAllocPriorityAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+  std::unique_ptr<RegAllocPriorityAdvisorProvider> &getProvider() override {
     Provider->setAnalyses(&getAnalysis<SlotIndexesWrapperPass>().getSI());
-    return Provider->getAdvisor(MF, RA);
+    return Provider;
   }
-
-  std::unique_ptr<DevelopmentModePriorityAdvisorProvider> Provider;
 };
 #endif //#ifdef LLVM_HAVE_TFLITE
 
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp
index de4e9b889ee1517..c53c7eb50d091e9 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.cpp
+++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp
@@ -2770,9 +2770,9 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
       getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().getProvider();
   EvictAdvisor = EvictAdvisorProvider.getAdvisor(*MF, *this, MBFI, Loops);
 
-  PriorityAdvisor =
-      getAnalysis<RegAllocPriorityAdvisorAnalysisLegacy>().getAdvisor(*MF,
-                                                                      *this);
+  PriorityAdvisor = getAnalysis<RegAllocPriorityAdvisorAnalysisLegacy>()
+                        .getProvider()
+                        ->getAdvisor(*MF, *this);
 
   VRAI = std::make_unique<VirtRegAuxInfo>(*MF, *LIS, *VRM, *Loops, *MBFI);
   SpillerInstance.reset(
diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
index 0bb1219333237ff..c66c9b1fd06d610 100644
--- a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
@@ -96,11 +96,11 @@ class DefaultPriorityAdvisorAnalysisLegacy final
     RegAllocPriorityAdvisorAnalysisLegacy::getAnalysisUsage(AU);
   }
 
-  std::unique_ptr<RegAllocPriorityAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+  std::unique_ptr<RegAllocPriorityAdvisorProvider> &getProvider() override {
     Provider->setAnalyses(&getAnalysis<SlotIndexesWrapperPass>().getSI());
-    return Provider->getAdvisor(MF, RA);
+    return Provider;
   }
+
   bool doInitialization(Module &M) override {
     Provider.reset(
         new DefaultPriorityAdvisorProvider(NotAsRequested, M.getContext()));
@@ -108,7 +108,7 @@ class DefaultPriorityAdvisorAnalysisLegacy final
   }
 
   const bool NotAsRequested;
-  std::unique_ptr<DefaultPriorityAdvisorProvider> Provider;
+  // std::unique_ptr<DefaultPriorityAdvisorProvider> Provider;
 };
 
 class DummyPriorityAdvisorAnalysis final

>From 4076367c6b9a3218e0f28f213bd6ebdc97e6168a Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Thu, 26 Dec 2024 09:31:53 +0000
Subject: [PATCH 3/6] Suggestions and remove createProvider methods

---
 .../llvm/CodeGen/RegAllocPriorityAdvisor.h    |  7 +---
 .../lib/CodeGen/MLRegAllocPriorityAdvisor.cpp | 29 +++++++++------
 llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp  | 37 ++++++++-----------
 3 files changed, 35 insertions(+), 38 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
index 600e63110ea915c..bc00fb0c85a027d 100644
--- a/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
+++ b/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
@@ -101,11 +101,6 @@ class RegAllocPriorityAdvisorProvider {
   const AdvisorMode Mode;
 };
 
-RegAllocPriorityAdvisorProvider *createReleaseModePriorityAdvisorProvider();
-
-RegAllocPriorityAdvisorProvider *
-createDevelopmentModePriorityAdvisorProvider(LLVMContext &Ctx);
-
 class RegAllocPriorityAdvisorAnalysis
     : public AnalysisInfoMixin<RegAllocPriorityAdvisorAnalysis> {
   static AnalysisKey Key;
@@ -128,6 +123,8 @@ class RegAllocPriorityAdvisorAnalysis
 
 private:
   void initializeProvider(LLVMContext &Ctx);
+  void initializeMLProvider(RegAllocPriorityAdvisorProvider::AdvisorMode Mode,
+                            LLVMContext &Ctx);
   std::unique_ptr<RegAllocPriorityAdvisorProvider> Provider;
 };
 
diff --git a/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
index dd598767f1bbc21..dedd20b31d1c022 100644
--- a/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
@@ -146,10 +146,6 @@ class ReleaseModePriorityAdvisorProvider final
   std::unique_ptr<MLModelRunner> Runner;
 };
 
-RegAllocPriorityAdvisorProvider *createReleaseModePriorityAdvisorProvider() {
-  return new ReleaseModePriorityAdvisorProvider();
-}
-
 class ReleaseModePriorityAdvisorAnalysisLegacy final
     : public RegAllocPriorityAdvisorAnalysisLegacy {
 public:
@@ -176,8 +172,6 @@ class ReleaseModePriorityAdvisorAnalysisLegacy final
     Provider = std::make_unique<ReleaseModePriorityAdvisorProvider>();
     return false;
   }
-
-  // std::unique_ptr<ReleaseModePriorityAdvisorProvider> Provider;
 };
 
 // ===================================
@@ -283,11 +277,6 @@ class DevelopmentModePriorityAdvisorProvider final
   std::unique_ptr<Logger> Log;
 };
 
-RegAllocPriorityAdvisorProvider *
-createDevelopmentModePriorityAdvisorProvider(LLVMContext &Ctx) {
-  return new DevelopmentModePriorityAdvisorProvider(Ctx);
-}
-
 class DevelopmentModePriorityAdvisorAnalysisLegacy final
     : public RegAllocPriorityAdvisorAnalysisLegacy {
 public:
@@ -410,3 +399,21 @@ DevelopmentModePriorityAdvisor::getPriority(const LiveInterval &LI) const {
 }
 
 #endif // #ifdef LLVM_HAVE_TFLITE
+
+void RegAllocPriorityAdvisorAnalysis::initializeMLProvider(
+    RegAllocPriorityAdvisorProvider::AdvisorMode Mode, LLVMContext &Ctx) {
+  if (Provider)
+    return;
+  switch (Mode) {
+  case RegAllocPriorityAdvisorProvider::AdvisorMode::Development:
+#if defined(LLVM_HAVE_TFLITE)
+    Provider.reset(new DevelopmentModePriorityAdvisorProvider(Ctx));
+#endif
+    break;
+  case RegAllocPriorityAdvisorProvider::AdvisorMode::Release:
+    Provider.reset(new ReleaseModePriorityAdvisorProvider());
+    break;
+  default:
+    break;
+  }
+}
\ No newline at end of file
diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
index c66c9b1fd06d610..d9a26186aca2553 100644
--- a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
@@ -129,10 +129,14 @@ class DummyPriorityAdvisorAnalysis final
     RegAllocPriorityAdvisorAnalysisLegacy::getAnalysisUsage(AU);
   }
 
-  std::unique_ptr<RegAllocPriorityAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
-    return std::make_unique<DummyPriorityAdvisor>(
-        MF, RA, &getAnalysis<SlotIndexesWrapperPass>().getSI());
+  std::unique_ptr<RegAllocPriorityAdvisorProvider>&
+  getProvider() override {
+    Provider->setAnalyses(&getAnalysis<SlotIndexesWrapperPass>().getSI());
+    return Provider;
+  }
+  bool doInitialization(Module &M) override {
+    Provider.reset(new DummyPriorityAdvisorProvider());
+    return false;
   }
 };
 
@@ -141,27 +145,16 @@ class DummyPriorityAdvisorAnalysis final
 void RegAllocPriorityAdvisorAnalysis::initializeProvider(LLVMContext &Ctx) {
   if (Provider)
     return;
-
-  switch (Mode) {
-  case RegAllocPriorityAdvisorProvider::AdvisorMode::Default:
-    Provider.reset(
-        new DefaultPriorityAdvisorProvider(/*NotAsRequested*/ false, Ctx));
-    break;
-  case RegAllocPriorityAdvisorProvider::AdvisorMode::Dummy:
+  if (Mode == RegAllocPriorityAdvisorProvider::AdvisorMode::Dummy)
     Provider.reset(new DummyPriorityAdvisorProvider());
-    break;
-  case RegAllocPriorityAdvisorProvider::AdvisorMode::Development:
-#if defined(LLVM_HAVE_TFLITE)
-    Provider.reset(createDevelopmentModePriorityAdvisorProvider(Ctx));
-#endif
-    break;
-  case RegAllocPriorityAdvisorProvider::AdvisorMode::Release:
-    Provider.reset(createReleaseModePriorityAdvisorProvider());
-    break;
-  }
+  else if (Mode == RegAllocPriorityAdvisorProvider::AdvisorMode::Default)
+    Provider.reset(
+        new DefaultPriorityAdvisorProvider(/*NotAsRequested=*/false, Ctx));
+  else
+    initializeMLProvider(Mode, Ctx);
   if (!Provider)
     Provider.reset(
-        new DefaultPriorityAdvisorProvider(/*NotAsRequested*/ true, Ctx));
+        new DefaultPriorityAdvisorProvider(/*NotAsRequested=*/true, Ctx));
 }
 
 AnalysisKey RegAllocPriorityAdvisorAnalysis::Key;

>From 47c5fa0e53d0e3e5b2305d02e40c6820bb292a46 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Thu, 26 Dec 2024 10:13:51 +0000
Subject: [PATCH 4/6] Apply suggestions and remove setAnalyses

---
 .../llvm/CodeGen/RegAllocPriorityAdvisor.h    | 12 ++++--------
 .../lib/CodeGen/MLRegAllocPriorityAdvisor.cpp | 16 ++++------------
 llvm/lib/CodeGen/RegAllocGreedy.cpp           |  2 +-
 llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp  | 19 ++++---------------
 4 files changed, 13 insertions(+), 36 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
index bc00fb0c85a027d..82232eb4732a310 100644
--- a/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
+++ b/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
@@ -85,18 +85,14 @@ class RegAllocPriorityAdvisorProvider {
   virtual ~RegAllocPriorityAdvisorProvider() = default;
 
   virtual void logRewardIfNeeded(const MachineFunction &MF,
-                                 llvm::function_ref<float()> GetReward) {};
+                                 function_ref<float()> GetReward) {};
 
   virtual std::unique_ptr<RegAllocPriorityAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
-
-  void setAnalyses(SlotIndexes *SI) { this->SI = SI; }
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
+             SlotIndexes *SI) = 0;
 
   AdvisorMode getAdvisorMode() const { return Mode; }
 
-protected:
-  SlotIndexes *SI;
-
 private:
   const AdvisorMode Mode;
 };
@@ -136,7 +132,7 @@ class RegAllocPriorityAdvisorAnalysisLegacy : public ImmutablePass {
   static char ID;
 
   /// Get an advisor for the given context (i.e. machine function, etc)
-  virtual std::unique_ptr<RegAllocPriorityAdvisorProvider> &getProvider() = 0;
+  RegAllocPriorityAdvisorProvider &getProvider() { return *Provider; }
   AdvisorMode getAdvisorMode() const { return Mode; }
   virtual void logRewardIfNeeded(const MachineFunction &MF,
                                  llvm::function_ref<float()> GetReward) {};
diff --git a/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
index dedd20b31d1c022..8b195b36a711faa 100644
--- a/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
@@ -127,7 +127,8 @@ class ReleaseModePriorityAdvisorProvider final
   ReleaseModePriorityAdvisorProvider()
       : RegAllocPriorityAdvisorProvider(AdvisorMode::Release) {}
   std::unique_ptr<RegAllocPriorityAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
+             SlotIndexes *SI) override {
     if (!Runner) {
       if (InteractiveChannelBaseName.empty())
         Runner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
@@ -163,11 +164,6 @@ class ReleaseModePriorityAdvisorAnalysisLegacy final
     RegAllocPriorityAdvisorAnalysisLegacy::getAnalysisUsage(AU);
   }
 
-  std::unique_ptr<RegAllocPriorityAdvisorProvider> &getProvider() override {
-    Provider->setAnalyses(&getAnalysis<SlotIndexesWrapperPass>().getSI());
-    return Provider;
-  }
-
   bool doInitialization(Module &M) override {
     Provider = std::make_unique<ReleaseModePriorityAdvisorProvider>();
     return false;
@@ -262,7 +258,8 @@ class DevelopmentModePriorityAdvisorProvider final
   }
 
   std::unique_ptr<RegAllocPriorityAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
+             SlotIndexes *SI) override {
     if (!Runner)
       return nullptr;
     if (Log) {
@@ -307,11 +304,6 @@ class DevelopmentModePriorityAdvisorAnalysisLegacy final
     return false;
     ;
   }
-
-  std::unique_ptr<RegAllocPriorityAdvisorProvider> &getProvider() override {
-    Provider->setAnalyses(&getAnalysis<SlotIndexesWrapperPass>().getSI());
-    return Provider;
-  }
 };
 #endif //#ifdef LLVM_HAVE_TFLITE
 
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp
index c53c7eb50d091e9..72c38ab8c7d07bc 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.cpp
+++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp
@@ -2772,7 +2772,7 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
 
   PriorityAdvisor = getAnalysis<RegAllocPriorityAdvisorAnalysisLegacy>()
                         .getProvider()
-                        ->getAdvisor(*MF, *this);
+                        .getAdvisor(*MF, *this, Indexes);
 
   VRAI = std::make_unique<VirtRegAuxInfo>(*MF, *LIS, *VRM, *Loops, *MBFI);
   SpillerInstance.reset(
diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
index d9a26186aca2553..97cd85eac61293c 100644
--- a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
@@ -57,7 +57,8 @@ class DefaultPriorityAdvisorProvider final
   }
 
   std::unique_ptr<RegAllocPriorityAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
+             SlotIndexes *SI) override {
     assert(SI && "SlotIndexes result must be set");
     return std::make_unique<DefaultPriorityAdvisor>(MF, RA, SI);
   }
@@ -72,7 +73,8 @@ class DummyPriorityAdvisorProvider final : public RegAllocPriorityAdvisorProvide
   }
 
   std::unique_ptr<RegAllocPriorityAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
+             SlotIndexes *SI) override {
     assert(SI && "SlotIndexes result must be set");
     return std::make_unique<DummyPriorityAdvisor>(MF, RA, SI);
   }
@@ -96,11 +98,6 @@ class DefaultPriorityAdvisorAnalysisLegacy final
     RegAllocPriorityAdvisorAnalysisLegacy::getAnalysisUsage(AU);
   }
 
-  std::unique_ptr<RegAllocPriorityAdvisorProvider> &getProvider() override {
-    Provider->setAnalyses(&getAnalysis<SlotIndexesWrapperPass>().getSI());
-    return Provider;
-  }
-
   bool doInitialization(Module &M) override {
     Provider.reset(
         new DefaultPriorityAdvisorProvider(NotAsRequested, M.getContext()));
@@ -108,7 +105,6 @@ class DefaultPriorityAdvisorAnalysisLegacy final
   }
 
   const bool NotAsRequested;
-  // std::unique_ptr<DefaultPriorityAdvisorProvider> Provider;
 };
 
 class DummyPriorityAdvisorAnalysis final
@@ -129,11 +125,6 @@ class DummyPriorityAdvisorAnalysis final
     RegAllocPriorityAdvisorAnalysisLegacy::getAnalysisUsage(AU);
   }
 
-  std::unique_ptr<RegAllocPriorityAdvisorProvider>&
-  getProvider() override {
-    Provider->setAnalyses(&getAnalysis<SlotIndexesWrapperPass>().getSI());
-    return Provider;
-  }
   bool doInitialization(Module &M) override {
     Provider.reset(new DummyPriorityAdvisorProvider());
     return false;
@@ -164,8 +155,6 @@ RegAllocPriorityAdvisorAnalysis::run(MachineFunction &MF,
                                      MachineFunctionAnalysisManager &MFAM) {
   // Lazily initialize the provider.
   initializeProvider(MF.getFunction().getContext());
-  // On each run, update the analysis for the provider.
-  Provider->setAnalyses(&MFAM.getResult<SlotIndexesAnalysis>(MF));
   // The requiring analysis will construct the advisor.
   return Result{Provider.get()};
 }

>From 035da90263baa655094dd60a4d1931a827ac47d8 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Tue, 7 Jan 2025 09:57:23 +0000
Subject: [PATCH 5/6] run clang format

---
 llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
index 97cd85eac61293c..e7cec6661bf7785 100644
--- a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
@@ -64,9 +64,11 @@ class DefaultPriorityAdvisorProvider final
   }
 };
 
-class DummyPriorityAdvisorProvider final : public RegAllocPriorityAdvisorProvider {
+class DummyPriorityAdvisorProvider final
+    : public RegAllocPriorityAdvisorProvider {
 public:
-  DummyPriorityAdvisorProvider() : RegAllocPriorityAdvisorProvider(AdvisorMode::Dummy) {}
+  DummyPriorityAdvisorProvider()
+      : RegAllocPriorityAdvisorProvider(AdvisorMode::Dummy) {}
 
   static bool classof(const RegAllocPriorityAdvisorProvider *R) {
     return R->getAdvisorMode() == AdvisorMode::Dummy;

>From 59f56139e89dd792e93734d48f0bf26ae742f8e9 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Thu, 23 Jan 2025 09:10:27 +0000
Subject: [PATCH 6/6] Add newline at EOF

---
 llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
index 8b195b36a711faa..20ae298d627aaae 100644
--- a/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
@@ -408,4 +408,4 @@ void RegAllocPriorityAdvisorAnalysis::initializeMLProvider(
   default:
     break;
   }
-}
\ No newline at end of file
+}



More information about the llvm-branch-commits mailing list