[PEI] Don't re-allocate a pre-allocated stack protector slot

[thegameg]

Summary

The LocalStackSlotPass pre-allocates a stack protector and makes sure that it comes before the local variables on the stack.

We need to make sure that later during PEI we don't re-allocate a new stack protector slot. If that happens, the new stack protector slot will end up being after the local variables that it should be protecting.

Therefore, we would have two slots assigned for two different stack protectors, one at the top of the stack, and one at the bottom. Since PEI will overwrite the assigned slot for the stack protector, the load that is used to compare the value of the stack protector will use the slot assigned by PEI, which is wrong.

For this, we need to check if the object is pre-allocated, and re-use that pre-allocated slot.

[efriedma]

llvm/lib/CodeGen/LocalStackSlotAllocation.cpp

202–211

Is this change a no-op?

llvm/lib/CodeGen/PrologEpilogInserter.cpp

932

We obviously don't want to re-allocate a slot for the stack protector, but what about the other stack objects? Do we guarantee that LocalStackSlotPass will allocate all obejcts where getObjectSSPLayout() is not SSPLK_None? If we do, we can just simplify this check to if (MFI.getUseLocalStackAllocationBlock()). Otherwise, we should still run this code for all the all the non-stack-protector frame indexes.

[thegameg]

llvm/lib/CodeGen/LocalStackSlotAllocation.cpp

202–211

Yes, it just adds the assert.

llvm/lib/CodeGen/PrologEpilogInserter.cpp

932

LocalStackSlotPass works in two steps:

1. Assign offsets to all stack objects except objects used for saving callee saved registers in MFI.LocalFrameObjects . These offsets will be potentially used by PEI by sliding them to an actual offset after allocating slots for other things like CSRs.
2. Replace FrameIndex operands with a vreg operand (itself defined by a FI + offset) + offset operand

1) always runs, but if 2) does not rewrite any FI operands, it does not set MFI.UseLocalStackAllocationBlock, and PEI will not use anything from MFI.LocalFrameObjects, and it will instead re-allocate everything in MFI.Objects as if the pass did not even run. In this specific case, 1) runs and sets MFI.StackProtectorIndex (which is "global"), but only assigns the offsets of local variables in MFI.LocalFrameObjects.

AdjustStackOffset from LocalStackSlotPass uses

MFI.mapLocalFrameObject(FrameIdx, LocalOffset);

AdjustStackOffset from PEI uses

MFI.setObjectOffset(FrameIdx, Offset)

[aadg]

llvm/lib/CodeGen/LocalStackSlotAllocation.cpp

202–211

Yes, this change is a no-op.

[efriedma]

llvm/lib/CodeGen/PrologEpilogInserter.cpp

932

I guess I commented too fast with my previous review comment; the check here doesn't simplify quite the way I suggested. It only simplifies to if (MFI.hasStackProtectorIndex() && !MFI.getUseLocalStackAllocationBlock()). And then we can get rid of the if (MFI.isObjectPreAllocated(i) && MFI.getUseLocalStackAllocationBlock()) check, because getUseLocalStackAllocationBlock() is never true inside the if statement.

The question I was trying to ask was, is the additional guard for the calls to AssignProtectedObjSet a no-op in practice? I don't think you addressed that at all.

[thegameg]

llvm/lib/CodeGen/PrologEpilogInserter.cpp

932

Yes, please ignore my previous answer, that was confusing on my part.

Do we guarantee that LocalStackSlotPass will allocate all obejcts where getObjectSSPLayout() is not SSPLK_None?
I'm not sure we guarantee it, but I believe this is what is happening right now. We could maybe assert if it's not true? FWIW, I think we should guarantee it, since the order actually matters based on the type, and if PEI allocates some of them, the objects won't be sorted as expected.

It only simplifies to if (MFI.hasStackProtectorIndex() && !MFI.getUseLocalStackAllocationBlock()).
Right, PreAllocated should not matter if MFI.UseLocalStackAllocationBlock is not set.

And then we can get rid of the if (MFI.isObjectPreAllocated(i) && MFI.getUseLocalStackAllocationBlock()) check
In that case, I agree, we could remove the check inside the loop.

is the additional guard for the calls to AssignProtectedObjSet a no-op in practice?
I would expect the sets to be empty if MFI.UseLocalStackAllocationBlock is set (if the assumptions from the first question are true).

[efriedma]

llvm/lib/CodeGen/PrologEpilogInserter.cpp

932

FWIW, I think we should guarantee it, since the order actually matters based on the type, and if PEI allocates some of them, the objects won't be sorted as expected.

This makes sense.

[thegameg]

• Move NFC changes into a separate commit
• In PEI, verify that LocalStackSlotPass assigned the stack protector and the protected stack objects

[john.brawn]

This will fix https://bugs.llvm.org/show_bug.cgi?id=25610, so it would be good if we could have a test for that as well (and to close that bug once this is committed).

This bug also affects the arm and powerpc targets (as they also return true for TargetRegisterInfo::requiresVirtualBaseRegisters), so it would be good to have tests for those as well. An example which affects all three targets (with -fomit-frame-pointer) is

