[CodeGen][NewPM] Port RegAllocPriorityAdvisor analysis to NPM

llvm/lib/CodeGen/RegAllocPriorityAdvisor.h → 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,20 +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,
+ function_ref<float()> GetReward) {};
+
+ virtual std::unique_ptr<RegAllocPriorityAdvisor>
+ getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
+ SlotIndexes &SI) = 0;
+
+ AdvisorMode getAdvisorMode() const { return Mode; }
+
+private:
+ const AdvisorMode Mode;
+};
+
+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);
+ void initializeMLProvider(RegAllocPriorityAdvisorProvider::AdvisorMode Mode,
+ 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)
- virtual std::unique_ptr<RegAllocPriorityAdvisor>
- getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
+ RegAllocPriorityAdvisorProvider &getProvider() { return *Provider; }
  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
@@ -90,18 +144,28 @@ class RegAllocPriorityAdvisorAnalysis : public ImmutablePass {
  AU.setPreservesAll();
  }
 
+ std::unique_ptr<RegAllocPriorityAdvisorProvider> Provider;
+
 private:
  StringRef getPassName() const override;
  const AdvisorMode Mode;
 };
 
 /// 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>();
+
+RegAllocPriorityAdvisorAnalysisLegacy *
+createReleaseModePriorityAdvisorAnalysis();
+
+RegAllocPriorityAdvisorAnalysisLegacy *
+createDevelopmentModePriorityAdvisorAnalysis();
 
-RegAllocPriorityAdvisorAnalysis *createReleaseModePriorityAdvisor();
+LLVM_ATTRIBUTE_RETURNS_NONNULL RegAllocPriorityAdvisorProvider *
+createReleaseModePriorityAdvisorProvider();
 
-RegAllocPriorityAdvisorAnalysis *createDevelopmentModePriorityAdvisor();
+LLVM_ATTRIBUTE_RETURNS_NONNULL RegAllocPriorityAdvisorProvider *
+createDevelopmentModePriorityAdvisorProvider(LLVMContext &Ctx);
 
 } // namespace llvm
 

llvm/include/llvm/InitializePasses.h

@@ -253,7 +253,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 &);

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())

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