[NFC][regalloc] Separate iteration from AllocationOrder

[mtrofin]

Summary

This separates the two concerns - encapsulation of traversal order; and iteration.

[wmi]

Is the intention to have multiple iterators live for the same AllocationOrder?

llvm/lib/CodeGen/AllocationOrder.h

44–45

Can you have a reference of AllocationOrder in Iterator instead of defining multiple fields? It looks a little clearer to access parent's fields directly.

[mtrofin]

llvm/lib/CodeGen/AllocationOrder.h

44–45

With the current way, I'm leveraging the ability of constructing a view over the Order ArrayRef. Also, we could author a test very easily now, avoiding the need for a Matrix, RegClassInfo, etc (what AllocationOrder needs) - albeit currently the .h is under lib/CodeGen; but we could as a next step.

[mtrofin]

In D88256#2293981, @wmi wrote:

Is the intention to have multiple iterators live for the same AllocationOrder?

The main goal is to avoid this:

void fooBar(AllocationOrder &Order) {

Order.rewind();
while (...Order.next()) {
 <lots of code>
}

}

So if the while loop is large, you need to go see that the order has actually been reset. Also, the state post-while loop should be reset. The alternative I'm proposing allows allocation orders be passed around without worrying they get surprisingly mutated - so easier to read and maintain.

[wmi]

Thanks for explaining the motivation of the change. That makes sense.

[wmi]

Since it changes the way how AllocationOrder will be used, add Quentin as reviewer.

[qcolombet]

Hi,

I like that we take the iterator path, but I don't think the current patch is clearer than the previous implementation.

First we duplicate some logic from AllocationOrder to AllocationOrder::Iterator (see my inline comment), second, I found the AllocationOrder::getIterator not iterator friendly.
What I mean here is that I would have expected a AllocationOrder::begin and AllocationOrder::end, so that we can use range based loop and so on.

As is I don't really see the value of that refactoring.

Cheers,
-Quentin

llvm/lib/CodeGen/AllocationOrder.h

64

I don't like that we duplicate the logic of AllocationOrder::isHint here.

[mtrofin]

In D88256#2295887, @qcolombet wrote:

Hi,

I like that we take the iterator path, but I don't think the current patch is clearer than the previous implementation.

First we duplicate some logic from AllocationOrder to AllocationOrder::Iterator (see my inline comment), second, I found the AllocationOrder::getIterator not iterator friendly.
What I mean here is that I would have expected a AllocationOrder::begin and AllocationOrder::end, so that we can use range based loop and so on.

As is I don't really see the value of that refactoring.

Cheers,
-Quentin

SGTM - I'll send first another patch with unittests for the current behavior; then refresh this one.

[mtrofin]

updated to follow iterator style (begin() / end())

[mtrofin]

gentle reminder - thanks!

[mtrofin]

Gentle reminder - thanks!

[qcolombet]

LGTM with nitpicks

llvm/lib/CodeGen/RegAllocGreedy.cpp

811

Could we just stick to assigning PhysReg at the beginning of the loop?
That way we don't have to modify the body of the loop.

1860

Could you call out the type here?
Put differently, I don't like the use of auto when the type is not immediately obvious.

2299

Ditto with the use of auto

2617

Ditto with the use of auto

[mtrofin]

feedback

[mtrofin]

llvm/lib/CodeGen/RegAllocGreedy.cpp

811

There's a slight nuance:

In the old code, suppose we never hit line 819 (the break). Then line 821 will have a 0 PhysReg (indicating loop completed, and no reg was found)

In the new code, if we just initialize PhysReg on line 812, then we'll never have a 0 PhysReg on line 826 (ok, not never - because it could be the loop is never entered).

We could have I in the outer scope, and check both on line 826; or - the path I preferred - we clarify that the loop searches for a value for PhysReg.

wdyt?

1860

Done - also changed the type of operator*, in light of our discussion - we could type it as MCPhysReg, but I think MCRegister would also be correct, and help reduce the concept count - wdyt?

[qcolombet]

llvm/lib/CodeGen/RegAllocGreedy.cpp

811

Sounds fine.
Let's stick to the version in that patch.

1860

I have a slight preference with MCPhysReg, but MCRegister is fine too.
That said, it is strange to have ArrayRef<MCPhysReg> getOrder returning a different type.
I would rather be consistent, but I can live with the iterator offering the richer type.

[mtrofin]

feedback

