[llvm] [CodeGen][NewPM] Port RegAllocEvictionAdvisor analysis to NPM (PR #117309)

Akshat Oke via llvm-commits llvm-commits at lists.llvm.org
Fri Nov 29 01:20:39 PST 2024


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

>From 0bef4ef6160dcaa5dd06b782326d1662037459cb Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Fri, 22 Nov 2024 09:31:50 +0000
Subject: [PATCH 1/2] [CodeGen][NewPM] Port RegAllocEvictionAdvisor analysis to
 NPM

---
 .../llvm}/CodeGen/RegAllocEvictionAdvisor.h   |  69 +++++++-
 llvm/include/llvm/InitializePasses.h          |   2 +-
 llvm/include/llvm/Passes/CodeGenPassBuilder.h |   6 +-
 llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp   | 167 +++++++++++++-----
 llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp  | 107 ++++++++---
 llvm/lib/CodeGen/RegAllocGreedy.cpp           |   9 +-
 llvm/lib/CodeGen/RegAllocGreedy.h             |   1 -
 llvm/lib/CodeGen/RegAllocPriorityAdvisor.h    |   2 +-
 llvm/lib/Passes/PassBuilder.cpp               |   1 +
 llvm/lib/Passes/PassRegistry.def              |   1 +
 10 files changed, 284 insertions(+), 81 deletions(-)
 rename llvm/{lib => include/llvm}/CodeGen/RegAllocEvictionAdvisor.h (75%)

diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
similarity index 75%
rename from llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
rename to llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
index 52dd946a685400a..847bf032235c1dd 100644
--- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
+++ b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
@@ -9,11 +9,13 @@
 #ifndef LLVM_CODEGEN_REGALLOCEVICTIONADVISOR_H
 #define LLVM_CODEGEN_REGALLOCEVICTIONADVISOR_H
 
+#include "llvm/ADT/Any.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/CodeGen/Register.h"
 #include "llvm/Config/llvm-config.h"
+#include "llvm/IR/PassManager.h"
 #include "llvm/MC/MCRegister.h"
 #include "llvm/Pass.h"
 
