[PEI, AArch64] Use empty spaces in stack area for local stack slot allocation.

[gberry]

Summary

If the target requests it, use emptry spaces in the fixed and
callee-save stack area to allocate local stack objects.

AArch64: Change last callee-save reg stack object alignment instead of
size to leave a gap to take advantage of above change.

[qcolombet]

Hi Geoff,

Couple of questions:

• Where are those empty slots? E.g., are they padding because of alignment, something else?
• Isn’t stack coloring doing already what you want?

Cheers,
-Quentin

[gberry]

Hi Quentin,

• In the case of AArch64, yes the empty slots are there because of alignment, specifically of the callee-save stack area. The stack pointer must always be 16-byte aligned on aarch64, so when allocating a callee-save area for say, 3 64-bit registers, there will be a 64-bit hole left. This change uses that hole to allocate a spill slot/alloca when possible.

• My understanding is that stack coloring is solving a different problem (technically two different passes solving similar problems). It tries to compact the stack space used by spill slots/allocas by re-using the same stack slot for different variables when their live ranges don't overlap. They run before PEI, so they don't even see the "holes" in the callee-save area that this patch is trying to eliminate.

[qcolombet]

Hi Geoff,

Thanks for your answer, we are on the same page then :).

I’ll have a look at the patch shortly.

Cheers,
-Quentin

[qcolombet]

lib/CodeGen/PrologEpilogInserter.cpp

593

resize(0) is weird when reading the code. Clear() looks more appropriate.

Also, if all the bytes are free, why do we stop?

601

use the actual type here.
It is important to know if it is a int an unsigned, etc.

766

Add a comment on what information we keep in FixedCSOffset, e.g., offset containing all the fixed sized and callee saved objects.

903

protecter -> protector

903

isn’t being used *and* the target...

906

We should gate that with the optimization level as well.

909

Allocated

910

Period.

913

Period.

922

A comment reminding that ObjOffset is negative would be nice :).

925

ObjStart

932

Could you move all that stuff in a helper function?

test/CodeGen/AArch64/aarch64-dynamic-stack-layout.ll

681

Why is the input of the test case changed?

[gberry]

Update to address Quentin's review feedback.

[gberry]

Hi Quentin,

Thanks for the review. I believe I've addressed all of the issues you raised, if not just let me know.

@t.p.northover Are you okay with the AArch64 specific changes?

lib/CodeGen/PrologEpilogInserter.cpp

925

Wow, that was a particularly bad typo :)

test/CodeGen/AArch64/aarch64-dynamic-stack-layout.ll

681

After my change, %tmp was getting placed in a hole in the Fixed/CS stack area, leading to no local stack space being allocated in the prologue, which ended up hiding the problem this test was checking for without actually fixing it. I made the change smaller in this revision by simply making %tmp bigger so it wouldn't fit in the unused part of the Fixed/CS stack area any more.

[qcolombet]

Thanks Geoff, the update looks good, just one missing thing.
(See my inlined comments.)

FWIW, the AArch64 changes looks good to me as well.

Q.

lib/CodeGen/PrologEpilogInserter.cpp

599

Putting back my comment from the other day: having the actual type (instead of auto) would be nice. At first glance, I have no idea if those are unsigned, int, etc.

622

The empty check is now also checked as part of the none check.
I would suggest to get rid of the empty check.

628

I guess we can get rid of the call to clear.
Maybe that was an optimization when StackBytesFree is not empty to have none being faster?
(I am fine keeping it.)

[gberry]

Thanks Quentin,

I'll address the issues you raised shortly.

lib/CodeGen/PrologEpilogInserter.cpp

599

Sorry, I didn't ignore you :) I think your original comment was on a different part of the original change and I didn't notice these needed fixing too. I'll double check all of the autos just in case.

622

Yeah, this is probably overkill.

628

Yes, my thinking was that later calls to empty()/none() would complete faster.

[gberry]

Update to address Quentin's latest comments.

[gberry]

lib/CodeGen/PrologEpilogInserter.cpp

