[CodeGen][NewPM] Port RegAllocGreedy to NPM

optimisan · optimisan · commit 36d5e74a5115 · 2024-12-11T11:26:16.000Z
diff --git a/llvm/include/llvm/CodeGen/InterferenceCache.h b/llvm/include/llvm/CodeGen/InterferenceCache.h
@@ -142,7 +142,7 @@ class LLVM_LIBRARY_VISIBILITY InterferenceCache {
 
   // Point to an entry for each physreg. The entry pointed to may not be up to
   // date, and it may have been reused for a different physreg.
-  unsigned char* PhysRegEntries = nullptr;
+  unsigned char *PhysRegEntries = nullptr;
   size_t PhysRegEntriesCount = 0;
 
   // Next round-robin entry to be picked.
@@ -158,9 +158,7 @@ class LLVM_LIBRARY_VISIBILITY InterferenceCache {
   InterferenceCache() = default;
   InterferenceCache &operator=(const InterferenceCache &other) = delete;
   InterferenceCache(const InterferenceCache &other) = delete;
-  ~InterferenceCache() {
-    free(PhysRegEntries);
-  }
+  ~InterferenceCache() { free(PhysRegEntries); }
 
   void reinitPhysRegEntries();
 
@@ -194,9 +192,7 @@ class LLVM_LIBRARY_VISIBILITY InterferenceCache {
     /// Cursor - Create a dangling cursor.
     Cursor() = default;
 
-    Cursor(const Cursor &O) {
-      setEntry(O.CacheEntry);
-    }
+    Cursor(const Cursor &O) { setEntry(O.CacheEntry); }
 
     Cursor &operator=(const Cursor &O) {
       setEntry(O.CacheEntry);
@@ -220,21 +216,15 @@ class LLVM_LIBRARY_VISIBILITY InterferenceCache {
     }
 
     /// hasInterference - Return true if the current block has any interference.
-    bool hasInterference() {
-      return Current->First.isValid();
-    }
+    bool hasInterference() { return Current->First.isValid(); }
 
     /// first - Return the starting index of the first interfering range in the
     /// current block.
-    SlotIndex first() {
-      return Current->First;
-    }
+    SlotIndex first() { return Current->First; }
 
     /// last - Return the ending index of the last interfering range in the
     /// current block.
-    SlotIndex last() {
-      return Current->Last;
-    }
+    SlotIndex last() { return Current->Last; }
   };
 };
 
diff --git a/llvm/include/llvm/CodeGen/MachineFunction.h b/llvm/include/llvm/CodeGen/MachineFunction.h
@@ -900,6 +900,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
@@ -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/CodeGen/RegAllocBase.h b/llvm/include/llvm/CodeGen/RegAllocBase.h
@@ -48,7 +48,7 @@ class LiveIntervals;
 class LiveRegMatrix;
 class MachineInstr;
 class MachineRegisterInfo;
-template<typename T> class SmallVectorImpl;
+template <typename T> class SmallVectorImpl;
 class Spiller;
 class TargetRegisterInfo;
 class VirtRegMap;
diff --git a/llvm/include/llvm/CodeGen/RegAllocGreedy.h b/llvm/include/llvm/CodeGen/RegAllocGreedy.h
@@ -12,26 +12,29 @@
 #ifndef LLVM_CODEGEN_REGALLOCGREEDY_H_
 #define LLVM_CODEGEN_REGALLOCGREEDY_H_
 
-#include "InterferenceCache.h"
-#include "RegAllocBase.h"
-#include "SplitKit.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/IndexedMap.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/CodeGen/CalcSpillWeights.h"
+#include "llvm/CodeGen/InterferenceCache.h"
 #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/RegAllocBase.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/SplitKit.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
+#include "llvm/IR/PassManager.h"
 #include <algorithm>
 #include <cstdint>
 #include <memory>
@@ -56,11 +59,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 +200,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 +308,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 +321,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 &,
@@ -451,5 +465,23 @@ class LLVM_LIBRARY_VISIBILITY RAGreedy : public MachineFunctionPass,
   /// Report the statistic for each loop.
   void reportStats();
 };
+
+class RAGreedyPass : public PassInfoMixin<RAGreedyPass> {
+  RegAllocFilterFunc Filter;
+
+public:
+  RAGreedyPass(RegAllocFilterFunc F = nullptr) : Filter(F) {}
+  PreservedAnalyses run(MachineFunction &F, MachineFunctionAnalysisManager &AM);
+
+  MachineFunctionProperties getRequiredProperties() const {
+    return MachineFunctionProperties().set(
+        MachineFunctionProperties::Property::NoPHIs);
+  }
+
+  MachineFunctionProperties getClearedProperties() const {
+    return MachineFunctionProperties().set(
+        MachineFunctionProperties::Property::IsSSA);
+  }
+};
 } // namespace llvm
 #endif // #ifndef LLVM_CODEGEN_REGALLOCGREEDY_H_
