[NFC][MLGO] ML Regalloc Priority Advisor

This patch introduces the priority analysis and the priority advisor,
the default implementation, and the scaffolding for introducing the
other implementations of the advisor.

Reviewed By: mtrofin

Differential Revision: https://reviews.llvm.org/D131220

Author: Flpha0830

llvm/lib/CodeGen/CMakeLists.txt

@@ -171,6 +171,7 @@ add_llvm_component_library(LLVMCodeGen
   RegAllocFast.cpp
   RegAllocGreedy.cpp
   RegAllocPBQP.cpp
+  RegAllocPriorityAdvisor.cpp
   RegAllocScore.cpp
   RegisterClassInfo.cpp
   RegisterCoalescer.cpp

llvm/lib/CodeGen/RegAllocGreedy.cpp

@@ -279,16 +279,28 @@ void RAGreedy::enqueueImpl(const LiveInterval *LI) { enqueue(Queue, LI); }
 void RAGreedy::enqueue(PQueue &CurQueue, const LiveInterval *LI) {
   // Prioritize live ranges by size, assigning larger ranges first.
   // The queue holds (size, reg) pairs.
-  const unsigned Size = LI->getSize();
   const Register Reg = LI->reg();
   assert(Reg.isVirtual() && "Can only enqueue virtual registers");
-  unsigned Prio;
 
   auto Stage = ExtraInfo->getOrInitStage(Reg);
   if (Stage == RS_New) {
     Stage = RS_Assign;
     ExtraInfo->setStage(Reg, Stage);
   }
+
+  unsigned Ret = PriorityAdvisor->getPriority(*LI);
+
+  // The virtual register number is a tie breaker for same-sized ranges.
+  // Give lower vreg numbers higher priority to assign them first.
+  CurQueue.push(std::make_pair(Ret, ~Reg));
+}
+
+unsigned DefaultPriorityAdvisor::getPriority(const LiveInterval &LI) const {
+  const unsigned Size = LI.getSize();
+  const Register Reg = LI.reg();
+  unsigned Prio;
+  LiveRangeStage Stage = RA.getExtraInfo().getStage(LI);
+
   if (Stage == RS_Split) {
     // Unsplit ranges that couldn't be allocated immediately are deferred until
     // everything else has been allocated.
@@ -309,18 +321,18 @@ void RAGreedy::enqueue(PQueue &CurQueue, const LiveInterval *LI) {
                            (2 * RegClassInfo.getNumAllocatableRegs(&RC));
     unsigned GlobalBit = 0;
 
-    if (Stage == RS_Assign && !ForceGlobal && !LI->empty() &&
-        LIS->intervalIsInOneMBB(*LI)) {
+    if (Stage == RS_Assign && !ForceGlobal && !LI.empty() &&
+        LIS->intervalIsInOneMBB(LI)) {
       // Allocate original local ranges in linear instruction order. Since they
       // are singly defined, this produces optimal coloring in the absence of
       // global interference and other constraints.
       if (!ReverseLocalAssignment)
-        Prio = LI->beginIndex().getInstrDistance(Indexes->getLastIndex());
+        Prio = LI.beginIndex().getInstrDistance(Indexes->getLastIndex());
       else {
         // Allocating bottom up may allow many short LRGs to be assigned first
         // to one of the cheap registers. This could be much faster for very
         // large blocks on targets with many physical registers.
-        Prio = Indexes->getZeroIndex().getInstrDistance(LI->endIndex());
+        Prio = Indexes->getZeroIndex().getInstrDistance(LI.endIndex());
       }
     } else {
       // Allocate global and split ranges in long->short order. Long ranges that
@@ -341,9 +353,8 @@ void RAGreedy::enqueue(PQueue &CurQueue, const LiveInterval *LI) {
     if (VRM->hasKnownPreference(Reg))
       Prio |= (1u << 30);
   }
-  // The virtual register number is a tie breaker for same-sized ranges.
-  // Give lower vreg numbers higher priority to assign them first.
-  CurQueue.push(std::make_pair(Prio, ~Reg));
+
+  return Prio;
 }
 
 const LiveInterval *RAGreedy::dequeue() { return dequeue(Queue); }
@@ -2554,6 +2565,7 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
   ExtraInfo.emplace();
   EvictAdvisor =
       getAnalysis<RegAllocEvictionAdvisorAnalysis>().getAdvisor(*MF, *this);
+  PriorityAdvisor = std::make_unique<DefaultPriorityAdvisor>(*MF, *this);
 
   VRAI = std::make_unique<VirtRegAuxInfo>(*MF, *LIS, *VRM, *Loops, *MBFI);
   SpillerInstance.reset(createInlineSpiller(*this, *MF, *VRM, *VRAI));

llvm/lib/CodeGen/RegAllocGreedy.h

@@ -15,6 +15,7 @@
 #include "InterferenceCache.h"
 #include "RegAllocBase.h"
 #include "RegAllocEvictionAdvisor.h"
+#include "RegAllocPriorityAdvisor.h"
 #include "SpillPlacement.h"
 #include "SplitKit.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -147,6 +148,16 @@ class LLVM_LIBRARY_VISIBILITY RAGreedy : public MachineFunctionPass,
   size_t getQueueSize() const { return Queue.size(); }
   // end (interface to eviction advisers)
 
+  // Interface to priority advisers
+  bool getRegClassPriorityTrumpsGlobalness() const {
+    return RegClassPriorityTrumpsGlobalness;
+  }
+  bool getReverseLocalAssignment() const { return ReverseLocalAssignment; }
+  // FIXME: this is unnecessary once priority advisers are created by an
+  // analysis pass, which can fetch the SlotIndexes analysis itself.
+  SlotIndexes *getIndexes() const { return Indexes; }
+  // end (interface to priority advisers)
+
 private:
   // Convenient shortcuts.
   using PQueue = std::priority_queue<std::pair<unsigned, unsigned>>;
@@ -180,6 +191,8 @@ class LLVM_LIBRARY_VISIBILITY RAGreedy : public MachineFunctionPass,
   Optional<ExtraRegInfo> ExtraInfo;
   std::unique_ptr<RegAllocEvictionAdvisor> EvictAdvisor;
 
+  std::unique_ptr<RegAllocPriorityAdvisor> PriorityAdvisor;
+
   // Enum CutOffStage to keep a track whether the register allocation failed
   // because of the cutoffs encountered in last chance recoloring.
   // Note: This is used as bitmask. New value should be next power of 2.

llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp

@@ -0,0 +1,29 @@
+//===- RegAllocPriorityAdvisor.cpp - live ranges priority advisor ---------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Implementation of the default priority advisor and of the Analysis pass.
+//
+//===----------------------------------------------------------------------===//
+
+#include "RegAllocPriorityAdvisor.h"
+#include "RegAllocGreedy.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/VirtRegMap.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+
+using namespace llvm;
+
+RegAllocPriorityAdvisor::RegAllocPriorityAdvisor(const MachineFunction &MF,
+                                                 const RAGreedy &RA)
+    : RA(RA), LIS(RA.getLiveIntervals()), VRM(RA.getVirtRegMap()),
+      MRI(&VRM->getRegInfo()), TRI(MF.getSubtarget().getRegisterInfo()),
+      RegClassInfo(RA.getRegClassInfo()), Indexes(RA.getIndexes()),
+      RegClassPriorityTrumpsGlobalness(
+          RA.getRegClassPriorityTrumpsGlobalness()),
+      ReverseLocalAssignment(RA.getReverseLocalAssignment()) {}

llvm/lib/CodeGen/RegAllocPriorityAdvisor.h

@@ -0,0 +1,58 @@
+//===- RegAllocPriorityAdvisor.h - live ranges priority advisor -*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_REGALLOCPRIORITYADVISOR_H
+#define LLVM_CODEGEN_REGALLOCPRIORITYADVISOR_H
+
+#include "RegAllocEvictionAdvisor.h"
+#include "llvm/CodeGen/SlotIndexes.h"
+#include "llvm/Pass.h"
+
+namespace llvm {
+
+class MachineFunction;
+class VirtRegMap;
+class RAGreedy;
+
+/// Interface to the priority advisor, which is responsible for prioritizing
+/// live ranges.
+class RegAllocPriorityAdvisor {
+public:
+  RegAllocPriorityAdvisor(const RegAllocPriorityAdvisor &) = delete;
+  RegAllocPriorityAdvisor(RegAllocPriorityAdvisor &&) = delete;
+  virtual ~RegAllocPriorityAdvisor() = default;
+
+  /// Find the priority value for a live range. A float value is used since ML
+  /// prefers it.
+  virtual unsigned getPriority(const LiveInterval &LI) const = 0;
+
+  RegAllocPriorityAdvisor(const MachineFunction &MF, const RAGreedy &RA);
+
+protected:
+  const RAGreedy &RA;
+  LiveIntervals *const LIS;
+  VirtRegMap *const VRM;
+  MachineRegisterInfo *const MRI;
+  const TargetRegisterInfo *const TRI;
+  const RegisterClassInfo &RegClassInfo;
+  SlotIndexes *const Indexes;
+  const bool RegClassPriorityTrumpsGlobalness;
+  const bool ReverseLocalAssignment;
+};
+
+class DefaultPriorityAdvisor : public RegAllocPriorityAdvisor {
+public:
+  DefaultPriorityAdvisor(const MachineFunction &MF, const RAGreedy &RA)
+      : RegAllocPriorityAdvisor(MF, RA) {}
+
+private:
+  unsigned getPriority(const LiveInterval &LI) const override;
+};
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_REGALLOCPRIORITYADVISOR_H