Add LLVM Segment Tree Register Allocator implementation

- Implement segment tree-based register allocator for LLVM
- Features coordinate compression and lazy propagation
- Add interference checking using interval tree queries
- Include performance benchmarking against RAGreedy
- Currently 6.7x slower than baseline (experimental research)

Author: steven-studio

llvm/include/llvm/CodeGen/RegAllocSegmentTree.h

@@ -31,6 +31,7 @@
 #include <memory>
 #include <utility>
 #include <algorithm>
+#include <queue>
 
 namespace llvm {
 
@@ -40,99 +41,73 @@ class MachineRegisterInfo;
 class TargetRegisterInfo;
 class VirtRegMap;
 
-/// SegmentTreeRegisterAllocator implements a register allocator by leveraging
-/// a segment tree to efficiently track the availability of physical registers
-/// across their live ranges.
-class RegAllocSegmentTree : public MachineFunctionPass, public RegAllocBase, private LiveRangeEdit::Delegate {
-
-  // 非常重要的：用於記錄每個物理暫存器當前的分配區間
-  // Key: 物理暫存器編號 (PhysReg)
-  // Value: 一個包含多個(LiveInterval, VirtReg)的列表，表示該物理暫存器上分配了哪些虛擬暫存器及其生命期
-  // 但注意：一個物理暫存器在同一時間只能被一個虛擬暫存器佔用，所以這實際上是一系列不重疊的區間
-  // 我們通常用一個按開始時間排序的區間列表來表示一個物理暫存器的使用情況
-  // 我們將為每個物理暫存器建立一個線段樹
-
+class RegAllocSegmentTree : public MachineFunctionPass, 
+                            public RegAllocBase, 
+                            private LiveRangeEdit::Delegate {
 public:
-  // 线段树节点结构体
-  struct SegmentTreeNode {
-    SlotIndex MaxEnd;  // 该节点覆盖区间内的最大结束时间
-    
-    SegmentTreeNode() : MaxEnd(SlotIndex()) {}
-  };
-
   static char ID; // 静态成员声明
 
   RegAllocSegmentTree();
 
-  /// 主要的分配接口，繼承自 MachineFunctionPass
+  // MachineFunctionPass interface
   bool runOnMachineFunction(MachineFunction &mf) override;
-
-  /// 獲取此分配器的名稱
   StringRef getPassName() const override { return "Segment Tree Register Allocator"; }
 
   /// 獲取分析依賴
-  void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.setPreservesCFG();
-    AU.addRequired<SlotIndexesWrapperPass>();
-    AU.addPreserved<SlotIndexesWrapperPass>();
-    AU.addRequired<LiveIntervalsWrapperPass>();
-    AU.addPreserved<LiveIntervalsWrapperPass>();
-    AU.addRequired<LiveStacksWrapperLegacy>();
-    AU.addPreserved<LiveStacksWrapperLegacy>();
-    AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
-    AU.addPreserved<MachineBlockFrequencyInfoWrapperPass>();
-    AU.addRequired<MachineDominatorTreeWrapperPass>();
-    AU.addPreserved<MachineDominatorTreeWrapperPass>();
-    AU.addRequired<MachineLoopInfoWrapperPass>();
-    AU.addPreserved<MachineLoopInfoWrapperPass>();
-    AU.addRequired<VirtRegMapWrapperLegacy>();
-    AU.addPreserved<VirtRegMapWrapperLegacy>();
-    AU.addRequired<LiveRegMatrixWrapperLegacy>();
-    AU.addPreserved<LiveRegMatrixWrapperLegacy>();
-    MachineFunctionPass::getAnalysisUsage(AU);
-  }
 
   // 我們需要覆蓋這個方法來初始化我們的自定義數據結構
   void init(VirtRegMap &vrm, LiveIntervals &lis, LiveRegMatrix &mat);
 
-  // 我們需要覆蓋這個方法來實現主要的分配邏輯
-  void allocatePhysRegs();
-
   // 你自己的 finalize（宣告）
   void finalizeAlloc(MachineFunction &MF, LiveIntervals &LIS, VirtRegMap &VRM) const;
 
+#ifndef NDEBUG
+  bool verifyAllocation(MachineFunction &MF, LiveIntervals &LIS, VirtRegMap &VRM);
+#endif
+
   // 你堅持要的同名接口（注意：**不要**寫 override）
   void postOptimization(Spiller &VRegSpiller, LiveIntervals &LIS);
-protected:
-  // 必须实现的纯虚函数
-  Spiller &spiller() override;
-  void enqueueImpl(const LiveInterval *LI) override;
-  const LiveInterval *dequeue() override;
-  MCRegister selectOrSplit(const LiveInterval &VirtReg,
-                           SmallVectorImpl<Register> &splitLVRs) override;
 
 private:
+  // Timer group constants
+  static const char TimerGroupName[];
+  static const char TimerGroupDescription[];
+  
+  // Performance statistics
+  mutable unsigned NumInterferenceChecks = 0;
+  mutable unsigned NumSegTreeUpdates = 0;
+  mutable unsigned NumCoordRebuilds = 0;
+  mutable unsigned NumAllocAttempts = 0;
+  mutable unsigned NumSpills = 0;
+
+  // Convenient shortcuts.
+  using PQueue = std::priority_queue<std::pair<unsigned, unsigned>>;
+  using SmallLISet = SmallSetVector<const LiveInterval *, 4>;
+
   void resetAllocatorState();  // ← 新增
+  void precomputeAllCoordinates();
   void performSegmentTreeSpecificOptimizations(LiveIntervals &LIS);
   void validatePostOptimizationState(LiveIntervals &LIS);
 
   MachineFunction *MF = nullptr;
+
+  // Shortcuts to some useful interface.
+  const TargetInstrInfo *TII = nullptr;
+
   LiveIntervals       *LIS = nullptr;
   VirtRegMap          *VRM = nullptr;
   LiveRegMatrix       *LRM = nullptr;
   MachineRegisterInfo *MRI = nullptr;
+  const TargetRegisterInfo *TRI = nullptr;
 
   RegisterClassInfo RCI;
 
-  std::unique_ptr<Spiller> VRegSpiller;  // Add this line
+  std::unique_ptr<Spiller> SpillerInstance;  // Add this line
 
   // 工作佇列與去重集合
-  llvm::SmallVector<const LiveInterval*, 64> WorkQ;
+  llvm::SmallVector<const LiveInterval*, 64> CurQueue;
   llvm::SmallDenseSet<unsigned, 64> InQ;  // 記 vreg id，避免重複排入
 
-  // Legacy demo per-physreg trees (kept for compatibility paths).
-  SmallVector<std::vector<SegmentTreeNode>> PhysRegSegmentTrees;
-
   // Bookkeeping of allocated intervals per physical register (for rebuild).
   std::vector<std::vector<std::pair<SlotIndex, SlotIndex>>> PhysRegIntervals;
 
@@ -166,22 +141,6 @@ class RegAllocSegmentTree : public MachineFunctionPass, public RegAllocBase, pri
   // 可放檢查（含別名）：任一段 maxCover>0 就不可放
   bool canPlaceOnPhysReg(unsigned PhysReg, const LiveInterval &LI) const;
 
-  // 线段树构建辅助函数（递归）
-  void buildSegmentTree(SegmentTreeNode *tree, unsigned idx,
-                        unsigned l, unsigned r,
-                        const std::vector<SlotIndex> &ends);
-
-  // 为物理寄存器构建线段树
-  void buildSegmentTreeForPhysReg(unsigned PhysReg, const std::vector<SlotIndex>& intervals);
-  
-  // 查询物理寄存器在指定区间是否可用
-  bool querySegmentTreeForRange(unsigned PhysReg, SlotIndex Start, SlotIndex End) const;
-
-  // 线段树查询辅助函数
-  bool querySegmentTree(const SegmentTreeNode *tree, unsigned idx,
-                        unsigned tree_l, unsigned tree_r,
-                        SlotIndex query_start, SlotIndex query_end) const;
-
   // 嘗試為給定的虛擬暫存器分配一個物理暫存器
   // 返回分配的物理暫存器，若失敗則返回0
   unsigned tryAllocateRegister(LiveInterval &VirtReg);
@@ -197,11 +156,27 @@ class RegAllocSegmentTree : public MachineFunctionPass, public RegAllocBase, pri
 
   // 更新物理寄存器的线段树（分配区间后）
   void updateSegmentTreeForPhysReg(unsigned PhysReg, SlotIndex Start, SlotIndex End);
+
+  // 清理失败的虚拟寄存器
+  void cleanupFailedVReg(Register FailedVReg, unsigned Depth,
+                         SmallVectorImpl<Register> &SplitRegs);
+
+  // 找到合適的拆分點
+  SlotIndex findSplitPoint(const LiveInterval &LI);
+public:
+  // RegAllocBase interface  
+  Spiller &spiller() override { return *SpillerInstance; };
+  void enqueueImpl(const LiveInterval *LI) override;
+  const LiveInterval *dequeue() override;
+  MCRegister selectOrSplit(const LiveInterval &VirtReg,
+                           SmallVectorImpl<Register> &splitLVRs) override;
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
 };
 
 // 創建SegmentTreeRegisterAllocator實例的函數
 FunctionPass *createRegAllocSegmentTree();
 
 } // end namespace llvm
 