diff --git a/llvm/include/llvm/CodeGen/SplitKit.h b/llvm/include/llvm/CodeGen/SplitKit.h
@@ -88,7 +88,6 @@ class LLVM_LIBRARY_VISIBILITY InsertPointAnalysis {
     }
     return Res;
   }
-
 };
 
 /// SplitAnalysis - Analyze a LiveInterval, looking for live range splitting
@@ -387,7 +386,7 @@ class LLVM_LIBRARY_VISIBILITY SplitEditor {
 
   /// removeBackCopies - Remove the copy instructions that defines the values
   /// in the vector in the complement interval.
-  void removeBackCopies(SmallVectorImpl<VNInfo*> &Copies);
+  void removeBackCopies(SmallVectorImpl<VNInfo *> &Copies);
 
   /// getShallowDominator - Returns the least busy dominator of MBB that is
   /// also dominated by DefMBB.  Busy is measured by loop depth.
@@ -430,8 +429,9 @@ class LLVM_LIBRARY_VISIBILITY SplitEditor {
   /// \p InsertBefore. This can be invoked with a \p LaneMask which may make it
   /// necessary to construct a sequence of copies to cover it exactly.
   SlotIndex buildCopy(Register FromReg, Register ToReg, LaneBitmask LaneMask,
-      MachineBasicBlock &MBB, MachineBasicBlock::iterator InsertBefore,
-      bool Late, unsigned RegIdx);
+                      MachineBasicBlock &MBB,
+                      MachineBasicBlock::iterator InsertBefore, bool Late,
+                      unsigned RegIdx);
 
   SlotIndex buildSingleSubRegCopy(Register FromReg, Register ToReg,
                                   MachineBasicBlock &MB,
@@ -448,7 +448,7 @@ class LLVM_LIBRARY_VISIBILITY SplitEditor {
               VirtRegAuxInfo &VRAI);
 
   /// reset - Prepare for a new split.
-  void reset(LiveRangeEdit&, ComplementSpillMode = SM_Partition);
+  void reset(LiveRangeEdit &, ComplementSpillMode = SM_Partition);
 
   /// Create a new virtual register and live interval.
   /// Return the interval index, starting from 1. Interval index 0 is the
@@ -533,9 +533,9 @@ class LLVM_LIBRARY_VISIBILITY SplitEditor {
   /// @param LeaveBefore When set, leave IntvIn before this point.
   /// @param IntvOut     Interval index leaving the block.
   /// @param EnterAfter  When set, enter IntvOut after this point.
-  void splitLiveThroughBlock(unsigned MBBNum,
-                             unsigned IntvIn, SlotIndex LeaveBefore,
-                             unsigned IntvOut, SlotIndex EnterAfter);
+  void splitLiveThroughBlock(unsigned MBBNum, unsigned IntvIn,
+                             SlotIndex LeaveBefore, unsigned IntvOut,
+                             SlotIndex EnterAfter);
 
   /// splitRegInBlock - Split CurLI in the given block such that it enters the
   /// block in IntvIn and leaves it on the stack (or not at all). Split points
@@ -545,8 +545,8 @@ class LLVM_LIBRARY_VISIBILITY SplitEditor {
   /// @param BI          Block descriptor.
   /// @param IntvIn      Interval index entering the block. Not 0.
   /// @param LeaveBefore When set, leave IntvIn before this point.
-  void splitRegInBlock(const SplitAnalysis::BlockInfo &BI,
-                       unsigned IntvIn, SlotIndex LeaveBefore);
+  void splitRegInBlock(const SplitAnalysis::BlockInfo &BI, unsigned IntvIn,
+                       SlotIndex LeaveBefore);
 
   /// splitRegOutBlock - Split CurLI in the given block such that it enters the
   /// block on the stack (or isn't live-in at all) and leaves it in IntvOut.
@@ -557,8 +557,8 @@ class LLVM_LIBRARY_VISIBILITY SplitEditor {
   /// @param BI          Block descriptor.
   /// @param IntvOut     Interval index leaving the block.
   /// @param EnterAfter  When set, enter IntvOut after this point.
-  void splitRegOutBlock(const SplitAnalysis::BlockInfo &BI,
-                        unsigned IntvOut, SlotIndex EnterAfter);
+  void splitRegOutBlock(const SplitAnalysis::BlockInfo &BI, unsigned IntvOut,
+                        SlotIndex EnterAfter);
 };
 
 } // end namespace llvm
diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
@@ -249,7 +249,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
@@ -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
@@ -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/InlineSpiller.cpp b/llvm/lib/CodeGen/InlineSpiller.cpp
@@ -11,7 +11,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "SplitKit.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
@@ -36,6 +35,7 @@
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/SlotIndexes.h"
 #include "llvm/CodeGen/Spiller.h"
+#include "llvm/CodeGen/SplitKit.h"
 #include "llvm/CodeGen/StackMaps.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
diff --git a/llvm/lib/CodeGen/InterferenceCache.cpp b/llvm/lib/CodeGen/InterferenceCache.cpp
@@ -10,7 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "InterferenceCache.h"
+#include "llvm/CodeGen/InterferenceCache.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -11,11 +11,11 @@
 //===----------------------------------------------------------------------===//
 
 #include "AllocationOrder.h"
-#include "RegAllocGreedy.h"
 #include "llvm/Analysis/InteractiveModelRunner.h"
 #include "llvm/Analysis/MLModelRunner.h"
 #include "llvm/Analysis/TensorSpec.h"
 #include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
+#include "llvm/CodeGen/RegAllocGreedy.h"
 #if defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL) || defined(LLVM_HAVE_TFLITE)
 #include "llvm/Analysis/ModelUnderTrainingRunner.h"
 #include "llvm/Analysis/NoInferenceModelRunner.h"
diff --git a/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp
@@ -11,7 +11,6 @@
 //===----------------------------------------------------------------------===//
 
 #include "AllocationOrder.h"
-#include "RegAllocGreedy.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/InteractiveModelRunner.h"
 #include "llvm/Analysis/MLModelRunner.h"
@@ -24,6 +23,7 @@
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/Passes.h"
+#include "llvm/CodeGen/RegAllocGreedy.h"
 #include "llvm/CodeGen/RegAllocPriorityAdvisor.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
 #include "llvm/CodeGen/SlotIndexes.h"
diff --git a/llvm/lib/CodeGen/RegAllocBase.cpp b/llvm/lib/CodeGen/RegAllocBase.cpp
@@ -11,7 +11,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "RegAllocBase.h"
+#include "llvm/CodeGen/RegAllocBase.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/CodeGen/LiveInterval.h"
diff --git a/llvm/lib/CodeGen/RegAllocBasic.cpp b/llvm/lib/CodeGen/RegAllocBasic.cpp
diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp
diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp
diff --git a/llvm/lib/CodeGen/SplitKit.cpp b/llvm/lib/CodeGen/SplitKit.cpp
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
