Enforce StackID definition in PEI

[sdesmalen]

Summary

There are various places in LLVM where the definition of StackID is not
properly honoured, for example in PEI where objects with a StackID > 0 are
allocated on the default stack (StackID0). This patch enforces that PEI
only considers allocating objects to StackID 0.

[arsenm]

lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

295–299

I'm curious how you ended up finding this?

[arsenm]

LGTM

[sdesmalen]

Thanks!

lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

295–299

Without this change I ended up with several unit test failures:

LLVM :: CodeGen/AMDGPU/control-flow-fastregalloc.ll
LLVM :: CodeGen/AMDGPU/si-spill-sgpr-stack.ll
LLVM :: CodeGen/AMDGPU/spill-m0.ll
LLVM :: CodeGen/AMDGPU/spill-wide-sgpr.ll

include/llvm/CodeGen/MachineFrameInfo.h

unsigned getObjectAlignment(int ObjectIdx) const {
     return Objects[ObjectIdx+NumFixedObjects].Alignment;
   }
 
   /// setObjectAlignment - Change the alignment of the specified stack object.
   void setObjectAlignment(int ObjectIdx, unsigned Align) {
     assert(unsigned(ObjectIdx+NumFixedObjects) < Objects.size() &&
            "Invalid Object Idx!");
     Objects[ObjectIdx+NumFixedObjects].Alignment = Align;
+
+    // Only ensure max alignment for the default stack.
+    if (getStackID(ObjectIdx) == 0)
       ensureMaxAlignment(Align);
   }
 
   /// Return the underlying Alloca of the specified
   /// stack object if it exists. Returns 0 if none exists.
   const AllocaInst* getObjectAllocation(int ObjectIdx) const {
     assert(unsigned(ObjectIdx+NumFixedObjects) < Objects.size() &&
            "Invalid Object Idx!");
     return Objects[ObjectIdx+NumFixedObjects].Alloca;
uint8_t getStackID(int ObjectIdx) const {
     return Objects[ObjectIdx+NumFixedObjects].StackID;
   }
 
   /// \see StackID
   void setStackID(int ObjectIdx, uint8_t ID) {
     assert(unsigned(ObjectIdx+NumFixedObjects) < Objects.size() &&
            "Invalid Object Idx!");
     Objects[ObjectIdx+NumFixedObjects].StackID = ID;
+    // If ID > 0, MaxAlignment may now be overly conservative.
+    // If ID == 0, MaxAlignment will need to be updated separately.
   }
 
   /// Returns true if the specified index corresponds to a dead object.
   bool isDeadObjectIndex(int ObjectIdx) const {
     assert(unsigned(ObjectIdx+NumFixedObjects) < Objects.size() &&
            "Invalid Object Idx!");
     return Objects[ObjectIdx+NumFixedObjects].Size == ~0ULL;
   }

lib/CodeGen/MIRParser/MIRParser.cpp

bool MIRParserImpl::initializeFrameInfo(PerFunctionMIParsingState &PFS,
   // Initialize the fixed frame objects.
   for (const auto &Object : YamlMF.FixedStackObjects) {
     int ObjectIdx;
     if (Object.Type != yaml::FixedMachineStackObject::SpillSlot)
       ObjectIdx = MFI.CreateFixedObject(Object.Size, Object.Offset,
                                         Object.IsImmutable, Object.IsAliased);
     else
       ObjectIdx = MFI.CreateFixedSpillStackObject(Object.Size, Object.Offset);
-    MFI.setObjectAlignment(ObjectIdx, Object.Alignment);
     MFI.setStackID(ObjectIdx, Object.StackID);
+    MFI.setObjectAlignment(ObjectIdx, Object.Alignment);
     if (!PFS.FixedStackObjectSlots.insert(std::make_pair(Object.ID.Value,
                                                          ObjectIdx))
              .second)
       return error(Object.ID.SourceRange.Start,
                    Twine("redefinition of fixed stack object '%fixed-stack.") +
                        Twine(Object.ID.Value) + "'");
     if (parseCalleeSavedRegister(PFS, CSIInfo, Object.CalleeSavedRegister,
                                  Object.CalleeSavedRestored, ObjectIdx))
if (!Name.Value.empty()) {
                          "' isn't defined in the function '" + F.getName() +
                          "'");
     }
     if (Object.Type == yaml::MachineStackObject::VariableSized)
       ObjectIdx = MFI.CreateVariableSizedObject(Object.Alignment, Alloca);
     else
       ObjectIdx = MFI.CreateStackObject(
           Object.Size, Object.Alignment,
-          Object.Type == yaml::MachineStackObject::SpillSlot, Alloca);
+          Object.Type == yaml::MachineStackObject::SpillSlot, Alloca,
+          Object.StackID);
     MFI.setObjectOffset(ObjectIdx, Object.Offset);
-    MFI.setStackID(ObjectIdx, Object.StackID);
 
     if (!PFS.StackObjectSlots.insert(std::make_pair(Object.ID.Value, ObjectIdx))
              .second)
       return error(Object.ID.SourceRange.Start,
                    Twine("redefinition of stack object '%stack.") +
                        Twine(Object.ID.Value) + "'");
     if (parseCalleeSavedRegister(PFS, CSIInfo, Object.CalleeSavedRegister,
                                  Object.CalleeSavedRestored, ObjectIdx))

lib/CodeGen/MachineFrameInfo.cpp

int MachineFrameInfo::CreateStackObject(uint64_t Size, unsigned Alignment,
                                         const AllocaInst *Alloca,
                                         uint8_t StackID) {
   assert(Size != 0 && "Cannot allocate zero size stack objects!");
   Alignment = clampStackAlignment(!StackRealignable, Alignment, StackAlignment);
   Objects.push_back(StackObject(Size, Alignment, 0, false, IsSpillSlot, Alloca,
                                 !IsSpillSlot, StackID));
   int Index = (int)Objects.size() - NumFixedObjects - 1;
   assert(Index >= 0 && "Bad frame index!");
+  if (StackID == 0)
     ensureMaxAlignment(Alignment);
   return Index;
 }
 
 int MachineFrameInfo::CreateSpillStackObject(uint64_t Size,
                                              unsigned Alignment) {
   Alignment = clampStackAlignment(!StackRealignable, Alignment, StackAlignment);
   CreateStackObject(Size, Alignment, true);
   int Index = (int)Objects.size() - NumFixedObjects - 1;
unsigned MachineFrameInfo::estimateStackSize(const MachineFunction &MF) const {
   unsigned MaxAlign = getMaxAlignment();
   int Offset = 0;
 
   // This code is very, very similar to PEI::calculateFrameObjectOffsets().
   // It really should be refactored to share code. Until then, changes
   // should keep in mind that there's tight coupling between the two.
 
   for (int i = getObjectIndexBegin(); i != 0; ++i) {
+    // Only estimate stack size of default stack.
+    if (getStackID(i))
+      continue;
     int FixedOff = -getObjectOffset(i);
     if (FixedOff > Offset) Offset = FixedOff;
   }
   for (unsigned i = 0, e = getObjectIndexEnd(); i != e; ++i) {
-    if (isDeadObjectIndex(i))
+    // Only estimate stack size of live objects on default stack.
+    if (isDeadObjectIndex(i) || getStackID(i))
       continue;
     Offset += getObjectSize(i);
     unsigned Align = getObjectAlignment(i);
     // Adjust to alignment boundary
     Offset = (Offset+Align-1)/Align*Align;
 
     MaxAlign = std::max(Align, MaxAlign);
   }

lib/CodeGen/PrologEpilogInserter.cpp

computeFreeStackSlots(MachineFrameInfo &MFI, bool StackGrowsDown,
   if (FixedCSEnd > std::numeric_limits<int>::max())
     return;
 
   StackBytesFree.resize(FixedCSEnd, true);
 
   SmallVector<int, 16> AllocatedFrameSlots;
   // Add fixed objects.
   for (int i = MFI.getObjectIndexBegin(); i != 0; ++i)
+    // StackSlot scavenging is only implemented for the default stack.
+    if (MFI.getStackID(i) == 0)
       AllocatedFrameSlots.push_back(i);
   // Add callee-save objects.
   for (int i = MinCSFrameIndex; i <= (int)MaxCSFrameIndex; ++i)
+    if (MFI.getStackID(i) == 0)
       AllocatedFrameSlots.push_back(i);
 
   for (int i : AllocatedFrameSlots) {
     // These are converted from int64_t, but they should always fit in int
     // because of the FixedCSEnd check above.
     int ObjOffset = MFI.getObjectOffset(i);
     int ObjSize = MFI.getObjectSize(i);
     int ObjStart, ObjEnd;
     if (StackGrowsDown) {
if (StackGrowsDown)
     LocalAreaOffset = -LocalAreaOffset;
   assert(LocalAreaOffset >= 0
          && "Local area offset should be in direction of stack growth");
   int64_t Offset = LocalAreaOffset;
 
   // Skew to be applied to alignment.
   unsigned Skew = TFI.getStackAlignmentSkew(MF);
 
+#ifdef EXPENSIVE_CHECKS
+  for (unsigned i = 0, e = MFI.getObjectIndexEnd(); i != e; ++i)
+    if (!MFI.isDeadObjectIndex(i) && MFI.getStackID(i) == 0)
+      assert(MFI.getObjectAlignment(i) <= MFI.getMaxAlignment() &&
+             "MaxAlignment is invalid");
+#endif
+
   // If there are fixed sized objects that are preallocated in the local area,
   // non-fixed objects can't be allocated right at the start of local area.
   // Adjust 'Offset' to point to the end of last fixed sized preallocated
   // object.
   for (int i = MFI.getObjectIndexBegin(); i != 0; ++i) {
+    if (MFI.getStackID(i)) // Only allocate objects on the default stack.
+      continue;
+
     int64_t FixedOff;
     if (StackGrowsDown) {
       // The maximum distance from the stack pointer is at lower address of
       // the object -- which is given by offset. For down growing stack
       // the offset is negative, so we negate the offset to get the distance.
       FixedOff = -MFI.getObjectOffset(i);
     } else {
       // The maximum distance from the start pointer is at the upper
       // address of the object.
       FixedOff = MFI.getObjectOffset(i) + MFI.getObjectSize(i);
     }
     if (FixedOff > Offset) Offset = FixedOff;
   }
 
   // First assign frame offsets to stack objects that are used to spill
   // callee saved registers.
   if (StackGrowsDown) {
     for (unsigned i = MinCSFrameIndex; i <= MaxCSFrameIndex; ++i) {
+      if (MFI.getStackID(i)) // Only allocate objects on the default stack.
+        continue;
+
       // If the stack grows down, we need to add the size to find the lowest
       // address of the object.
       Offset += MFI.getObjectSize(i);
 
       unsigned Align = MFI.getObjectAlignment(i);
       // Adjust to alignment boundary
       Offset = alignTo(Offset, Align, Skew);
 
       LLVM_DEBUG(dbgs() << "alloc FI(" << i << ") at SP[" << -Offset << "]\n");
       MFI.setObjectOffset(i, -Offset);        // Set the computed offset
     }
   } else if (MaxCSFrameIndex >= MinCSFrameIndex) {
     // Be careful about underflow in comparisons agains MinCSFrameIndex.
     for (unsigned i = MaxCSFrameIndex; i != MinCSFrameIndex - 1; --i) {
+      if (MFI.getStackID(i)) // Only allocate objects on the default stack.
+        continue;
+
       if (MFI.isDeadObjectIndex(i))
         continue;
 
       unsigned Align = MFI.getObjectAlignment(i);
       // Adjust to alignment boundary
       Offset = alignTo(Offset, Align, Skew);
 
       LLVM_DEBUG(dbgs() << "alloc FI(" << i << ") at SP[" << Offset << "]\n");
for (unsigned i = 0, e = MFI.getObjectIndexEnd(); i != e; ++i) {
         continue;
       if (RS && RS->isScavengingFrameIndex((int)i))
         continue;
       if (MFI.isDeadObjectIndex(i))
         continue;
       if (MFI.getStackProtectorIndex() == (int)i ||
           EHRegNodeFrameIndex == (int)i)
         continue;
+      if (MFI.getStackID(i)) // Only allocate objects on the default stack.
+        continue;
 
       switch (MFI.getObjectSSPLayout(i)) {
       case MachineFrameInfo::SSPLK_None:
         continue;
       case MachineFrameInfo::SSPLK_SmallArray:
         SmallArrayObjs.insert(i);
         continue;
       case MachineFrameInfo::SSPLK_AddrOf:
if (RS && RS->isScavengingFrameIndex((int)i))
       continue;
     if (MFI.isDeadObjectIndex(i))
       continue;
     if (MFI.getStackProtectorIndex() == (int)i ||
         EHRegNodeFrameIndex == (int)i)
       continue;
     if (ProtectedObjs.count(i))
       continue;
+    if (MFI.getStackID(i)) // Only allocate objects on the default stack.
+      continue;
 
     // Add the objects that we need to allocate to our working set.
     ObjectsToAllocate.push_back(i);
   }
 
   // Allocate the EH registration node first if one is present.
   if (EHRegNodeFrameIndex != std::numeric_limits<int>::max())
     AdjustStackOffset(MFI, EHRegNodeFrameIndex, StackGrowsDown, Offset,

lib/Target/AMDGPU/SIFrameLowering.cpp

for (MachineBasicBlock &MBB : MF) {
             bool Spilled = TRI.eliminateSGPRToVGPRSpillFrameIndex(MI, FI, RS);
             (void)Spilled;
             assert(Spilled && "failed to spill SGPR to VGPR when allocated");
           } else
             AllSGPRSpilledToVGPRs = false;
         }
       }
     }
+  }
 
-    FuncInfo->removeSGPRToVGPRFrameIndices(MFI);
-  }
 
   // FIXME: The other checks should be redundant with allStackObjectsAreDead,
   // but currently hasNonSpillStackObjects is set only from source
   // allocas. Stack temps produced from legalization are not counted currently.
   if (FuncInfo->hasNonSpillStackObjects() || FuncInfo->hasSpilledVGPRs() ||
       !AllSGPRSpilledToVGPRs || !allStackObjectsAreDead(MFI)) {
     assert(RS && "RegScavenger required if spilling");
 

lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

bool SIMachineFunctionInfo::allocateSGPRSpillToVGPR(MachineFunction &MF,
   }
 
   return true;
 }
 
 void SIMachineFunctionInfo::removeSGPRToVGPRFrameIndices(MachineFrameInfo &MFI) {
   for (auto &R : SGPRToVGPRSpills)
     MFI.RemoveStackObject(R.first);
+  // All other SPGRs must be allocated on the default stack, so reset
+  // the stack ID.
+  for (unsigned i = MFI.getObjectIndexBegin(), e = MFI.getObjectIndexEnd();
+       i != e; ++i)
+    MFI.setStackID(i, 0);
 }
 
 
 /// \returns VGPR used for \p Dim' work item ID.
 unsigned SIMachineFunctionInfo::getWorkItemIDVGPR(unsigned Dim) const {
   switch (Dim) {
   case 0:
     assert(hasWorkItemIDX());

test/CodeGen/AArch64/stack-id-pei-alloc.mir

+# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass=prologepilog %s -o - | FileCheck %s
+...
+# Ensure that objects with StackID > 0 are not allocated on the default stack
+# (will not be allocated an offset) and are not considered in the calculation of
+# the StackSize.
+# CHECK:      name: test_allocate
+# CHECK:      stackSize: 16
+# CHECK:      stack:
+# CHECK:        id: 0, name: '', type: default, offset: -8, size: 8, alignment: 8,
+# CHECK-NEXT:   stack-id: 0
+# CHECK:        id: 1, name: '', type: default, offset: -16, size: 8, alignment: 8,
+# CHECK-NEXT:   stack-id: 0
+# CHECK:        id: 2, name: '', type: default, offset: 0, size: 8, alignment: 8,
+# CHECK-NEXT:   stack-id: 42
+name:            test_allocate
+frameInfo:
+  maxAlignment:  16
+stack:
+  - { id: 0, stack-id: 0, size: 8, alignment: 8, offset: 0 }
+  - { id: 1, stack-id: 0, size: 8, alignment: 8, offset: 0 }
+  - { id: 2, stack-id: 42, size: 8, alignment: 8, offset: 0 }
+body:             |
+  bb.0.entry:
+    RET_ReallyLR
+---
+...
+# Ensure MaxAlignment becomes '32' even though we also have an object
+# with alignment of 64. MaxAlignment only pertains to the default stack
+# (StackID 0), so objects associated with a different StackID should
+# not be considered.
+#
+# CHECK: name: test_maxalign
+# CHECK: maxAlignment: 32
+name:            test_maxalign
+frameInfo:
+  maxAlignment:  16
+stack:
+  - { id: 0, stack-id: 0, size: 16, alignment: 32 }
+  - { id: 1, stack-id: 42, size: 16, alignment: 64 }
+body:             |
+  bb.0.entry:
+    RET_ReallyLR
+---
+...
+# CHECK: name: test_maxalign_fixedstack
+# CHECK: maxAlignment: 32
+name:            test_maxalign_fixedstack
+frameInfo:
+  maxAlignment:  16
+fixedStack:
+  - { id: 0, stack-id: 0, size: 16, alignment: 32 }
+  - { id: 1, stack-id: 42, size: 16, alignment: 64 }
+body:             |
+  bb.0.entry:
+    RET_ReallyLR
+---

test/CodeGen/AArch64/stack-id-stackslot-scavenging.mir

+# RUN: llc -mtriple=aarch64-none-linux-gnu -start-before=greedy -stop-after=prologepilog %s -o - | FileCheck %s
+...
+# Ensure that an object with a different stack-id is not allocated to a
+# callee-save slot using stack-slot scavenging. This test saves X28 which
+# creates a hole in the CSR stack region, but it should not be saved to.
+# Instead of saving to SP + 1 (which would be the hole in the region), it
+# should save to SP + 2 (since AArch64 codegen currently does not support
+# (and thus allocate) objects with a stack-id > 0).
+name:            test_no_stackslot_scavenging
+# CHECK: name: test_no_stackslot_scavenging
+# CHECK: STRXui $x0, $sp, 2
+tracksRegLiveness: true
+frameInfo:
+  maxAlignment:  16
+stack:
+  - { id: 0, stack-id: 42, size: 8, alignment: 8 }
+body:             |
+  bb.0.entry:
+    liveins: $x0
+    STRXui $x0, %stack.0, 0
+    ; Force preserve a CSR to create a hole in the CSR stack region.
+    $x28 = IMPLICIT_DEF
+    RET_ReallyLR
+---