llvm/lib/CodeGen/RegAllocGreedy.cpp

1860

ack - fixed here and at line 2298. Left at 2617, which before was Register, and also in RegAllocBasic, same reason.

llvm/lib/CodeGen/AllocationOrder.h

 // hints and target hooks. The AllocationOrder class encapsulates all of that.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_CODEGEN_ALLOCATIONORDER_H
 #define LLVM_LIB_CODEGEN_ALLOCATIONORDER_H
 
 #include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/MC/MCRegister.h"
 
 namespace llvm {
 
 class RegisterClassInfo;
 class VirtRegMap;
 class LiveRegMatrix;
 
 class LLVM_LIBRARY_VISIBILITY AllocationOrder {
   const SmallVector<MCPhysReg, 16> Hints;
   ArrayRef<MCPhysReg> Order;
-  int Pos = 0;
+  // How far into the Order we can iterate. This is 0 if the AllocationOrder is
+  // constructed with HardHints = true, Order.size() otherwise. While
+  // technically a size_t, it will participate in comparisons with the
+  // Iterator's Pos, which must be signed, so it's typed here as signed, too, to
+  // avoid warnings and under the assumption that the size of Order is
+  // relatively small.
+  // IterationLimit defines an invalid iterator position.
+  const int IterationLimit;
 
-  // If HardHints is true, *only* Hints will be returned.
-  const bool HardHints;
+public:
+  /// Forward iterator for an AllocationOrder.
+  class Iterator final {
+    const AllocationOrder &AO;
+    int Pos = 0;
 
   public:
+    Iterator(const AllocationOrder &AO, int Pos) : AO(AO), Pos(Pos) {}
+
+    /// Return true if the curent position is that of a preferred register.
+    bool isHint() const { return Pos < 0; }
+
+    /// Return the next physical register in the allocation order.
+    unsigned operator*() const {
+      if (Pos < 0)
+        return AO.Hints.end()[Pos];
+      assert(Pos < AO.IterationLimit);
+      return AO.Order[Pos];
+    }
+
+    /// Advance the iterator to the next position. If that's past the Hints
+    /// list, advance to the first value that's not also in the Hints list.
+    Iterator &operator++() {
+      if (Pos < AO.IterationLimit)
+        ++Pos;
+      while (Pos >= 0 && Pos < AO.IterationLimit && AO.isHint(AO.Order[Pos]))
+        ++Pos;
+      return *this;
+    }
+
+    bool operator==(const Iterator &Other) const {
+      assert(&AO == &Other.AO);
+      return Pos == Other.Pos;
+    }
+
+    bool operator!=(const Iterator &Other) const { return !(*this == Other); }
+  };
 
   /// Create a new AllocationOrder for VirtReg.
   /// @param VirtReg      Virtual register to allocate for.
   /// @param VRM          Virtual register map for function.
   /// @param RegClassInfo Information about reserved and allocatable registers.
   static AllocationOrder create(unsigned VirtReg, const VirtRegMap &VRM,
                                 const RegisterClassInfo &RegClassInfo,
                                 const LiveRegMatrix *Matrix);
 
   /// Create an AllocationOrder given the Hits, Order, and HardHits values.
   /// Use the create method above - the ctor is for unittests.
   AllocationOrder(SmallVector<MCPhysReg, 16> &&Hints, ArrayRef<MCPhysReg> Order,
                   bool HardHints)
       : Hints(std::move(Hints)), Order(Order),
-        Pos(-static_cast<int>(this->Hints.size())), HardHints(HardHints) {}
-
-  /// Get the allocation order without reordered hints.
-  ArrayRef<MCPhysReg> getOrder() const { return Order; }
+        IterationLimit(HardHints ? 0 : static_cast<int>(Order.size())) {}
 
-  /// Return the next physical register in the allocation order, or 0.
-  /// It is safe to call next() again after it returned 0, it will keep
-  /// returning 0 until rewind() is called.
-  MCPhysReg next(unsigned Limit = 0) {
-    if (Pos < 0)
-      return Hints.end()[Pos++];
-    if (HardHints)
-      return 0;
-    if (!Limit)
-      Limit = Order.size();
-    while (Pos < int(Limit)) {
-      unsigned Reg = Order[Pos++];
-      if (!isHint(Reg))
-        return Reg;
-    }
-    return 0;
+  Iterator begin() const {
+    return Iterator(*this, -(static_cast<int>(Hints.size())));
   }
 
-  /// Start over from the beginning.
-  void rewind() { Pos = -int(Hints.size()); }
+  Iterator end() const { return Iterator(*this, IterationLimit); }
+
+  Iterator getOrderLimitEnd(unsigned OrderLimit) const {
+    assert(OrderLimit <= Order.size());
+    if (OrderLimit == 0)
+      return end();
+    Iterator Ret(*this,
+                 std::min(static_cast<int>(OrderLimit) - 1, IterationLimit));
+    return ++Ret;
+  }
 
-  /// Return true if the last register returned from next() was a preferred register.
-  bool isHint() const { return Pos <= 0; }
+  /// Get the allocation order without reordered hints.
+  ArrayRef<MCPhysReg> getOrder() const { return Order; }
 
   /// Return true if PhysReg is a preferred register.
   bool isHint(unsigned PhysReg) const { return is_contained(Hints, PhysReg); }
 };
 
 } // end namespace llvm
 
 #endif

llvm/lib/CodeGen/RegAllocBasic.cpp

 Register RABasic::selectOrSplit(LiveInterval &VirtReg,
                                 SmallVectorImpl<Register> &SplitVRegs) {
   // Populate a list of physical register spill candidates.
   SmallVector<Register, 8> PhysRegSpillCands;
 
   // Check for an available register in this class.
   auto Order =
       AllocationOrder::create(VirtReg.reg(), *VRM, RegClassInfo, Matrix);
-  while (Register PhysReg = Order.next()) {
+  for (Register PhysReg : Order) {
+    assert(PhysReg);
     // Check for interference in PhysReg
     switch (Matrix->checkInterference(VirtReg, PhysReg)) {
     case LiveRegMatrix::IK_Free:
       // PhysReg is available, allocate it.
       return PhysReg;
 
     case LiveRegMatrix::IK_VirtReg:
       // Only virtual registers in the way, we may be able to spill them.

llvm/lib/CodeGen/RegAllocGreedy.cpp

 //                            Direct Assignment
 //===----------------------------------------------------------------------===//
 
 /// tryAssign - Try to assign VirtReg to an available register.
 Register RAGreedy::tryAssign(LiveInterval &VirtReg,
                              AllocationOrder &Order,
                              SmallVectorImpl<Register> &NewVRegs,
                              const SmallVirtRegSet &FixedRegisters) {
-  Order.rewind();
   Register PhysReg;
-  while ((PhysReg = Order.next()))
-    if (!Matrix->checkInterference(VirtReg, PhysReg))
-      break;
-  if (!PhysReg || Order.isHint())
+  for (auto I = Order.begin(), E = Order.end(); I != E && !PhysReg; ++I) {
+    assert(*I);
+    if (!Matrix->checkInterference(VirtReg, *I)) {
+      if (I.isHint())
+        return *I;
+      else
+        PhysReg = *I;
+    }
+  }
+  if (!PhysReg)
     return PhysReg;
 
   // PhysReg is available, but there may be a better choice.
 
   // If we missed a simple hint, try to cheaply evict interference from the
   // preferred register.
   if (Register Hint = MRI->getSimpleHint(VirtReg.reg()))
     if (Order.isHint(Hint)) {

 //===----------------------------------------------------------------------===//
 //                         Interference eviction
 //===----------------------------------------------------------------------===//
 
 Register RAGreedy::canReassign(LiveInterval &VirtReg, Register PrevReg) {
   auto Order =
       AllocationOrder::create(VirtReg.reg(), *VRM, RegClassInfo, Matrix);
   Register PhysReg;
-  while ((PhysReg = Order.next())) {
-    if (PhysReg == PrevReg)
+  for (auto I = Order.begin(), E = Order.end(); I != E && !PhysReg; ++I) {
+    if (*I == PrevReg)
       continue;
 
-    MCRegUnitIterator Units(PhysReg, TRI);
+    MCRegUnitIterator Units(*I, TRI);
     for (; Units.isValid(); ++Units) {
       // Instantiate a "subquery", not to be confused with the Queries array.
       LiveIntervalUnion::Query subQ(VirtReg, Matrix->getLiveUnions()[*Units]);
       if (subQ.checkInterference())
         break;
     }
     // If no units have interference, break out with the current PhysReg.
     if (!Units.isValid())
-      break;
+      PhysReg = *I;
   }
   if (PhysReg)
     LLVM_DEBUG(dbgs() << "can reassign: " << VirtReg << " from "
                       << printReg(PrevReg, TRI) << " to "
                       << printReg(PhysReg, TRI) << '\n');
   return PhysReg;
 }
 
if (CostPerUseLimit < ~0u) {
     // the same cost. We don't need to look at them if they're too expensive.
     if (TRI->getCostPerUse(Order.getOrder().back()) >= CostPerUseLimit) {
       OrderLimit = RegClassInfo.getLastCostChange(RC);
       LLVM_DEBUG(dbgs() << "Only trying the first " << OrderLimit
                         << " regs.\n");
     }
   }
 
-  Order.rewind();
-  while (MCRegister PhysReg = Order.next(OrderLimit)) {
+  for (auto I = Order.begin(), E = Order.getOrderLimitEnd(OrderLimit); I != E;
+       ++I) {
+    MCRegister PhysReg = *I;
+    assert(PhysReg);
     if (TRI->getCostPerUse(PhysReg) >= CostPerUseLimit)
       continue;
     // The first use of a callee-saved register in a function has cost 1.
     // Don't start using a CSR when the CostPerUseLimit is low.
     if (CostPerUseLimit == 1 && isUnusedCalleeSavedReg(PhysReg)) {
       LLVM_DEBUG(
           dbgs() << printReg(PhysReg, TRI) << " would clobber CSR "
                  << printReg(RegClassInfo.getLastCalleeSavedAlias(PhysReg), TRI)
                  << '\n');
       continue;
     }
 
     if (!canEvictInterference(VirtReg, PhysReg, false, BestCost,
                               FixedRegisters))
       continue;
 
     // Best so far.
     BestPhys = PhysReg;
 
     // Stop if the hint can be used.
-    if (Order.isHint())
+    if (I.isHint())
       break;
   }
 
   if (!BestPhys)
     return 0;
 
   evictInterference(VirtReg, BestPhys, NewVRegs);
   return BestPhys;

 }
 
 unsigned RAGreedy::calculateRegionSplitCost(LiveInterval &VirtReg,
                                             AllocationOrder &Order,
                                             BlockFrequency &BestCost,
                                             unsigned &NumCands, bool IgnoreCSR,
                                             bool *CanCauseEvictionChain) {
   unsigned BestCand = NoCand;
-  Order.rewind();
-  while (unsigned PhysReg = Order.next()) {
+  for (auto PhysReg : Order) {
+    assert(PhysReg);
     if (IgnoreCSR && isUnusedCalleeSavedReg(PhysReg))
       continue;
 
     // Discard bad candidates before we run out of interference cache cursors.
     // This will only affect register classes with a lot of registers (>32).
     if (NumCands == IntfCache.getMaxCursors()) {
       unsigned WorstCount = ~0u;
       unsigned Worst = 0;
unsigned RAGreedy::tryLocalSplit(LiveInterval &VirtReg, AllocationOrder &Order,
   unsigned BestAfter = 0;
   float BestDiff = 0;
 
   const float blockFreq =
     SpillPlacer->getBlockFrequency(BI.MBB->getNumber()).getFrequency() *
     (1.0f / MBFI->getEntryFreq());
   SmallVector<float, 8> GapWeight;
 
-  Order.rewind();
-  while (unsigned PhysReg = Order.next()) {
+  for (auto PhysReg : Order) {
+    assert(PhysReg);
     // Keep track of the largest spill weight that would need to be evicted in
     // order to make use of PhysReg between UseSlots[I] and UseSlots[I + 1].
     calcGapWeights(PhysReg, GapWeight);
 
     // Remove any gaps with regmask clobbers.
     if (Matrix->checkRegMaskInterference(VirtReg, PhysReg))
       for (unsigned I = 0, E = RegMaskGaps.size(); I != E; ++I)
         GapWeight[RegMaskGaps[I]] = huge_valf;
unsigned RAGreedy::tryLastChanceRecoloring(LiveInterval &VirtReg,
   // the recoloring fails.
   DenseMap<Register, Register> VirtRegToPhysReg;
   // Mark VirtReg as fixed, i.e., it will not be recolored pass this point in
   // this recoloring "session".
   assert(!FixedRegisters.count(VirtReg.reg()));
   FixedRegisters.insert(VirtReg.reg());
   SmallVector<Register, 4> CurrentNewVRegs;
 
-  Order.rewind();
-  while (Register PhysReg = Order.next()) {
+  for (auto PhysReg : Order) {
+    assert(PhysReg);
     LLVM_DEBUG(dbgs() << "Try to assign: " << VirtReg << " to "
                       << printReg(PhysReg, TRI) << '\n');
     RecoloringCandidates.clear();
     VirtRegToPhysReg.clear();
     CurrentNewVRegs.clear();
 
     // It is only possible to recolor virtual register interference.
     if (Matrix->checkInterference(VirtReg, PhysReg) >

llvm/unittests/CodeGen/AllocationOrderTest.cpp

 //===- llvm/unittest/CodeGen/AllocationOrderTest.cpp - AllocationOrder tests =//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "../lib/CodeGen/AllocationOrder.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 namespace {
-std::vector<MCPhysReg> loadOrder(AllocationOrder &O, unsigned Limit = 0) {
+std::vector<MCPhysReg> loadOrder(const AllocationOrder &O, unsigned Limit = 0) {
   std::vector<MCPhysReg> Ret;
-  O.rewind();
-  while (auto R = O.next(Limit))
+  if (Limit == 0)
+    for (auto R : O)
       Ret.push_back(R);
+  else
+    for (auto I = O.begin(), E = O.getOrderLimitEnd(Limit); I != E; ++I)
+      Ret.push_back(*I);
   return Ret;
 }
 } // namespace
 
 TEST(AllocationOrderTest, Basic) {
   SmallVector<MCPhysReg, 16> Hints = {1, 2, 3};
   SmallVector<MCPhysReg, 16> Order = {4, 5, 6, 7};
   AllocationOrder O(std::move(Hints), Order, false);

 }
 
 TEST(AllocationOrderTest, LimitsBasic) {
   SmallVector<MCPhysReg, 16> Hints = {1, 2, 3};
   SmallVector<MCPhysReg, 16> Order = {4, 5, 6, 7};
   AllocationOrder O(std::move(Hints), Order, false);
   EXPECT_EQ((std::vector<MCPhysReg>{1, 2, 3, 4, 5, 6, 7}), loadOrder(O, 0));
   EXPECT_EQ((std::vector<MCPhysReg>{1, 2, 3, 4}), loadOrder(O, 1));
+  EXPECT_EQ(O.end(), O.getOrderLimitEnd(0));
 }
 
 TEST(AllocationOrderTest, LimitsDuplicates) {
   SmallVector<MCPhysReg, 16> Hints = {1, 2, 3};
   SmallVector<MCPhysReg, 16> Order = {4, 1, 5, 6};
   AllocationOrder O(std::move(Hints), Order, false);
   EXPECT_EQ((std::vector<MCPhysReg>{1, 2, 3, 4}), loadOrder(O, 1));
   EXPECT_EQ((std::vector<MCPhysReg>{1, 2, 3, 4}), loadOrder(O, 2));
TEST(AllocationOrderTest, NoHints) {
   EXPECT_EQ((std::vector<MCPhysReg>{1, 2}), loadOrder(O, 2));
   EXPECT_EQ((std::vector<MCPhysReg>{1, 2, 3}), loadOrder(O, 3));
 }
 
 TEST(AllocationOrderTest, IsHintTest) {
   SmallVector<MCPhysReg, 16> Hints = {1, 2, 3};
   SmallVector<MCPhysReg, 16> Order = {4, 1, 5, 6};
   AllocationOrder O(std::move(Hints), Order, false);
-  O.rewind();
-  auto V = O.next();
-  EXPECT_TRUE(O.isHint());
+  auto I = O.begin();
+  auto V = *I;
+  EXPECT_TRUE(I.isHint());
   EXPECT_EQ(V, 1U);
-  O.next();
-  EXPECT_TRUE(O.isHint());
-  O.next();
-  EXPECT_TRUE(O.isHint());
-  V = O.next();
-  EXPECT_FALSE(O.isHint());
+  ++I;
+  EXPECT_TRUE(I.isHint());
+  ++I;
+  EXPECT_TRUE(I.isHint());
+  V = *(++I);
+  EXPECT_FALSE(I.isHint());
   EXPECT_EQ(V, 4U);
-  V = O.next();
+  V = *(++I);
   EXPECT_TRUE(O.isHint(1));
-  EXPECT_FALSE(O.isHint());
+  EXPECT_FALSE(I.isHint());
   EXPECT_EQ(V, 5U);
 }