[StackColoring] Change the StackColoring logic + enables it to handle spills

llvm/include/llvm/CodeGen/LiveStacks.h

@@ -40,49 +40,46 @@ class LiveStacks {
   ///
   VNInfo::Allocator VNInfoAllocator;
 
-  /// S2IMap - Stack slot indices to live interval mapping.
-  using SS2IntervalMap = std::unordered_map<int, LiveInterval>;
-  SS2IntervalMap S2IMap;
-
-  /// S2RCMap - Stack slot indices to register class mapping.
-  std::map<int, const TargetRegisterClass *> S2RCMap;
+  int StartIdx = -1;
+  SmallVector<LiveInterval*> S2LI;
+  SmallVector<const TargetRegisterClass *> S2RC;
 
 public:
-  using iterator = SS2IntervalMap::iterator;
-  using const_iterator = SS2IntervalMap::const_iterator;
+  using iterator = SmallVector<LiveInterval*>::iterator;
+  using const_iterator = SmallVector<LiveInterval*>::const_iterator;
 
-  const_iterator begin() const { return S2IMap.begin(); }
-  const_iterator end() const { return S2IMap.end(); }
-  iterator begin() { return S2IMap.begin(); }
-  iterator end() { return S2IMap.end(); }
+  const_iterator begin() const { return S2LI.begin(); }
+  const_iterator end() const { return S2LI.end(); }
+  iterator begin() { return S2LI.begin(); }
+  iterator end() { return S2LI.end(); }
 
-  unsigned getNumIntervals() const { return (unsigned)S2IMap.size(); }
+  unsigned getStartIdx() const {
+    assert(StartIdx != -1);
+    return StartIdx;
+  }
+  unsigned getNumIntervals() const { return (unsigned)S2LI.size(); }
 
   LiveInterval &getOrCreateInterval(int Slot, const TargetRegisterClass *RC);
 
   LiveInterval &getInterval(int Slot) {
     assert(Slot >= 0 && "Spill slot indice must be >= 0");
-    SS2IntervalMap::iterator I = S2IMap.find(Slot);
-    assert(I != S2IMap.end() && "Interval does not exist for stack slot");
-    return I->second;
+    return *S2LI[Slot - StartIdx];
   }
 
   const LiveInterval &getInterval(int Slot) const {
     assert(Slot >= 0 && "Spill slot indice must be >= 0");
-    SS2IntervalMap::const_iterator I = S2IMap.find(Slot);
-    assert(I != S2IMap.end() && "Interval does not exist for stack slot");
-    return I->second;
+    return *S2LI[Slot - StartIdx];
   }
 
-  bool hasInterval(int Slot) const { return S2IMap.count(Slot); }
+  bool hasInterval(int Slot) const {
+    if (Slot < StartIdx || StartIdx == -1)
+      return false;
+    return !getInterval(Slot).empty();
+  }
 
   const TargetRegisterClass *getIntervalRegClass(int Slot) const {
     assert(Slot >= 0 && "Spill slot indice must be >= 0");
-    std::map<int, const TargetRegisterClass *>::const_iterator I =
-        S2RCMap.find(Slot);
-    assert(I != S2RCMap.end() &&
-           "Register class info does not exist for stack slot");
-    return I->second;
+    return S2RC[Slot - StartIdx];
   }
 
   VNInfo::Allocator &getVNInfoAllocator() { return VNInfoAllocator; }

llvm/include/llvm/CodeGen/MachineFrameInfo.h

@@ -120,12 +120,18 @@ class MachineFrameInfo {
                       ///< triggered protection.  3rd closest to the protector.
   };
 
-private:
+  static constexpr int NoUnderlyingSlot = std::numeric_limits<int>::min();
+  static constexpr int IsUnderlyingSlot = std::numeric_limits<int>::min() + 1;
+
+  private:
   // Represent a single object allocated on the stack.
   struct StackObject {
     // The offset of this object from the stack pointer on entry to
     // the function.  This field has no meaning for a variable sized element.
-    int64_t SPOffset;
+    // After getting placed this is relative to SP
+    // If UnderlyingSlot is not NoUnderlyingSlot, this is relative to the start
+    // of the UnderlyingSlot
+    int64_t Offset;
 
     // The size of this object on the stack. 0 means a variable sized object,
     // ~0ULL means a dead object.
@@ -134,6 +140,10 @@ class MachineFrameInfo {
     // The required alignment of this stack slot.
     Align Alignment;
 
+    // If not NoUnderlyingSlot, it Indicate that this slot should be placed
+    // at Offset, into the slot UnderlyingSlot
+    int UnderlyingSlot = NoUnderlyingSlot;
+
     // If true, the value of the stack object is set before
     // entering the function and is not modified inside the function. By
     // default, fixed objects are immutable unless marked otherwise.
@@ -183,10 +193,10 @@ class MachineFrameInfo {
 
     uint8_t SSPLayout = SSPLK_None;
 
-    StackObject(uint64_t Size, Align Alignment, int64_t SPOffset,
+    StackObject(uint64_t Size, Align Alignment, int64_t Offset,
                 bool IsImmutable, bool IsSpillSlot, const AllocaInst *Alloca,
                 bool IsAliased, uint8_t StackID = 0)
-        : SPOffset(SPOffset), Size(Size), Alignment(Alignment),
+        : Offset(Offset), Size(Size), Alignment(Alignment),
           isImmutable(IsImmutable), isSpillSlot(IsSpillSlot), StackID(StackID),
           Alloca(Alloca), isAliased(IsAliased) {}
   };
@@ -532,7 +542,7 @@ class MachineFrameInfo {
            "Invalid Object Idx!");
     assert(!isDeadObjectIndex(ObjectIdx) &&
            "Getting frame offset for a dead object?");
-    return Objects[ObjectIdx+NumFixedObjects].SPOffset;
+    return Objects[ObjectIdx+NumFixedObjects].Offset;
   }
 
   bool isObjectZExt(int ObjectIdx) const {
@@ -561,12 +571,12 @@ class MachineFrameInfo {
 
   /// Set the stack frame offset of the specified object. The
   /// offset is relative to the stack pointer on entry to the function.
-  void setObjectOffset(int ObjectIdx, int64_t SPOffset) {
+  void setObjectOffset(int ObjectIdx, int64_t Offset) {
     assert(unsigned(ObjectIdx+NumFixedObjects) < Objects.size() &&
            "Invalid Object Idx!");
     assert(!isDeadObjectIndex(ObjectIdx) &&
            "Setting frame offset for a dead object?");
-    Objects[ObjectIdx+NumFixedObjects].SPOffset = SPOffset;
+    Objects[ObjectIdx+NumFixedObjects].Offset = Offset;
   }
 
   SSPLayoutKind getObjectSSPLayout(int ObjectIdx) const {
@@ -762,6 +772,18 @@ class MachineFrameInfo {
     // If ID == 0, MaxAlignment will need to be updated separately.
   }
 
+  int getUnderlyingSlot(int ObjectIdx) {
+    assert(unsigned(ObjectIdx + NumFixedObjects) < Objects.size() &&
+           "Invalid Object Idx!");
+    return Objects[ObjectIdx + NumFixedObjects].UnderlyingSlot;
+  }
+
+  void setUnderlyingSlot(int ObjectIdx, int Underlying) {
+    assert(unsigned(ObjectIdx + NumFixedObjects) < Objects.size() &&
+           "Invalid Object Idx!");
+    Objects[ObjectIdx + NumFixedObjects].UnderlyingSlot = Underlying;
+  }
+
   /// Returns true if the specified index corresponds to a dead object.
   bool isDeadObjectIndex(int ObjectIdx) const {
     assert(unsigned(ObjectIdx+NumFixedObjects) < Objects.size() &&

llvm/include/llvm/CodeGen/MachineInstr.h

@@ -1340,7 +1340,9 @@ class MachineInstr
   }
 
   // True if the instruction represents a position in the function.
-  bool isPosition() const { return isLabel() || isCFIInstruction(); }
+  bool isPosition() const {
+    return isLifetimeMarker() || isLabel() || isCFIInstruction();
+  }
 
   bool isNonListDebugValue() const {
     return getOpcode() == TargetOpcode::DBG_VALUE;

llvm/lib/CodeGen/LiveStacks.cpp

@@ -37,10 +37,12 @@ void LiveStacksWrapperLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
 }
 
 void LiveStacks::releaseMemory() {
+  for (int Idx = 0; Idx < (int)S2LI.size(); ++Idx)
+    S2LI[Idx]->~LiveInterval();
   // Release VNInfo memory regions, VNInfo objects don't need to be dtor'd.
   VNInfoAllocator.Reset();
-  S2IMap.clear();
-  S2RCMap.clear();
+  S2LI.clear();
+  S2RC.clear();
 }
 
 void LiveStacks::init(MachineFunction &MF) {
@@ -52,20 +54,22 @@ void LiveStacks::init(MachineFunction &MF) {
 LiveInterval &
 LiveStacks::getOrCreateInterval(int Slot, const TargetRegisterClass *RC) {
   assert(Slot >= 0 && "Spill slot indice must be >= 0");
-  SS2IntervalMap::iterator I = S2IMap.find(Slot);
-  if (I == S2IMap.end()) {
-    I = S2IMap
-            .emplace(
-                std::piecewise_construct, std::forward_as_tuple(Slot),
-                std::forward_as_tuple(Register::index2StackSlot(Slot), 0.0F))
-            .first;
-    S2RCMap.insert(std::make_pair(Slot, RC));
+  if (StartIdx == -1)
+    StartIdx = Slot;
+
+  int Idx = Slot - StartIdx;
+  assert(Idx >= 0 && "Slot not in order ?");
+  if (Idx < (int)S2LI.size()) {
+    S2RC[Idx] = TRI->getCommonSubClass(S2RC[Idx], RC);
   } else {
-    // Use the largest common subclass register class.
-    const TargetRegisterClass *&OldRC = S2RCMap[Slot];
-    OldRC = TRI->getCommonSubClass(OldRC, RC);
+    S2RC.resize(Idx + 1);
+    S2LI.resize(Idx + 1);
+    S2LI[Idx] = this->VNInfoAllocator.Allocate<LiveInterval>();
+    new (S2LI[Idx]) LiveInterval(Register::index2StackSlot(Slot), 0.0F);
+    S2RC[Idx] = RC;
   }
-  return I->second;
+  assert(S2RC.size() == S2LI.size());
+  return *S2LI[Idx];
 }
 
 AnalysisKey LiveStacksAnalysis::Key;
@@ -96,13 +100,12 @@ void LiveStacksWrapperLegacy::print(raw_ostream &OS, const Module *) const {
 }
 
 /// print - Implement the dump method.
-void LiveStacks::print(raw_ostream &OS, const Module*) const {
+void LiveStacks::print(raw_ostream &OS, const Module *) const {
 
   OS << "********** INTERVALS **********\n";
-  for (const_iterator I = begin(), E = end(); I != E; ++I) {
-    I->second.print(OS);
-    int Slot = I->first;
-    const TargetRegisterClass *RC = getIntervalRegClass(Slot);
+  for (int Idx = 0; Idx < (int)S2LI.size(); ++Idx) {
+    S2LI[Idx]->print(OS);
+    const TargetRegisterClass *RC = S2RC[Idx];
     if (RC)
       OS << " [" << TRI->getRegClassName(RC) << "]\n";
     else

llvm/lib/CodeGen/MachineFrameInfo.cpp

@@ -13,6 +13,7 @@
 #include "llvm/CodeGen/MachineFrameInfo.h"
 
 #include "llvm/ADT/BitVector.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetFrameLowering.h"
@@ -80,7 +81,7 @@ int MachineFrameInfo::CreateVariableSizedObject(Align Alignment,
   return (int)Objects.size()-NumFixedObjects-1;
 }
 
-int MachineFrameInfo::CreateFixedObject(uint64_t Size, int64_t SPOffset,
+int MachineFrameInfo::CreateFixedObject(uint64_t Size, int64_t Offset,
                                         bool IsImmutable, bool IsAliased) {
   assert(Size != 0 && "Cannot allocate zero size fixed stack objects!");
   // The alignment of the frame index can be determined from its offset from
@@ -90,23 +91,23 @@ int MachineFrameInfo::CreateFixedObject(uint64_t Size, int64_t SPOffset,
   // stack needs realignment, we can't assume that the stack will in fact be
   // aligned.
   Align Alignment =
-      commonAlignment(ForcedRealign ? Align(1) : StackAlignment, SPOffset);
+      commonAlignment(ForcedRealign ? Align(1) : StackAlignment, Offset);
   Alignment = clampStackAlignment(!StackRealignable, Alignment, StackAlignment);
   Objects.insert(Objects.begin(),
-                 StackObject(Size, Alignment, SPOffset, IsImmutable,
+                 StackObject(Size, Alignment, Offset, IsImmutable,
                              /*IsSpillSlot=*/false, /*Alloca=*/nullptr,
                              IsAliased));
   return -++NumFixedObjects;
 }
 
 int MachineFrameInfo::CreateFixedSpillStackObject(uint64_t Size,
-                                                  int64_t SPOffset,
+                                                  int64_t Offset,
                                                   bool IsImmutable) {
   Align Alignment =
-      commonAlignment(ForcedRealign ? Align(1) : StackAlignment, SPOffset);
+      commonAlignment(ForcedRealign ? Align(1) : StackAlignment, Offset);
   Alignment = clampStackAlignment(!StackRealignable, Alignment, StackAlignment);
   Objects.insert(Objects.begin(),
-                 StackObject(Size, Alignment, SPOffset, IsImmutable,
+                 StackObject(Size, Alignment, Offset, IsImmutable,
                              /*IsSpillSlot=*/true, /*Alloca=*/nullptr,
                              /*IsAliased=*/false));
   return -++NumFixedObjects;
@@ -221,6 +222,12 @@ void MachineFrameInfo::print(const MachineFunction &MF, raw_ostream &OS) const{
     if (SO.StackID != 0)
       OS << "id=" << static_cast<unsigned>(SO.StackID) << ' ';
 
+    if (SO.Alloca)
+      OS << "alloca=" << SO.Alloca->getName() << ' ';
+
+    if (SO.isSpillSlot)
+      OS << "spill ";
+
     if (SO.Size == ~0ULL) {
       OS << "dead\n";
       continue;
@@ -233,8 +240,13 @@ void MachineFrameInfo::print(const MachineFunction &MF, raw_ostream &OS) const{
 
     if (i < NumFixedObjects)
       OS << ", fixed";
-    if (i < NumFixedObjects || SO.SPOffset != -1) {
-      int64_t Off = SO.SPOffset - ValOffset;
+    if (SO.UnderlyingSlot == MachineFrameInfo::IsUnderlyingSlot)
+      OS << ", underlying";
+    if (SO.UnderlyingSlot > MachineFrameInfo::IsUnderlyingSlot) {
+      OS << ", placed=" << "fi#" << (int)(SO.UnderlyingSlot - NumFixedObjects)
+         << "+" << SO.Offset;
+    } else if (i < NumFixedObjects || SO.Offset != -1) {
+      int64_t Off = SO.Offset - ValOffset;
       OS << ", at location [SP";
       if (Off > 0)
         OS << "+" << Off;

llvm/lib/CodeGen/MachineInstr.cpp

@@ -1417,10 +1417,6 @@ bool MachineInstr::isDead(const MachineRegisterInfo &MRI,
   if (isInlineAsm())
     return false;
 
-  // FIXME: See issue #105950 for why LIFETIME markers are considered dead here.
-  if (isLifetimeMarker())
-    return true;
-
   // If there are no defs with uses, then we call the instruction dead so long
   // as we do not suspect it may have sideeffects.
   return wouldBeTriviallyDead();

llvm/lib/CodeGen/PrologEpilogInserter.cpp

@@ -694,6 +694,13 @@ void PEIImpl::spillCalleeSavedRegs(MachineFunction &MF) {
   }
 }
 
+static inline void UpdateOffset(MachineFrameInfo &MFI, int FrameIdx,
+                                int64_t Offset) {
+    LLVM_DEBUG(dbgs() << "alloc FI(" << FrameIdx << ") at SP[" << Offset
+                      << "]\n");
+    MFI.setObjectOffset(FrameIdx, Offset); // Set the computed offset
+}
+
 /// AdjustStackOffset - Helper function used to adjust the stack frame offset.
 static inline void AdjustStackOffset(MachineFrameInfo &MFI, int FrameIdx,
                                      bool StackGrowsDown, int64_t &Offset,
@@ -712,13 +719,9 @@ static inline void AdjustStackOffset(MachineFrameInfo &MFI, int FrameIdx,
   Offset = alignTo(Offset, Alignment);
 
   if (StackGrowsDown) {
-    LLVM_DEBUG(dbgs() << "alloc FI(" << FrameIdx << ") at SP[" << -Offset
-                      << "]\n");
-    MFI.setObjectOffset(FrameIdx, -Offset); // Set the computed offset
+    UpdateOffset(MFI, FrameIdx, -Offset);
   } else {
-    LLVM_DEBUG(dbgs() << "alloc FI(" << FrameIdx << ") at SP[" << Offset
-                      << "]\n");
-    MFI.setObjectOffset(FrameIdx, Offset);
+    UpdateOffset(MFI, FrameIdx, Offset);
     Offset += MFI.getObjectSize(FrameIdx);
   }
 }
@@ -1044,6 +1047,7 @@ void PEIImpl::calculateFrameObjectOffsets(MachineFunction &MF) {
   }
 
   SmallVector<int, 8> ObjectsToAllocate;
+  SmallVector<int, 8> UpdateOffsetAfterAllocate;
 
   // Then prepare to assign frame offsets to stack objects that are not used to
   // spill callee saved registers.
@@ -1064,6 +1068,11 @@ void PEIImpl::calculateFrameObjectOffsets(MachineFunction &MF) {
     if (MFI.getStackID(i) != TargetStackID::Default)
       continue;
 
+    if (MFI.getUnderlyingSlot(i) > MachineFrameInfo::IsUnderlyingSlot) {
+      UpdateOffsetAfterAllocate.push_back(i);
+      continue;
+    }
+
     // Add the objects that we need to allocate to our working set.
     ObjectsToAllocate.push_back(i);
   }
@@ -1104,6 +1113,14 @@ void PEIImpl::calculateFrameObjectOffsets(MachineFunction &MF) {
       AdjustStackOffset(MFI, SFI, StackGrowsDown, Offset, MaxAlign);
   }
 
+  for (int FrameIdx : UpdateOffsetAfterAllocate) {
+    int UnderlyingSlot = MFI.getUnderlyingSlot(FrameIdx);
+    int64_t ObjOffset =
+        MFI.getObjectOffset(UnderlyingSlot) + MFI.getObjectOffset(FrameIdx);
+    UpdateOffset(MFI, FrameIdx, ObjOffset);
+    MFI.setUnderlyingSlot(FrameIdx, MachineFrameInfo::NoUnderlyingSlot);
+  }
+
   if (!TFI.targetHandlesStackFrameRounding()) {
     // If we have reserved argument space for call sites in the function
     // immediately on entry to the current function, count it as part of the