[llvm-branch-commits] [llvm] [RegAlloc][NewPM] Plug Greedy RA in codegen pipeline (PR #120557)

Akshat Oke via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Tue Dec 31 22:55:58 PST 2024


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

>From 4496358f0a2d81710b0d3e451ccb93962c8978a7 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Wed, 1 Jan 2025 06:33:26 +0000
Subject: [PATCH 01/22] [Support] Recycler: Enforce minimum allocation size

Recycler uses reinterpret_cast to an internal structure of size 8.
Invalid write occurs if Recycler is used for objects with sizes less
than 8.
---
 llvm/include/llvm/Support/Recycler.h    |  1 +
 llvm/unittests/Support/CMakeLists.txt   |  1 +
 llvm/unittests/Support/RecyclerTest.cpp | 46 +++++++++++++++++++++++++
 3 files changed, 48 insertions(+)
 create mode 100644 llvm/unittests/Support/RecyclerTest.cpp

diff --git a/llvm/include/llvm/Support/Recycler.h b/llvm/include/llvm/Support/Recycler.h
index bbd9ae321ae30c..2bf30b0b3d5247 100644
--- a/llvm/include/llvm/Support/Recycler.h
+++ b/llvm/include/llvm/Support/Recycler.h
@@ -85,6 +85,7 @@ class Recycler {
                   "Recycler allocation alignment is less than object align!");
     static_assert(sizeof(SubClass) <= Size,
                   "Recycler allocation size is less than object size!");
+    static_assert(Size >= 8 && "Recycler size must be atleast 8");
     return FreeList ? reinterpret_cast<SubClass *>(pop_val())
                     : static_cast<SubClass *>(Allocator.Allocate(Size, Align));
   }
diff --git a/llvm/unittests/Support/CMakeLists.txt b/llvm/unittests/Support/CMakeLists.txt
index d64f89847aa8e7..6de81658264420 100644
--- a/llvm/unittests/Support/CMakeLists.txt
+++ b/llvm/unittests/Support/CMakeLists.txt
@@ -69,6 +69,7 @@ add_llvm_unittest(SupportTests
   PerThreadBumpPtrAllocatorTest.cpp
   ProcessTest.cpp
   ProgramTest.cpp
+  RecyclerTest.cpp
   RegexTest.cpp
   ReverseIterationTest.cpp
   ReplaceFileTest.cpp
diff --git a/llvm/unittests/Support/RecyclerTest.cpp b/llvm/unittests/Support/RecyclerTest.cpp
new file mode 100644
index 00000000000000..8cd763c0b83f8a
--- /dev/null
+++ b/llvm/unittests/Support/RecyclerTest.cpp
@@ -0,0 +1,46 @@
+//===--- unittest/Support/RecyclerTest.cpp --------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Support/Recycler.h"
+#include "llvm/Support/AllocatorBase.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+struct Object1 {
+  char Data[1];
+};
+
+class DecoratedMallocAllocator : public MallocAllocator {
+public:
+  int DeallocCount = 0;
+
+  template <typename T> void Deallocate(T *Ptr) {
+    DeallocCount++;
+    MallocAllocator::Deallocate(Ptr);
+  }
+};
+
+TEST(RecyclerTest, RecycleAllocation) {
+  DecoratedMallocAllocator Allocator;
+  // Recycler needs size to be atleast 8 bytes.
+  Recycler<Object1, 8, 8> R;
+  Object1 *A1 = R.Allocate(Allocator);
+  Object1 *A2 = R.Allocate(Allocator);
+  R.Deallocate(Allocator, A2);
+  Object1 *A3 = R.Allocate(Allocator);
+  EXPECT_EQ(A2, A3); // reuse the deallocated object.
+  R.Deallocate(Allocator, A1);
+  R.Deallocate(Allocator, A3);
+  R.clear(Allocator); // Should deallocate A1 and A3.
+  EXPECT_EQ(Allocator.DeallocCount, 2);
+}
+
+} // end anonymous namespace

>From a55eb363197ba7f3bc7c175c4a77d6bc62dbfaa8 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Thu, 19 Dec 2024 06:57:46 +0000
Subject: [PATCH 02/22] [Support] Recycler: Implement move constructor

---
 llvm/include/llvm/Support/Recycler.h | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/llvm/include/llvm/Support/Recycler.h b/llvm/include/llvm/Support/Recycler.h
index 2bf30b0b3d5247..19e5faa535edd7 100644
--- a/llvm/include/llvm/Support/Recycler.h
+++ b/llvm/include/llvm/Support/Recycler.h
@@ -60,6 +60,10 @@ class Recycler {
     // clear() before deleting the Recycler.
     assert(!FreeList && "Non-empty recycler deleted!");
   }
+  Recycler(const Recycler &) = delete;
+  Recycler(Recycler &&Other)
+      : FreeList(std::exchange(Other.FreeList, nullptr)) {}
+  Recycler() = default;
 
   /// clear - Release all the tracked allocations to the allocator. The
   /// recycler must be free of any tracked allocations before being

>From 30e38c215d219164d2257023416a2f584f312d56 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Wed, 1 Jan 2025 06:35:09 +0000
Subject: [PATCH 03/22] Add test

---
 llvm/unittests/Support/RecyclerTest.cpp | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/llvm/unittests/Support/RecyclerTest.cpp b/llvm/unittests/Support/RecyclerTest.cpp
index 8cd763c0b83f8a..0ef9c5aa0ec82b 100644
--- a/llvm/unittests/Support/RecyclerTest.cpp
+++ b/llvm/unittests/Support/RecyclerTest.cpp
@@ -18,6 +18,10 @@ struct Object1 {
   char Data[1];
 };
 
+struct Object8 {
+  char Data[8];
+};
+
 class DecoratedMallocAllocator : public MallocAllocator {
 public:
   int DeallocCount = 0;
@@ -43,4 +47,19 @@ TEST(RecyclerTest, RecycleAllocation) {
   EXPECT_EQ(Allocator.DeallocCount, 2);
 }
 
+TEST(RecyclerTest, MoveConstructor) {
+  DecoratedMallocAllocator Allocator;
+  Recycler<Object8> R;
+  Object8 *A1 = R.Allocate(Allocator);
+  Object8 *A2 = R.Allocate(Allocator);
+  R.Deallocate(Allocator, A1);
+  R.Deallocate(Allocator, A2);
+  Recycler<Object8> R2(std::move(R));
+  Object8 *A3 = R2.Allocate(Allocator);
+  R2.Deallocate(Allocator, A3);
+  R.clear(Allocator); // Should not deallocate anything as it was moved from.
+  EXPECT_EQ(Allocator.DeallocCount, 0);
+  R2.clear(Allocator);
+  EXPECT_EQ(Allocator.DeallocCount, 2);
+}
 } // end anonymous namespace

>From 91998f98bcd3eaa309d4b4a79effb5a1d9c16959 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 04/22] [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 52dd946a685400..847bf032235c1d 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 1cb9013bc48cc5..8915bf258005a7 100644
--- a/llvm/include/llvm/InitializePasses.h
+++ b/llvm/include/llvm/InitializePasses.h
@@ -253,7 +253,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 d2e9e8185a2b90..aa6437c1d2f017 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 9c6487b40d6061..1eca5c92e2fac4 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"
@@ -114,7 +114,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);
@@ -372,11 +372,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)
@@ -386,17 +387,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>
@@ -411,12 +412,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;
 };
 
 // ===================================
@@ -451,11 +487,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)
@@ -477,7 +516,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;
   }
 
@@ -497,14 +536,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 {
@@ -550,14 +588,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
@@ -1094,8 +1171,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(
@@ -1161,18 +1243,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 a1f441ebd0d5e4..09fb5df259a97e 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 4fa2bc76b38b4b..7743de321716ec 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 "SplitKit.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -46,6 +45,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(SpillPlacementWrapperLegacy)
 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<SpillPlacementWrapperLegacy>();
   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 594c481826cf09..001a1e86a96914 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 "SplitKit.h"
 #include "llvm/ADT/ArrayRef.h"
diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
index 1e9fa967214cc8..2d42a43c4c6372 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 aa5b9077376e05..132492ccbdd55e 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -123,6 +123,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 ecc9f554937103..3ca94e47de7821 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -33,6 +33,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 15c98ffff9da190cd3fcc6022ad78aee9d9f876a 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 05/22] 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 847bf032235c1d..8261f6d1cf0bbd 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 aa6437c1d2f017..bcd335f2f99427 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 5a4e79d7225db1..faa7c68ad2acc6 100644
--- a/llvm/include/llvm/Passes/MachinePassRegistry.def
+++ b/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())
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 1eca5c92e2fac4..9523183745d119 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -376,8 +376,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)
@@ -391,15 +391,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) {
@@ -412,6 +403,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);
   }
@@ -419,8 +412,6 @@ class ReleaseModeEvictionAdvisorProvider final
 private:
   std::vector<TensorSpec> InputFeatures;
   std::unique_ptr<MLModelRunner> Runner;
-  MachineBlockFrequencyInfo *MBFI;
-  MachineLoopInfo *Loops;
 };
 
 class ReleaseModeEvictionAdvisorAnalysisLegacy final
@@ -433,12 +424,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) {
@@ -452,7 +450,7 @@ class ReleaseModeEvictionAdvisorAnalysisLegacy final
   }
 
 private:
-  ReleaseModeEvictionAdvisorProvider Provider;
+  std::unique_ptr<ReleaseModeEvictionAdvisorProvider> Provider;
 };
 
 // ===================================
@@ -490,11 +488,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)
@@ -514,43 +509,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);
@@ -559,15 +521,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()))
@@ -579,7 +541,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>
@@ -600,8 +583,6 @@ class DevelopmentModeEvictionAdvisorProvider final
 
   std::unique_ptr<MLModelRunner> Runner;
   std::unique_ptr<Logger> Log;
-  const MachineBlockFrequencyInfo *MBFI;
-  const MachineLoopInfo *Loops;
 };
 
 class DevelopmentModeEvictionAdvisorAnalysisLegacy final
@@ -611,15 +592,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.
@@ -634,7 +622,7 @@ class DevelopmentModeEvictionAdvisorAnalysisLegacy final
   }
 
 private:
-  DevelopmentModeEvictionAdvisorProvider Provider;
+  std::unique_ptr<DevelopmentModeEvictionAdvisorProvider> Provider;
 };
 
 #endif // #ifdef LLVM_HAVE_TFLITE
@@ -1171,8 +1159,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 *
@@ -1251,8 +1240,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 09fb5df259a97e..4c7528ebffbfd5 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 3ca94e47de7821..ecc9f554937103 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -33,7 +33,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())
 

>From 82db37e37c7c140158d9bf6d83492c8a97eca027 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Tue, 3 Dec 2024 06:17:46 +0000
Subject: [PATCH 06/22] Remove parenthesis

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

diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 9523183745d119..2eab8fccceced9 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -571,7 +571,7 @@ class DevelopmentModeEvictionAdvisorProvider final
       return nullptr;
     if (Log)
       Log->switchContext(MF.getName());