-#endif // LLVM_CODEGEN_REGALLOCSEGMENTTREE_H
+#endif // LLVM_CODEGEN_REGALLOCSEGMENTTREE_H

llvm/include/llvm/InitializePasses.h

@@ -271,6 +271,7 @@ initializeRegAllocEvictionAdvisorAnalysisLegacyPass(PassRegistry &);
 LLVM_ABI void initializeRegAllocFastPass(PassRegistry &);
 LLVM_ABI void
 initializeRegAllocPriorityAdvisorAnalysisLegacyPass(PassRegistry &);
+LLVM_ABI void initializeRegAllocSegmentTreePass(PassRegistry&);
 LLVM_ABI void initializeRegAllocScoringPass(PassRegistry &);
 LLVM_ABI void initializeRegBankSelectPass(PassRegistry &);
 LLVM_ABI void initializeRegToMemWrapperPassPass(PassRegistry &);

llvm/lib/CodeGen/CodeGen.cpp

@@ -113,6 +113,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeRABasicPass(Registry);
   initializeRAGreedyLegacyPass(Registry);
   initializeRegAllocFastPass(Registry);
+  initializeRegAllocSegmentTreePass(Registry);  // Add this line
   initializeRegUsageInfoCollectorLegacyPass(Registry);
   initializeRegUsageInfoPropagationLegacyPass(Registry);
   initializeRegisterCoalescerLegacyPass(Registry);