[NFC][regalloc] Pass RAGreedy to eviction adviser

[mtrofin]

Summary

This patch simplifies the interface between RAGreedy and the eviction
adviser by passing the allocator to the adviser, which allows the latter
to extract needed information as needed, rather than requiring it be passed
piecemeal at construction time (which would also complicate later
evolution).

Part of this, the patch also moves ExtraRegInfo back to RAGreedy. We
keep the encapsulation of ExtraRegInfo because it has benefits (e.g.
improved readability by abstracting access to the cascade info) and also
simpler re-initialization at regalloc pass re-entry time (we just flush
the Optional).

[mtrofin]

Gentle reminder - thanks!

llvm/lib/CodeGen/RegAllocEvictionAdvisor.h

struct EvictionCost {
   void setBrokenHints(unsigned NHints) { BrokenHints = NHints; }
 
   bool operator<(const EvictionCost &O) const {
     return std::tie(BrokenHints, MaxWeight) <
            std::tie(O.BrokenHints, O.MaxWeight);
   }
 };
 
-/// Track allocation stage and eviction loop prevention during allocation.
-// TODO(mtrofin): Consider exposing RAGreedy in a header instead, and folding
-// this back into it.
-class ExtraRegInfo final {
-  // RegInfo - Keep additional information about each live range.
-  struct RegInfo {
-    LiveRangeStage Stage = RS_New;
-
-    // Cascade - Eviction loop prevention. See
-    // canEvictInterferenceBasedOnCost().
-    unsigned Cascade = 0;
-
-    RegInfo() = default;
-  };
-
-  IndexedMap<RegInfo, VirtReg2IndexFunctor> Info;
-  unsigned NextCascade = 1;
-
-public:
-  ExtraRegInfo() = default;
-  ExtraRegInfo(const ExtraRegInfo &) = delete;
-
-  LiveRangeStage getStage(Register Reg) const { return Info[Reg].Stage; }
-
-  LiveRangeStage getStage(const LiveInterval &VirtReg) const {
-    return getStage(VirtReg.reg());
-  }
-
-  void setStage(Register Reg, LiveRangeStage Stage) {
-    Info.grow(Reg.id());
-    Info[Reg].Stage = Stage;
-  }
-
-  void setStage(const LiveInterval &VirtReg, LiveRangeStage Stage) {
-    setStage(VirtReg.reg(), Stage);
-  }
-
-  /// Return the current stage of the register, if present, otherwise initialize
-  /// it and return that.
-  LiveRangeStage getOrInitStage(Register Reg) {
-    Info.grow(Reg.id());
-    return getStage(Reg);
-  }
-
-  unsigned getCascade(Register Reg) const { return Info[Reg].Cascade; }
-
-  void setCascade(Register Reg, unsigned Cascade) {
-    Info.grow(Reg.id());
-    Info[Reg].Cascade = Cascade;
-  }
-
-  unsigned getOrAssignNewCascade(Register Reg) {
-    unsigned Cascade = getCascade(Reg);
-    if (!Cascade) {
-      Cascade = NextCascade++;
-      setCascade(Reg, Cascade);
-    }
-    return Cascade;
-  }
-
-  unsigned getCascadeOrCurrentNext(Register Reg) const {
-    unsigned Cascade = getCascade(Reg);
-    if (!Cascade)
-      Cascade = NextCascade;
-    return Cascade;
-  }
-
-  template <typename Iterator>
-  void setStage(Iterator Begin, Iterator End, LiveRangeStage NewStage) {
-    for (; Begin != End; ++Begin) {
-      Register Reg = *Begin;
-      Info.grow(Reg.id());
-      if (Info[Reg].Stage == RS_New)
-        Info[Reg].Stage = NewStage;
-    }
-  }
-  void LRE_DidCloneVirtReg(Register New, Register Old);
-};
-
 /// Interface to the eviction advisor, which is responsible for making a
 /// decision as to which live ranges should be evicted (if any).