-    assert((MBFI && Loops) &&
+    assert(MBFI && Loops &&
            "Invalid provider state: must have analysis available");
     return std::make_unique<DevelopmentModeEvictAdvisor>(
         MF, RA, Runner.get(), *MBFI, *Loops, Log.get());

>From c62ac13c52a6620a263346682776f64072fa64a9 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Wed, 4 Dec 2024 06:44:24 +0000
Subject: [PATCH 07/22] Add TODO to remove provider, remove paren

---
 llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h | 1 +
 llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp         | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
index 8261f6d1cf0bbd..9ebe5c22f47399 100644
--- a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
+++ b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
@@ -199,6 +199,7 @@ class RegAllocEvictionAdvisorAnalysisLegacy : public ImmutablePass {
 /// 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.
+/// TODO: Coalesce this with the NPM analysis when legacy PM is removed.
 class RegAllocEvictionAdvisorProvider {
 public:
   enum class AdvisorMode : int { Default, Release, Development };
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 2eab8fccceced9..534237713f0cb5 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -403,7 +403,7 @@ class ReleaseModeEvictionAdvisorProvider final
             InteractiveChannelBaseName + ".out",
             InteractiveChannelBaseName + ".in");
     }
-    assert((MBFI && Loops) &&
+    assert(MBFI && Loops &&
            "Invalid provider state: must have analysis available");
     return std::make_unique<MLEvictAdvisor>(MF, RA, Runner.get(), *MBFI,
                                             *Loops);

>From 4ca73fe4fe488c285690731053e719b99282c231 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Mon, 9 Dec 2024 08:28:04 +0000
Subject: [PATCH 08/22] Remove createProvider*(), add new method in ML.cpp

---
 .../llvm/CodeGen/RegAllocEvictionAdvisor.h    | 13 +++++-----
 llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp   | 24 ++++++++++++-------
 llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp  | 20 +++++-----------
 3 files changed, 29 insertions(+), 28 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
index 9ebe5c22f47399..3e6127351911bd 100644
--- a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
+++ b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
@@ -233,11 +233,6 @@ class RegAllocEvictionAdvisorProvider {
   const AdvisorMode Mode;
 };
 
-RegAllocEvictionAdvisorProvider *
-createReleaseModeAdvisorProvider(LLVMContext &Ctx);
-RegAllocEvictionAdvisorProvider *
-createDevelopmentModeAdvisorProvider(LLVMContext &Ctx);
-
 /// 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
@@ -276,7 +271,13 @@ class RegAllocEvictionAdvisorAnalysis
   Result run(MachineFunction &MF, MachineFunctionAnalysisManager &MAM);
 
 private:
-  void initializeProvider(LLVMContext &Ctx);
+  void
+  initializeProvider(RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode Mode,
+                     LLVMContext &Ctx);
+  void
+  initializeMLProvider(RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode Mode,
+                       LLVMContext &Ctx);
+
   std::unique_ptr<RegAllocEvictionAdvisorProvider> Provider;
 };
 
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 534237713f0cb5..26bd241fc6ae37 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -1159,11 +1159,6 @@ void llvm::extractMBBFrequency(
 // Development mode-specific implementations
 #ifdef LLVM_HAVE_TFLITE
 
-RegAllocEvictionAdvisorProvider *
-llvm::createDevelopmentModeAdvisorProvider(LLVMContext &Ctx) {
-  return new DevelopmentModeEvictionAdvisorProvider(Ctx);
-}
-
 RegAllocEvictionAdvisorAnalysisLegacy *
 llvm::createDevelopmentModeAdvisorAnalysisLegacy() {
   return new DevelopmentModeEvictionAdvisorAnalysisLegacy();
@@ -1240,9 +1235,22 @@ bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
 }
 #endif // #ifdef LLVM_HAVE_TFLITE
 
-RegAllocEvictionAdvisorProvider *
-llvm::createReleaseModeAdvisorProvider(LLVMContext &Ctx) {
-  return new ReleaseModeEvictionAdvisorProvider(Ctx);
+void RegAllocEvictionAdvisorAnalysis::initializeMLProvider(
+    RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode Mode, LLVMContext &Ctx) {
+  if (Provider)
+    return;
+  switch (Mode) {
+  case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Development:
+#if defined(LLVM_HAVE_TFLITE)
+    Provider.reset(new DevelopmentModeEvictionAdvisorProvider(Ctx));
+#endif
+    break;
+  case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Release:
+    Provider.reset(new ReleaseModeEvictionAdvisorProvider(Ctx));
+    break;
+  default:
+    break;
+  }
 }
 
 RegAllocEvictionAdvisorAnalysisLegacy *
diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
index 4c7528ebffbfd5..1307ed2651a91a 100644
--- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
@@ -120,23 +120,15 @@ class DefaultEvictionAdvisorAnalysisLegacy final
 
 AnalysisKey RegAllocEvictionAdvisorAnalysis::Key;
 
-void RegAllocEvictionAdvisorAnalysis::initializeProvider(LLVMContext &Ctx) {
+void RegAllocEvictionAdvisorAnalysis::initializeProvider(
+    RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode Mode, LLVMContext &Ctx) {
   if (Provider)
     return;
-  switch (Mode) {
-  case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default:
+  if (Mode == RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default)
     Provider.reset(
         new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ false, Ctx));
-    break;
-  case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Development:
-#if defined(LLVM_HAVE_TFLITE)
-    Provider.reset(createDevelopmentModeAdvisorProvider(Ctx));
-#endif
-    break;
-  case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Release:
-    Provider.reset(createReleaseModeAdvisorProvider(Ctx));
-    break;
-  }
+  else
+    initializeMLProvider(Mode, Ctx);
   if (!Provider)
     Provider.reset(
         new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ true, Ctx));
@@ -146,7 +138,7 @@ RegAllocEvictionAdvisorAnalysis::Result
 RegAllocEvictionAdvisorAnalysis::run(MachineFunction &MF,
                                      MachineFunctionAnalysisManager &MFAM) {
   // Lazy initialization of the provider.
-  initializeProvider(MF.getFunction().getContext());
+  initializeProvider(::Mode, MF.getFunction().getContext());
   auto *MBFI = &MFAM.getResult<MachineBlockFrequencyAnalysis>(MF);
   auto *Loops = &MFAM.getResult<MachineLoopAnalysis>(MF);
   Provider->setAnalyses(MBFI, Loops);

>From 96b42b033a5d55b1cc1f329a623c35013cdb6d9a Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Wed, 11 Dec 2024 09:08:06 +0000
Subject: [PATCH 09/22] Use the provider as analysis result

---
 .../llvm/CodeGen/RegAllocEvictionAdvisor.h    | 81 ++++++++++---------
 llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp   | 16 ++--
 llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp  |  7 +-
 llvm/lib/CodeGen/RegAllocGreedy.cpp           |  3 +-
 4 files changed, 53 insertions(+), 54 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
index 3e6127351911bd..25df101ee512e9 100644
--- a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
+++ b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
@@ -153,6 +153,44 @@ class RegAllocEvictionAdvisor {
   const bool EnableLocalReassign;
 };
 
+/// 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.
+/// TODO: Coalesce this with the NPM analysis when legacy PM is removed.
+class RegAllocEvictionAdvisorProvider {
+public:
+  enum class AdvisorMode : int { Default, Release, Development };
+  RegAllocEvictionAdvisorProvider(AdvisorMode Mode, LLVMContext &Ctx)
+      : Ctx(Ctx), Mode(Mode) {}
+
+  virtual ~RegAllocEvictionAdvisorProvider() = default;
+
+  virtual void logRewardIfNeeded(const MachineFunction &MF,
+                                 llvm::function_ref<float()> GetReward) {}
+
+  virtual std::unique_ptr<RegAllocEvictionAdvisor>
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
+
+  /// 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;
+};
+
 /// ImmutableAnalysis abstraction for fetching the Eviction Advisor. We model it
 /// as an analysis to decouple the user from the implementation insofar as
 /// dependencies on other analyses goes. The motivation for it being an
@@ -177,8 +215,8 @@ class RegAllocEvictionAdvisorAnalysisLegacy : public ImmutablePass {
   static char ID;
 
   /// Get an advisor for the given context (i.e. machine function, etc)
-  virtual std::unique_ptr<RegAllocEvictionAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
+  virtual std::unique_ptr<RegAllocEvictionAdvisorProvider>&
+  getProvider() = 0;
   AdvisorMode getAdvisorMode() const { return Mode; }
   virtual void logRewardIfNeeded(const MachineFunction &MF,
                                  llvm::function_ref<float()> GetReward) {};
@@ -189,50 +227,13 @@ class RegAllocEvictionAdvisorAnalysisLegacy : public ImmutablePass {
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
   }
+  std::unique_ptr<RegAllocEvictionAdvisorProvider> Provider;
 
 private:
   StringRef getPassName() const override;
   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.
-/// TODO: Coalesce this with the NPM analysis when legacy PM is removed.
-class RegAllocEvictionAdvisorProvider {
-public:
-  enum class AdvisorMode : int { Default, Release, Development };
-  RegAllocEvictionAdvisorProvider(AdvisorMode Mode, LLVMContext &Ctx)
-      : Ctx(Ctx), Mode(Mode) {}
-
-  virtual ~RegAllocEvictionAdvisorProvider() = default;
-
-  virtual void logRewardIfNeeded(const MachineFunction &MF,
-                                 llvm::function_ref<float()> GetReward) {}
-
-  virtual std::unique_ptr<RegAllocEvictionAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
-
-  /// 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;
-};
-
 /// 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
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 26bd241fc6ae37..368bf592ebac90 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -420,12 +420,12 @@ class ReleaseModeEvictionAdvisorAnalysisLegacy final
   ReleaseModeEvictionAdvisorAnalysisLegacy()
       : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Release) {}
 
-  std::unique_ptr<RegAllocEvictionAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+  std::unique_ptr<RegAllocEvictionAdvisorProvider>&
+  getProvider() override {
     auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
     auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
     Provider->setAnalyses(MBFI, Loops);
-    return Provider->getAdvisor(MF, RA);
+    return Provider;
   }
 
   void logRewardIfNeeded(const MachineFunction &MF,
@@ -450,7 +450,7 @@ class ReleaseModeEvictionAdvisorAnalysisLegacy final
   }
 
 private:
-  std::unique_ptr<ReleaseModeEvictionAdvisorProvider> Provider;
+  // std::unique_ptr<ReleaseModeEvictionAdvisorProvider> Provider;
 };
 
 // ===================================
@@ -602,12 +602,12 @@ class DevelopmentModeEvictionAdvisorAnalysisLegacy final
     Provider->logRewardIfNeeded(MF, GetReward);
   }
 
-  std::unique_ptr<RegAllocEvictionAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+  std::unique_ptr<RegAllocEvictionAdvisorProvider>&
+  getProvider() override {
     auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
     auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
     Provider->setAnalyses(MBFI, Loops);
-    return Provider->getAdvisor(MF, RA);
+    return Provider;
   }
 
   // support for isa<> and dyn_cast.
@@ -622,7 +622,7 @@ class DevelopmentModeEvictionAdvisorAnalysisLegacy final
   }
 
 private:
-  std::unique_ptr<DevelopmentModeEvictionAdvisorProvider> Provider;
+  // std::unique_ptr<DevelopmentModeEvictionAdvisorProvider> Provider;
 };
 
 #endif // #ifdef LLVM_HAVE_TFLITE
diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
index 1307ed2651a91a..e00b81d8dcf61d 100644
--- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
@@ -95,10 +95,10 @@ class DefaultEvictionAdvisorAnalysisLegacy final
       : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Default),
         NotAsRequested(NotAsRequested) {}
 
-  std::unique_ptr<RegAllocEvictionAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+  std::unique_ptr<RegAllocEvictionAdvisorProvider>&
+  getProvider() override {
     // MBFI and Loops not required here.
-    return Provider->getAdvisor(MF, RA);
+    return Provider;
   }
 
   bool doInitialization(Module &M) override {
@@ -113,7 +113,6 @@ class DefaultEvictionAdvisorAnalysisLegacy final
   }
 
 private:
-  std::unique_ptr<DefaultEvictionAdvisorProvider> Provider;
   const bool NotAsRequested;
 };
 } // namespace
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp
index 7743de321716ec..4571e982e24183 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.cpp
+++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp
@@ -2748,8 +2748,7 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
 
   ExtraInfo.emplace();
   EvictAdvisor =
-      getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().getAdvisor(*MF,
-                                                                      *this);
+      getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().getProvider()->getAdvisor(*MF, *this);
   PriorityAdvisor =
       getAnalysis<RegAllocPriorityAdvisorAnalysis>().getAdvisor(*MF, *this);
 

>From 364cfd1ba4d761b06b2bcfb08da07c59594f39dd Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Thu, 26 Dec 2024 08:25:11 +0000
Subject: [PATCH 10/22] Remove setAnalyses(); pass analysis in getAdvisor

---
 .../llvm/CodeGen/RegAllocEvictionAdvisor.h    | 19 +++++---------
 llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp   | 25 +++----------------
 llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp  | 16 +++---------
 llvm/lib/CodeGen/RegAllocGreedy.cpp           |  7 ++++--
 4 files changed, 19 insertions(+), 48 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
index 25df101ee512e9..781afbc571abbf 100644
--- a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
+++ b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
@@ -170,22 +170,13 @@ class RegAllocEvictionAdvisorProvider {
                                  llvm::function_ref<float()> GetReward) {}
 
   virtual std::unique_ptr<RegAllocEvictionAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
-
-  /// 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;
-  }
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
+             MachineBlockFrequencyInfo *MBFI, MachineLoopInfo *Loops) = 0;
 
   AdvisorMode getAdvisorMode() const { return Mode; }
 
 protected:
   LLVMContext &Ctx;
-  MachineBlockFrequencyInfo *MBFI;
-  MachineLoopInfo *Loops;
 
 private:
   const AdvisorMode Mode;
@@ -215,8 +206,10 @@ class RegAllocEvictionAdvisorAnalysisLegacy : public ImmutablePass {
   static char ID;
 
   /// Get an advisor for the given context (i.e. machine function, etc)
-  virtual std::unique_ptr<RegAllocEvictionAdvisorProvider>&
-  getProvider() = 0;
+  std::unique_ptr<RegAllocEvictionAdvisorProvider> &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/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 368bf592ebac90..fa315146090b8a 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -392,7 +392,8 @@ class ReleaseModeEvictionAdvisorProvider final
   }
 
   std::unique_ptr<RegAllocEvictionAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
+             MachineBlockFrequencyInfo *MBFI, MachineLoopInfo *Loops) override {
     if (!Runner) {
       if (InteractiveChannelBaseName.empty())
         Runner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
@@ -420,14 +421,6 @@ class ReleaseModeEvictionAdvisorAnalysisLegacy final
   ReleaseModeEvictionAdvisorAnalysisLegacy()
       : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Release) {}
 
-  std::unique_ptr<RegAllocEvictionAdvisorProvider>&
-  getProvider() override {
-    auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
-    auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
-    Provider->setAnalyses(MBFI, Loops);
-    return Provider;
-  }
-
   void logRewardIfNeeded(const MachineFunction &MF,
                          llvm::function_ref<float()> GetReward) override {
     // No-op in release mode
@@ -448,9 +441,6 @@ class ReleaseModeEvictionAdvisorAnalysisLegacy final
     AU.addRequired<MachineLoopInfoWrapperPass>();
     RegAllocEvictionAdvisorAnalysisLegacy::getAnalysisUsage(AU);
   }
-
-private:
-  // std::unique_ptr<ReleaseModeEvictionAdvisorProvider> Provider;
 };
 
 // ===================================
@@ -566,7 +556,8 @@ class DevelopmentModeEvictionAdvisorProvider final
   }
 
   std::unique_ptr<RegAllocEvictionAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
+             MachineBlockFrequencyInfo *MBFI, MachineLoopInfo *Loops) override {
     if (!Runner)
       return nullptr;
     if (Log)
@@ -602,14 +593,6 @@ class DevelopmentModeEvictionAdvisorAnalysisLegacy final
     Provider->logRewardIfNeeded(MF, GetReward);
   }
 
-  std::unique_ptr<RegAllocEvictionAdvisorProvider>&
-  getProvider() override {
-    auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
-    auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
-    Provider->setAnalyses(MBFI, Loops);
-    return Provider;
-  }
-
   // support for isa<> and dyn_cast.
   static bool classof(const RegAllocEvictionAdvisorAnalysisLegacy *R) {
     return R->getAdvisorMode() == AdvisorMode::Development;
diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
index e00b81d8dcf61d..73de74c5aad925 100644
--- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
@@ -83,7 +83,8 @@ class DefaultEvictionAdvisorProvider final
   }
 
   std::unique_ptr<RegAllocEvictionAdvisor>
-  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
+             MachineBlockFrequencyInfo *, MachineLoopInfo *) override {
     return std::make_unique<DefaultEvictionAdvisor>(MF, RA);
   }
 };
