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.

Author: optimisan

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

llvm/include/llvm/Passes/CodeGenPassBuilder.h

@@ -1065,10 +1065,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());
 }
 

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -114,6 +114,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("spill-code-placement", SpillPlacementAnalysis())
 MACHINE_FUNCTION_ANALYSIS("virtregmap", VirtRegMapAnalysis())

llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp

@@ -393,8 +393,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)
@@ -408,15 +408,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) {
@@ -429,15 +420,15 @@ 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);
   }
 
 private:
   std::vector<TensorSpec> InputFeatures;
   std::unique_ptr<MLModelRunner> Runner;
-  MachineBlockFrequencyInfo *MBFI;
-  MachineLoopInfo *Loops;
 };
 
 class ReleaseModeEvictionAdvisorAnalysisLegacy final
@@ -450,12 +441,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) {
@@ -469,7 +467,7 @@ class ReleaseModeEvictionAdvisorAnalysisLegacy final
   }
 
 private:
-  ReleaseModeEvictionAdvisorProvider Provider;
+  std::unique_ptr<ReleaseModeEvictionAdvisorProvider> Provider;
 };
 
 // ===================================
@@ -507,11 +505,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)
@@ -531,43 +526,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);
@@ -576,15 +538,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()))
@@ -596,7 +558,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>
@@ -617,8 +600,6 @@ class DevelopmentModeEvictionAdvisorProvider final
 
   std::unique_ptr<MLModelRunner> Runner;
   std::unique_ptr<Logger> Log;
-  const MachineBlockFrequencyInfo *MBFI;
-  const MachineLoopInfo *Loops;
 };
 
 class DevelopmentModeEvictionAdvisorAnalysisLegacy final
@@ -628,15 +609,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.
@@ -651,7 +639,7 @@ class DevelopmentModeEvictionAdvisorAnalysisLegacy final
   }
 
 private:
-  DevelopmentModeEvictionAdvisorProvider Provider;
+  std::unique_ptr<DevelopmentModeEvictionAdvisorProvider> Provider;
 };
 
 #endif // #ifdef LLVM_HAVE_TFLITE
@@ -1204,8 +1192,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 *
@@ -1284,8 +1273,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 *