+class RAGreedy;
 class RegAllocEvictionAdvisor {
 public:
   RegAllocEvictionAdvisor(const RegAllocEvictionAdvisor &) = delete;
   RegAllocEvictionAdvisor(RegAllocEvictionAdvisor &&) = delete;
   virtual ~RegAllocEvictionAdvisor() = default;
 
   /// Find a physical register that can be freed by evicting the FixedRegisters,
   /// or return NoRegister. The eviction decision is assumed to be correct (i.e.
public:
   canEvictHintInterference(LiveInterval &VirtReg, MCRegister PhysReg,
                            const SmallVirtRegSet &FixedRegisters) const = 0;
 
   /// Returns true if the given \p PhysReg is a callee saved register and has
   /// not been used for allocation yet.
   bool isUnusedCalleeSavedReg(MCRegister PhysReg) const;
 
 protected:
-  RegAllocEvictionAdvisor(const MachineFunction &MF, LiveRegMatrix *Matrix,
-                          LiveIntervals *LIS, VirtRegMap *VRM,
-                          const RegisterClassInfo &RegClassInfo,
-                          ExtraRegInfo *ExtraInfo);
+  RegAllocEvictionAdvisor(const MachineFunction &MF, const RAGreedy &RA);
 
   Register canReassign(LiveInterval &VirtReg, Register PrevReg) const;
 
   const MachineFunction &MF;
+  const RAGreedy &RA;
   LiveRegMatrix *const Matrix;
   LiveIntervals *const LIS;
   VirtRegMap *const VRM;
   MachineRegisterInfo *const MRI;
   const TargetRegisterInfo *const TRI;
   const RegisterClassInfo &RegClassInfo;
   const ArrayRef<uint8_t> RegCosts;
-  ExtraRegInfo *const ExtraInfo;
 
   /// Run or not the local reassignment heuristic. This information is
   /// obtained from the TargetSubtargetInfo.
   const bool EnableLocalReassign;
 
 private:
   unsigned NextCascade = 1;
 };
public:
   enum class AdvisorMode : int { Default, Release, Development };
 
   RegAllocEvictionAdvisorAnalysis(AdvisorMode Mode)
       : ImmutablePass(ID), Mode(Mode){};
   static char ID;
 
   /// Get an advisor for the given context (i.e. machine function, etc)
   virtual std::unique_ptr<RegAllocEvictionAdvisor>
-  getAdvisor(const MachineFunction &MF, LiveRegMatrix *Matrix,
-             LiveIntervals *LIS, VirtRegMap *VRM,
-             const RegisterClassInfo &RegClassInfo,
-             ExtraRegInfo *ExtraInfo) = 0;
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
   AdvisorMode getAdvisorMode() const { return Mode; }
 
 private:
   // This analysis preserves everything, and subclasses may have additional
   // requirements.
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
   }

 #ifdef LLVM_HAVE_TF_API
 RegAllocEvictionAdvisorAnalysis *createDevelopmentModeAdvisor();
 #endif
 
 // TODO: move to RegAllocEvictionAdvisor.cpp when we move implementation
 // out of RegAllocGreedy.cpp
 class DefaultEvictionAdvisor : public RegAllocEvictionAdvisor {
 public:
-  DefaultEvictionAdvisor(const MachineFunction &MF, LiveRegMatrix *Matrix,
-                         LiveIntervals *LIS, VirtRegMap *VRM,
-                         const RegisterClassInfo &RegClassInfo,
-                         ExtraRegInfo *ExtraInfo)
-      : RegAllocEvictionAdvisor(MF, Matrix, LIS, VRM, RegClassInfo, ExtraInfo) {
-  }
+  DefaultEvictionAdvisor(const MachineFunction &MF, const RAGreedy &RA)
+      : RegAllocEvictionAdvisor(MF, RA) {}
 
 private:
   MCRegister tryFindEvictionCandidate(LiveInterval &, const AllocationOrder &,
                                       uint8_t,
                                       const SmallVirtRegSet &) const override;
   bool canEvictHintInterference(LiveInterval &, MCRegister,
                                 const SmallVirtRegSet &) const override;
   bool canEvictInterferenceBasedOnCost(LiveInterval &, MCRegister, bool,
                                        EvictionCost &,
                                        const SmallVirtRegSet &) const;
   bool shouldEvict(LiveInterval &A, bool, LiveInterval &B, bool) const;
 };
 } // namespace llvm
 
 #endif // LLVM_CODEGEN_REGALLOCEVICTIONADVISOR_H

llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp

 //===- RegAllocEvictionAdvisor.cpp - eviction 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 eviction advisor and of the Analysis pass.
 //
 //===----------------------------------------------------------------------===//
 
 #include "RegAllocEvictionAdvisor.h"
+#include "RegAllocGreedy.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
 #include "llvm/CodeGen/VirtRegMap.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/PassRegistry.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ErrorHandling.h"
public:
 
   // support for isa<> and dyn_cast.
   static bool classof(const RegAllocEvictionAdvisorAnalysis *R) {
     return R->getAdvisorMode() == AdvisorMode::Default;
   }
 
 private:
   std::unique_ptr<RegAllocEvictionAdvisor>
-  getAdvisor(const MachineFunction &MF, LiveRegMatrix *Matrix,
-             LiveIntervals *LIS, VirtRegMap *VRM,
-             const RegisterClassInfo &RegClassInfo,
-             ExtraRegInfo *ExtraInfo) override {
-    return std::make_unique<DefaultEvictionAdvisor>(MF, Matrix, LIS, VRM,
-                                                    RegClassInfo, ExtraInfo);
+  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 be created. Using default");
     return RegAllocEvictionAdvisorAnalysis::doInitialization(M);
   }
   const bool NotAsRequested;
StringRef RegAllocEvictionAdvisorAnalysis::getPassName() const {
   case AdvisorMode::Release:
     return "Release mode Regalloc Eviction Advisor";
   case AdvisorMode::Development:
     return "Development mode Regalloc Eviction Advisor";
   }
   llvm_unreachable("Unknown advisor kind");
 }
 
-RegAllocEvictionAdvisor::RegAllocEvictionAdvisor(
-    const MachineFunction &MF, LiveRegMatrix *Matrix, LiveIntervals *LIS,
-    VirtRegMap *VRM, const RegisterClassInfo &RegClassInfo,
-    ExtraRegInfo *ExtraInfo)
-    : MF(MF), Matrix(Matrix), LIS(LIS), VRM(VRM), MRI(&VRM->getRegInfo()),
-      TRI(MF.getSubtarget().getRegisterInfo()), RegClassInfo(RegClassInfo),
-      RegCosts(TRI->getRegisterCosts(MF)), ExtraInfo(ExtraInfo),
+RegAllocEvictionAdvisor::RegAllocEvictionAdvisor(const MachineFunction &MF,
+                                                 const RAGreedy &RA)
+    : MF(MF), RA(RA), Matrix(RA.getInterferenceMatrix()),
+      LIS(RA.getLiveIntervals()), VRM(RA.getVirtRegMap()),
+      MRI(&VRM->getRegInfo()), TRI(MF.getSubtarget().getRegisterInfo()),
+      RegClassInfo(RA.getRegClassInfo()), RegCosts(TRI->getRegisterCosts(MF)),
       EnableLocalReassign(EnableLocalReassignment ||
                           MF.getSubtarget().enableRALocalReassignment(
                               MF.getTarget().getOptLevel())) {}