587

This check and the corresponding change of ObjOffest and ObjSize below to 'int' are the only significant changes in this revision.

[Gerolf]

Since I couldn't spot this in the code: in the scenario where there is one empty reallocate both of them to the empty slot or give up after it reallocated the first object?

Thanks
Gerolf

lib/CodeGen/PrologEpilogInserter.cpp

632

Could this be function like FreeStart = FindSlot(...) -- line 632 - 657

[gberry]

Hi Gerolf,

I'm not sure I understand exactly what you're asking, but this change will allocate multiple stack slots to free space in the fixed/CS area. It will assign later slots to this free area even if earlier slots were not assigned there. It will not assign multiple slots to overlapping regions of the fixed/CS area.

lib/CodeGen/PrologEpilogInserter.cpp

632

I suppose, but I don't think it would make the loop any clearer since I don't think it would simplify the control flow of the loop. Also, this function isn't really doing much else.

[Gerolf]

Thanks, Geoff. I take back my previous comment. I was worried that when A and B are allocated at stack offset x and this optimization changes x, it would only change it for A but not for B. But looking deeper into the code this is not the case.

lib/CodeGen/PrologEpilogInserter.cpp

632

I agree with your functionality argument. A function would just improve readability.

[qcolombet]

Hi Geoff,

LGTM.

Did not mark it as accepted as you asked Tim to check the AArch64 part + I am not sure what you'd agreed with Gerolf :).

Cheers,
-Quentin

lib/CodeGen/PrologEpilogInserter.cpp

599

No worries, I figured that was what happened :). Phab can be quite confusing after a few revisions.

[gberry]

I'm going to go ahead and commit this based on Quentin's review. @t.p.northover if you do take a look and have any objections, just let me know and I'll resolve them.

[mcrosier]

Sounds reasonable to me.

include/llvm/Target/TargetFrameLowering.h

virtual bool targetHandlesStackFrameRounding() const {
     return false;
   }
 
   /// Returns true if the target will correctly handle shrink wrapping.
   virtual bool enableShrinkWrapping(const MachineFunction &MF) const {
     return false;
   }
 
+  /// Returns true if the stack slot holes in the fixed and callee-save stack
+  /// area should be used when allocating other stack locations to reduce stack
+  /// size.
+  virtual bool enableStackSlotScavenging(const MachineFunction &MF) const {
+    return false;
+  }
+
   /// emitProlog/emitEpilog - These methods insert prolog and epilog code into
   /// the function.
   virtual void emitPrologue(MachineFunction &MF,
                             MachineBasicBlock &MBB) const = 0;
   virtual void emitEpilogue(MachineFunction &MF,
                             MachineBasicBlock &MBB) const = 0;
 
   /// Replace a StackProbe stub (if any) with the actual probe code inline

lib/CodeGen/PrologEpilogInserter.cpp

if (StackGrowsDown) {
     MFI->setObjectOffset(FrameIdx, -Offset); // Set the computed offset
   } else {
     DEBUG(dbgs() << "alloc FI(" << FrameIdx << ") at SP[" << Offset << "]\n");
     MFI->setObjectOffset(FrameIdx, Offset);
     Offset += MFI->getObjectSize(FrameIdx);
   }
 }
 