@@ -95,12 +96,6 @@ class DefaultEvictionAdvisorAnalysisLegacy final
       : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Default),
         NotAsRequested(NotAsRequested) {}
 
-  std::unique_ptr<RegAllocEvictionAdvisorProvider>&
-  getProvider() override {
-    // MBFI and Loops not required here.
-    return Provider;
-  }
-
   bool doInitialization(Module &M) override {
     Provider.reset(
         new DefaultEvictionAdvisorProvider(NotAsRequested, M.getContext()));
@@ -125,12 +120,12 @@ void RegAllocEvictionAdvisorAnalysis::initializeProvider(
     return;
   if (Mode == RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default)
     Provider.reset(
-        new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ false, Ctx));
+        new DefaultEvictionAdvisorProvider(/*NotAsRequested=*/false, Ctx));
   else
     initializeMLProvider(Mode, Ctx);
   if (!Provider)
     Provider.reset(
-        new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ true, Ctx));
+        new DefaultEvictionAdvisorProvider(/*NotAsRequested=*/true, Ctx));
 }
 
 RegAllocEvictionAdvisorAnalysis::Result
@@ -138,9 +133,6 @@ RegAllocEvictionAdvisorAnalysis::run(MachineFunction &MF,
                                      MachineFunctionAnalysisManager &MFAM) {
   // Lazy initialization of the provider.
   initializeProvider(::Mode, MF.getFunction().getContext());
-  auto *MBFI = &MFAM.getResult<MachineBlockFrequencyAnalysis>(MF);
-  auto *Loops = &MFAM.getResult<MachineLoopAnalysis>(MF);
-  Provider->setAnalyses(MBFI, Loops);
   return Result{Provider.get()};
 }
 
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp
index 4571e982e24183..847a6ad29a1df2 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.cpp
+++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp
@@ -2747,8 +2747,11 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
                                : TRI->reverseLocalAssignment();
 
   ExtraInfo.emplace();
-  EvictAdvisor =
-      getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().getProvider()->getAdvisor(*MF, *this);
+
+  auto &EvictAdvisorProvider =
+      getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().getProvider();
+  EvictAdvisor = EvictAdvisorProvider->getAdvisor(*MF, *this, MBFI, Loops);
+
   PriorityAdvisor =
       getAnalysis<RegAllocPriorityAdvisorAnalysis>().getAdvisor(*MF, *this);
 

>From 40f515a8dc5036f419242ec7938434a1b84a67ac 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 11/22] [CodeGen][NewPM] Port RegAllocPriorityAdvisor analysis
 to NPM

---
 .../llvm}/CodeGen/RegAllocPriorityAdvisor.h   |  79 +++++++-
 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  | 120 +++++++++---
 llvm/lib/Passes/PassBuilder.cpp               |   1 +
 10 files changed, 294 insertions(+), 112 deletions(-)
 rename llvm/{lib => include/llvm}/CodeGen/RegAllocPriorityAdvisor.h (53%)

diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
similarity index 53%
rename from llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
rename to llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
index 2d42a43c4c6372..bddfe15bf17751 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 {
@@ -56,12 +58,73 @@ class DefaultPriorityAdvisor : 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 };
 
-  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:
+  enum class AdvisorMode : int { Default, Release, Development };
+
+  RegAllocPriorityAdvisorAnalysisLegacy(AdvisorMode Mode)
+      : ImmutablePass(ID), Mode(Mode) {};
   static char ID;
 
   /// Get an advisor for the given context (i.e. machine function, etc)
@@ -69,7 +132,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
@@ -85,11 +148,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 8915bf258005a7..e74b85c0de886f 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 faa7c68ad2acc6..c23e4af1a342bd 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 fa315146090b8a..8370bd974f9f4e 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -115,7 +115,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);
   }
@@ -1212,8 +1212,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 9638df81770c1e..dfc06eaa0bfe69 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 73de74c5aad925..23e76c4363f7b1 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 847a6ad29a1df2..64ae6ef660eb04 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);
 }
 
@@ -2753,7 +2753,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(createInlineSpiller(*this, *MF, *VRM, *VRAI));
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.h b/llvm/lib/CodeGen/RegAllocGreedy.h
index 001a1e86a96914..89ceefd37795bc 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 0650aaff56ea00..09078fc57b1bf1 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,76 +20,138 @@
 
 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")));
 
-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:
-  DefaultPriorityAdvisorAnalysis(bool NotAsRequested)
-      : RegAllocPriorityAdvisorAnalysis(AdvisorMode::Default),
+  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 DefaultPriorityAdvisorAnalysisLegacy final
+    : public RegAllocPriorityAdvisorAnalysisLegacy {
+public:
+  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;
 };
 } // namespace
 
-template <> Pass *llvm::callDefaultCtor<RegAllocPriorityAdvisorAnalysis>() {
+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::Development:
+#if defined(LLVM_HAVE_TFLITE)
+    Provider.reset(createDevelopmentModePriorityAdvisorProvider(Ctx));
+#endif
+    break;
+  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 RegAllocPriorityAdvisorAnalysis::AdvisorMode::Default:
-    Ret = new DefaultPriorityAdvisorAnalysis(/*NotAsRequested*/ false);
+  case RegAllocPriorityAdvisorProvider::AdvisorMode::Default:
+    Ret = new DefaultPriorityAdvisorAnalysisLegacy(/*NotAsRequested*/ false);
     break;
-  case RegAllocPriorityAdvisorAnalysis::AdvisorMode::Development:
+  case RegAllocPriorityAdvisorProvider::AdvisorMode::Development:
 #if defined(LLVM_HAVE_TFLITE)
-    Ret = createDevelopmentModePriorityAdvisor();
+    Ret = createDevelopmentModePriorityAdvisorAnalysis();
 #endif
     break;
-  case RegAllocPriorityAdvisorAnalysis::AdvisorMode::Release:
-    Ret = createReleaseModePriorityAdvisor();
+  case RegAllocPriorityAdvisorProvider::AdvisorMode::Release:
+    Ret = createReleaseModePriorityAdvisorAnalysis();
     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 132492ccbdd55e..5394e06282e50c 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -125,6 +125,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 9dd5454205bccbc9b3a35b51de6bb21f4f662a7e 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 12/22] 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 bddfe15bf17751..f557165692a46a 100644
--- a/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
+++ b/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
@@ -128,8 +128,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) {};
@@ -141,6 +140,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 dfc06eaa0bfe69..dd598767f1bbc2 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 64ae6ef660eb04..cd264695287f48 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.cpp
+++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp
@@ -2752,9 +2752,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(createInlineSpiller(*this, *MF, *VRM, *VRAI));
diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
index 09078fc57b1bf1..6699da052de5cf 100644
--- a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
@@ -78,11 +78,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()));
@@ -90,7 +90,7 @@ class DefaultPriorityAdvisorAnalysisLegacy final
   }
 
   const bool NotAsRequested;
-  std::unique_ptr<DefaultPriorityAdvisorProvider> Provider;
+  // std::unique_ptr<DefaultPriorityAdvisorProvider> Provider;
 };
 } // namespace
 

>From 94f324099e7974ec9c92df4cf1e03f7298479dbd 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 13/22] Suggestions and remove createProvider methods

---
 .../llvm/CodeGen/RegAllocPriorityAdvisor.h    |  7 ++---
 .../lib/CodeGen/MLRegAllocPriorityAdvisor.cpp | 29 ++++++++++++-------
 llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp  | 20 ++++---------
 3 files changed, 25 insertions(+), 31 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
index f557165692a46a..3aba88fb802583 100644
--- a/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
+++ b/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
@@ -89,11 +89,6 @@ class RegAllocPriorityAdvisorProvider {
   const AdvisorMode Mode;
 };
 
-RegAllocPriorityAdvisorProvider *createReleaseModePriorityAdvisorProvider();
-
-RegAllocPriorityAdvisorProvider *
-createDevelopmentModePriorityAdvisorProvider(LLVMContext &Ctx);
-
 class RegAllocPriorityAdvisorAnalysis
     : public AnalysisInfoMixin<RegAllocPriorityAdvisorAnalysis> {
   static AnalysisKey Key;
@@ -116,6 +111,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 dd598767f1bbc2..dedd20b31d1c02 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 6699da052de5cf..407b8f8a283ca1 100644
--- a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
@@ -97,24 +97,14 @@ class DefaultPriorityAdvisorAnalysisLegacy final
 void RegAllocPriorityAdvisorAnalysis::initializeProvider(LLVMContext &Ctx) {
   if (Provider)
     return;
-
-  switch (Mode) {
-  case RegAllocPriorityAdvisorProvider::AdvisorMode::Default:
+  if (Mode == RegAllocPriorityAdvisorProvider::AdvisorMode::Default)
     Provider.reset(
-        new DefaultPriorityAdvisorProvider(/*NotAsRequested*/ false, Ctx));
-    break;
-  case RegAllocPriorityAdvisorProvider::AdvisorMode::Development:
-#if defined(LLVM_HAVE_TFLITE)
-    Provider.reset(createDevelopmentModePriorityAdvisorProvider(Ctx));
-#endif
-    break;
-  case RegAllocPriorityAdvisorProvider::AdvisorMode::Release:
-    Provider.reset(createReleaseModePriorityAdvisorProvider());
-    break;
-  }
+        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 d8714650324ad27c75a3edaf892c2a89a545e492 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 14/22] Apply suggestions and remove setAnalyses

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

diff --git a/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
index 3aba88fb802583..c9ab4a4ee8399f 100644
--- a/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
+++ b/llvm/include/llvm/CodeGen/RegAllocPriorityAdvisor.h
@@ -73,18 +73,15 @@ 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;
 };
@@ -125,7 +122,9 @@ 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;
+  std::unique_ptr<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 dedd20b31d1c02..8b195b36a711fa 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 cd264695287f48..466c3f69c6acda 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.cpp
+++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp
@@ -2754,7 +2754,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(createInlineSpiller(*this, *MF, *VRM, *VRAI));
diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
index 407b8f8a283ca1..05444189c12f8d 100644
--- a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
@@ -54,7 +54,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);
   }
@@ -78,11 +79,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()));
@@ -90,7 +86,6 @@ class DefaultPriorityAdvisorAnalysisLegacy final
   }
 
   const bool NotAsRequested;
-  // std::unique_ptr<DefaultPriorityAdvisorProvider> Provider;
 };
 } // namespace
 
@@ -114,8 +109,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 8151817211d10ae9cf8f9a931aa6e2416134bda8 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Mon, 9 Dec 2024 07:58:48 +0000
Subject: [PATCH 15/22] Spiller: Deatach legacy pass and supply analyses
 instead

---
 llvm/include/llvm/CodeGen/Spiller.h | 16 +++++++++++--
 llvm/lib/CodeGen/InlineSpiller.cpp  | 36 +++++++++++------------------
 llvm/lib/CodeGen/RegAllocBasic.cpp  | 16 +++++++++----
 llvm/lib/CodeGen/RegAllocGreedy.cpp |  4 +++-
 llvm/lib/CodeGen/RegAllocPBQP.cpp   |  5 +++-
 5 files changed, 46 insertions(+), 31 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/Spiller.h b/llvm/include/llvm/CodeGen/Spiller.h
index 51ad36bc6b1f8b..3132cefeb6c68a 100644
--- a/llvm/include/llvm/CodeGen/Spiller.h
+++ b/llvm/include/llvm/CodeGen/Spiller.h
@@ -19,6 +19,10 @@ class MachineFunction;
 class MachineFunctionPass;
 class VirtRegMap;
 class VirtRegAuxInfo;
+class LiveIntervals;
+class LiveStacks;
+class MachineDominatorTree;
+class MachineBlockFrequencyInfo;
 
 /// Spiller interface.
 ///
@@ -41,12 +45,20 @@ class Spiller {
   virtual ArrayRef<Register> getReplacedRegs() = 0;
 
   virtual void postOptimization() {}
+
+  struct RequiredAnalyses {
+    LiveIntervals &LIS;
+    LiveStacks &LSS;
+    MachineDominatorTree &MDT;
+    const MachineBlockFrequencyInfo &MBFI;
+  };
 };
 
 /// Create and return a spiller that will insert spill code directly instead
 /// of deferring though VirtRegMap.
-Spiller *createInlineSpiller(MachineFunctionPass &Pass, MachineFunction &MF,
-                             VirtRegMap &VRM, VirtRegAuxInfo &VRAI);
+Spiller *createInlineSpiller(const Spiller::RequiredAnalyses &Analyses,
+                             MachineFunction &MF, VirtRegMap &VRM,
+                             VirtRegAuxInfo &VRAI);
 
 } // end namespace llvm
 