#include <string.h>
#include <stdio.h>

int fn(const char *str) {
  char buffer[65536];
  strcpy(buffer, str);
  puts(buffer);
  return buffer[65535];
}

[efriedma]

LGTM

[thegameg]

• Add common test for ARM, AArch64 and PowerPC in previous commit
• Update tests with the modifications coming from this patch

[thegameg]

Add test from PR25610.

[john.brawn]

Tests look good to me.

[aadg]

LGTM as well.

llvm/lib/CodeGen/LocalStackSlotAllocation.cpp

void LocalStackSlotPass::calculateFrameObjectOffsets(MachineFunction &Fn) {
   bool StackGrowsDown =
     TFI.getStackGrowthDirection() == TargetFrameLowering::StackGrowsDown;
   int64_t Offset = 0;
   unsigned MaxAlign = 0;
 
   // Make sure that the stack protector comes before the local variables on the
   // stack.
   SmallSet<int, 16> ProtectedObjs;
   if (MFI.hasStackProtectorIndex()) {
     int StackProtectorFI = MFI.getStackProtectorIndex();
+
+    // We need to make sure we didn't pre-allocate the stack protector when
+    // doing this.
+    // If we already have a stack protector, this will re-assign it to a slot
+    // that is **not** covering the protected objects.
+    assert(!MFI.isObjectPreAllocated(StackProtectorFI) &&
+           "Stack protector pre-allocated in LocalStackSlotAllocation");
+
     StackObjSet LargeArrayObjs;
     StackObjSet SmallArrayObjs;
     StackObjSet AddrOfObjs;
 
     AdjustStackOffset(MFI, StackProtectorFI, Offset, StackGrowsDown, MaxAlign);
 
     // Assign large stack objects first.
     for (unsigned i = 0, e = MFI.getObjectIndexEnd(); i != e; ++i) {

llvm/lib/CodeGen/PrologEpilogInserter.cpp

#endif
   if (const WinEHFuncInfo *FuncInfo = MF.getWinEHFuncInfo())
     EHRegNodeFrameIndex = FuncInfo->EHRegNodeFrameIndex;
 
   // Make sure that the stack protector comes before the local variables on the
   // stack.
   SmallSet<int, 16> ProtectedObjs;
   if (MFI.hasStackProtectorIndex()) {
     int StackProtectorFI = MFI.getStackProtectorIndex();
     StackObjSet LargeArrayObjs;
     StackObjSet SmallArrayObjs;
     StackObjSet AddrOfObjs;
 
+    // If we need a stack protector, we need to make sure that
+    // LocalStackSlotPass didn't already allocate a slot for it.
+    // If we are told to use the LocalStackAllocationBlock, the stack protector
+    // is expected to be already pre-allocated.
+    if (!MFI.getUseLocalStackAllocationBlock())
       AdjustStackOffset(MFI, StackProtectorFI, StackGrowsDown, Offset, MaxAlign,
                         Skew);
+    else if (!MFI.isObjectPreAllocated(MFI.getStackProtectorIndex()))
+      llvm_unreachable(
+          "Stack protector not pre-allocated by LocalStackSlotPass.");
 
     // Assign large stack objects first.
     for (unsigned i = 0, e = MFI.getObjectIndexEnd(); i != e; ++i) {
       if (MFI.isObjectPreAllocated(i) && MFI.getUseLocalStackAllocationBlock())
         continue;
       if (i >= MinCSFrameIndex && i <= MaxCSFrameIndex)
         continue;
       if (RS && RS->isScavengingFrameIndex((int)i))
for (unsigned i = 0, e = MFI.getObjectIndexEnd(); i != e; ++i) {
         continue;
       case MachineFrameInfo::SSPLK_LargeArray:
         LargeArrayObjs.insert(i);
         continue;
       }
       llvm_unreachable("Unexpected SSPLayoutKind.");
     }
 
+    // We expect **all** the protected stack objects to be pre-allocated by
+    // LocalStackSlotPass. If it turns out that PEI still has to allocate some
+    // of them, we may end up messing up the expected order of the objects.
+    if (MFI.getUseLocalStackAllocationBlock() &&
+        !(LargeArrayObjs.empty() && SmallArrayObjs.empty() &&
+          AddrOfObjs.empty()))
+      llvm_unreachable("Found protected stack objects not pre-allocated by "
+                       "LocalStackSlotPass.");
+
     AssignProtectedObjSet(LargeArrayObjs, ProtectedObjs, MFI, StackGrowsDown,
                           Offset, MaxAlign, Skew);
     AssignProtectedObjSet(SmallArrayObjs, ProtectedObjs, MFI, StackGrowsDown,
                           Offset, MaxAlign, Skew);
     AssignProtectedObjSet(AddrOfObjs, ProtectedObjs, MFI, StackGrowsDown,
                           Offset, MaxAlign, Skew);
   }
 

llvm/test/CodeGen/AArch64/stack-guard-reassign.ll

 ; RUN: llc -O0 --frame-pointer=all -mtriple=aarch64-- -o - %S/../Inputs/stack-guard-reassign.ll | FileCheck %s
 
 ; Verify that the offset assigned to the stack protector is at the top of the
 ; frame, covering the locals.
 ; CHECK-LABEL: fn:
-; CHECK:      add x8, sp, #24
+; CHECK:      sub x8, x29, #24
 ; CHECK-NEXT: adrp x9, __stack_chk_guard
 ; CHECK-NEXT: ldr x9, [x9, :lo12:__stack_chk_guard]
 ; CHECK-NEXT: str x9, [x8]

llvm/test/CodeGen/AArch64/stack-guard-reassign.mir

 # RUN: llc -mtriple=arm64-apple-ios -start-before=localstackalloc -stop-after=prologepilog -o - %s | FileCheck %s
 
 --- |
   @__stack_chk_guard = external global i8*
   define i32 @main(i32, i8**) {
     %StackGuardSlot = alloca i8*
     unreachable
   }
 ...
 ---
 name:            main
 tracksRegLiveness: true
 frameInfo:
-# CHECK: stackSize: 560
+# CHECK: stackSize: 544
   stackProtector:  '%stack.0.StackGuardSlot'
 stack:
   - { id: 0, name: StackGuardSlot, size: 8, alignment: 8, stack-id: default }
 # Verify that the offset assigned to the stack protector is at the top of the
 # frame, covering the locals.
-# CHECK:       - { id: 0, name: StackGuardSlot, type: default, offset: -552, size: 8,
+# CHECK:       - { id: 0, name: StackGuardSlot, type: default, offset: -24, size: 8,
 # CHECK-NEXT:      alignment: 8, stack-id: default, callee-saved-register: '', callee-saved-restored: true,
 # CHECK-NEXT:      local-offset: -8, debug-info-variable: '', debug-info-expression: '',
 # CHECK-NEXT:      debug-info-location: '' }
   - { id: 1, size: 512, alignment: 1, stack-id: default }
   - { id: 2, size: 4, alignment: 4, stack-id: default }
 body:             |
   bb.0:
     %25:gpr64common = LOAD_STACK_GUARD :: (dereferenceable invariant load 8 from @__stack_chk_guard)
     STRXui killed %25, %stack.0.StackGuardSlot, 0 :: (volatile store 8 into %stack.0.StackGuardSlot)
     %28:gpr64 = LDRXui %stack.0.StackGuardSlot, 0 :: (volatile load 8 from %stack.0.StackGuardSlot)
     %29:gpr64common = LOAD_STACK_GUARD :: (dereferenceable invariant load 8 from @__stack_chk_guard)
     RET_ReallyLR implicit undef $w0, implicit killed %28, implicit killed %29
 
 ...

llvm/test/CodeGen/ARM/stack-guard-reassign.ll

 ; RUN: llc -O0 --frame-pointer=none -mtriple=arm-- -o - %S/../Inputs/stack-guard-reassign.ll | FileCheck %s
 
 ; Verify that the offset assigned to the stack protector is at the top of the
 ; frame, covering the locals.
 ; CHECK-LABEL: fn:
-; CHECK:      sub sp, sp, #40
+; CHECK:      sub sp, sp, #32
 ; CHECK-NEXT: sub sp, sp, #65536
-; CHECK-NEXT: add r1, sp, #28
+; CHECK-NEXT: add lr, sp, #65536
+; CHECK-NEXT: add r1, lr, #28
 ; CHECK-NEXT: ldr r2, .LCPI0_0
 ; CHECK-NEXT: ldr r3, [r2]
 ; CHECK-NEXT: str r3, [r1]
-; CHECK-NEXT: str r0, [sp, #32]
 ; CHECK: .LCPI0_0:
 ; CHECK-NEXT: .long __stack_chk_guard

llvm/test/CodeGen/PowerPC/stack-guard-reassign.ll

 ; RUN: llc -O0 --frame-pointer=none -mtriple=powerpc-- -o - %S/../Inputs/stack-guard-reassign.ll | FileCheck %s
 
 ; Verify that the offset assigned to the stack protector is at the top of the
 ; frame, covering the locals.
 ; CHECK-LABEL: fn:
 ; CHECK:      mflr 0
 ; CHECK-NEXT: stw 0, 4(1)
 ; CHECK-NEXT: lis 0, -2
 ; CHECK-NEXT: ori 0, 0, 65488
 ; CHECK-NEXT: stwux 1, 1, 0
 ; CHECK-NEXT: subf 0, 0, 1
-; CHECK-NEXT: addi 4, 1, 36
+; CHECK-NEXT: lis 4, 1
+; CHECK-NEXT: ori 4, 4, 44
+; CHECK-NEXT: add 4, 1, 4
 ; CHECK-NEXT: lis 5, __stack_chk_guard@ha
 ; CHECK-NEXT: lwz 6, __stack_chk_guard@l(5)
 ; CHECK-NEXT: stw 6, 0(4)