+/// Compute which bytes of fixed and callee-save stack area are unused and keep
+/// track of them in StackBytesFree.
+///
+static inline void
+computeFreeStackSlots(MachineFrameInfo *MFI, bool StackGrowsDown,
+                      unsigned MinCSFrameIndex, unsigned MaxCSFrameIndex,
+                      int64_t FixedCSEnd, BitVector &StackBytesFree) {
+  StackBytesFree.resize(FixedCSEnd, true);
+
+  SmallVector<int, 16> AllocatedFrameSlots;
+  // Add fixed objects.
+  for (int i = MFI->getObjectIndexBegin(); i != 0; ++i)
+    AllocatedFrameSlots.push_back(i);
+  // Add callee-save objects.
+  for (int i = MinCSFrameIndex; i <= (int)MaxCSFrameIndex; ++i)
+    AllocatedFrameSlots.push_back(i);
+
+  for (int i : AllocatedFrameSlots) {
+    auto ObjOffset = MFI->getObjectOffset(i);
+    auto ObjSize = MFI->getObjectSize(i);
+    int ObjStart, ObjEnd;
+    if (StackGrowsDown) {
+      // ObjOffset is negative when StackGrowsDown is true.
+      ObjStart = -ObjOffset - ObjSize;
+      ObjEnd = -ObjOffset;
+    } else {
+      ObjStart = ObjOffset;
+      ObjEnd = ObjOffset + ObjSize;
+    }
+    // Ignore fixed holes that are in the previous stack frame.
+    if (ObjEnd > 0)
+      StackBytesFree.reset(ObjStart, ObjEnd);
+  }
+}
+
+/// Assign frame object to an unused portion of the stack in the fixed stack
+/// object range.  Return true if the allocation was successful.
+///
+static inline bool scavengeStackSlot(MachineFrameInfo *MFI, int FrameIdx,
+                                     bool StackGrowsDown, unsigned MaxAlign,
+                                     BitVector &StackBytesFree) {
+  if (StackBytesFree.empty())
+    return false;
+
+  if (MFI->isVariableSizedObjectIndex(FrameIdx))
+    return false;
+
+  if (StackBytesFree.none()) {
+    StackBytesFree.clear();
+    return false;
+  }
+
+  unsigned ObjAlign = MFI->getObjectAlignment(FrameIdx);
+  if (ObjAlign > MaxAlign)
+    return false;
+
+  int64_t ObjSize = MFI->getObjectSize(FrameIdx);
+  int FreeStart;
+  for (FreeStart = StackBytesFree.find_first(); FreeStart != -1;
+       FreeStart = StackBytesFree.find_next(FreeStart)) {
+
+    // Check that free space has suitable alignment.
+    unsigned ObjStart = StackGrowsDown ? FreeStart + ObjSize : FreeStart;
+    if (alignTo(ObjStart, ObjAlign) != ObjStart)
+      continue;
+
+    if (FreeStart + ObjSize > StackBytesFree.size())
+      return false;
+
+    bool AllBytesFree = true;
+    for (unsigned Byte = 0; Byte < ObjSize; ++Byte)
+      if (!StackBytesFree.test(FreeStart + Byte)) {
+        AllBytesFree = false;
+        break;
+      }
+    if (AllBytesFree)
+      break;
+  }
+
+  if (FreeStart == -1)
+    return false;
+
+  if (StackGrowsDown) {
+    int ObjStart = -(FreeStart + ObjSize);
+    DEBUG(dbgs() << "alloc FI(" << FrameIdx << ") scavenged at SP[" << ObjStart
+                 << "]\n");
+    MFI->setObjectOffset(FrameIdx, ObjStart);
+  } else {
+    DEBUG(dbgs() << "alloc FI(" << FrameIdx << ") scavenged at SP[" << FreeStart
+                 << "]\n");
+    MFI->setObjectOffset(FrameIdx, FreeStart);
+  }
+
+  StackBytesFree.reset(FreeStart, FreeStart + ObjSize);
+  return true;
+}
+
 /// AssignProtectedObjSet - Helper function to assign large stack objects (i.e.,
 /// those required to be close to the Stack Protector) to stack offsets.
 static void
 AssignProtectedObjSet(const StackObjSet &UnassignedObjs,
                       SmallSet<int, 16> &ProtectedObjs,
                       MachineFrameInfo *MFI, bool StackGrowsDown,
                       int64_t &Offset, unsigned &MaxAlign, unsigned Skew) {
 
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(Fn);
 
   // 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.
-  // We currently don't support filling in holes in between fixed sized
-  // objects, so we adjust 'Offset' to point to the end of last fixed sized
-  // preallocated object.
+  // Adjust 'Offset' to point to the end of last fixed sized preallocated
+  // object.
   for (int i = MFI->getObjectIndexBegin(); i != 0; ++i) {
     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 {
for (unsigned i = MaxCSFI; i >= MinCSFI; --i) {
       Offset = alignTo(Offset, Align, Skew);
 
       DEBUG(dbgs() << "alloc FI(" << i << ") at SP[" << Offset << "]\n");
       MFI->setObjectOffset(i, Offset);
       Offset += MFI->getObjectSize(i);
     }
   }
 
+  // FixedCSEnd is the stack offset to the end of the fixed and callee-save
+  // stack area.
+  int64_t FixedCSEnd = Offset;
   unsigned MaxAlign = MFI->getMaxAlignment();
 
   // Make sure the special register scavenging spill slot is closest to the
   // incoming stack pointer if a frame pointer is required and is closer
   // to the incoming rather than the final stack pointer.
   const TargetRegisterInfo *RegInfo = Fn.getSubtarget().getRegisterInfo();
   bool EarlyScavengingSlots = (TFI.hasFP(Fn) &&
                                TFI.isFPCloseToIncomingSP() &&
void PEI::calculateFrameObjectOffsets(MachineFunction &Fn) {
   if (EHRegNodeFrameIndex != INT_MAX)
     AdjustStackOffset(MFI, EHRegNodeFrameIndex, StackGrowsDown, Offset,
                       MaxAlign, Skew);
 
   // Give the targets a chance to order the objects the way they like it.
   if (Fn.getTarget().getOptLevel() != CodeGenOpt::None &&
       Fn.getTarget().Options.StackSymbolOrdering)
     TFI.orderFrameObjects(Fn, ObjectsToAllocate);
-  
+  // Keep track of which bytes in the fixed and callee-save range are used so we
+  // can use the holes when allocating later stack objects.  Only do this if
+  // stack protector isn't being used and the target requests it and we're
+  // optimizing.
+  BitVector StackBytesFree;
+  if (!ObjectsToAllocate.empty() &&
+      Fn.getTarget().getOptLevel() != CodeGenOpt::None &&
+      MFI->getStackProtectorIndex() < 0 && TFI.enableStackSlotScavenging(Fn))
+    computeFreeStackSlots(MFI, StackGrowsDown, MinCSFrameIndex, MaxCSFrameIndex,
+                          FixedCSEnd, StackBytesFree);
+
   // Now walk the objects and actually assign base offsets to them.
   for (auto &Object : ObjectsToAllocate)
+    if (!scavengeStackSlot(MFI, Object, StackGrowsDown, MaxAlign,
+                           StackBytesFree))
       AdjustStackOffset(MFI, Object, StackGrowsDown, Offset, MaxAlign, Skew);
 
   // Make sure the special register scavenging spill slot is closest to the
   // stack pointer.
   if (RS && !EarlyScavengingSlots) {
     SmallVector<int, 2> SFIs;
     RS->getScavengingFrameIndices(SFIs);
     for (SmallVectorImpl<int>::iterator I = SFIs.begin(),
            IE = SFIs.end(); I != IE; ++I)
       AdjustStackOffset(MFI, *I, StackGrowsDown, Offset, MaxAlign, Skew);
   }
 
   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
     // overall stack size.
     if (MFI->adjustsStack() && TFI.hasReservedCallFrame(Fn))
       Offset += MFI->getMaxCallFrameSize();
 
     // Round up the size to a multiple of the alignment.  If the function has
     // any calls or alloca's, align to the target's StackAlignment value to
     // ensure that the callee's frame or the alloca data is suitably aligned;
     // otherwise, for leaf functions, align to the TransientStackAlignment
     // value.
     unsigned StackAlign;

lib/Target/AArch64/AArch64FrameLowering.h

public:
   void determineCalleeSaves(MachineFunction &MF, BitVector &SavedRegs,
                             RegScavenger *RS) const override;
 
   /// Returns true if the target will correctly handle shrink wrapping.
   bool enableShrinkWrapping(const MachineFunction &MF) const override {
     return true;
   }
 
+  bool enableStackSlotScavenging(const MachineFunction &MF) const override;
+
 private:
   bool shouldCombineCSRLocalStackBump(MachineFunction &MF,
                                       unsigned StackBumpBytes) const;
 };
 
 } // End llvm namespace
 
 #endif

lib/Target/AArch64/AArch64FrameLowering.cpp

for (unsigned i = 0; i < Count; ++i) {
 
     RPI.FrameIdx = CSI[i].getFrameIdx();
 
     if (Count * 8 != AFI->getCalleeSavedStackSize() && !RPI.isPaired()) {
       // Round up size of non-pair to pair size if we need to pad the
       // callee-save area to ensure 16-byte alignment.
       Offset -= 16;
       assert(MFI->getObjectAlignment(RPI.FrameIdx) <= 16);
-      MFI->setObjectSize(RPI.FrameIdx, 16);
+      MFI->setObjectAlignment(RPI.FrameIdx, 16);
+      AFI->setCalleeSaveStackHasFreeSpace(true);
     } else
       Offset -= RPI.isPaired() ? 16 : 8;
     assert(Offset % 8 == 0);
     RPI.Offset = Offset / 8;
     assert((RPI.Offset >= -64 && RPI.Offset <= 63) &&
            "Offset out of bounds for LDP/STP immediate");
 
     RegPairs.push_back(RPI);
if (!ExtraCSSpill) {
                    << " as the emergency spill slot.\n");
     }
   }
 
   // Round up to register pair alignment to avoid additional SP adjustment
   // instructions.
   AFI->setCalleeSavedStackSize(alignTo(8 * NumRegsSpilled, 16));
 }
