Enforce StackID definition in PEI

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.

Reviewers: arsenm, thegameg, MatzeB

Reviewed By: arsenm

Differential Revision: https://reviews.llvm.org/D60062

llvm-svn: 357460

llvm/include/llvm/CodeGen/MachineFrameInfo.h

@@ -470,7 +470,10 @@ public:
     assert(unsigned(ObjectIdx+NumFixedObjects) < Objects.size() &&
            "Invalid Object Idx!");
     Objects[ObjectIdx+NumFixedObjects].Alignment = Align;
-    ensureMaxAlignment(Align);
+
+    // Only ensure max alignment for the default stack.
+    if (getStackID(ObjectIdx) == 0)
+      ensureMaxAlignment(Align);
   }
 
   /// Return the underlying Alloca of the specified
@@ -697,6 +700,8 @@ public:
     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.

llvm/lib/CodeGen/MIRParser/MIRParser.cpp

@@ -620,8 +620,8 @@ bool MIRParserImpl::initializeFrameInfo(PerFunctionMIParsingState &PFS,
                                         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)
@@ -654,9 +654,9 @@ bool MIRParserImpl::initializeFrameInfo(PerFunctionMIParsingState &PFS,
     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)

llvm/lib/CodeGen/MachineFrameInfo.cpp

@@ -56,7 +56,8 @@ int MachineFrameInfo::CreateStackObject(uint64_t Size, unsigned Alignment,
                                 !IsSpillSlot, StackID));
   int Index = (int)Objects.size() - NumFixedObjects - 1;
   assert(Index >= 0 && "Bad frame index!");
-  ensureMaxAlignment(Alignment);
+  if (StackID == 0)
+    ensureMaxAlignment(Alignment);
   return Index;
 }
 
@@ -141,11 +142,15 @@ unsigned MachineFrameInfo::estimateStackSize(const MachineFunction &MF) const {
   // 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);

llvm/lib/CodeGen/PrologEpilogInserter.cpp

@@ -661,10 +661,13 @@ computeFreeStackSlots(MachineFrameInfo &MFI, bool StackGrowsDown,
   SmallVector<int, 16> AllocatedFrameSlots;
   // Add fixed objects.
   for (int i = MFI.getObjectIndexBegin(); i != 0; ++i)
-    AllocatedFrameSlots.push_back(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)
-    AllocatedFrameSlots.push_back(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
@@ -786,11 +789,21 @@ void PEI::calculateFrameObjectOffsets(MachineFunction &MF) {
   // 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
@@ -809,6 +822,9 @@ void PEI::calculateFrameObjectOffsets(MachineFunction &MF) {
   // 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);
@@ -823,6 +839,9 @@ void PEI::calculateFrameObjectOffsets(MachineFunction &MF) {
   } 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;
 
@@ -913,6 +932,8 @@ void PEI::calculateFrameObjectOffsets(MachineFunction &MF) {
       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:
@@ -956,6 +977,8 @@ void PEI::calculateFrameObjectOffsets(MachineFunction &MF) {
       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);

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

@@ -699,10 +699,10 @@ void SIFrameLowering::processFunctionBeforeFrameFinalized(
         }
       }
     }
-
-    FuncInfo->removeSGPRToVGPRFrameIndices(MFI);
   }
 
+  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.

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

@@ -293,6 +293,11 @@ bool SIMachineFunctionInfo::allocateSGPRSpillToVGPR(MachineFunction &MF,
 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);
 }
 
 

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

@@ -0,0 +1,56 @@
+# 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
+---

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

@@ -0,0 +1,24 @@
+# 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
+---