@@ -164,12 +166,12 @@ class RegAllocEvictionAdvisor {
 ///
 /// Because we need to offer additional services in 'development' mode, the
 /// implementations of this analysis need to implement RTTI support.
-class RegAllocEvictionAdvisorAnalysis : public ImmutablePass {
+class RegAllocEvictionAdvisorAnalysisLegacy : public ImmutablePass {
 public:
   enum class AdvisorMode : int { Default, Release, Development };
 
-  RegAllocEvictionAdvisorAnalysis(AdvisorMode Mode)
-      : ImmutablePass(ID), Mode(Mode){};
+  RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode Mode)
+      : ImmutablePass(ID), Mode(Mode) {};
   static char ID;
 
   /// Get an advisor for the given context (i.e. machine function, etc)
@@ -177,7 +179,7 @@ class RegAllocEvictionAdvisorAnalysis : 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
@@ -191,13 +193,66 @@ class RegAllocEvictionAdvisorAnalysis : public ImmutablePass {
   const AdvisorMode Mode;
 };
 
+/// Common provider for legacy and new pass managers.
+/// This keeps the state for logging, and sets up and holds the provider.
+/// The legacy pass itself used to keep the logging state and provider,
+/// so this extraction helps the NPM analysis to reuse the logic.
+class RegAllocEvictionAdvisorProvider {
+public:
+  enum class AdvisorMode : int { Default, Release, Development };
+  RegAllocEvictionAdvisorProvider(AdvisorMode Mode) : Mode(Mode) {}
+
+  virtual ~RegAllocEvictionAdvisorProvider() = default;
+
+  virtual bool doInitialization(Module &M) { return false; }
+
+  virtual void logRewardIfNeeded(const MachineFunction &MF,
+                                 llvm::function_ref<float()> GetReward) {}
+
+  virtual std::unique_ptr<RegAllocEvictionAdvisor>
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
+
+  /// Set the analyses that the advisor needs to use as they might not be
+  /// available before the advisor is created.
+  virtual void setAnalyses(std::initializer_list<llvm::Any> AnalysisP) {}
+
+  AdvisorMode getAdvisorMode() const { return Mode; }
+
+private:
+  const AdvisorMode Mode;
+};
+
+RegAllocEvictionAdvisorProvider *createReleaseModeAdvisorProvider();
+RegAllocEvictionAdvisorProvider *createDevelopmentModeAdvisorProvider();
+
+/// A Module analysis for fetching the Eviction Advisor. This is not a
+/// MachineFunction analysis for two reasons:
+/// - in the ML implementation case, the evaluator is stateless but (especially
+/// in the development mode) expensive to set up. With a Module Analysis, we
+/// `require` it and set it up once.
+/// - in the 'development' mode ML case, we want to capture the training log
+/// during allocation (this is a log of features encountered and decisions
+/// made), and then measure a score, potentially a few steps after allocation
+/// completes. So we need a Module analysis to keep the logger state around
+/// until we can make that measurement.
+class RegAllocEvictionAdvisorAnalysis
+    : public AnalysisInfoMixin<RegAllocEvictionAdvisorAnalysis> {
+  static AnalysisKey Key;
+  friend AnalysisInfoMixin<RegAllocEvictionAdvisorAnalysis>;
+
+public:
+  using Result = std::unique_ptr<RegAllocEvictionAdvisorProvider>;
+  Result run(Module &MF, ModuleAnalysisManager &MAM);
+};
+
 /// Specialization for the API used by the analysis infrastructure to create
 /// an instance of the eviction advisor.
-template <> Pass *callDefaultCtor<RegAllocEvictionAdvisorAnalysis>();
+template <> Pass *callDefaultCtor<RegAllocEvictionAdvisorAnalysisLegacy>();
 
-RegAllocEvictionAdvisorAnalysis *createReleaseModeAdvisor();
+RegAllocEvictionAdvisorAnalysisLegacy *createReleaseModeAdvisorAnalysisLegacy();
 
-RegAllocEvictionAdvisorAnalysis *createDevelopmentModeAdvisor();
+RegAllocEvictionAdvisorAnalysisLegacy *
+createDevelopmentModeAdvisorAnalysisLegacy();
 
 // TODO: move to RegAllocEvictionAdvisor.cpp when we move implementation
 // out of RegAllocGreedy.cpp
diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
index e883aae2758688d..75212cb2a777ab8 100644
--- a/llvm/include/llvm/InitializePasses.h
+++ b/llvm/include/llvm/InitializePasses.h
@@ -252,7 +252,7 @@ void initializePseudoProbeInserterPass(PassRegistry &);
 void initializeRAGreedyPass(PassRegistry &);
 void initializeReachingDefAnalysisPass(PassRegistry &);
 void initializeReassociateLegacyPassPass(PassRegistry &);
-void initializeRegAllocEvictionAdvisorAnalysisPass(PassRegistry &);
+void initializeRegAllocEvictionAdvisorAnalysisLegacyPass(PassRegistry &);
 void initializeRegAllocFastPass(PassRegistry &);
 void initializeRegAllocPriorityAdvisorAnalysisPass(PassRegistry &);
 void initializeRegAllocScoringPass(PassRegistry &);
diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index d2e9e8185a2b90c..aa6437c1d2f0172 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -54,6 +54,7 @@
 #include "llvm/CodeGen/PHIElimination.h"
 #include "llvm/CodeGen/PeepholeOptimizer.h"
 #include "llvm/CodeGen/PreISelIntrinsicLowering.h"
+#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
 #include "llvm/CodeGen/RegAllocFast.h"
 #include "llvm/CodeGen/RegUsageInfoCollector.h"
 #include "llvm/CodeGen/RegUsageInfoPropagate.h"
@@ -1061,9 +1062,10 @@ void CodeGenPassBuilder<Derived, TargetMachineT>::addMachineSSAOptimization(
 template <typename Derived, typename TargetMachineT>
 void CodeGenPassBuilder<Derived, TargetMachineT>::addTargetRegisterAllocator(
     AddMachinePass &addPass, bool Optimized) const {
-  if (Optimized)
+  if (Optimized) {
+    addPass(RequireAnalysis<RegAllocEvictionAdvisorAnalysis>());
     addPass(RAGreedyPass());
-  else
+  } else
     addPass(RegAllocFastPass());
 }
 
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index d099544c2a49187..9c15c327bb202c2 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -11,11 +11,11 @@
 //===----------------------------------------------------------------------===//
 
 #include "AllocationOrder.h"
-#include "RegAllocEvictionAdvisor.h"
 #include "RegAllocGreedy.h"
 #include "llvm/Analysis/InteractiveModelRunner.h"
 #include "llvm/Analysis/MLModelRunner.h"
 #include "llvm/Analysis/TensorSpec.h"
+#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
 #if defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL) || defined(LLVM_HAVE_TFLITE)
 #include "llvm/Analysis/ModelUnderTrainingRunner.h"
 #include "llvm/Analysis/NoInferenceModelRunner.h"
@@ -108,7 +108,7 @@ class RegAllocScoring : public MachineFunctionPass {
   /// RegAllocReward analysis usage.
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
-    AU.addRequired<RegAllocEvictionAdvisorAnalysis>();
+    AU.addRequired<RegAllocEvictionAdvisorAnalysisLegacy>();
     AU.addRequired<RegAllocPriorityAdvisorAnalysis>();
     AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
     MachineFunctionPass::getAnalysisUsage(AU);
@@ -366,11 +366,12 @@ class MLEvictAdvisor : public RegAllocEvictionAdvisor {
 // ===================================
 // Release (AOT) - specifics
 // ===================================
-class ReleaseModeEvictionAdvisorAnalysis final
-    : public RegAllocEvictionAdvisorAnalysis {
+/// Common provider for legacy and new pass managers.
+class ReleaseModeEvictionAdvisorProvider final
+    : public RegAllocEvictionAdvisorProvider {
 public:
-  ReleaseModeEvictionAdvisorAnalysis()
-      : RegAllocEvictionAdvisorAnalysis(AdvisorMode::Release) {
+  ReleaseModeEvictionAdvisorProvider()
+      : RegAllocEvictionAdvisorProvider(AdvisorMode::Release) {
     if (EnableDevelopmentFeatures) {
       InputFeatures = {RA_EVICT_FEATURES_LIST(
           _DECL_FEATURES) RA_EVICT_FIRST_DEVELOPMENT_FEATURE(_DECL_FEATURES)
@@ -380,17 +381,17 @@ class ReleaseModeEvictionAdvisorAnalysis final
     }
   }
   // support for isa<> and dyn_cast.
-  static bool classof(const RegAllocEvictionAdvisorAnalysis *R) {
+  static bool classof(const RegAllocEvictionAdvisorProvider *R) {
     return R->getAdvisorMode() == AdvisorMode::Release;
   }
 
-private:
-  std::vector<TensorSpec> InputFeatures;
-
-  void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
-    AU.addRequired<MachineLoopInfoWrapperPass>();
-    RegAllocEvictionAdvisorAnalysis::getAnalysisUsage(AU);
+  void setAnalyses(std::initializer_list<llvm::Any> AnalysisList) override {
+    for (auto Analysis : AnalysisList) {
+      if (auto **MBFI = llvm::any_cast<MachineBlockFrequencyInfo *>(&Analysis))
+        this->MBFI = *MBFI;
+      if (auto **Loops = llvm::any_cast<MachineLoopInfo *>(&Analysis))
+        this->Loops = *Loops;
+    }
   }
 
   std::unique_ptr<RegAllocEvictionAdvisor>
@@ -405,12 +406,47 @@ class ReleaseModeEvictionAdvisorAnalysis final
             InteractiveChannelBaseName + ".out",
             InteractiveChannelBaseName + ".in");
     }
-    return std::make_unique<MLEvictAdvisor>(
-        MF, RA, Runner.get(),
-        getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI(),
-        getAnalysis<MachineLoopInfoWrapperPass>().getLI());
+    return std::make_unique<MLEvictAdvisor>(MF, RA, Runner.get(), *MBFI,
+                                            *Loops);
   }
+
+private:
+  std::vector<TensorSpec> InputFeatures;
   std::unique_ptr<MLModelRunner> Runner;
+  MachineBlockFrequencyInfo *MBFI;
+  MachineLoopInfo *Loops;
+};
+
+class ReleaseModeEvictionAdvisorAnalysisLegacy final
+    : public RegAllocEvictionAdvisorAnalysisLegacy {
+public:
+  ReleaseModeEvictionAdvisorAnalysisLegacy()
+      : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Release) {}
+
+  std::unique_ptr<RegAllocEvictionAdvisor>
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+    auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
+    auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
+    Provider.setAnalyses({MBFI, Loops});
+    return Provider.getAdvisor(MF, RA);
+  }
+
+  bool doInitialization(Module &M) override {
+    return Provider.doInitialization(M);
+  }
+
+  static bool classof(const RegAllocEvictionAdvisorAnalysisLegacy *R) {
+    return R->getAdvisorMode() == AdvisorMode::Release;
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
+    AU.addRequired<MachineLoopInfoWrapperPass>();
+    RegAllocEvictionAdvisorAnalysisLegacy::getAnalysisUsage(AU);
+  }
+
+private:
+  ReleaseModeEvictionAdvisorProvider Provider;
 };
 
 // ===================================
@@ -445,11 +481,14 @@ class DevelopmentModeEvictAdvisor : public MLEvictAdvisor {
   Logger *const Log;
 };
 
-class DevelopmentModeEvictionAdvisorAnalysis final
-    : public RegAllocEvictionAdvisorAnalysis {
+class DevelopmentModeEvictionAdvisorProvider final
+    : public RegAllocEvictionAdvisorProvider {
 public:
-  DevelopmentModeEvictionAdvisorAnalysis()
-      : RegAllocEvictionAdvisorAnalysis(AdvisorMode::Development) {
+  DevelopmentModeEvictionAdvisorProvider(
+      MachineBlockFrequencyInfo *MBFI = nullptr,
+      MachineLoopInfo *Loops = nullptr)
+      : RegAllocEvictionAdvisorProvider(AdvisorMode::Development), MBFI(MBFI),
+        Loops(Loops) {
     if (EnableDevelopmentFeatures) {
       InputFeatures = {RA_EVICT_FEATURES_LIST(
           _DECL_FEATURES) RA_EVICT_FIRST_DEVELOPMENT_FEATURE(_DECL_FEATURES)
@@ -471,7 +510,7 @@ class DevelopmentModeEvictionAdvisorAnalysis final
     }
   }
   // support for isa<> and dyn_cast.
-  static bool classof(const RegAllocEvictionAdvisorAnalysis *R) {
+  static bool classof(const RegAllocEvictionAdvisorProvider *R) {
     return R->getAdvisorMode() == AdvisorMode::Development;
   }
 
@@ -491,14 +530,13 @@ class DevelopmentModeEvictionAdvisorAnalysis final
       Log->logReward<float>(GetReward());
   }
 
-private:
-  std::vector<TensorSpec> InputFeatures;
-  std::vector<TensorSpec> TrainingInputFeatures;
-
-  void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
-    AU.addRequired<MachineLoopInfoWrapperPass>();
-    RegAllocEvictionAdvisorAnalysis::getAnalysisUsage(AU);
+  void setAnalyses(std::initializer_list<llvm::Any> AnalysisList) override {
+    for (auto Analysis : AnalysisList) {
+      if (auto **MBFI = llvm::any_cast<MachineBlockFrequencyInfo *>(&Analysis))
+        this->MBFI = *MBFI;
+      if (auto **Loops = llvm::any_cast<MachineLoopInfo *>(&Analysis))
+        this->Loops = *Loops;
+    }
   }
 
   bool doInitialization(Module &M) override {
@@ -544,14 +582,53 @@ class DevelopmentModeEvictionAdvisorAnalysis final
       return nullptr;
     if (Log)
       Log->switchContext(MF.getName());
+    assert((MBFI && Loops) &&
+           "Invalid provider state: must have analysis available");
     return std::make_unique<DevelopmentModeEvictAdvisor>(
-        MF, RA, Runner.get(),
-        getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI(),
-        getAnalysis<MachineLoopInfoWrapperPass>().getLI(), Log.get());
+        MF, RA, Runner.get(), *MBFI, *Loops, Log.get());
   }
 
+private:
+  std::vector<TensorSpec> InputFeatures;
+  std::vector<TensorSpec> TrainingInputFeatures;
+
   std::unique_ptr<MLModelRunner> Runner;
   std::unique_ptr<Logger> Log;
+  const MachineBlockFrequencyInfo *MBFI;
+  const MachineLoopInfo *Loops;
+};
+
+class DevelopmentModeEvictionAdvisorAnalysisLegacy final
+    : public RegAllocEvictionAdvisorAnalysisLegacy {
+public:
+  DevelopmentModeEvictionAdvisorAnalysisLegacy()
+      : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Development) {}
+
+  bool doInitialization(Module &M) override {
+    auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
+    auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
+    Provider.setAnalyses({MBFI, Loops});
+    return Provider.doInitialization(M);
+  }
+
+  std::unique_ptr<RegAllocEvictionAdvisor>
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+    return Provider.getAdvisor(MF, RA);
+  }
+
+  // support for isa<> and dyn_cast.
+  static bool classof(const RegAllocEvictionAdvisorAnalysisLegacy *R) {
+    return R->getAdvisorMode() == AdvisorMode::Development;
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
+    AU.addRequired<MachineLoopInfoWrapperPass>();
+    RegAllocEvictionAdvisorAnalysisLegacy::getAnalysisUsage(AU);
+  }
+
+private:
+  DevelopmentModeEvictionAdvisorProvider Provider;
 };
 
 #endif //#ifdef LLVM_HAVE_TFLITE
@@ -1079,8 +1156,13 @@ void llvm::extractMBBFrequency(
 // Development mode-specific implementations
 #ifdef LLVM_HAVE_TFLITE
 
-RegAllocEvictionAdvisorAnalysis *llvm::createDevelopmentModeAdvisor() {
-  return new DevelopmentModeEvictionAdvisorAnalysis();
+RegAllocEvictionAdvisorProvider *llvm::createDevelopmentModeAdvisorProvider() {
+  return new DevelopmentModeEvictionAdvisorProvider();
+}
+
+RegAllocEvictionAdvisorAnalysisLegacy *
+llvm::createDevelopmentModeAdvisorAnalysisLegacy() {
+  return new DevelopmentModeEvictionAdvisorAnalysisLegacy();
 }
 
 int64_t DevelopmentModeEvictAdvisor::tryFindEvictionCandidatePosition(
@@ -1146,18 +1228,23 @@ bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
     return *CachedReward;
   };
 
-  getAnalysis<RegAllocEvictionAdvisorAnalysis>().logRewardIfNeeded(MF,
-                                                                   GetReward);
+  getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().logRewardIfNeeded(
+      MF, GetReward);
   getAnalysis<RegAllocPriorityAdvisorAnalysis>().logRewardIfNeeded(MF,
                                                                    GetReward);
   return false;
 }
 #endif // #ifdef LLVM_HAVE_TFLITE
 
-RegAllocEvictionAdvisorAnalysis *llvm::createReleaseModeAdvisor() {
+RegAllocEvictionAdvisorProvider *llvm::createReleaseModeAdvisorProvider() {
+  return new ReleaseModeEvictionAdvisorProvider();
+}
+
+RegAllocEvictionAdvisorAnalysisLegacy *
+llvm::createReleaseModeAdvisorAnalysisLegacy() {
   return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
                  !InteractiveChannelBaseName.empty()
-             ? new ReleaseModeEvictionAdvisorAnalysis()
+             ? new ReleaseModeEvictionAdvisorAnalysisLegacy()
              : nullptr;
 }
 
diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
index a1f441ebd0d5e41..09fb5df259a97e9 100644
--- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
@@ -10,9 +10,10 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "RegAllocEvictionAdvisor.h"
+#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
 #include "AllocationOrder.h"
 #include "RegAllocGreedy.h"
+#include "RegAllocPriorityAdvisor.h"
 #include "llvm/CodeGen/LiveRegMatrix.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
@@ -26,17 +27,18 @@
 
 using namespace llvm;
 
-static cl::opt<RegAllocEvictionAdvisorAnalysis::AdvisorMode> Mode(
+static cl::opt<RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode> Mode(
     "regalloc-enable-advisor", cl::Hidden,
-    cl::init(RegAllocEvictionAdvisorAnalysis::AdvisorMode::Default),
+    cl::init(RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default),
     cl::desc("Enable regalloc advisor mode"),
     cl::values(
-        clEnumValN(RegAllocEvictionAdvisorAnalysis::AdvisorMode::Default,
+        clEnumValN(RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default,
                    "default", "Default"),
-        clEnumValN(RegAllocEvictionAdvisorAnalysis::AdvisorMode::Release,
+        clEnumValN(RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Release,
                    "release", "precompiled"),
-        clEnumValN(RegAllocEvictionAdvisorAnalysis::AdvisorMode::Development,
-                   "development", "for training")));
+        clEnumValN(
+            RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Development,
+            "development", "for training")));
 
 static cl::opt<bool> EnableLocalReassignment(
     "enable-local-reassign", cl::Hidden,
@@ -59,59 +61,114 @@ cl::opt<unsigned> EvictInterferenceCutoff(
 #define LLVM_HAVE_TF_AOT
 #endif
 
-char RegAllocEvictionAdvisorAnalysis::ID = 0;
-INITIALIZE_PASS(RegAllocEvictionAdvisorAnalysis, "regalloc-evict",
+char RegAllocEvictionAdvisorAnalysisLegacy::ID = 0;
+INITIALIZE_PASS(RegAllocEvictionAdvisorAnalysisLegacy, "regalloc-evict",
                 "Regalloc eviction policy", false, true)
 
 namespace {
-class DefaultEvictionAdvisorAnalysis final
-    : public RegAllocEvictionAdvisorAnalysis {
+class DefaultEvictionAdvisorProvider final
+    : public RegAllocEvictionAdvisorProvider {
 public:
-  DefaultEvictionAdvisorAnalysis(bool NotAsRequested)
-      : RegAllocEvictionAdvisorAnalysis(AdvisorMode::Default),
+  DefaultEvictionAdvisorProvider(bool NotAsRequested)
+      : RegAllocEvictionAdvisorProvider(AdvisorMode::Default),
         NotAsRequested(NotAsRequested) {}
 
   // support for isa<> and dyn_cast.
-  static bool classof(const RegAllocEvictionAdvisorAnalysis *R) {
+  static bool classof(const RegAllocEvictionAdvisorProvider *R) {
     return R->getAdvisorMode() == AdvisorMode::Default;
   }
 
-private:
   std::unique_ptr<RegAllocEvictionAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
     return std::make_unique<DefaultEvictionAdvisor>(MF, RA);
   }
+
   bool doInitialization(Module &M) override {
     if (NotAsRequested)
       M.getContext().emitError("Requested regalloc eviction advisor analysis "
                                "could not be created. Using default");
-    return RegAllocEvictionAdvisorAnalysis::doInitialization(M);
+    return RegAllocEvictionAdvisorProvider::doInitialization(M);
+  }
+
+private:
+  const bool NotAsRequested;
+};
+
+class DefaultEvictionAdvisorAnalysisLegacy final
+    : public RegAllocEvictionAdvisorAnalysisLegacy {
+public:
+  DefaultEvictionAdvisorAnalysisLegacy(bool NotAsRequested)
+      : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Default),
+        NotAsRequested(NotAsRequested) {}
+
+  std::unique_ptr<RegAllocEvictionAdvisor>
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+    return Provider->getAdvisor(MF, RA);
+  }
+
+  bool doInitialization(Module &M) override {
+    Provider.reset(new DefaultEvictionAdvisorProvider(NotAsRequested));
+    return Provider->doInitialization(M);
   }
+
+  // support for isa<> and dyn_cast.
+  static bool classof(const RegAllocEvictionAdvisorAnalysisLegacy *R) {
+    return R->getAdvisorMode() == AdvisorMode::Default;
+  }
+
+private:
+  std::unique_ptr<DefaultEvictionAdvisorProvider> Provider;
   const bool NotAsRequested;
 };
 } // namespace
 
-template <> Pass *llvm::callDefaultCtor<RegAllocEvictionAdvisorAnalysis>() {
+AnalysisKey RegAllocEvictionAdvisorAnalysis::Key;
+
+RegAllocEvictionAdvisorAnalysis::Result
+RegAllocEvictionAdvisorAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
+  std::unique_ptr<RegAllocEvictionAdvisorProvider> Provider;
+  switch (Mode) {
+  case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default:
+    Provider.reset(
+        new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ false));
+    break;
+  case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Development:
+#if defined(LLVM_HAVE_TFLITE)
+    Provider.reset(createDevelopmentModeAdvisorProvider());
+#endif
+    break;
+  case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Release:
+    Provider.reset(createReleaseModeAdvisorProvider());
+    break;
+  }
+  if (!Provider)
+    Provider.reset(new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ true));
+  Provider->doInitialization(M);
+  return Provider;
+}
+
+template <>
+Pass *llvm::callDefaultCtor<RegAllocEvictionAdvisorAnalysisLegacy>() {
   Pass *Ret = nullptr;
   switch (Mode) {
-  case RegAllocEvictionAdvisorAnalysis::AdvisorMode::Default:
-    Ret = new DefaultEvictionAdvisorAnalysis(/*NotAsRequested*/ false);
+  case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default:
+    Ret = new DefaultEvictionAdvisorAnalysisLegacy(/*NotAsRequested*/ false);
     break;
-  case RegAllocEvictionAdvisorAnalysis::AdvisorMode::Development:
+  case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Development:
 #if defined(LLVM_HAVE_TFLITE)
-    Ret = createDevelopmentModeAdvisor();
+    Ret = createDevelopmentModeAdvisorAnalysisLegacy();
 #endif
     break;
-  case RegAllocEvictionAdvisorAnalysis::AdvisorMode::Release:
-    Ret = createReleaseModeAdvisor();
+  case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Release:
+    Ret = createReleaseModeAdvisorAnalysisLegacy();
     break;
   }
   if (Ret)
     return Ret;
-  return new DefaultEvictionAdvisorAnalysis(/*NotAsRequested*/ true);
+  return new DefaultEvictionAdvisorAnalysisLegacy(/*NotAsRequested*/ true);
 }
 
-StringRef RegAllocEvictionAdvisorAnalysis::getPassName() const {
+StringRef RegAllocEvictionAdvisorAnalysisLegacy::getPassName() const {
   switch (getAdvisorMode()) {
   case AdvisorMode::Default:
     return "Default Regalloc Eviction Advisor";
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp
index 3542bfe18af46f9..8f9a7b46691f0d7 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 "RegAllocEvictionAdvisor.h"
 #include "RegAllocPriorityAdvisor.h"
 #include "SpillPlacement.h"
 #include "SplitKit.h"
@@ -47,6 +46,7 @@
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
 #include "llvm/CodeGen/RegAllocRegistry.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
 #include "llvm/CodeGen/SlotIndexes.h"
@@ -164,7 +164,7 @@ INITIALIZE_PASS_DEPENDENCY(LiveRegMatrixWrapperLegacy)
 INITIALIZE_PASS_DEPENDENCY(EdgeBundlesWrapperLegacy)
 INITIALIZE_PASS_DEPENDENCY(SpillPlacement)
 INITIALIZE_PASS_DEPENDENCY(MachineOptimizationRemarkEmitterPass)
-INITIALIZE_PASS_DEPENDENCY(RegAllocEvictionAdvisorAnalysis)
+INITIALIZE_PASS_DEPENDENCY(RegAllocEvictionAdvisorAnalysisLegacy)
 INITIALIZE_PASS_DEPENDENCY(RegAllocPriorityAdvisorAnalysis)
 INITIALIZE_PASS_END(RAGreedy, "greedy",
                 "Greedy Register Allocator", false, false)
@@ -219,7 +219,7 @@ void RAGreedy::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequired<EdgeBundlesWrapperLegacy>();
   AU.addRequired<SpillPlacement>();
   AU.addRequired<MachineOptimizationRemarkEmitterPass>();
-  AU.addRequired<RegAllocEvictionAdvisorAnalysis>();
+  AU.addRequired<RegAllocEvictionAdvisorAnalysisLegacy>();
   AU.addRequired<RegAllocPriorityAdvisorAnalysis>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
@@ -2748,7 +2748,8 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
 
   ExtraInfo.emplace();
   EvictAdvisor =
-      getAnalysis<RegAllocEvictionAdvisorAnalysis>().getAdvisor(*MF, *this);
+      getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().getAdvisor(*MF,
+                                                                      *this);
   PriorityAdvisor =
       getAnalysis<RegAllocPriorityAdvisorAnalysis>().getAdvisor(*MF, *this);
 
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.h b/llvm/lib/CodeGen/RegAllocGreedy.h
index 2e7608a53e9cec1..62752a7962db804 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.h
+++ b/llvm/lib/CodeGen/RegAllocGreedy.h
@@ -14,7 +14,6 @@
 
 #include "InterferenceCache.h"
 #include "RegAllocBase.h"
-#include "RegAllocEvictionAdvisor.h"
 #include "RegAllocPriorityAdvisor.h"
 #include "SpillPlacement.h"
 #include "SplitKit.h"
diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
index 1e9fa967214cc87..2d42a43c4c63724 100644
--- a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
+++ b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
@@ -9,7 +9,7 @@
 #ifndef LLVM_CODEGEN_REGALLOCPRIORITYADVISOR_H
 #define LLVM_CODEGEN_REGALLOCPRIORITYADVISOR_H
 
-#include "RegAllocEvictionAdvisor.h"
+#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
 #include "llvm/CodeGen/SlotIndexes.h"
 #include "llvm/Pass.h"
 
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index cf7ceed63607a67..af2c2e0058d1c7b 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -121,6 +121,7 @@
 #include "llvm/CodeGen/PHIElimination.h"
 #include "llvm/CodeGen/PeepholeOptimizer.h"
 #include "llvm/CodeGen/PreISelIntrinsicLowering.h"
+#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
 #include "llvm/CodeGen/RegAllocFast.h"
 #include "llvm/CodeGen/RegUsageInfoCollector.h"
 #include "llvm/CodeGen/RegUsageInfoPropagate.h"
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 7c3798f6462a46d..2a2d2ad674680b7 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -32,6 +32,7 @@ MODULE_ANALYSIS("no-op-module", NoOpModuleAnalysis())
 MODULE_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis(PIC))
 MODULE_ANALYSIS("profile-summary", ProfileSummaryAnalysis())
 MODULE_ANALYSIS("reg-usage", PhysicalRegisterUsageAnalysis())
+MODULE_ANALYSIS("regalloc-evict", RegAllocEvictionAdvisorAnalysis())
 MODULE_ANALYSIS("stack-safety", StackSafetyGlobalAnalysis())
 MODULE_ANALYSIS("verify", VerifierAnalysis())
 

>From 6faee3f719576c5dbd448a246aece2eae22db9f6 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Fri, 29 Nov 2024 09:18:59 +0000
Subject: [PATCH 2/2] Remove Provider::doInit, change to MF analysis

init is in the constructor. The new MF analysis lazily intiailizes the
provider. The (wrapped) provider is the analysis result.
---
 .../llvm/CodeGen/RegAllocEvictionAdvisor.h    |  58 ++++++--
 llvm/include/llvm/Passes/CodeGenPassBuilder.h |   5 +-
 .../llvm/Passes/MachinePassRegistry.def       |   1 +
 llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp   | 136 ++++++++----------
 llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp  |  56 ++++----
 llvm/lib/Passes/PassRegistry.def              |   1 -
 6 files changed, 143 insertions(+), 114 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
index 847bf032235c1dd..8261f6d1cf0bbde 100644
--- a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
+++ b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
@@ -13,6 +13,8 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
+#include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/Register.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/IR/PassManager.h"
@@ -200,33 +202,43 @@ class RegAllocEvictionAdvisorAnalysisLegacy : public ImmutablePass {
 class RegAllocEvictionAdvisorProvider {
 public:
   enum class AdvisorMode : int { Default, Release, Development };
-  RegAllocEvictionAdvisorProvider(AdvisorMode Mode) : Mode(Mode) {}
+  RegAllocEvictionAdvisorProvider(AdvisorMode Mode, LLVMContext &Ctx)
+      : Ctx(Ctx), Mode(Mode) {}
 
   virtual ~RegAllocEvictionAdvisorProvider() = default;
 
-  virtual bool doInitialization(Module &M) { return false; }
-
   virtual void logRewardIfNeeded(const MachineFunction &MF,
                                  llvm::function_ref<float()> GetReward) {}
 
   virtual std::unique_ptr<RegAllocEvictionAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
 
-  /// Set the analyses that the advisor needs to use as they might not be
-  /// available before the advisor is created.
-  virtual void setAnalyses(std::initializer_list<llvm::Any> AnalysisP) {}
+  /// We create this provider in doInitialization which doesn't have these
+  /// analyses. For NPM, we do have them in run(MachineFunction&)
+  virtual void setAnalyses(MachineBlockFrequencyInfo *MBFI,
+                           MachineLoopInfo *Loops) {
+    this->MBFI = MBFI;
+    this->Loops = Loops;
+  }
 
   AdvisorMode getAdvisorMode() const { return Mode; }
 
+protected:
+  LLVMContext &Ctx;
+  MachineBlockFrequencyInfo *MBFI;
+  MachineLoopInfo *Loops;
+
 private:
   const AdvisorMode Mode;
 };
 
-RegAllocEvictionAdvisorProvider *createReleaseModeAdvisorProvider();
-RegAllocEvictionAdvisorProvider *createDevelopmentModeAdvisorProvider();
+RegAllocEvictionAdvisorProvider *
+createReleaseModeAdvisorProvider(LLVMContext &Ctx);
+RegAllocEvictionAdvisorProvider *
+createDevelopmentModeAdvisorProvider(LLVMContext &Ctx);
 
-/// A Module analysis for fetching the Eviction Advisor. This is not a
-/// MachineFunction analysis for two reasons:
+/// A MachineFunction analysis for fetching the Eviction Advisor.
+/// This sets up the Provider lazily and caches it.
 /// - in the ML implementation case, the evaluator is stateless but (especially
 /// in the development mode) expensive to set up. With a Module Analysis, we
 /// `require` it and set it up once.
@@ -241,8 +253,30 @@ class RegAllocEvictionAdvisorAnalysis
   friend AnalysisInfoMixin<RegAllocEvictionAdvisorAnalysis>;
 
 public:
-  using Result = std::unique_ptr<RegAllocEvictionAdvisorProvider>;
-  Result run(Module &MF, ModuleAnalysisManager &MAM);
+  struct Result {
+    // owned by this analysis
+    RegAllocEvictionAdvisorProvider *Provider;
+
+    bool invalidate(MachineFunction &MF, const PreservedAnalyses &PA,
+                    MachineFunctionAnalysisManager::Invalidator &Inv) {
+      auto PAC = PA.getChecker<RegAllocEvictionAdvisorAnalysis>();
+      // If we are in default mode, the provider is always valid.
+      if (Provider->getAdvisorMode() ==
+          RegAllocEvictionAdvisorProvider::AdvisorMode::Default)
+        return !PAC.preservedWhenStateless();
+      // MBFI and Loops are used in release and development modes, so check
+      // those.
+      return !PAC.preservedWhenStateless() ||
+             Inv.invalidate<MachineBlockFrequencyAnalysis>(MF, PA) ||
+             Inv.invalidate<MachineLoopAnalysis>(MF, PA);
+    }
+  };
+
+  Result run(MachineFunction &MF, MachineFunctionAnalysisManager &MAM);
+
+private:
+  void initializeProvider(LLVMContext &Ctx);
+  std::unique_ptr<RegAllocEvictionAdvisorProvider> Provider;
 };
 
 /// Specialization for the API used by the analysis infrastructure to create
diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index aa6437c1d2f0172..bcd335f2f99427e 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -1062,10 +1062,9 @@ void CodeGenPassBuilder<Derived, TargetMachineT>::addMachineSSAOptimization(
 template <typename Derived, typename TargetMachineT>
 void CodeGenPassBuilder<Derived, TargetMachineT>::addTargetRegisterAllocator(
     AddMachinePass &addPass, bool Optimized) const {
-  if (Optimized) {
-    addPass(RequireAnalysis<RegAllocEvictionAdvisorAnalysis>());
+  if (Optimized)
     addPass(RAGreedyPass());
-  } else
+  else
     addPass(RegAllocFastPass());
 }
 
diff --git a/llvm/include/llvm/Passes/MachinePassRegistry.def b/llvm/include/llvm/Passes/MachinePassRegistry.def
index cb1c295d824787b..e84f2d19f2bd621 100644
--- a/llvm/include/llvm/Passes/MachinePassRegistry.def
+++ b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -112,6 +112,7 @@ MACHINE_FUNCTION_ANALYSIS("machine-post-dom-tree",
                           MachinePostDominatorTreeAnalysis())
 MACHINE_FUNCTION_ANALYSIS("machine-trace-metrics", MachineTraceMetricsAnalysis())
 MACHINE_FUNCTION_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis(PIC))
+MACHINE_FUNCTION_ANALYSIS("regalloc-evict", RegAllocEvictionAdvisorAnalysis())
 MACHINE_FUNCTION_ANALYSIS("slot-indexes", SlotIndexesAnalysis())
 MACHINE_FUNCTION_ANALYSIS("virtregmap", VirtRegMapAnalysis())
 // MACHINE_FUNCTION_ANALYSIS("live-stacks", LiveStacksPass())
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 9c15c327bb202c2..76011b0a6e4bfc5 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -370,8 +370,8 @@ class MLEvictAdvisor : public RegAllocEvictionAdvisor {
 class ReleaseModeEvictionAdvisorProvider final
     : public RegAllocEvictionAdvisorProvider {
 public:
-  ReleaseModeEvictionAdvisorProvider()
-      : RegAllocEvictionAdvisorProvider(AdvisorMode::Release) {
+  ReleaseModeEvictionAdvisorProvider(LLVMContext &Ctx)
+      : RegAllocEvictionAdvisorProvider(AdvisorMode::Release, Ctx) {
     if (EnableDevelopmentFeatures) {
       InputFeatures = {RA_EVICT_FEATURES_LIST(
           _DECL_FEATURES) RA_EVICT_FIRST_DEVELOPMENT_FEATURE(_DECL_FEATURES)
@@ -385,15 +385,6 @@ class ReleaseModeEvictionAdvisorProvider final
     return R->getAdvisorMode() == AdvisorMode::Release;
   }
 
-  void setAnalyses(std::initializer_list<llvm::Any> AnalysisList) override {
-    for (auto Analysis : AnalysisList) {
-      if (auto **MBFI = llvm::any_cast<MachineBlockFrequencyInfo *>(&Analysis))
-        this->MBFI = *MBFI;
-      if (auto **Loops = llvm::any_cast<MachineLoopInfo *>(&Analysis))
-        this->Loops = *Loops;
-    }
-  }
-
   std::unique_ptr<RegAllocEvictionAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
     if (!Runner) {
@@ -406,6 +397,8 @@ class ReleaseModeEvictionAdvisorProvider final
             InteractiveChannelBaseName + ".out",
             InteractiveChannelBaseName + ".in");
     }
+    assert((MBFI && Loops) &&
+           "Invalid provider state: must have analysis available");
     return std::make_unique<MLEvictAdvisor>(MF, RA, Runner.get(), *MBFI,
                                             *Loops);
   }
@@ -413,8 +406,6 @@ class ReleaseModeEvictionAdvisorProvider final
 private:
   std::vector<TensorSpec> InputFeatures;
   std::unique_ptr<MLModelRunner> Runner;
-  MachineBlockFrequencyInfo *MBFI;
-  MachineLoopInfo *Loops;
 };
 
 class ReleaseModeEvictionAdvisorAnalysisLegacy final
@@ -427,12 +418,19 @@ class ReleaseModeEvictionAdvisorAnalysisLegacy final
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
     auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
     auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
-    Provider.setAnalyses({MBFI, Loops});
-    return Provider.getAdvisor(MF, RA);
+    Provider->setAnalyses(MBFI, Loops);
+    return Provider->getAdvisor(MF, RA);
+  }
+
+  void logRewardIfNeeded(const MachineFunction &MF,
+                         llvm::function_ref<float()> GetReward) override {
+    // No-op in release mode
   }
 
   bool doInitialization(Module &M) override {
-    return Provider.doInitialization(M);
+    Provider =
+        std::make_unique<ReleaseModeEvictionAdvisorProvider>(M.getContext());
+    return false;
   }
 
   static bool classof(const RegAllocEvictionAdvisorAnalysisLegacy *R) {
@@ -446,7 +444,7 @@ class ReleaseModeEvictionAdvisorAnalysisLegacy final
   }
 
 private:
-  ReleaseModeEvictionAdvisorProvider Provider;
+  std::unique_ptr<ReleaseModeEvictionAdvisorProvider> Provider;
 };
 
 // ===================================
@@ -484,11 +482,8 @@ class DevelopmentModeEvictAdvisor : public MLEvictAdvisor {
 class DevelopmentModeEvictionAdvisorProvider final
     : public RegAllocEvictionAdvisorProvider {
 public:
-  DevelopmentModeEvictionAdvisorProvider(
-      MachineBlockFrequencyInfo *MBFI = nullptr,
-      MachineLoopInfo *Loops = nullptr)
-      : RegAllocEvictionAdvisorProvider(AdvisorMode::Development), MBFI(MBFI),
-        Loops(Loops) {
+  DevelopmentModeEvictionAdvisorProvider(LLVMContext &Ctx)
+      : RegAllocEvictionAdvisorProvider(AdvisorMode::Development, Ctx) {
     if (EnableDevelopmentFeatures) {
       InputFeatures = {RA_EVICT_FEATURES_LIST(
           _DECL_FEATURES) RA_EVICT_FIRST_DEVELOPMENT_FEATURE(_DECL_FEATURES)
@@ -508,43 +503,10 @@ class DevelopmentModeEvictionAdvisorProvider final
           TensorSpec::createSpec<int32_t>("action_step_type", {1}),
           TensorSpec::createSpec<float>("action_reward", {1})};
     }
-  }
-  // support for isa<> and dyn_cast.
-  static bool classof(const RegAllocEvictionAdvisorProvider *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());
-  }
-
-  void setAnalyses(std::initializer_list<llvm::Any> AnalysisList) override {
-    for (auto Analysis : AnalysisList) {
-      if (auto **MBFI = llvm::any_cast<MachineBlockFrequencyInfo *>(&Analysis))
-        this->MBFI = *MBFI;
-      if (auto **Loops = llvm::any_cast<MachineLoopInfo *>(&Analysis))
-        this->Loops = *Loops;
-    }
-  }
-
-  bool doInitialization(Module &M) override {
-    LLVMContext &Ctx = M.getContext();
     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);
@@ -553,15 +515,15 @@ class DevelopmentModeEvictionAdvisorProvider 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()))
@@ -573,7 +535,28 @@ class DevelopmentModeEvictionAdvisorProvider final
 
     Log = std::make_unique<Logger>(std::move(OS), LFS, Reward,
                                    /*IncludeReward*/ true);
-    return false;
+    return;
+  }
+
+  // support for isa<> and dyn_cast.
+  static bool classof(const RegAllocEvictionAdvisorProvider *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());
   }
 
   std::unique_ptr<RegAllocEvictionAdvisor>
@@ -594,8 +577,6 @@ class DevelopmentModeEvictionAdvisorProvider final
 
   std::unique_ptr<MLModelRunner> Runner;
   std::unique_ptr<Logger> Log;
-  const MachineBlockFrequencyInfo *MBFI;
-  const MachineLoopInfo *Loops;
 };
 
 class DevelopmentModeEvictionAdvisorAnalysisLegacy final
@@ -605,15 +586,22 @@ class DevelopmentModeEvictionAdvisorAnalysisLegacy final
       : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Development) {}
 
   bool doInitialization(Module &M) override {
-    auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
-    auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
-    Provider.setAnalyses({MBFI, Loops});
-    return Provider.doInitialization(M);
+    Provider = std::make_unique<DevelopmentModeEvictionAdvisorProvider>(
+        M.getContext());
+    return false;
+  }
+
+  void logRewardIfNeeded(const MachineFunction &MF,
+                         llvm::function_ref<float()> GetReward) override {
+    Provider->logRewardIfNeeded(MF, GetReward);
   }
 
   std::unique_ptr<RegAllocEvictionAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
-    return Provider.getAdvisor(MF, RA);
+    auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
+    auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
+    Provider->setAnalyses(MBFI, Loops);
+    return Provider->getAdvisor(MF, RA);
   }
 
   // support for isa<> and dyn_cast.
@@ -628,7 +616,7 @@ class DevelopmentModeEvictionAdvisorAnalysisLegacy final
   }
 
 private:
-  DevelopmentModeEvictionAdvisorProvider Provider;
+  std::unique_ptr<DevelopmentModeEvictionAdvisorProvider> Provider;
 };
 
 #endif //#ifdef LLVM_HAVE_TFLITE
@@ -1156,8 +1144,9 @@ void llvm::extractMBBFrequency(
 // Development mode-specific implementations
 #ifdef LLVM_HAVE_TFLITE
 
-RegAllocEvictionAdvisorProvider *llvm::createDevelopmentModeAdvisorProvider() {
-  return new DevelopmentModeEvictionAdvisorProvider();
+RegAllocEvictionAdvisorProvider *
+llvm::createDevelopmentModeAdvisorProvider(LLVMContext &Ctx) {
+  return new DevelopmentModeEvictionAdvisorProvider(Ctx);
 }
 
 RegAllocEvictionAdvisorAnalysisLegacy *
@@ -1236,8 +1225,9 @@ bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
 }
 #endif // #ifdef LLVM_HAVE_TFLITE
 
-RegAllocEvictionAdvisorProvider *llvm::createReleaseModeAdvisorProvider() {
-  return new ReleaseModeEvictionAdvisorProvider();
+RegAllocEvictionAdvisorProvider *
+llvm::createReleaseModeAdvisorProvider(LLVMContext &Ctx) {
+  return new ReleaseModeEvictionAdvisorProvider(Ctx);
 }
 
 RegAllocEvictionAdvisorAnalysisLegacy *
diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
index 09fb5df259a97e9..4c7528ebffbfd55 100644
--- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
@@ -9,13 +9,14 @@
 // Implementation of the default eviction advisor and of the Analysis pass.
 //
 //===----------------------------------------------------------------------===//
-
 #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/RegisterClassInfo.h"
 #include "llvm/CodeGen/VirtRegMap.h"
 #include "llvm/IR/Module.h"
@@ -69,9 +70,12 @@ namespace {
 class DefaultEvictionAdvisorProvider final
     : public RegAllocEvictionAdvisorProvider {
 public:
-  DefaultEvictionAdvisorProvider(bool NotAsRequested)
-      : RegAllocEvictionAdvisorProvider(AdvisorMode::Default),
-        NotAsRequested(NotAsRequested) {}
+  DefaultEvictionAdvisorProvider(bool NotAsRequested, LLVMContext &Ctx)
+      : RegAllocEvictionAdvisorProvider(AdvisorMode::Default, Ctx) {
+    if (NotAsRequested)
+      Ctx.emitError("Requested regalloc eviction advisor analysis "
+                    "could not be created. Using default");
+  }
 
   // support for isa<> and dyn_cast.
   static bool classof(const RegAllocEvictionAdvisorProvider *R) {
@@ -82,16 +86,6 @@ class DefaultEvictionAdvisorProvider final
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
     return std::make_unique<DefaultEvictionAdvisor>(MF, RA);
   }
-
-  bool doInitialization(Module &M) override {
-    if (NotAsRequested)
-      M.getContext().emitError("Requested regalloc eviction advisor analysis "
-                               "could not be created. Using default");
-    return RegAllocEvictionAdvisorProvider::doInitialization(M);
-  }
-
-private:
-  const bool NotAsRequested;
 };
 
 class DefaultEvictionAdvisorAnalysisLegacy final
@@ -103,12 +97,14 @@ class DefaultEvictionAdvisorAnalysisLegacy final
 
   std::unique_ptr<RegAllocEvictionAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+    // MBFI and Loops not required here.
     return Provider->getAdvisor(MF, RA);
   }
 
   bool doInitialization(Module &M) override {
-    Provider.reset(new DefaultEvictionAdvisorProvider(NotAsRequested));
-    return Provider->doInitialization(M);
+    Provider.reset(
+        new DefaultEvictionAdvisorProvider(NotAsRequested, M.getContext()));
+    return false;
   }
 
   // support for isa<> and dyn_cast.
@@ -124,27 +120,37 @@ class DefaultEvictionAdvisorAnalysisLegacy final
 
 AnalysisKey RegAllocEvictionAdvisorAnalysis::Key;
 
-RegAllocEvictionAdvisorAnalysis::Result
-RegAllocEvictionAdvisorAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
-  std::unique_ptr<RegAllocEvictionAdvisorProvider> Provider;
+void RegAllocEvictionAdvisorAnalysis::initializeProvider(LLVMContext &Ctx) {
+  if (Provider)
+    return;
   switch (Mode) {
   case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default:
     Provider.reset(
-        new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ false));
+        new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ false, Ctx));
     break;
   case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Development:
 #if defined(LLVM_HAVE_TFLITE)
-    Provider.reset(createDevelopmentModeAdvisorProvider());
+    Provider.reset(createDevelopmentModeAdvisorProvider(Ctx));
 #endif
     break;
   case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Release:
-    Provider.reset(createReleaseModeAdvisorProvider());
+    Provider.reset(createReleaseModeAdvisorProvider(Ctx));
     break;
   }
   if (!Provider)
-    Provider.reset(new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ true));
-  Provider->doInitialization(M);
-  return Provider;
+    Provider.reset(
+        new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ true, Ctx));
+}
+
+RegAllocEvictionAdvisorAnalysis::Result
+RegAllocEvictionAdvisorAnalysis::run(MachineFunction &MF,
+                                     MachineFunctionAnalysisManager &MFAM) {
+  // Lazy initialization of the provider.
+  initializeProvider(MF.getFunction().getContext());
+  auto *MBFI = &MFAM.getResult<MachineBlockFrequencyAnalysis>(MF);
+  auto *Loops = &MFAM.getResult<MachineLoopAnalysis>(MF);
+  Provider->setAnalyses(MBFI, Loops);
+  return Result{Provider.get()};
 }
 
 template <>
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 2a2d2ad674680b7..7c3798f6462a46d 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -32,7 +32,6 @@ MODULE_ANALYSIS("no-op-module", NoOpModuleAnalysis())
 MODULE_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis(PIC))
 MODULE_ANALYSIS("profile-summary", ProfileSummaryAnalysis())
 MODULE_ANALYSIS("reg-usage", PhysicalRegisterUsageAnalysis())
-MODULE_ANALYSIS("regalloc-evict", RegAllocEvictionAdvisorAnalysis())
 MODULE_ANALYSIS("stack-safety", StackSafetyGlobalAnalysis())
 MODULE_ANALYSIS("verify", VerifierAnalysis())
 



More information about the llvm-commits mailing list