+
+bool AArch64FrameLowering::enableStackSlotScavenging(
+    const MachineFunction &MF) const {
+  const AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
+  return AFI->hasCalleeSaveStackFreeSpace();
+}

lib/Target/AArch64/AArch64MachineFunctionInfo.h

class AArch64FunctionInfo : public MachineFunctionInfo {
   /// True if this function has a subset of CSRs that is handled explicitly via
   /// copies.
   bool IsSplitCSR;
 
   /// True when the stack gets realigned dynamically because the size of stack
   /// frame is unknown at compile time. e.g., in case of VLAs.
   bool StackRealigned;
 
+  /// True when the callee-save stack area has unused gaps that may be used for
+  /// other stack allocations.
+  bool CalleeSaveStackHasFreeSpace;
+
 public:
   AArch64FunctionInfo()
       : BytesInStackArgArea(0), ArgumentStackToRestore(0), HasStackFrame(false),
         NumLocalDynamicTLSAccesses(0), VarArgsStackIndex(0), VarArgsGPRIndex(0),
         VarArgsGPRSize(0), VarArgsFPRIndex(0), VarArgsFPRSize(0),
-        IsSplitCSR(false), StackRealigned(false) {}
+        IsSplitCSR(false), StackRealigned(false),
+        CalleeSaveStackHasFreeSpace(false) {}
 