llvm/lib/CodeGen/RegAllocGreedy.h

 #include <queue>
 #include <tuple>
 #include <utility>
 
 namespace llvm {
 class LLVM_LIBRARY_VISIBILITY RAGreedy : public MachineFunctionPass,
                                          public RegAllocBase,
                                          private LiveRangeEdit::Delegate {
+  // Interface to eviction advisers
+public:
+  /// Track allocation stage and eviction loop prevention during allocation.
+  class ExtraRegInfo final {
+    // RegInfo - Keep additional information about each live range.
+    struct RegInfo {
+      LiveRangeStage Stage = RS_New;
+
+      // Cascade - Eviction loop prevention. See
+      // canEvictInterferenceBasedOnCost().
+      unsigned Cascade = 0;
+
+      RegInfo() = default;
+    };
+
+    IndexedMap<RegInfo, VirtReg2IndexFunctor> Info;
+    unsigned NextCascade = 1;
+
+  public:
+    ExtraRegInfo() = default;
+    ExtraRegInfo(const ExtraRegInfo &) = delete;
+
+    LiveRangeStage getStage(Register Reg) const { return Info[Reg].Stage; }
+
+    LiveRangeStage getStage(const LiveInterval &VirtReg) const {
+      return getStage(VirtReg.reg());
+    }
+
+    void setStage(Register Reg, LiveRangeStage Stage) {
+      Info.grow(Reg.id());
+      Info[Reg].Stage = Stage;
+    }
+
+    void setStage(const LiveInterval &VirtReg, LiveRangeStage Stage) {
+      setStage(VirtReg.reg(), Stage);
+    }
+
+    /// Return the current stage of the register, if present, otherwise
+    /// initialize it and return that.
+    LiveRangeStage getOrInitStage(Register Reg) {
+      Info.grow(Reg.id());
+      return getStage(Reg);
+    }
+
+    unsigned getCascade(Register Reg) const { return Info[Reg].Cascade; }
+
+    void setCascade(Register Reg, unsigned Cascade) {
+      Info.grow(Reg.id());
+      Info[Reg].Cascade = Cascade;
+    }
+
+    unsigned getOrAssignNewCascade(Register Reg) {
+      unsigned Cascade = getCascade(Reg);
+      if (!Cascade) {
+        Cascade = NextCascade++;
+        setCascade(Reg, Cascade);
+      }
+      return Cascade;
+    }
+
+    unsigned getCascadeOrCurrentNext(Register Reg) const {
+      unsigned Cascade = getCascade(Reg);
+      if (!Cascade)
+        Cascade = NextCascade;
+      return Cascade;
+    }
+
+    template <typename Iterator>
+    void setStage(Iterator Begin, Iterator End, LiveRangeStage NewStage) {
+      for (; Begin != End; ++Begin) {
+        Register Reg = *Begin;
+        Info.grow(Reg.id());
+        if (Info[Reg].Stage == RS_New)
+          Info[Reg].Stage = NewStage;
+      }
+    }
+    void LRE_DidCloneVirtReg(Register New, Register Old);
+  };
+
+  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; }
+  // end (interface to eviction advisers)
+
+private:
   // Convenient shortcuts.
   using PQueue = std::priority_queue<std::pair<unsigned, unsigned>>;
   using SmallLISet = SmallPtrSet<LiveInterval *, 4>;
 
   // context
   MachineFunction *MF;
 
   // Shortcuts to some useful interface.

llvm/lib/CodeGen/RegAllocGreedy.cpp

void RAGreedy::LRE_WillShrinkVirtReg(Register VirtReg) {
   Matrix->unassign(LI);
   RegAllocBase::enqueue(&LI);
 }
 
 void RAGreedy::LRE_DidCloneVirtReg(Register New, Register Old) {
   ExtraInfo->LRE_DidCloneVirtReg(New, Old);
 }
 
-void ExtraRegInfo::LRE_DidCloneVirtReg(Register New, Register Old) {
+void RAGreedy::ExtraRegInfo::LRE_DidCloneVirtReg(Register New, Register Old) {
   // Cloning a register we haven't even heard about yet?  Just ignore it.
   if (!Info.inBounds(Old))
     return;
 
   // LRE may clone a virtual register because dead code elimination causes it to
   // be split into connected components. The new components are much smaller
   // than the original, so they should get a new chance at being assigned.
   // same stage as the parent.

 /// @param A          The live range to be assigned.
 /// @param IsHint     True when A is about to be assigned to its preferred
 ///                   register.
 /// @param B          The live range to be evicted.
 /// @param BreaksHint True when B is already assigned to its preferred register.
 bool DefaultEvictionAdvisor::shouldEvict(LiveInterval &A, bool IsHint,
                                          LiveInterval &B,
                                          bool BreaksHint) const {
-  bool CanSplit = ExtraInfo->getStage(B) < RS_Spill;
+  bool CanSplit = RA.getExtraInfo().getStage(B) < RS_Spill;
 
   // Be fairly aggressive about following hints as long as the evictee can be
   // split.
   if (CanSplit && IsHint && !BreaksHint)
     return true;
 
   if (A.weight() > B.weight()) {
     LLVM_DEBUG(dbgs() << "should evict: " << B << " w= " << B.weight() << '\n');
bool DefaultEvictionAdvisor::canEvictInterferenceBasedOnCost(
 
   // Find VirtReg's cascade number. This will be unassigned if VirtReg was never
   // involved in an eviction before. If a cascade number was assigned, deny
   // evicting anything with the same or a newer cascade number. This prevents
   // infinite eviction loops.
   //
   // This works out so a register without a cascade number is allowed to evict
   // anything, and it can be evicted by anything.
-  unsigned Cascade = ExtraInfo->getCascadeOrCurrentNext(VirtReg.reg());
+  unsigned Cascade = RA.getExtraInfo().getCascadeOrCurrentNext(VirtReg.reg());
 
   EvictionCost Cost;
   for (MCRegUnitIterator Units(PhysReg, TRI); Units.isValid(); ++Units) {
     LiveIntervalUnion::Query &Q = Matrix->query(VirtReg, *Units);
     // If there is 10 or more interferences, chances are one is heavier.
     const auto &Interferences = Q.interferingVRegs(10);
     if (Interferences.size() >= 10)
       return false;
 
     // Check if any interfering live range is heavier than MaxWeight.
     for (LiveInterval *Intf : reverse(Interferences)) {
       assert(Register::isVirtualRegister(Intf->reg()) &&
              "Only expecting virtual register interference from query");
 
       // Do not allow eviction of a virtual register if we are in the middle
       // of last-chance recoloring and this virtual register is one that we
       // have scavenged a physical register for.
       if (FixedRegisters.count(Intf->reg()))
         return false;
 
       // Never evict spill products. They cannot split or spill.
-      if (ExtraInfo->getStage(*Intf) == RS_Done)
+      if (RA.getExtraInfo().getStage(*Intf) == RS_Done)
         return false;
       // Once a live range becomes small enough, it is urgent that we find a
       // register for it. This is indicated by an infinite spill weight. These
       // urgent live ranges get to evict almost anything.
       //
       // Also allow urgent evictions of unspillable ranges from a strictly
       // larger allocation order.
       bool Urgent =
           !VirtReg.isSpillable() &&
           (Intf->isSpillable() ||
            RegClassInfo.getNumAllocatableRegs(MRI->getRegClass(VirtReg.reg())) <
                RegClassInfo.getNumAllocatableRegs(
                    MRI->getRegClass(Intf->reg())));
       // Only evict older cascades or live ranges without a cascade.
-      unsigned IntfCascade = ExtraInfo->getCascade(Intf->reg());
+      unsigned IntfCascade = RA.getExtraInfo().getCascade(Intf->reg());
       if (Cascade <= IntfCascade) {
         if (!Urgent)
           return false;
         // We permit breaking cascades for urgent evictions. It should be the
         // last resort, though, so make it really expensive.
         Cost.BrokenHints += 10;
       }
       // Would this break a satisfied hint?
bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
 
   VRAI->calculateSpillWeightsAndHints();
 
   LLVM_DEBUG(LIS->dump());
 
   SA.reset(new SplitAnalysis(*VRM, *LIS, *Loops));
   SE.reset(new SplitEditor(*SA, *AA, *LIS, *VRM, *DomTree, *MBFI, *VRAI));
   ExtraInfo.emplace();
-  EvictAdvisor = getAnalysis<RegAllocEvictionAdvisorAnalysis>().getAdvisor(
-      *MF, Matrix, LIS, VRM, RegClassInfo, &*ExtraInfo);
+  EvictAdvisor =
+      getAnalysis<RegAllocEvictionAdvisorAnalysis>().getAdvisor(*MF, *this);
   IntfCache.init(MF, Matrix->getLiveUnions(), Indexes, LIS, TRI);
   GlobalCand.resize(32);  // This will grow as needed.
   SetOfBrokenHints.clear();
   LastEvicted.clear();
 
   allocatePhysRegs();
   tryHintsRecoloring();
 
   if (VerifyEnabled)
     MF->verify(this, "Before post optimization");
   postOptimization();
   reportStats();
 
   releaseMemory();
   return true;
 }