Add new StackColoring algo

Ralender · Ralender · commit fb19fd014669 · 2025-06-22T13:31:39.000+02:00
diff --git a/llvm/include/llvm/CodeGen/MachineFrameInfo.h b/llvm/include/llvm/CodeGen/MachineFrameInfo.h
@@ -120,12 +120,18 @@ class MachineFrameInfo {
                       ///< triggered protection.  3rd closest to the protector.
   };
 
+  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() &&
diff --git a/llvm/lib/CodeGen/MachineFrameInfo.cpp b/llvm/lib/CodeGen/MachineFrameInfo.cpp
@@ -81,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
@@ -91,23 +91,22 @@ 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,
+int MachineFrameInfo::CreateFixedSpillStackObject(uint64_t Size, 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;
@@ -240,8 +239,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;
diff --git a/llvm/lib/CodeGen/PrologEpilogInserter.cpp b/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
diff --git a/llvm/lib/CodeGen/StackColoring.cpp b/llvm/lib/CodeGen/StackColoring.cpp
