[NFC][MLGO] ML Regalloc Priority Advisor

[Flpha0830]

Summary

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

[Flpha0830]

This comment was removed by Flpha0830.

[mtrofin]

Add a patch description: what is this patch trying to achieve?

llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp

1 ↗

(On Diff #450897)

This should come in a different patch.

llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp

30

new line

llvm/lib/CodeGen/RegAllocPriorityAdvisor.h

26

I'd prefer this were introduced in a subsequent change. This current change can stay focused at factoring out the prioritization logic behind the RegAllocPriorityAdvisor interface.

59

add a new line

[Flpha0830]

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

[Flpha0830]

Add LICENSE and remove unused code

[barannikov88]

llvm/lib/CodeGen/RegAllocPriorityAdvisor.h

12

Header guard is still missing

[mtrofin]

llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp

2

file header

llvm/lib/CodeGen/RegAllocPriorityAdvisor.h

2

add a header (see other .h files)

[MatzeB]

llvm/lib/CodeGen/RegAllocPriorityAdvisor.h

2

This comment seems off/wrong.

7–17

We really should have some comments here, either on the file header or the classes to have at least some sentences explaining what this is doing / used for.

[Flpha0830]

Add header guard and some comments

[mtrofin]

llvm/lib/CodeGen/RegAllocGreedy.cpp

357

rename to getPriority, and it should just take LIas parameter. Then, its implementation is all the stuff above.

llvm/lib/CodeGen/RegAllocPriorityAdvisor.h

35

remove logReward. this is something that only matters in the training scenario, so it doesn't belong on the abstraction.

38

do you need MF and RA here? IIUC, you just need the LiveInterval (passed in); maybe have the advisor take just RAGreedy, so it can get to ExtraInfo and LIS from it. At this stage, too, the ctor can be public.

45

remove the analysis stuff, it'll come in with a subsequent patch

83

remove logReward.

[Flpha0830]

Remove the analysis stuff and logReward.

[Flpha0830]

Remove MF

[mtrofin]

Nice! This is almost ready to go, left a few comments that should be easy to address.

llvm/lib/CodeGen/CMakeLists.txt

171

nit: place it in alphabetical order, I think this goes under PBQP

llvm/lib/CodeGen/RegAllocGreedy.cpp

298

make it const LiveInterval &LI: it's never expected to be nullptr, so pass-by-reference strongly suggests that to one reading the code.

also: remove the LiveRangeStage parameter. You can get it from RA.getExtraInfo().getStage(LI). This makes the interface very simple: you pass in a LI, you get a priority, the rest of the values that may factor in the calculation are up to the advisor implementation to get from wherever - other analyses, or the RA.

llvm/lib/CodeGen/RegAllocGreedy.h

152

this API almost feels like could be removed and we just pass RegClassPriorityTrumpsGlobalness via the ctor to the advisor, but we can easily do that once the full design is in, so lgtm

153

you can put a FIXME comment above to say this can be removed once we have an analysis, because the analysis can getAnalysis<SlotIndexes> and pass this. See elsewhere in the codebase how these comments look like (basically //FIXME: this is unnecessary once priority advisers are created by an analysis pass, which can fetch the SlotIndexes analysis itself.)

llvm/lib/CodeGen/RegAllocPriorityAdvisor.h

45

SlotIndexes * const Indexes;
const bool RegClassPriorityTrumpsGlobalness;

because they are ctor-time init-ed and then never change.

[mtrofin]

In D131220#3707921, @Flpha0830 wrote:

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

Please remember to:

• add this to the git commit (i.e. git commit --amend and make sure this is the description)
• also add this to the review message (from Edit Revision button on the right)

[Flpha0830]

Remove LiveRangeStage parameter and Pass LiveInterval by reference.

[MatzeB]

llvm/lib/CodeGen/RegAllocGreedy.cpp

357

Did you consider the precision loss that can happen when casting a 32bit unsigned number to 32bit float?

[MatzeB]

llvm/lib/CodeGen/RegAllocPriorityAdvisor.h

2

You marked this comment as "done", but this still reads "interference resolution" which I am pretty sure makes no sense.

[Flpha0830]

Leave the value as unsigned

[Flpha0830]

llvm/lib/CodeGen/RegAllocGreedy.cpp

357

Thanks for pointing this out. Because ML prefers small, bounded float values, we probably need float values. We can address this subsequently.

llvm/lib/CodeGen/RegAllocPriorityAdvisor.h

2

Thanks for the catch. This time it's fixed.

[Flpha0830]

Fix conflicts

llvm/lib/CodeGen/CMakeLists.txt

add_llvm_component_library(LLVMCodeGen
   PseudoSourceValue.cpp
   RDFGraph.cpp
   RDFLiveness.cpp
   RDFRegisters.cpp
   ReachingDefAnalysis.cpp
   RegAllocBase.cpp
   RegAllocBasic.cpp
   RegAllocEvictionAdvisor.cpp
   RegAllocFast.cpp
   RegAllocGreedy.cpp
   RegAllocPBQP.cpp
+  RegAllocPriorityAdvisor.cpp
   RegAllocScore.cpp
   RegisterClassInfo.cpp
   RegisterCoalescer.cpp
   RegisterPressure.cpp
   RegisterScavenging.cpp
   RemoveRedundantDebugValues.cpp
   RenameIndependentSubregs.cpp
   MachineStableHash.cpp

llvm/lib/CodeGen/RegAllocGreedy.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CODEGEN_REGALLOCGREEDY_H_
 #define LLVM_CODEGEN_REGALLOCGREEDY_H_
 
 #include "InterferenceCache.h"
 #include "RegAllocBase.h"
 #include "RegAllocEvictionAdvisor.h"
+#include "RegAllocPriorityAdvisor.h"
 #include "SpillPlacement.h"
 #include "SplitKit.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/IndexedMap.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
public:
   LiveRegMatrix *getInterferenceMatrix() const { return Matrix; }
   LiveIntervals *getLiveIntervals() const { return LIS; }
   VirtRegMap *getVirtRegMap() const { return VRM; }
   const RegisterClassInfo &getRegClassInfo() const { return RegClassInfo; }
   const ExtraRegInfo &getExtraInfo() const { return *ExtraInfo; }
   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>>;
   using SmallLISet = SmallSetVector<const LiveInterval *, 4>;
 
   // We need to track all tentative recolorings so we can roll back any
   // successful and unsuccessful recoloring attempts.
   using RecoloringStack =
private:
 
   // state
   std::unique_ptr<Spiller> SpillerInstance;
   PQueue Queue;
   std::unique_ptr<VirtRegAuxInfo> VRAI;
   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.
   enum CutOffStage {
     // No cutoffs encountered
     CO_None = 0,
 
     // lcr-max-depth cutoff encountered

llvm/lib/CodeGen/RegAllocGreedy.cpp

void RAGreedy::releaseMemory() {
   GlobalCand.clear();
 }
 
 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.
     Prio = Size;
   } else if (Stage == RS_Memory) {
     // Memory operand should be considered last.
     // Change the priority such that Memory operand are assigned in
     // the reverse order that they came in.
     // TODO: Make this a member variable and probably do something about hints.
     static unsigned MemOp = 0;
     Prio = MemOp++;
   } else {
     // Giant live ranges fall back to the global assignment heuristic, which
     // prevents excessive spilling in pathological cases.
     const TargetRegisterClass &RC = *MRI->getRegClass(Reg);
     bool ForceGlobal = !ReverseLocalAssignment &&
                        (Size / SlotIndex::InstrDist) >
                            (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
       // don't fit should be spilled (or split) ASAP so they don't create
       // interference.  Mark a bit to prioritize global above local ranges.
       Prio = Size;
       GlobalBit = 1;
     }
     if (RegClassPriorityTrumpsGlobalness)
       Prio |= RC.AllocationPriority << 25 | GlobalBit << 24;
     else
       Prio |= GlobalBit << 29 | RC.AllocationPriority << 24;
 
     // Mark a higher bit to prioritize global and local above RS_Split.
     Prio |= (1u << 31);
 
     // Boost ranges that have a physical register hint.
     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); }
 
 const LiveInterval *RAGreedy::dequeue(PQueue &CurQueue) {
   if (CurQueue.empty())
     return nullptr;
   LiveInterval *LI = &LIS->getInterval(~CurQueue.top().second);
bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
 
   ReverseLocalAssignment = GreedyReverseLocalAssignment.getNumOccurrences()
                                ? GreedyReverseLocalAssignment
                                : TRI->reverseLocalAssignment();
 
   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));
 
   VRAI->calculateSpillWeightsAndHints();
 
   LLVM_DEBUG(LIS->dump());
 

llvm/lib/CodeGen/RegAllocPriorityAdvisor.h

+//===- 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

llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp

+//===- 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()) {}