   explicit AArch64FunctionInfo(MachineFunction &MF)
       : BytesInStackArgArea(0), ArgumentStackToRestore(0), HasStackFrame(false),
         NumLocalDynamicTLSAccesses(0), VarArgsStackIndex(0), VarArgsGPRIndex(0),
         VarArgsGPRSize(0), VarArgsFPRIndex(0), VarArgsFPRSize(0),
-        IsSplitCSR(false), StackRealigned(false) {
+        IsSplitCSR(false), StackRealigned(false),
+        CalleeSaveStackHasFreeSpace(false) {
     (void)MF;
   }
 
   unsigned getBytesInStackArgArea() const { return BytesInStackArgArea; }
   void setBytesInStackArgArea(unsigned bytes) { BytesInStackArgArea = bytes; }
 
   unsigned getArgumentStackToRestore() const { return ArgumentStackToRestore; }
   void setArgumentStackToRestore(unsigned bytes) {
     ArgumentStackToRestore = bytes;
   }
 
   bool hasStackFrame() const { return HasStackFrame; }
   void setHasStackFrame(bool s) { HasStackFrame = s; }
 
   bool isStackRealigned() const { return StackRealigned; }
   void setStackRealigned(bool s) { StackRealigned = s; }
 
+  bool hasCalleeSaveStackFreeSpace() const {
+    return CalleeSaveStackHasFreeSpace;
+  }
+  void setCalleeSaveStackHasFreeSpace(bool s) {
+    CalleeSaveStackHasFreeSpace = s;
+  }
+
   bool isSplitCSR() const { return IsSplitCSR; }
   void setIsSplitCSR(bool s) { IsSplitCSR = s; }
 
   void setLocalStackSize(unsigned Size) { LocalStackSize = Size; }
   unsigned getLocalStackSize() const { return LocalStackSize; }
 
   void setCalleeSavedStackSize(unsigned Size) { CalleeSavedStackSize = Size; }
   unsigned getCalleeSavedStackSize() const { return CalleeSavedStackSize; }

test/CodeGen/AArch64/aarch64-dynamic-stack-layout.ll

 ; CHECK:  sub  [[REG:x[01-9]+]], sp, #64
 ; CHECK:  and  sp, [[REG]], #0xffffffffffffffe0
 ; CHECK:  .[[LABEL]]:
 ; CHECK:  ret
 
 
 define void @realign_conditional2(i1 %b) {
 entry:
-  %tmp = alloca i8, i32 4
+  %tmp = alloca i8, i32 16
   br i1 %b, label %bb0, label %bb1
 
 bb0:
   %MyAlloca = alloca i8, i64 64, align 32
   br label %bb1
 
 bb1:
   ret void
 }
 
 ; CHECK-LABEL: realign_conditional2
 ; Extra realignment in the prologue (performance issue).

test/CodeGen/AArch64/arm64-hello.ll

 ; RUN: llc < %s -mtriple=arm64-apple-ios7.0 -disable-post-ra -disable-fp-elim | FileCheck %s
 ; RUN: llc < %s -mtriple=arm64-linux-gnu -disable-post-ra | FileCheck %s --check-prefix=CHECK-LINUX
 
 ; CHECK-LABEL: main:
 ; CHECK:  sub sp, sp, #32
 ; CHECK-NEXT: stp x29, x30, [sp, #16]
 ; CHECK-NEXT: add x29, sp, #16
 ; CHECK-NEXT: stur  wzr, [x29, #-4]
 ; CHECK:  adrp  x0, L_.str@PAGE
 ; CHECK:  add x0, x0, L_.str@PAGEOFF
 ; CHECK-NEXT: bl  _puts
 ; CHECK-NEXT: ldp x29, x30, [sp, #16]
 ; CHECK-NEXT: add sp, sp, #32
 ; CHECK-NEXT: ret
 
 ; CHECK-LINUX-LABEL: main:
-; CHECK-LINUX:  sub sp, sp, #32
-; CHECK-LINUX-NEXT: str x30, [sp, #16]
+; CHECK-LINUX:  str x30, [sp, #-16]!
 ; CHECK-LINUX-NEXT: str wzr, [sp, #12]
 ; CHECK-LINUX:  adrp  x0, .L.str
 ; CHECK-LINUX:  add x0, x0, :lo12:.L.str
 ; CHECK-LINUX-NEXT: bl  puts
-; CHECK-LINUX-NEXT: ldr x30, [sp, #16]
-; CHECK-LINUX-NEXT: add sp, sp, #32
+; CHECK-LINUX-NEXT: ldr x30, [sp], #16
 ; CHECK-LINUX-NEXT: ret
 
 @.str = private unnamed_addr constant [7 x i8] c"hello\0A\00"
 
 define i32 @main() nounwind ssp {
 entry:
   %retval = alloca i32, align 4
   store i32 0, i32* %retval
   %call = call i32 @puts(i8* getelementptr inbounds ([7 x i8], [7 x i8]* @.str, i32 0, i32 0))
   ret i32 %call
 }
 
 declare i32 @puts(i8*)