diff --git a/llvm/lib/CodeGen/InlineSpiller.cpp b/llvm/lib/CodeGen/InlineSpiller.cpp
index 64f290f5930a1b..b9768d5c63a5d1 100644
--- a/llvm/lib/CodeGen/InlineSpiller.cpp
+++ b/llvm/lib/CodeGen/InlineSpiller.cpp
@@ -75,7 +75,6 @@ RestrictStatepointRemat("restrict-statepoint-remat",
                        cl::desc("Restrict remat for statepoint operands"));
 
 namespace {
-
 class HoistSpillHelper : private LiveRangeEdit::Delegate {
   MachineFunction &MF;
   LiveIntervals &LIS;
@@ -128,15 +127,11 @@ class HoistSpillHelper : private LiveRangeEdit::Delegate {
                       DenseMap<MachineBasicBlock *, unsigned> &SpillsToIns);
 
 public:
-  HoistSpillHelper(MachineFunctionPass &pass, MachineFunction &mf,
-                   VirtRegMap &vrm)
-      : MF(mf), LIS(pass.getAnalysis<LiveIntervalsWrapperPass>().getLIS()),
-        LSS(pass.getAnalysis<LiveStacksWrapperLegacy>().getLS()),
-        MDT(pass.getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree()),
+  HoistSpillHelper(const Spiller::RequiredAnalyses &Analyses,
+                   MachineFunction &mf, VirtRegMap &vrm)
+      : MF(mf), LIS(Analyses.LIS), LSS(Analyses.LSS), MDT(Analyses.MDT),
         VRM(vrm), MRI(mf.getRegInfo()), TII(*mf.getSubtarget().getInstrInfo()),
-        TRI(*mf.getSubtarget().getRegisterInfo()),
-        MBFI(
-            pass.getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI()),
+        TRI(*mf.getSubtarget().getRegisterInfo()), MBFI(Analyses.MBFI),
         IPA(LIS, mf.getNumBlockIDs()) {}
 
   void addToMergeableSpills(MachineInstr &Spill, int StackSlot,
@@ -190,16 +185,12 @@ class InlineSpiller : public Spiller {
   ~InlineSpiller() override = default;
 
 public:
-  InlineSpiller(MachineFunctionPass &Pass, MachineFunction &MF, VirtRegMap &VRM,
-                VirtRegAuxInfo &VRAI)
-      : MF(MF), LIS(Pass.getAnalysis<LiveIntervalsWrapperPass>().getLIS()),
-        LSS(Pass.getAnalysis<LiveStacksWrapperLegacy>().getLS()),
-        MDT(Pass.getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree()),
+  InlineSpiller(const Spiller::RequiredAnalyses &Analyses, MachineFunction &MF,
+                VirtRegMap &VRM, VirtRegAuxInfo &VRAI)
+      : MF(MF), LIS(Analyses.LIS), LSS(Analyses.LSS), MDT(Analyses.MDT),
         VRM(VRM), MRI(MF.getRegInfo()), TII(*MF.getSubtarget().getInstrInfo()),
-        TRI(*MF.getSubtarget().getRegisterInfo()),
-        MBFI(
-            Pass.getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI()),
-        HSpiller(Pass, MF, VRM), VRAI(VRAI) {}
+        TRI(*MF.getSubtarget().getRegisterInfo()), MBFI(Analyses.MBFI),
+        HSpiller(Analyses, MF, VRM), VRAI(VRAI) {}
 
   void spill(LiveRangeEdit &) override;
   ArrayRef<Register> getSpilledRegs() override { return RegsToSpill; }
@@ -237,10 +228,11 @@ Spiller::~Spiller() = default;
 
 void Spiller::anchor() {}
 
-Spiller *llvm::createInlineSpiller(MachineFunctionPass &Pass,
-                                   MachineFunction &MF, VirtRegMap &VRM,
-                                   VirtRegAuxInfo &VRAI) {
-  return new InlineSpiller(Pass, MF, VRM, VRAI);
+Spiller *
+llvm::createInlineSpiller(const InlineSpiller::RequiredAnalyses &Analyses,
+                          MachineFunction &MF, VirtRegMap &VRM,
+                          VirtRegAuxInfo &VRAI) {
+  return new InlineSpiller(Analyses, MF, VRM, VRAI);
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/llvm/lib/CodeGen/RegAllocBasic.cpp b/llvm/lib/CodeGen/RegAllocBasic.cpp
index c05aa1e40e4779..f3f34f890be11e 100644
--- a/llvm/lib/CodeGen/RegAllocBasic.cpp
+++ b/llvm/lib/CodeGen/RegAllocBasic.cpp
@@ -22,6 +22,7 @@
 #include "llvm/CodeGen/LiveRegMatrix.h"
 #include "llvm/CodeGen/LiveStacks.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
+#include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/Passes.h"
@@ -187,6 +188,7 @@ void RABasic::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequired<ProfileSummaryInfoWrapperPass>();
   AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
   AU.addPreserved<MachineBlockFrequencyInfoWrapperPass>();
+  AU.addRequired<MachineDominatorTreeWrapperPass>();
   AU.addRequiredID(MachineDominatorsID);
   AU.addPreservedID(MachineDominatorsID);
   AU.addRequired<MachineLoopInfoWrapperPass>();
@@ -310,16 +312,20 @@ bool RABasic::runOnMachineFunction(MachineFunction &mf) {
                     << "********** Function: " << mf.getName() << '\n');
 
   MF = &mf;
+  auto &MBFI = getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
+  auto &LiveStks = getAnalysis<LiveStacksWrapperLegacy>().getLS();
+  auto &MDT = getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
+
   RegAllocBase::init(getAnalysis<VirtRegMapWrapperLegacy>().getVRM(),
                      getAnalysis<LiveIntervalsWrapperPass>().getLIS(),
                      getAnalysis<LiveRegMatrixWrapperLegacy>().getLRM());
-  VirtRegAuxInfo VRAI(
-      *MF, *LIS, *VRM, getAnalysis<MachineLoopInfoWrapperPass>().getLI(),
-      getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI(),
-      &getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI());
+  VirtRegAuxInfo VRAI(*MF, *LIS, *VRM,
+                      getAnalysis<MachineLoopInfoWrapperPass>().getLI(), MBFI,
+                      &getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI());
   VRAI.calculateSpillWeightsAndHints();
 
-  SpillerInstance.reset(createInlineSpiller(*this, *MF, *VRM, VRAI));
+  SpillerInstance.reset(
+      createInlineSpiller({*LIS, LiveStks, MDT, MBFI}, *MF, *VRM, VRAI));
 
   allocatePhysRegs();
   postOptimization();
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp
index 466c3f69c6acda..5a7282809f5f7d 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.cpp
+++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp
@@ -2733,6 +2733,7 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
   Bundles = &getAnalysis<EdgeBundlesWrapperLegacy>().getEdgeBundles();
   SpillPlacer = &getAnalysis<SpillPlacementWrapperLegacy>().getResult();
   DebugVars = &getAnalysis<LiveDebugVariablesWrapperLegacy>().getLDV();
+  auto &LSS = getAnalysis<LiveStacksWrapperLegacy>().getLS();
 
   initializeCSRCost();
 
@@ -2757,7 +2758,8 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
                         ->getAdvisor(*MF, *this, Indexes);
 
   VRAI = std::make_unique<VirtRegAuxInfo>(*MF, *LIS, *VRM, *Loops, *MBFI);
-  SpillerInstance.reset(createInlineSpiller(*this, *MF, *VRM, *VRAI));
+  SpillerInstance.reset(
+      createInlineSpiller({*LIS, LSS, *DomTree, *MBFI}, *MF, *VRM, *VRAI));
 
   VRAI->calculateSpillWeightsAndHints();
 
diff --git a/llvm/lib/CodeGen/RegAllocPBQP.cpp b/llvm/lib/CodeGen/RegAllocPBQP.cpp
index 696c312e4ba00a..e230a1be95c9fa 100644
--- a/llvm/lib/CodeGen/RegAllocPBQP.cpp
+++ b/llvm/lib/CodeGen/RegAllocPBQP.cpp
@@ -794,6 +794,9 @@ bool RegAllocPBQP::runOnMachineFunction(MachineFunction &MF) {
   MachineBlockFrequencyInfo &MBFI =
       getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
 
+  auto &LiveStks = getAnalysis<LiveStacksWrapperLegacy>().getLS();
+  auto &MDT = getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
+
   VirtRegMap &VRM = getAnalysis<VirtRegMapWrapperLegacy>().getVRM();
 
   PBQPVirtRegAuxInfo VRAI(
@@ -807,7 +810,7 @@ bool RegAllocPBQP::runOnMachineFunction(MachineFunction &MF) {
   VirtRegAuxInfo DefaultVRAI(
       MF, LIS, VRM, getAnalysis<MachineLoopInfoWrapperPass>().getLI(), MBFI);
   std::unique_ptr<Spiller> VRegSpiller(
-      createInlineSpiller(*this, MF, VRM, DefaultVRAI));
+      createInlineSpiller({LIS, LiveStks, MDT, MBFI}, MF, VRM, DefaultVRAI));
 
   MF.getRegInfo().freezeReservedRegs();
 

>From 757e3d26e4e0f5bdb76afd483a9e5b40e38bfeff Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Wed, 11 Dec 2024 08:51:55 +0000
Subject: [PATCH 16/22] [CodeGen][NewPM] Port RegAllocGreedy to NPM

---
 llvm/include/llvm/CodeGen/MachineFunction.h   |   1 +
 llvm/include/llvm/CodeGen/Passes.h            |   2 +-
 llvm/include/llvm/InitializePasses.h          |   2 +-
 .../llvm/Passes/MachinePassRegistry.def       |   9 +
 llvm/lib/CodeGen/CodeGen.cpp                  |   2 +-
 llvm/lib/CodeGen/RegAllocGreedy.cpp           | 185 ++++++++++++++----
 llvm/lib/CodeGen/RegAllocGreedy.h             |  57 +++---
 llvm/lib/Passes/PassBuilder.cpp               |   1 +
 8 files changed, 196 insertions(+), 63 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/MachineFunction.h b/llvm/include/llvm/CodeGen/MachineFunction.h
index d696add8a1af53..662272e5e09618 100644
--- a/llvm/include/llvm/CodeGen/MachineFunction.h
+++ b/llvm/include/llvm/CodeGen/MachineFunction.h
@@ -901,6 +901,7 @@ class LLVM_ABI MachineFunction {
 
   /// Run the current MachineFunction through the machine code verifier, useful
   /// for debugger use.
+  /// TODO: Add the param LiveStks
   /// \returns true if no problems were found.
   bool verify(LiveIntervals *LiveInts, SlotIndexes *Indexes,
               const char *Banner = nullptr, raw_ostream *OS = nullptr,
diff --git a/llvm/include/llvm/CodeGen/Passes.h b/llvm/include/llvm/CodeGen/Passes.h
index d1fac4a304cffe..1096c34b307f9b 100644
--- a/llvm/include/llvm/CodeGen/Passes.h
+++ b/llvm/include/llvm/CodeGen/Passes.h
@@ -167,7 +167,7 @@ namespace llvm {
   extern char &LiveRangeShrinkID;
 
   /// Greedy register allocator.
-  extern char &RAGreedyID;
+  extern char &RAGreedyLegacyID;
 
   /// Basic register allocator.
   extern char &RABasicID;
diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
index e74b85c0de886f..afe0aa6113dd21 100644
--- a/llvm/include/llvm/InitializePasses.h
+++ b/llvm/include/llvm/InitializePasses.h
@@ -250,7 +250,7 @@ void initializeProfileSummaryInfoWrapperPassPass(PassRegistry &);
 void initializePromoteLegacyPassPass(PassRegistry &);
 void initializeRABasicPass(PassRegistry &);
 void initializePseudoProbeInserterPass(PassRegistry &);
-void initializeRAGreedyPass(PassRegistry &);
+void initializeRAGreedyLegacyPass(PassRegistry &);
 void initializeReachingDefAnalysisPass(PassRegistry &);
 void initializeReassociateLegacyPassPass(PassRegistry &);
 void initializeRegAllocEvictionAdvisorAnalysisLegacyPass(PassRegistry &);
diff --git a/llvm/include/llvm/Passes/MachinePassRegistry.def b/llvm/include/llvm/Passes/MachinePassRegistry.def
index c23e4af1a342bd..bf12a2c9aca90b 100644
--- a/llvm/include/llvm/Passes/MachinePassRegistry.def
+++ b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -187,6 +187,15 @@ MACHINE_FUNCTION_PASS_WITH_PARAMS(
       return parseRegAllocFastPassOptions(*PB, Params);
     },
     "filter=reg-filter;no-clear-vregs")
+
+MACHINE_FUNCTION_PASS_WITH_PARAMS(
+    "regallocgreedy", "RAGreedy",
+    [](RegAllocFilterFunc F) { return RAGreedyPass(F); },
+    [PB = this](StringRef Params) {
+      // TODO: parseRegAllocFilter(*PB, Params);
+      return Expected<RegAllocFilterFunc>(nullptr);
+    }, ""
+)
 #undef MACHINE_FUNCTION_PASS_WITH_PARAMS
 
 // After a pass is converted to new pass manager, its entry should be moved from
diff --git a/llvm/lib/CodeGen/CodeGen.cpp b/llvm/lib/CodeGen/CodeGen.cpp
index 8efe540770913a..6acff9cd21134b 100644
--- a/llvm/lib/CodeGen/CodeGen.cpp
+++ b/llvm/lib/CodeGen/CodeGen.cpp
@@ -111,7 +111,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializePreISelIntrinsicLoweringLegacyPassPass(Registry);
   initializeProcessImplicitDefsPass(Registry);
   initializeRABasicPass(Registry);
-  initializeRAGreedyPass(Registry);
+  initializeRAGreedyLegacyPass(Registry);
   initializeRegAllocFastPass(Registry);
   initializeRegUsageInfoCollectorLegacyPass(Registry);
   initializeRegUsageInfoPropagationLegacyPass(Registry);
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp
index 5a7282809f5f7d..3a96325144efbc 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.cpp
+++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp
@@ -43,8 +43,10 @@
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
+#include "llvm/CodeGen/MachinePassManager.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
+#include "llvm/CodeGen/RegAllocGreedyPass.h"
 #include "llvm/CodeGen/RegAllocPriorityAdvisor.h"
 #include "llvm/CodeGen/RegAllocRegistry.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
@@ -55,6 +57,7 @@
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/CodeGen/VirtRegMap.h"
+#include "llvm/IR/Analysis.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/LLVMContext.h"
@@ -146,11 +149,134 @@ static cl::opt<unsigned> SplitThresholdForRegWithHint(
 static RegisterRegAlloc greedyRegAlloc("greedy", "greedy register allocator",
                                        createGreedyRegisterAllocator);
 
-char RAGreedy::ID = 0;
-char &llvm::RAGreedyID = RAGreedy::ID;
+namespace {
+class RAGreedyLegacy : public MachineFunctionPass {
+  RegAllocFilterFunc F;
 
-INITIALIZE_PASS_BEGIN(RAGreedy, "greedy",
-                "Greedy Register Allocator", false, false)
+public:
+  RAGreedyLegacy(const RegAllocFilterFunc F = nullptr);
+
+  static char ID;
+  /// Return the pass name.
+  StringRef getPassName() const override { return "Greedy Register Allocator"; }
+
+  /// RAGreedy analysis usage.
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
+  /// Perform register allocation.
+  bool runOnMachineFunction(MachineFunction &mf) override;
+
+  MachineFunctionProperties getRequiredProperties() const override {
+    return MachineFunctionProperties().set(
+        MachineFunctionProperties::Property::NoPHIs);
+  }
+
+  MachineFunctionProperties getClearedProperties() const override {
+    return MachineFunctionProperties().set(
+        MachineFunctionProperties::Property::IsSSA);
+  }
+};
+
+} // end anonymous namespace
+
+RAGreedyLegacy::RAGreedyLegacy(const RegAllocFilterFunc F)
+    : MachineFunctionPass(ID), F(F) {
+  initializeRAGreedyLegacyPass(*PassRegistry::getPassRegistry());
+}
+
+RAGreedy::RAGreedy(const RegAllocFilterFunc F) : RegAllocBase(F) {}
+
+void RAGreedy::setAnalyses(RequiredAnalyses &Analyses) {
+  VRM = Analyses.VRM;
+  LIS = Analyses.LIS;
+  Matrix = Analyses.LRM;
+  Indexes = Analyses.Indexes;
+  MBFI = Analyses.MBFI;
+  DomTree = Analyses.DomTree;
+  Loops = Analyses.Loops;
+  ORE = Analyses.ORE;
+  Bundles = Analyses.Bundles;
+  SpillPlacer = Analyses.SpillPlacer;
+  DebugVars = Analyses.DebugVars;
+  LSS = Analyses.LSS;
+  EvictProvider = Analyses.EvictProvider;
+  PriorityProvider = Analyses.PriorityProvider;
+}
+
+PreservedAnalyses RAGreedyPass::run(MachineFunction &MF,
+                                    MachineFunctionAnalysisManager &MFAM) {
+  MFPropsModifier _(*this, MF);
+
+  RAGreedy Impl(Filter);
+  RAGreedy::RequiredAnalyses Analyses;
+
+  Analyses.VRM = &MFAM.getResult<VirtRegMapAnalysis>(MF);
+  Analyses.LIS = &MFAM.getResult<LiveIntervalsAnalysis>(MF);
+  Analyses.LRM = &MFAM.getResult<LiveRegMatrixAnalysis>(MF);
+  Analyses.LSS = &MFAM.getResult<LiveStacksAnalysis>(MF);
+  Analyses.Indexes = &MFAM.getResult<SlotIndexesAnalysis>(MF);
+  Analyses.MBFI = &MFAM.getResult<MachineBlockFrequencyAnalysis>(MF);
+  Analyses.DomTree = &MFAM.getResult<MachineDominatorTreeAnalysis>(MF);
+  Analyses.ORE = &MFAM.getResult<MachineOptimizationRemarkEmitterAnalysis>(MF);
+  Analyses.Loops = &MFAM.getResult<MachineLoopAnalysis>(MF);
+  Analyses.Bundles = &MFAM.getResult<EdgeBundlesAnalysis>(MF);
+  Analyses.SpillPlacer = &MFAM.getResult<SpillPlacementAnalysis>(MF);
+  Analyses.DebugVars = &MFAM.getResult<LiveDebugVariablesAnalysis>(MF);
+  Analyses.EvictProvider =
+      MFAM.getResult<RegAllocEvictionAdvisorAnalysis>(MF).Provider;
+  Analyses.PriorityProvider =
+      MFAM.getResult<RegAllocPriorityAdvisorAnalysis>(MF).Provider;
+
+  Impl.setAnalyses(Analyses);
+  bool Changed = Impl.run(MF);
+  if (!Changed)
+    return PreservedAnalyses::all();
+  auto PA = getMachineFunctionPassPreservedAnalyses();
+  PA.preserveSet<CFGAnalyses>();
+  PA.preserve<MachineBlockFrequencyAnalysis>();
+  PA.preserve<LiveIntervalsAnalysis>();
+  PA.preserve<SlotIndexesAnalysis>();
+  PA.preserve<LiveDebugVariablesAnalysis>();
+  PA.preserve<LiveStacksAnalysis>();
+  PA.preserve<MachineDominatorTreeAnalysis>();
+  PA.preserve<MachineLoopAnalysis>();
+  PA.preserve<VirtRegMapAnalysis>();
+  PA.preserve<LiveRegMatrixAnalysis>();
+  return PA;
+}
+
+bool RAGreedyLegacy::runOnMachineFunction(MachineFunction &MF) {
+  RAGreedy Impl(F);
+
+  RAGreedy::RequiredAnalyses Analyses;
+  Analyses.VRM = &getAnalysis<VirtRegMapWrapperLegacy>().getVRM();
+  Analyses.LIS = &getAnalysis<LiveIntervalsWrapperPass>().getLIS();
+  Analyses.LSS = &getAnalysis<LiveStacksWrapperLegacy>().getLS();
+  Analyses.LRM = &getAnalysis<LiveRegMatrixWrapperLegacy>().getLRM();
+  Analyses.Indexes = &getAnalysis<SlotIndexesWrapperPass>().getSI();
+  Analyses.MBFI =
+      &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
+  Analyses.DomTree =
+      &getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
+  Analyses.ORE = &getAnalysis<MachineOptimizationRemarkEmitterPass>().getORE();
+  Analyses.Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
+  Analyses.Bundles = &getAnalysis<EdgeBundlesWrapperLegacy>().getEdgeBundles();
+  Analyses.SpillPlacer =
+      &getAnalysis<SpillPlacementWrapperLegacy>().getResult();
+  Analyses.DebugVars = &getAnalysis<LiveDebugVariablesWrapperLegacy>().getLDV();
+  Analyses.EvictProvider =
+      getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().getProvider().get();
+  Analyses.PriorityProvider =
+      getAnalysis<RegAllocPriorityAdvisorAnalysisLegacy>().getProvider().get();
+
+  Impl.setAnalyses(Analyses);
+  return Impl.run(MF);
+}
+
+char RAGreedyLegacy::ID = 0;
+char &llvm::RAGreedyLegacyID = RAGreedyLegacy::ID;
+
+INITIALIZE_PASS_BEGIN(RAGreedyLegacy, "greedy", "Greedy Register Allocator",
+                      false, false)
 INITIALIZE_PASS_DEPENDENCY(LiveDebugVariablesWrapperLegacy)
 INITIALIZE_PASS_DEPENDENCY(SlotIndexesWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LiveIntervalsWrapperPass)
@@ -166,8 +292,8 @@ INITIALIZE_PASS_DEPENDENCY(SpillPlacementWrapperLegacy)
 INITIALIZE_PASS_DEPENDENCY(MachineOptimizationRemarkEmitterPass)
 INITIALIZE_PASS_DEPENDENCY(RegAllocEvictionAdvisorAnalysisLegacy)
 INITIALIZE_PASS_DEPENDENCY(RegAllocPriorityAdvisorAnalysisLegacy)
-INITIALIZE_PASS_END(RAGreedy, "greedy",
-                "Greedy Register Allocator", false, false)
+INITIALIZE_PASS_END(RAGreedyLegacy, "greedy", "Greedy Register Allocator",
+                    false, false)
 
 #ifndef NDEBUG
 const char *const RAGreedy::StageName[] = {
@@ -186,17 +312,14 @@ const char *const RAGreedy::StageName[] = {
 const float Hysteresis = (2007 / 2048.0f); // 0.97998046875
 
 FunctionPass* llvm::createGreedyRegisterAllocator() {
-  return new RAGreedy();
+  return new RAGreedyLegacy();
 }
 
 FunctionPass *llvm::createGreedyRegisterAllocator(RegAllocFilterFunc Ftor) {
-  return new RAGreedy(Ftor);
+  return new RAGreedyLegacy(Ftor);
 }
 
-RAGreedy::RAGreedy(RegAllocFilterFunc F)
-    : MachineFunctionPass(ID), RegAllocBase(F) {}
-
-void RAGreedy::getAnalysisUsage(AnalysisUsage &AU) const {
+void RAGreedyLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesCFG();
   AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
   AU.addPreserved<MachineBlockFrequencyInfoWrapperPass>();
@@ -1051,7 +1174,8 @@ void RAGreedy::splitAroundRegion(LiveRangeEdit &LREdit,
   }
 
   if (VerifyEnabled)
-    MF->verify(this, "After splitting live range around region", &errs());
+    MF->verify(LIS, Indexes, "After splitting live range around region",
+               &errs());
 }
 
 MCRegister RAGreedy::tryRegionSplit(const LiveInterval &VirtReg,
@@ -1320,7 +1444,8 @@ unsigned RAGreedy::tryBlockSplit(const LiveInterval &VirtReg,
   }
 
   if (VerifyEnabled)
-    MF->verify(this, "After splitting live range around basic blocks", &errs());
+    MF->verify(LIS, Indexes, "After splitting live range around basic blocks",
+               &errs());
   return 0;
 }
 
@@ -2507,7 +2632,7 @@ MCRegister RAGreedy::selectOrSplitImpl(const LiveInterval &VirtReg,
       DebugVars->splitRegister(r, LRE.regs(), *LIS);
 
     if (VerifyEnabled)
-      MF->verify(this, "After spilling", &errs());
+      MF->verify(LIS, Indexes, "After spilling", &errs());
   }
 
   // The live virtual register requesting allocation was spilled, so tell
@@ -2703,7 +2828,7 @@ bool RAGreedy::hasVirtRegAlloc() {
   return false;
 }
 
-bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
+bool RAGreedy::run(MachineFunction &mf) {
   LLVM_DEBUG(dbgs() << "********** GREEDY REGISTER ALLOCATION **********\n"
                     << "********** Function: " << mf.getName() << '\n');
 
@@ -2711,29 +2836,18 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
   TII = MF->getSubtarget().getInstrInfo();
 
   if (VerifyEnabled)
-    MF->verify(this, "Before greedy register allocator", &errs());
+    MF->verify(LIS, Indexes, "Before greedy register allocator", &errs());
 
-  RegAllocBase::init(getAnalysis<VirtRegMapWrapperLegacy>().getVRM(),
-                     getAnalysis<LiveIntervalsWrapperPass>().getLIS(),
-                     getAnalysis<LiveRegMatrixWrapperLegacy>().getLRM());
+  RegAllocBase::init(*this->VRM, *this->LIS, *this->Matrix);
 
   // Early return if there is no virtual register to be allocated to a
   // physical register.
   if (!hasVirtRegAlloc())
     return false;
 
-  Indexes = &getAnalysis<SlotIndexesWrapperPass>().getSI();
   // Renumber to get accurate and consistent results from
   // SlotIndexes::getApproxInstrDistance.
   Indexes->packIndexes();
-  MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
-  DomTree = &getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
-  ORE = &getAnalysis<MachineOptimizationRemarkEmitterPass>().getORE();
-  Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
-  Bundles = &getAnalysis<EdgeBundlesWrapperLegacy>().getEdgeBundles();
-  SpillPlacer = &getAnalysis<SpillPlacementWrapperLegacy>().getResult();
-  DebugVars = &getAnalysis<LiveDebugVariablesWrapperLegacy>().getLDV();
-  auto &LSS = getAnalysis<LiveStacksWrapperLegacy>().getLS();
 
   initializeCSRCost();
 
@@ -2749,17 +2863,12 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
 
   ExtraInfo.emplace();
 
-  auto &EvictAdvisorProvider =
-      getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().getProvider();
-  EvictAdvisor = EvictAdvisorProvider->getAdvisor(*MF, *this, MBFI, Loops);
-
-  PriorityAdvisor = getAnalysis<RegAllocPriorityAdvisorAnalysisLegacy>()
-                        .getProvider()
-                        ->getAdvisor(*MF, *this, Indexes);
+  EvictAdvisor = EvictProvider->getAdvisor(*MF, *this, MBFI, Loops);
+  PriorityAdvisor = PriorityProvider->getAdvisor(*MF, *this, Indexes);
 
   VRAI = std::make_unique<VirtRegAuxInfo>(*MF, *LIS, *VRM, *Loops, *MBFI);
   SpillerInstance.reset(
-      createInlineSpiller({*LIS, LSS, *DomTree, *MBFI}, *MF, *VRM, *VRAI));
+      createInlineSpiller({*LIS, *LSS, *DomTree, *MBFI}, *MF, *VRM, *VRAI));
 
   VRAI->calculateSpillWeightsAndHints();
 
@@ -2776,7 +2885,7 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
   tryHintsRecoloring();
 
   if (VerifyEnabled)
-    MF->verify(this, "Before post optimization", &errs());
+    MF->verify(LIS, Indexes, "Before post optimization", &errs());
   postOptimization();
   reportStats();
 
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.h b/llvm/lib/CodeGen/RegAllocGreedy.h
index 89ceefd37795bc..278ffd0ed906b5 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.h
+++ b/llvm/lib/CodeGen/RegAllocGreedy.h
@@ -25,13 +25,15 @@
 #include "llvm/CodeGen/LiveDebugVariables.h"
 #include "llvm/CodeGen/LiveInterval.h"
 #include "llvm/CodeGen/LiveRangeEdit.h"
+#include "llvm/CodeGen/LiveStacks.h"
 #include "llvm/CodeGen/MachineFunction.h"
-#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
 #include "llvm/CodeGen/RegAllocPriorityAdvisor.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
 #include "llvm/CodeGen/SpillPlacement.h"
 #include "llvm/CodeGen/Spiller.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
+#include "llvm/IR/PassManager.h"
 #include <algorithm>
 #include <cstdint>
 #include <memory>
@@ -56,11 +58,30 @@ class SlotIndexes;
 class TargetInstrInfo;
 class VirtRegMap;
 
-class LLVM_LIBRARY_VISIBILITY RAGreedy : public MachineFunctionPass,
-                                         public RegAllocBase,
+class LLVM_LIBRARY_VISIBILITY RAGreedy : public RegAllocBase,
                                          private LiveRangeEdit::Delegate {
-  // Interface to eviction advisers
 public:
+  struct RequiredAnalyses {
+    VirtRegMap *VRM = nullptr;
+    LiveIntervals *LIS = nullptr;
+    LiveRegMatrix *LRM = nullptr;
+    SlotIndexes *Indexes = nullptr;
+    MachineBlockFrequencyInfo *MBFI = nullptr;
+    MachineDominatorTree *DomTree = nullptr;
+    MachineLoopInfo *Loops = nullptr;
+    MachineOptimizationRemarkEmitter *ORE = nullptr;
+    EdgeBundles *Bundles = nullptr;
+    SpillPlacement *SpillPlacer = nullptr;
+    LiveDebugVariables *DebugVars = nullptr;
+
+    // Used by InlineSpiller
+    LiveStacks *LSS;
+    // Proxies for eviction and priority advisors
+    RegAllocEvictionAdvisorProvider *EvictProvider;
+    RegAllocPriorityAdvisorProvider *PriorityProvider;
+  };
+
+  // Interface to eviction advisers
   /// Track allocation stage and eviction loop prevention during allocation.
   class ExtraRegInfo final {
     // RegInfo - Keep additional information about each live range.
@@ -178,6 +199,10 @@ class LLVM_LIBRARY_VISIBILITY RAGreedy : public MachineFunctionPass,
   EdgeBundles *Bundles = nullptr;
   SpillPlacement *SpillPlacer = nullptr;
   LiveDebugVariables *DebugVars = nullptr;
+  LiveStacks *LSS = nullptr; // Used by InlineSpiller
+  // Proxy for the advisors
+  RegAllocEvictionAdvisorProvider *EvictProvider = nullptr;
+  RegAllocPriorityAdvisorProvider *PriorityProvider = nullptr;
 
   // state
   std::unique_ptr<Spiller> SpillerInstance;
@@ -282,13 +307,11 @@ class LLVM_LIBRARY_VISIBILITY RAGreedy : public MachineFunctionPass,
 
 public:
   RAGreedy(const RegAllocFilterFunc F = nullptr);
+  // Evict and priority advisors use this object, so we can construct those
+  // first and pass them here.
+  // Not required once legacy PM is removed.
+  void setAnalyses(RequiredAnalyses &Analyses);
 
-  /// Return the pass name.
-  StringRef getPassName() const override { return "Greedy Register Allocator"; }
-
-  /// RAGreedy analysis usage.
-  void getAnalysisUsage(AnalysisUsage &AU) const override;
-  void releaseMemory() override;
   Spiller &spiller() override { return *SpillerInstance; }
   void enqueueImpl(const LiveInterval *LI) override;
   const LiveInterval *dequeue() override;
@@ -297,19 +320,9 @@ class LLVM_LIBRARY_VISIBILITY RAGreedy : public MachineFunctionPass,
   void aboutToRemoveInterval(const LiveInterval &) override;
 
   /// Perform register allocation.
-  bool runOnMachineFunction(MachineFunction &mf) override;
-
-  MachineFunctionProperties getRequiredProperties() const override {
-    return MachineFunctionProperties().set(
-        MachineFunctionProperties::Property::NoPHIs);
-  }
-
-  MachineFunctionProperties getClearedProperties() const override {
-    return MachineFunctionProperties().set(
-        MachineFunctionProperties::Property::IsSSA);
-  }
+  bool run(MachineFunction &mf);
 
-  static char ID;
+  void releaseMemory();
 
 private:
   MCRegister selectOrSplitImpl(const LiveInterval &,
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 5394e06282e50c..76359f369946c4 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -125,6 +125,7 @@
 #include "llvm/CodeGen/PreISelIntrinsicLowering.h"
 #include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
 #include "llvm/CodeGen/RegAllocFast.h"
+#include "llvm/CodeGen/RegAllocGreedyPass.h"
 #include "llvm/CodeGen/RegAllocPriorityAdvisor.h"
 #include "llvm/CodeGen/RegUsageInfoCollector.h"
 #include "llvm/CodeGen/RegUsageInfoPropagate.h"

>From cb8aa70a4eb17ce78c08078302ad32bb1cb20610 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Fri, 13 Dec 2024 10:59:04 +0000
Subject: [PATCH 17/22] Move pass to a new header, add options

---
 .../include/llvm/CodeGen/RegAllocGreedyPass.h | 43 +++++++++++++++++++
 llvm/include/llvm/Passes/CodeGenPassBuilder.h |  1 +
 .../llvm/Passes/MachinePassRegistry.def       | 11 +++--
 llvm/lib/CodeGen/RegAllocGreedy.cpp           |  7 ++-
 4 files changed, 55 insertions(+), 7 deletions(-)
 create mode 100644 llvm/include/llvm/CodeGen/RegAllocGreedyPass.h

diff --git a/llvm/include/llvm/CodeGen/RegAllocGreedyPass.h b/llvm/include/llvm/CodeGen/RegAllocGreedyPass.h
new file mode 100644
index 00000000000000..f325224c5384cb
--- /dev/null
+++ b/llvm/include/llvm/CodeGen/RegAllocGreedyPass.h
@@ -0,0 +1,43 @@
+//==- RegAllocGreedyPass.h --- greedy register allocator pass ------*-C++-*-==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/RegAllocCommon.h"
+#include "llvm/CodeGen/RegAllocFast.h"
+#include "llvm/IR/PassManager.h"
+
+using namespace llvm;
+
+class RAGreedyPass : public PassInfoMixin<RAGreedyPass> {
+
+public:
+  struct Options {
+    RegAllocFilterFunc Filter;
+    StringRef FilterName;
+    Options(RegAllocFilterFunc F = nullptr, StringRef FN = "all")
+        : Filter(F), FilterName(FN) {};
+  };
+
+  RAGreedyPass(Options Opts = Options()) : Opts(Opts) {}
+  PreservedAnalyses run(MachineFunction &F, MachineFunctionAnalysisManager &AM);
+
+  MachineFunctionProperties getRequiredProperties() const {
+    return MachineFunctionProperties().set(
+        MachineFunctionProperties::Property::NoPHIs);
+  }
+
+  MachineFunctionProperties getClearedProperties() const {
+    return MachineFunctionProperties().set(
+        MachineFunctionProperties::Property::IsSSA);
+  }
+
+  void printPipeline(raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) const;
+  static bool isRequired() { return true; }
+
+private:
+  Options Opts;
+};
diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index bcd335f2f99427..28768a72c83fa3 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -56,6 +56,7 @@
 #include "llvm/CodeGen/PreISelIntrinsicLowering.h"
 #include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
 #include "llvm/CodeGen/RegAllocFast.h"
+#include "llvm/CodeGen/RegAllocGreedyPass.h"
 #include "llvm/CodeGen/RegUsageInfoCollector.h"
 #include "llvm/CodeGen/RegUsageInfoPropagate.h"
 #include "llvm/CodeGen/RegisterUsageInfo.h"
diff --git a/llvm/include/llvm/Passes/MachinePassRegistry.def b/llvm/include/llvm/Passes/MachinePassRegistry.def
index bf12a2c9aca90b..35314473caa610 100644
--- a/llvm/include/llvm/Passes/MachinePassRegistry.def
+++ b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -189,12 +189,12 @@ MACHINE_FUNCTION_PASS_WITH_PARAMS(
     "filter=reg-filter;no-clear-vregs")
 
 MACHINE_FUNCTION_PASS_WITH_PARAMS(
-    "regallocgreedy", "RAGreedy",
-    [](RegAllocFilterFunc F) { return RAGreedyPass(F); },
+    "regallocgreedy", "RAGreedyPass",
+    [](RAGreedyPass::Options Opts) { return RAGreedyPass(Opts); },
     [PB = this](StringRef Params) {
-      // TODO: parseRegAllocFilter(*PB, Params);
-      return Expected<RegAllocFilterFunc>(nullptr);
-    }, ""
+      // TODO: parseRegAllocGreedyFilterFunc(*PB, Params);
+      return Expected<RAGreedyPass::Options>(RAGreedyPass::Options{});
+    }, "reg-filter"
 )
 #undef MACHINE_FUNCTION_PASS_WITH_PARAMS
 
@@ -266,7 +266,6 @@ DUMMY_MACHINE_FUNCTION_PASS("processimpdefs", ProcessImplicitDefsPass)
 DUMMY_MACHINE_FUNCTION_PASS("prologepilog", PrologEpilogInserterPass)
 DUMMY_MACHINE_FUNCTION_PASS("prologepilog-code", PrologEpilogCodeInserterPass)
 DUMMY_MACHINE_FUNCTION_PASS("ra-basic", RABasicPass)
-DUMMY_MACHINE_FUNCTION_PASS("ra-greedy", RAGreedyPass)
 DUMMY_MACHINE_FUNCTION_PASS("ra-pbqp", RAPBQPPass)
 DUMMY_MACHINE_FUNCTION_PASS("regalloc", RegAllocPass)
 DUMMY_MACHINE_FUNCTION_PASS("regallocscoringpass", RegAllocScoringPass)
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp
index 3a96325144efbc..a07406fef76952 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.cpp
+++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp
@@ -202,11 +202,16 @@ void RAGreedy::setAnalyses(RequiredAnalyses &Analyses) {
   PriorityProvider = Analyses.PriorityProvider;
 }
 
+void RAGreedyPass::printPipeline(raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) const {
+  StringRef FilterName = Opts.FilterName.empty() ? "all" : Opts.FilterName;
+  OS << "regallocgreedy<" << FilterName << ">";
+}
+
 PreservedAnalyses RAGreedyPass::run(MachineFunction &MF,
                                     MachineFunctionAnalysisManager &MFAM) {
   MFPropsModifier _(*this, MF);
 
-  RAGreedy Impl(Filter);
+  RAGreedy Impl(Opts.Filter);
   RAGreedy::RequiredAnalyses Analyses;
 
   Analyses.VRM = &MFAM.getResult<VirtRegMapAnalysis>(MF);

>From f18aa9e22d886a6574c8f1d06986156bcff836b9 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Mon, 16 Dec 2024 09:34:30 +0000
Subject: [PATCH 18/22] Move VRM after LiveIntervals

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

diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp
index a07406fef76952..cad49d2cc11bcb 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.cpp
+++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp
@@ -214,7 +214,6 @@ PreservedAnalyses RAGreedyPass::run(MachineFunction &MF,
   RAGreedy Impl(Opts.Filter);
   RAGreedy::RequiredAnalyses Analyses;
 
-  Analyses.VRM = &MFAM.getResult<VirtRegMapAnalysis>(MF);
   Analyses.LIS = &MFAM.getResult<LiveIntervalsAnalysis>(MF);
   Analyses.LRM = &MFAM.getResult<LiveRegMatrixAnalysis>(MF);
   Analyses.LSS = &MFAM.getResult<LiveStacksAnalysis>(MF);
@@ -230,6 +229,7 @@ PreservedAnalyses RAGreedyPass::run(MachineFunction &MF,
       MFAM.getResult<RegAllocEvictionAdvisorAnalysis>(MF).Provider;
   Analyses.PriorityProvider =
       MFAM.getResult<RegAllocPriorityAdvisorAnalysis>(MF).Provider;
+  Analyses.VRM = &MFAM.getResult<VirtRegMapAnalysis>(MF);
 
   Impl.setAnalyses(Analyses);
   bool Changed = Impl.run(MF);

>From c1f4df989c63c9b9107ec2aceead3d935deae7cf Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Wed, 11 Dec 2024 10:57:21 +0000
Subject: [PATCH 19/22] [RegAlloc][NewPM] Plug Greedy RA in codegen pipeline

---
 llvm/include/llvm/Passes/CodeGenPassBuilder.h | 18 +++++++++++++++-
 .../llvm/Passes/MachinePassRegistry.def       |  4 ++--
 .../include/llvm/Target/CGPassBuilderOption.h |  2 +-
 llvm/lib/Passes/PassBuilder.cpp               | 13 ++++++++++++
 ...plicit-def-remat-requires-impdef-check.mir |  1 +
 ...implicit-def-with-impdef-greedy-assert.mir |  1 +
 llvm/test/CodeGen/AArch64/pr51516.mir         |  1 +
 llvm/test/CodeGen/AArch64/spill-fold.mir      |  2 ++
 .../extend-phi-subrange-not-in-parent.mir     |  1 +
 llvm/test/CodeGen/MIR/Generic/runPass.mir     |  1 +
 .../SystemZ/clear-liverange-spillreg.mir      |  1 +
 llvm/test/CodeGen/Thumb/high-reg-clobber.mir  |  1 +
 llvm/test/CodeGen/X86/limit-split-cost.mir    |  1 +
 .../test/tools/llc/new-pm/regalloc-amdgpu.mir | 17 +++++++++------
 llvm/tools/llc/NewPMDriver.cpp                | 21 +++++++++++++++----
 15 files changed, 71 insertions(+), 14 deletions(-)

diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index 28768a72c83fa3..4135abdad90ff6 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -1059,7 +1059,7 @@ void CodeGenPassBuilder<Derived, TargetMachineT>::addMachineSSAOptimization(
 ///
 /// A target that uses the standard regalloc pass order for fast or optimized
 /// allocation may still override this for per-target regalloc
-/// selection. But -regalloc=... always takes precedence.
+/// selection. But -regalloc-npm=... always takes precedence.
 template <typename Derived, typename TargetMachineT>
 void CodeGenPassBuilder<Derived, TargetMachineT>::addTargetRegisterAllocator(
     AddMachinePass &addPass, bool Optimized) const {
@@ -1076,6 +1076,22 @@ template <typename Derived, typename TargetMachineT>
 void CodeGenPassBuilder<Derived, TargetMachineT>::addRegAllocPass(
     AddMachinePass &addPass, bool Optimized) const {
   // TODO: Parse Opt.RegAlloc to add register allocator.
+  // Use the specified -regalloc-npm={basic|greedy|fast|pbqp}
+  if (Opt.RegAlloc > RegAllocType::Default) {
+    switch (Opt.RegAlloc) {
+      case RegAllocType::Fast:
+        addPass(RegAllocFastPass());
+        break;
+      case RegAllocType::Greedy:
+        addPass(RAGreedyPass());
+        break;
+      default:
+        llvm_unreachable("Register allocator not supported yet.");
+    }
+    return;
+  }
+  // -regalloc=default or unspecified, so pick based on the optimization level.
+  derived().addTargetRegisterAllocator(addPass, Optimized);
 }
 
 template <typename Derived, typename TargetMachineT>
diff --git a/llvm/include/llvm/Passes/MachinePassRegistry.def b/llvm/include/llvm/Passes/MachinePassRegistry.def
index 35314473caa610..61c3c688eda52f 100644
--- a/llvm/include/llvm/Passes/MachinePassRegistry.def
+++ b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -188,12 +188,12 @@ MACHINE_FUNCTION_PASS_WITH_PARAMS(
     },
     "filter=reg-filter;no-clear-vregs")
 
+// 'all' is the default filter
 MACHINE_FUNCTION_PASS_WITH_PARAMS(
     "regallocgreedy", "RAGreedyPass",
     [](RAGreedyPass::Options Opts) { return RAGreedyPass(Opts); },
     [PB = this](StringRef Params) {
-      // TODO: parseRegAllocGreedyFilterFunc(*PB, Params);
-      return Expected<RAGreedyPass::Options>(RAGreedyPass::Options{});
+      return parseRegAllocGreedyFilterFunc(*PB, Params);
     }, "reg-filter"
 )
 #undef MACHINE_FUNCTION_PASS_WITH_PARAMS
diff --git a/llvm/include/llvm/Target/CGPassBuilderOption.h b/llvm/include/llvm/Target/CGPassBuilderOption.h
index 29bdb9c1746d3c..05c84c1bdec851 100644
--- a/llvm/include/llvm/Target/CGPassBuilderOption.h
+++ b/llvm/include/llvm/Target/CGPassBuilderOption.h
@@ -53,7 +53,7 @@ struct CGPassBuilderOption {
   bool RequiresCodeGenSCCOrder = false;
 
   RunOutliner EnableMachineOutliner = RunOutliner::TargetDefault;
-  StringRef RegAlloc = "default";
+  RegAllocType RegAlloc = RegAllocType::Default;
   std::optional<GlobalISelAbortMode> EnableGlobalISelAbort;
   std::string FSProfileFile;
   std::string FSRemappingFile;
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 76359f369946c4..a733df8c4ba2e8 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -1312,6 +1312,19 @@ parseBoundsCheckingOptions(StringRef Params) {
   return Mode;
 }
 
+Expected<RAGreedyPass::Options> parseRegAllocGreedyFilterFunc(PassBuilder &PB, StringRef Params) {
+  if (Params.empty() || Params == "all") {
+    return RAGreedyPass::Options();
+  }
+  std::optional<RegAllocFilterFunc> Filter = PB.parseRegAllocFilter(Params);
+  if (!Filter) {
+    return make_error<StringError>(
+        formatv("invalid regallocgreedy register filter '{0}' ", Params).str(),
+        inconvertibleErrorCode());
+  }
+  return RAGreedyPass::Options{*Filter, Params};
+}
+
 } // namespace
 
 /// Tests whether a pass name starts with a valid prefix for a default pipeline
diff --git a/llvm/test/CodeGen/AArch64/implicit-def-remat-requires-impdef-check.mir b/llvm/test/CodeGen/AArch64/implicit-def-remat-requires-impdef-check.mir
index 47aa34e3c01156..a168c2891c7d6f 100644
--- a/llvm/test/CodeGen/AArch64/implicit-def-remat-requires-impdef-check.mir
+++ b/llvm/test/CodeGen/AArch64/implicit-def-remat-requires-impdef-check.mir
@@ -1,5 +1,6 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 4
 # RUN: llc -mtriple=arm64-apple-macosx -mcpu=apple-m1 -stress-regalloc=4 -verify-regalloc -run-pass=greedy -o - %s | FileCheck %s
+# RUN: llc -mtriple=arm64-apple-macosx -mcpu=apple-m1 -stress-regalloc=4 -verify-regalloc -passes=regallocgreedy -o - %s | FileCheck %s
 
 --- |
   define void @inst_stores_to_dead_spill_implicit_def_impdef() {
diff --git a/llvm/test/CodeGen/AArch64/implicit-def-with-impdef-greedy-assert.mir b/llvm/test/CodeGen/AArch64/implicit-def-with-impdef-greedy-assert.mir
index a5d74ef75f0a0a..d9edda47638a3f 100644
--- a/llvm/test/CodeGen/AArch64/implicit-def-with-impdef-greedy-assert.mir
+++ b/llvm/test/CodeGen/AArch64/implicit-def-with-impdef-greedy-assert.mir
@@ -1,5 +1,6 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 3
 # RUN: llc -mtriple=arm64-apple-ios -run-pass=greedy -o - %s | FileCheck %s
+# RUN: llc -mtriple=arm64-apple-ios -passes=regallocgreedy -o - %s | FileCheck %s
 
 ---
 name:            widget
diff --git a/llvm/test/CodeGen/AArch64/pr51516.mir b/llvm/test/CodeGen/AArch64/pr51516.mir
index 910bfb858b50f6..e84f0ca2015ce5 100644
--- a/llvm/test/CodeGen/AArch64/pr51516.mir
+++ b/llvm/test/CodeGen/AArch64/pr51516.mir
@@ -1,4 +1,5 @@
 # RUN: llc -mtriple=aarch64-unknown-fuchsia -run-pass=greedy -verify-machineinstrs -o - %s | FileCheck %s
+# RUN: llc -mtriple=aarch64-unknown-fuchsia -passes=regallocgreedy -verify-machineinstrs -o - %s | FileCheck %s
 
 # Check that we spill %31 and do not rematerialize it since the use operand
 # of ADDXri is killed by the STRXui in this block.
diff --git a/llvm/test/CodeGen/AArch64/spill-fold.mir b/llvm/test/CodeGen/AArch64/spill-fold.mir
index b1e7ebe3a7e82b..2773b5f19618a9 100644
--- a/llvm/test/CodeGen/AArch64/spill-fold.mir
+++ b/llvm/test/CodeGen/AArch64/spill-fold.mir
@@ -1,5 +1,7 @@
 # RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass greedy -verify-machineinstrs  -o - %s | FileCheck %s
 # RUN: llc -mtriple=aarch64_be-none-linux-gnu -run-pass greedy -verify-machineinstrs  -o - %s | FileCheck %s
+# RUN: llc -mtriple=aarch64-none-linux-gnu -passes=regallocgreedy -o - %s | FileCheck %s
+# RUN: llc -mtriple=aarch64_be-none-linux-gnu -passes=regallocgreedy -o - %s | FileCheck %s
 --- |
   define i64 @test_subreg_spill_fold() { ret i64 0 }
   define i64 @test_subreg_spill_fold2() { ret i64 0 }
diff --git a/llvm/test/CodeGen/AMDGPU/extend-phi-subrange-not-in-parent.mir b/llvm/test/CodeGen/AMDGPU/extend-phi-subrange-not-in-parent.mir
index 760ae6032230f5..42bba4d1504013 100644
--- a/llvm/test/CodeGen/AMDGPU/extend-phi-subrange-not-in-parent.mir
+++ b/llvm/test/CodeGen/AMDGPU/extend-phi-subrange-not-in-parent.mir
@@ -1,5 +1,6 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx90a -verify-regalloc -run-pass=greedy -o - %s | FileCheck %s
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx90a -verify-regalloc -passes=regallocgreedy -o - %s | FileCheck %s
 
 # Initially %2 starts out with 2 subranges (one for sub0, and one for
 # the rest of the lanes). After %2 is split, after refineSubRanges the
diff --git a/llvm/test/CodeGen/MIR/Generic/runPass.mir b/llvm/test/CodeGen/MIR/Generic/runPass.mir
index 75763c5389b09e..41dd98ff909b0c 100644
--- a/llvm/test/CodeGen/MIR/Generic/runPass.mir
+++ b/llvm/test/CodeGen/MIR/Generic/runPass.mir
@@ -2,6 +2,7 @@
 # RUN: llc -run-pass=regallocbasic -debug-pass=Arguments -o - %s | FileCheck %s
 # RUN: llc -run-pass=regallocfast -debug-pass=Arguments -o - %s | FileCheck %s
 # RUN: llc -passes=regallocfast -o - %s | FileCheck %s
+# RUN: llc -passes=regallocgreedy -o - %s | FileCheck %s
 
 # Check that passes are initialized correctly, so that it's possible to
 # use -run-pass.
diff --git a/llvm/test/CodeGen/SystemZ/clear-liverange-spillreg.mir b/llvm/test/CodeGen/SystemZ/clear-liverange-spillreg.mir
index 197c3d8551fc38..de0db97f14bf3c 100644
--- a/llvm/test/CodeGen/SystemZ/clear-liverange-spillreg.mir
+++ b/llvm/test/CodeGen/SystemZ/clear-liverange-spillreg.mir
@@ -1,4 +1,5 @@
 #RUN: llc -o - %s -mtriple=s390x-ibm-linux -run-pass=greedy
+#RUN: llc -o - %s -mtriple=s390x-ibm-linux -passes=regallocgreedy
 #PR34502. Check HoistSpill works properly after the live range of spilled
 #virtual register is cleared.
 --- |
diff --git a/llvm/test/CodeGen/Thumb/high-reg-clobber.mir b/llvm/test/CodeGen/Thumb/high-reg-clobber.mir
index 1402c7c2cbca36..e085e38ae5fe31 100644
--- a/llvm/test/CodeGen/Thumb/high-reg-clobber.mir
+++ b/llvm/test/CodeGen/Thumb/high-reg-clobber.mir
@@ -3,6 +3,7 @@
 # RUN: llc -mtriple thumbv6m-arm-none-eabi -run-pass greedy %s -o - | FileCheck %s
 # RUN: llc -mtriple thumbv6m-arm-none-eabi -run-pass regallocfast %s -o - | FileCheck %s --check-prefix=FAST
 # RUN: llc -mtriple thumbv6m-arm-none-eabi -passes=regallocfast %s -o - | FileCheck %s --check-prefix=FAST
+# RUN: llc -mtriple thumbv6m-arm-none-eabi -passes=regallocgreedy %s -o - | FileCheck %s
 
 ...
 ---
diff --git a/llvm/test/CodeGen/X86/limit-split-cost.mir b/llvm/test/CodeGen/X86/limit-split-cost.mir
index 7ec0404e0f737c..eec8a3939151a6 100644
--- a/llvm/test/CodeGen/X86/limit-split-cost.mir
+++ b/llvm/test/CodeGen/X86/limit-split-cost.mir
@@ -1,5 +1,6 @@
 # REQUIRES: asserts
 # RUN: llc -mtriple=x86_64-- -run-pass=greedy %s -debug-only=regalloc -huge-size-for-split=0 -o /dev/null 2>&1 | FileCheck %s
+# RUN: llc -mtriple=x86_64-- -passes=regallocgreedy %s -debug-only=regalloc -huge-size-for-split=0 -o /dev/null 2>&1 | FileCheck %s
 # Check no global region split is needed because the live range to split is trivially rematerializable.
 # CHECK-NOT: Compact region bundles
 --- |
diff --git a/llvm/test/tools/llc/new-pm/regalloc-amdgpu.mir b/llvm/test/tools/llc/new-pm/regalloc-amdgpu.mir
index 07f2d350ffd9c0..66c9d8942f3da4 100644
--- a/llvm/test/tools/llc/new-pm/regalloc-amdgpu.mir
+++ b/llvm/test/tools/llc/new-pm/regalloc-amdgpu.mir
@@ -1,12 +1,17 @@
 # REQUIRES: amdgpu-registered-target
-# RUN: llc -mtriple=amdgcn --passes='regallocfast<filter=sgpr>,regallocfast<filter=wwm>,regallocfast<filter=vgpr>' --print-pipeline-passes --filetype=null %s | FileCheck %s --check-prefix=PASS
-# RUN: not llc -mtriple=amdgcn --passes='regallocfast<filter=bad-filter>' --print-pipeline-passes --filetype=null %s 2>&1 | FileCheck %s --check-prefix=BAD-FILTER
+# RUN: llc -mtriple=amdgcn --passes='regallocfast<filter=sgpr>,regallocfast<filter=wwm>,regallocfast<filter=vgpr>' --print-pipeline-passes --filetype=null %s | FileCheck %s --check-prefix=RAFAST
+# RUN: not llc -mtriple=amdgcn --passes='regallocfast<filter=bad-filter>' --print-pipeline-passes --filetype=null %s 2>&1 | FileCheck %s --check-prefix=RAFAST-BAD-FILTER
 
-# PASS: regallocfast<filter=sgpr>
-# PASS: regallocfast<filter=wwm>
-# PASS: regallocfast<filter=vgpr>
-# BAD-FILTER: invalid regallocfast register filter 'bad-filter'
+# RUN: llc -mtriple=amdgcn -passes='regallocgreedy<sgpr>' --print-pipeline-passes --filetype=null %s | FileCheck %s --check-prefix=RAGREEDY
+# RUN: not llc -mtriple=amdgcn -passes='regallocgreedy<bad-filter>' --print-pipeline-passes --filetype=null %s 2>&1 | FileCheck %s --check-prefix=RAGREEDY-BAD-FILTER
 
+# RAFAST: regallocfast<filter=sgpr>
+# RAFAST: regallocfast<filter=wwm>
+# RAFAST: regallocfast<filter=vgpr>
+# RAFAST-BAD-FILTER: invalid regallocfast register filter 'bad-filter'
+
+# RAGREEDY: regallocgreedy<sgpr>
+# RAGREEDY-BAD-FILTER: invalid regallocgreedy register filter 'bad-filter'
 ---
 name: f
 ...
diff --git a/llvm/tools/llc/NewPMDriver.cpp b/llvm/tools/llc/NewPMDriver.cpp
index 3892fbb8c74f78..06b7fa7ac351b0 100644
--- a/llvm/tools/llc/NewPMDriver.cpp
+++ b/llvm/tools/llc/NewPMDriver.cpp
@@ -48,10 +48,23 @@
 
 using namespace llvm;
 
-static cl::opt<std::string>
-    RegAlloc("regalloc-npm",
-             cl::desc("Register allocator to use for new pass manager"),
-             cl::Hidden, cl::init("default"));
+// static cl::opt<std::string>
+//     RegAlloc("regalloc-npm",
+//              cl::desc("Register allocator to use for new pass manager"),
+//              cl::Hidden, cl::init("default"));
+
+// create option for RegAllocType enum
+static cl::opt<RegAllocType> RegAlloc(
+    "regalloc-npm", cl::desc("Register allocator to use for new pass manager"),
+    cl::Hidden, cl::init(RegAllocType::Default),
+    cl::values(
+        clEnumValN(RegAllocType::Default, "default",
+                   "Default register allocator"),
+        clEnumValN(RegAllocType::PBQP, "pbqp", "PBQP register allocator"),
+        clEnumValN(RegAllocType::Fast, "fast", "Fast register allocator"),
+        clEnumValN(RegAllocType::Basic, "basic", "Basic register allocator"),
+        clEnumValN(RegAllocType::Greedy, "greedy",
+                   "Greedy register allocator")));
 
 static cl::opt<bool>
     DebugPM("debug-pass-manager", cl::Hidden,

>From 9d7cf7471f36883c32242317e81673b369812201 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Thu, 19 Dec 2024 11:00:46 +0000
Subject: [PATCH 20/22] Error out on AMDGPU for regalloc-npm flag

---
 llvm/include/llvm/Passes/CodeGenPassBuilder.h | 29 ++++++++++---------
 .../include/llvm/Target/CGPassBuilderOption.h |  4 +--
 .../lib/Target/AMDGPU/AMDGPUTargetMachine.cpp | 23 +++++++++++++++
 llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h  |  2 ++
 .../CodeGen/AMDGPU/sgpr-regalloc-flags.ll     |  3 ++
 llvm/tools/llc/NewPMDriver.cpp                |  2 +-
 6 files changed, 47 insertions(+), 16 deletions(-)

diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index 4135abdad90ff6..31fb8af7576c7b 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -1086,7 +1086,7 @@ void CodeGenPassBuilder<Derived, TargetMachineT>::addRegAllocPass(
         addPass(RAGreedyPass());
         break;
       default:
-        llvm_unreachable("Register allocator not supported yet.");
+        report_fatal_error("Register allocator not supported yet.", false);
     }
     return;
   }
@@ -1162,20 +1162,23 @@ void CodeGenPassBuilder<Derived, TargetMachineT>::addOptimizedRegAlloc(
   // PreRA instruction scheduling.
   addPass(MachineSchedulerPass());
 
-  if (derived().addRegAssignmentOptimized(addPass)) {
-    // Allow targets to expand pseudo instructions depending on the choice of
-    // registers before MachineCopyPropagation.
-    derived().addPostRewrite(addPass);
+  if (auto E = derived().addRegAssignmentOptimized(addPass)) {
+    // addRegAssignmentOptimized did not add a reg alloc pass, so do nothing.
+    // FIXME: This is not really an error.
+    return;
+  }
+  // Allow targets to expand pseudo instructions depending on the choice of
+  // registers before MachineCopyPropagation.
+  derived().addPostRewrite(addPass);
 
-    // Copy propagate to forward register uses and try to eliminate COPYs that
-    // were not coalesced.
-    addPass(MachineCopyPropagationPass());
+  // Copy propagate to forward register uses and try to eliminate COPYs that
+  // were not coalesced.
+  addPass(MachineCopyPropagationPass());
 
-    // Run post-ra machine LICM to hoist reloads / remats.
-    //
-    // FIXME: can this move into MachineLateOptimization?
-    addPass(MachineLICMPass());
-  }
+  // Run post-ra machine LICM to hoist reloads / remats.
+  //
+  // FIXME: can this move into MachineLateOptimization?
+  addPass(MachineLICMPass());
 }
 
 //===---------------------------------------------------------------------===//
diff --git a/llvm/include/llvm/Target/CGPassBuilderOption.h b/llvm/include/llvm/Target/CGPassBuilderOption.h
index 05c84c1bdec851..f146cc7b8028d8 100644
--- a/llvm/include/llvm/Target/CGPassBuilderOption.h
+++ b/llvm/include/llvm/Target/CGPassBuilderOption.h
@@ -20,7 +20,7 @@
 namespace llvm {
 
 enum class RunOutliner { TargetDefault, AlwaysOutline, NeverOutline };
-enum class RegAllocType { Default, Basic, Fast, Greedy, PBQP };
+enum class RegAllocType { Unset, Default, Basic, Fast, Greedy, PBQP };
 
 // Not one-on-one but mostly corresponding to commandline options in
 // TargetPassConfig.cpp.
@@ -53,7 +53,7 @@ struct CGPassBuilderOption {
   bool RequiresCodeGenSCCOrder = false;
 
   RunOutliner EnableMachineOutliner = RunOutliner::TargetDefault;
-  RegAllocType RegAlloc = RegAllocType::Default;
+  RegAllocType RegAlloc = RegAllocType::Unset;
   std::optional<GlobalISelAbortMode> EnableGlobalISelAbort;
   std::string FSProfileFile;
   std::string FSRemappingFile;
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
index 7256eec89008a5..952bf479827f48 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -68,6 +68,7 @@
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Support/FormatVariadic.h"
+#include "llvm/Target/CGPassBuilderOption.h"
 #include "llvm/Transforms/HipStdPar/HipStdPar.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/IPO/AlwaysInliner.h"
@@ -2099,6 +2100,28 @@ void AMDGPUCodeGenPassBuilder::addMachineSSAOptimization(
   addPass(SIShrinkInstructionsPass());
 }
 
+static const char RegAllocNPMNotSupportedMessage[] =
+    "-regalloc-npm not supported with amdgcn. Use -sgpr-regalloc-npm, "
+    "-wwm-regalloc-npm, and -vgpr-regalloc-npm";
+
+Error AMDGPUCodeGenPassBuilder::addRegAssignmentOptimized(
+    AddMachinePass &addPass) const {
+  if (Opt.RegAlloc != RegAllocType::Unset)
+    report_fatal_error(RegAllocNPMNotSupportedMessage, false);
+
+  return make_error<StringError>("not implemented yet",
+                                 inconvertibleErrorCode());
+}
+
+Error AMDGPUCodeGenPassBuilder::addRegAssignmentFast(
+    AddMachinePass &addPass) const {
+  if (Opt.RegAlloc != RegAllocType::Unset)
+    report_fatal_error(RegAllocNPMNotSupportedMessage, false);
+
+  return make_error<StringError>("not implemented yet",
+                                 inconvertibleErrorCode());
+}
+
 bool AMDGPUCodeGenPassBuilder::isPassEnabled(const cl::opt<bool> &Opt,
                                              CodeGenOptLevel Level) const {
   if (Opt.getNumOccurrences())
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h
index 5ba58a92621edb..197476a0f80574 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h
@@ -176,6 +176,8 @@ class AMDGPUCodeGenPassBuilder
   void addAsmPrinter(AddMachinePass &, CreateMCStreamer) const;
   Error addInstSelector(AddMachinePass &) const;
   void addMachineSSAOptimization(AddMachinePass &) const;
+  Error addRegAssignmentOptimized(AddMachinePass &) const;
+  Error addRegAssignmentFast(AddMachinePass &) const;
 
   /// Check if a pass is enabled given \p Opt option. The option always
   /// overrides defaults if explicitly used. Otherwise its default will be used
diff --git a/llvm/test/CodeGen/AMDGPU/sgpr-regalloc-flags.ll b/llvm/test/CodeGen/AMDGPU/sgpr-regalloc-flags.ll
index 52ad7e5355207d..a54fee3a0f964c 100644
--- a/llvm/test/CodeGen/AMDGPU/sgpr-regalloc-flags.ll
+++ b/llvm/test/CodeGen/AMDGPU/sgpr-regalloc-flags.ll
@@ -12,8 +12,11 @@
 ; RUN: not --crash llc -verify-machineinstrs=0 -regalloc=basic -mtriple=amdgcn-amd-amdhsa -debug-pass=Structure -o /dev/null %s 2>&1 | FileCheck -check-prefix=REGALLOC %s
 ; RUN: not --crash llc -verify-machineinstrs=0 -regalloc=fast -O0 -mtriple=amdgcn-amd-amdhsa -debug-pass=Structure -o /dev/null %s 2>&1 | FileCheck -check-prefix=REGALLOC %s
 
+; RUN: not llc -enable-new-pm -verify-machineinstrs=0 -regalloc-npm=fast -O0 -mtriple=amdgcn-amd-amdhsa -o /dev/null %s 2>&1 | FileCheck -check-prefix=REGALLOC-NPM %s
+; RUN: not llc -enable-new-pm -verify-machineinstrs=0 -regalloc-npm=basic -O3 -mtriple=amdgcn-amd-amdhsa -o /dev/null %s 2>&1 | FileCheck -check-prefix=REGALLOC-NPM %s
 
 ; REGALLOC: -regalloc not supported with amdgcn. Use -sgpr-regalloc, -wwm-regalloc, and -vgpr-regalloc
+; REGALLOC-NPM: -regalloc-npm not supported with amdgcn. Use -sgpr-regalloc-npm, -wwm-regalloc-npm, and -vgpr-regalloc-npm
 
 ; DEFAULT: Greedy Register Allocator
 ; DEFAULT-NEXT: Virtual Register Rewriter
diff --git a/llvm/tools/llc/NewPMDriver.cpp b/llvm/tools/llc/NewPMDriver.cpp
index 06b7fa7ac351b0..de15b3717da699 100644
--- a/llvm/tools/llc/NewPMDriver.cpp
+++ b/llvm/tools/llc/NewPMDriver.cpp
@@ -56,7 +56,7 @@ using namespace llvm;
 // create option for RegAllocType enum
 static cl::opt<RegAllocType> RegAlloc(
     "regalloc-npm", cl::desc("Register allocator to use for new pass manager"),
-    cl::Hidden, cl::init(RegAllocType::Default),
+    cl::Hidden, cl::init(RegAllocType::Unset),
     cl::values(
         clEnumValN(RegAllocType::Default, "default",
                    "Default register allocator"),

>From 5ae7794e5d19435eecc47cd7df6f4cdf633e3cca Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Thu, 19 Dec 2024 14:24:53 +0000
Subject: [PATCH 21/22] Remove the TODO as it's implemented

---
 llvm/include/llvm/Passes/CodeGenPassBuilder.h | 1 -
 llvm/tools/llc/NewPMDriver.cpp                | 6 ------
 2 files changed, 7 deletions(-)

diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index 31fb8af7576c7b..d74939d9efb8f4 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -1075,7 +1075,6 @@ void CodeGenPassBuilder<Derived, TargetMachineT>::addTargetRegisterAllocator(
 template <typename Derived, typename TargetMachineT>
 void CodeGenPassBuilder<Derived, TargetMachineT>::addRegAllocPass(
     AddMachinePass &addPass, bool Optimized) const {
-  // TODO: Parse Opt.RegAlloc to add register allocator.
   // Use the specified -regalloc-npm={basic|greedy|fast|pbqp}
   if (Opt.RegAlloc > RegAllocType::Default) {
     switch (Opt.RegAlloc) {
diff --git a/llvm/tools/llc/NewPMDriver.cpp b/llvm/tools/llc/NewPMDriver.cpp
index de15b3717da699..0f7aa6284962a2 100644
--- a/llvm/tools/llc/NewPMDriver.cpp
+++ b/llvm/tools/llc/NewPMDriver.cpp
@@ -48,12 +48,6 @@
 
 using namespace llvm;
 
-// static cl::opt<std::string>
-//     RegAlloc("regalloc-npm",
-//              cl::desc("Register allocator to use for new pass manager"),
-//              cl::Hidden, cl::init("default"));
-
-// create option for RegAllocType enum
 static cl::opt<RegAllocType> RegAlloc(
     "regalloc-npm", cl::desc("Register allocator to use for new pass manager"),
     cl::Hidden, cl::init(RegAllocType::Unset),

>From 9ec74b7c6761ba98dcae413bcc79516e384b1e35 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Thu, 19 Dec 2024 14:49:58 +0000
Subject: [PATCH 22/22] clang format files

---
 llvm/include/llvm/Passes/CodeGenPassBuilder.h | 16 ++++++++--------
 llvm/lib/Passes/PassBuilder.cpp               |  3 ++-
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index d74939d9efb8f4..2dbcd7122f6aa8 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -1078,14 +1078,14 @@ void CodeGenPassBuilder<Derived, TargetMachineT>::addRegAllocPass(
   // Use the specified -regalloc-npm={basic|greedy|fast|pbqp}
   if (Opt.RegAlloc > RegAllocType::Default) {
     switch (Opt.RegAlloc) {
-      case RegAllocType::Fast:
-        addPass(RegAllocFastPass());
-        break;
-      case RegAllocType::Greedy:
-        addPass(RAGreedyPass());
-        break;
-      default:
-        report_fatal_error("Register allocator not supported yet.", false);
+    case RegAllocType::Fast:
+      addPass(RegAllocFastPass());
+      break;
+    case RegAllocType::Greedy:
+      addPass(RAGreedyPass());
+      break;
+    default:
+      report_fatal_error("Register allocator not supported yet.", false);
     }
     return;
   }
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index a733df8c4ba2e8..17a5a6379f061f 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -1312,7 +1312,8 @@ parseBoundsCheckingOptions(StringRef Params) {
   return Mode;
 }
 
-Expected<RAGreedyPass::Options> parseRegAllocGreedyFilterFunc(PassBuilder &PB, StringRef Params) {
+Expected<RAGreedyPass::Options>
+parseRegAllocGreedyFilterFunc(PassBuilder &PB, StringRef Params) {
   if (Params.empty() || Params == "all") {
     return RAGreedyPass::Options();
   }



More information about the llvm-branch-commits mailing list