Update selection heristics to avoid code-size regression in average

Author: Ralender

llvm/include/llvm/CodeGen/MachineFrameInfo.h

@@ -772,7 +772,7 @@ class MachineFrameInfo {
     // If ID == 0, MaxAlignment will need to be updated separately.
   }
 
-  int getUnderlyingSlot(int ObjectIdx) {
+  int getUnderlyingSlot(int ObjectIdx) const {
     assert(unsigned(ObjectIdx + NumFixedObjects) < Objects.size() &&
            "Invalid Object Idx!");
     return Objects[ObjectIdx + NumFixedObjects].UnderlyingSlot;

llvm/lib/CodeGen/StackColoring.cpp

@@ -35,6 +35,7 @@
 #include "llvm/CodeGen/PseudoSourceValueManager.h"
 #include "llvm/CodeGen/SlotIndexes.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/CodeGen/WinEHFuncInfo.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/IR/Constants.h"
@@ -99,7 +100,6 @@ static cl::opt<unsigned> MaxCandidatesOpt(
         "Max number of candidates that will be evaluated, 0 means no limit"));
 
 STATISTIC(NumMarkerSeen, "Number of lifetime markers found.");
-STATISTIC(GeneratedWorse, "Number of times worse layout were generated");
 STATISTIC(StackSpaceSaved, "Number of bytes saved due to merging slots.");
 STATISTIC(StackSlotMerged, "Number of stack slot merged.");
 STATISTIC(EscapedAllocas, "Number of allocas that escaped the lifetime region");
@@ -400,7 +400,9 @@ class StackColoring {
     // Use to make overlap queries faster
     SmallVector<unsigned, 4> StartLiveness;
 
-    uint64_t SlotPriority = 0;
+    int64_t SlotPriority = 0;
+
+    unsigned UseCount = 0;
 
     unsigned Offset = InvalidIdx;
 
@@ -653,9 +655,11 @@ StackColoring::SlotInfo::dump(const StackColoring *State) const {
   dbgs() << ":";
   if (Offset != InvalidIdx)
     dbgs() << " offset=" << Offset;
+  dbgs() << " uses=" << UseCount;
+  dbgs() << " prio=" << SlotPriority;
   if (State) {
     if (State->MFI->getObjectAllocation(Slot))
-      dbgs() << " \"" << State->MFI->getObjectAllocation(Slot)->getName()
+      dbgs() << " alloca=\"" << State->MFI->getObjectAllocation(Slot)->getName()
              << "\"";
     if (State->MFI->isSpillSlotObjectIndex(Slot))
       dbgs() << " spill";
@@ -803,6 +807,7 @@ unsigned StackColoring::collectMarkers(unsigned NumSlot) {
           int Slot = MO.getIndex();
           if (Slot < 0)
             continue;
+          Slot2Info[Slot].UseCount++;
           if (!BetweenStartEnd.test(Slot)) {
             ConservativeSlots.set(Slot);
           }
@@ -1525,35 +1530,24 @@ unsigned StackColoring::doMerging(unsigned NumSlots) {
   if (SlotStack.size() <= 1)
     return InvalidIdx;
 
-  // This Whole block is only used to try and order the stack, such that the
-  // Slots are processed in an order that helps getting good packing
-  {
-    // Find how much usage of every livepoint there is.
-    SmallVector<unsigned> CumulatedUsage;
-    CumulatedUsage.resize(LivenessSize, 0);
-
-    for (unsigned Idx = 0; Idx < SlotStack.size(); Idx++) {
-      SlotInfo &Info = Slot2Info[SlotStack[Idx]];
-      for (unsigned Pt : Info.Liveness.set_bits()) {
-        CumulatedUsage[Pt] += Info.Size;
-      }
-    }
+  // This logic is optimized for x86_64, it probably needs to be adapted to
+  // other targets to get good code-size/stack-size balance.
+  // Its inspired from X86FrameLowering::orderFrameObjects, but modified weight
+  // in alignments helping with stack size
+  auto IsLower = [&](unsigned Lhs, unsigned Rhs) {
+    SlotInfo &L = Slot2Info[Lhs];
+    SlotInfo &R = Slot2Info[Rhs];
+    uint64_t DensityLScaled = static_cast<uint64_t>(L.UseCount) *
+                              static_cast<uint64_t>(R.Size + Log2(R.Align));
+    uint64_t DensityRScaled = static_cast<uint64_t>(R.UseCount) *
+                              static_cast<uint64_t>(L.Size + Log2(L.Align));
+    return DensityLScaled < DensityRScaled;
+  };
+  std::stable_sort(SlotStack.begin(), SlotStack.end(), IsLower);
 
-    for (unsigned Idx = 0; Idx < SlotStack.size(); Idx++) {
-      SlotInfo &Info = Slot2Info[SlotStack[Idx]];
-      for (unsigned Pt : Info.Liveness.set_bits()) {
-        // Since the goal is to minimize the max usage, blocks that are in high
-        // contention areas are given more priority
-        Info.SlotPriority +=
-            (uint64_t)CumulatedUsage[Pt] * (uint64_t)CumulatedUsage[Pt] +
-            (uint64_t)Info.Size * (uint64_t)Info.Align.value();
-      }
-    }
-    std::stable_sort(
-        SlotStack.begin(), SlotStack.end(), [&](unsigned Lhs, unsigned Rhs) {
-          return Slot2Info[Lhs].SlotPriority < Slot2Info[Rhs].SlotPriority;
-        });
-  }
+  int Prio = 0;
+  for (int Slot : SlotStack)
+    Slot2Info[Slot].SlotPriority = Prio++;
 
   SlotInfo *LastQueryLhs = nullptr;
   SlotInfo *LastQueryRhs = nullptr;
@@ -1666,24 +1660,27 @@ unsigned StackColoring::doMerging(unsigned NumSlots) {
 
       Offset = alignTo(Offset, Info.Align);
 
-      LLVM_DEBUG(dbgs() << "fi#" << Candidates[K] << "@" << Offset << "->";
-                 if (PrevSlot == InvalidIdx) dbgs() << "bottom";
-                 else dbgs() << "fi#" << PrevSlot; dbgs() << ", ";);
+      LLVM_DEBUG({
+        dbgs() << "fi#" << Candidates[K] << "@" << Offset;
+        if (PrevSlot != InvalidIdx)
+          dbgs() << "->" << "fi#" << PrevSlot;
+        dbgs() << ", ";
+      });
 
       bool IsBetter = [&] {
+        if (BestIdx == InvalidIdx)
+          return true;
+        SlotInfo &Best = Slot2Info[Candidates[BestIdx]];
         if (BestOffset != Offset)
           return BestOffset > Offset;
-        SlotInfo &Other = Slot2Info[Candidates[K]];
-        if (Other.Align != Info.Align)
-          return Other.Align < Info.Align;
-        if (Other.Size != Info.Size)
-          return Other.Size < Info.Size;
-        if (Other.SlotPriority != Info.SlotPriority)
-          return Other.SlotPriority < Info.SlotPriority;
+        if (Best.SlotPriority != Info.SlotPriority)
+          return Best.SlotPriority < Info.SlotPriority;
+        if (Best.Align != Info.Align)
+          return Best.Align < Info.Align;
 
         // Both are always stored in Slot2Info, so this is equivalent to
         // FrameIndex comparaison
-        return &Other < &Info;
+        return &Best < &Info;
       }();
 
       if (IsBetter) {
@@ -1726,7 +1723,6 @@ unsigned StackColoring::doMerging(unsigned NumSlots) {
   LLVM_DEBUG(dbgs() << "MergedSize=" << FinalSize << " OrigPesSize="
                     << OrigPesSize << " OrigOptSize" << OrigOptSize << "\n");
   if (FinalSize >= OrigPesSize) {
-    GeneratedWorse++;
     return InvalidIdx;
   }
 
@@ -1774,6 +1770,7 @@ bool StackColoring::run(MachineFunction &Func) {
   Intervals.reserve(NumSlots);
   LiveStarts.resize(NumSlots);
 
+  Slot2Info.resize(NumSlots);
   unsigned NumMarkers = collectMarkers(NumSlots);
 
   unsigned TotalSize = 0;
@@ -1792,7 +1789,6 @@ bool StackColoring::run(MachineFunction &Func) {
     return removeAllMarkers();
   }
 
-  Slot2Info.resize(NumSlots);
   for (unsigned i = 0; i < NumSlots; ++i) {
     std::unique_ptr<LiveRange> LI(new LiveRange());
     LI->getNextValue(Indexes->getZeroIndex(), VNInfoAllocator);