Fix AArch64 prologue for empty frame with dynamic allocas

[eugenis]

Summary

Fixes PR23804: assertion failure in emitPrologue in the case of a function with an empty frame and a dynamic alloca that needs stack realignment.
This is a typical case for AddressSanitizer.

[t.p.northover]

I think there's more going on here that needs fixing too. For example if I add a simple "%tmp = alloca i8, i32 4" to your entry block I get the same assertion failure David mentioned in the PR (even with your patch).

To me it looks like NeedsRealignment is misdefined, though that might not be the end of it. Elsewhere we used RegInfo->needsStackRealignment(MF), which looks like it would be an improvement, but doesn't actually fix the issue.

Another possible problem in this area is that if NumBytes != 0, we end up realigning the stack twice: once in the prologue and once at the alloca. That's arguably a separate performance problem though.

Cheers.

Tim.

[eugenis]

In D10826#197173, @t.p.northover wrote:

I think there's more going on here that needs fixing too. For example if I add a simple "%tmp = alloca i8, i32 4" to your entry block I get the same assertion failure David mentioned in the PR (even with your patch).

To me it looks like NeedsRealignment is misdefined, though that might not be the end of it. Elsewhere we used RegInfo->needsStackRealignment(MF), which looks like it would be an improvement, but doesn't actually fix the issue.

Another possible problem in this area is that if NumBytes != 0, we end up realigning the stack twice: once in the prologue and once at the alloca. That's arguably a separate performance problem though.

Good point, thanks. This particular case is easy to fix by reusing the base register for scratchSPReg. This matches register reservation code in processFunctionBeforeCalleeSavedScan.

RegInfo->needsStackRealignment is also true in this case, and it has roughly the same meaning as getMaxAlignment, i.e. it does not take alloca scopes into account. This problem (extra re-alignment) is also present on i686 and x86_64 targets.

[eugenis]

PTAL

[kristof.beyls]

I think the logic of this patch is roughly right, I only have a few minor comments.
After addressing those comments, this LGTM.

lib/Target/AArch64/AArch64FrameLowering.cpp

352

This should become

const bool NeedsRealignment = RegInfo->needsStackRealignment(MF);

That's in line with what other backends do and respects the DRY principle more.

356–358

I think slightly nicer/better logic is

if (RegInfo->hasBasePointer(MF))
  scratchSPReg = RegInfo->getBaseRegister();
else
  scratchSPReg = AArch64::X9;
assert(MF.getRegInfo().isPhysRegUsed(scratchSPReg) &&
       "No scratch register to align SP!");

With the above code, a few small tweaks to the tests in test/CodeGen/AArch64/aarch64-dynamic-stack-layout.ll are needed (to use the base pointer X19 instead of X9 for a few tests). When both X9 and X19 are available as scratch registers it's an arbitrary decision which one to use.

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

1 ↗

(On Diff #28822)

I don't think there is a need to create a separate test file just for pr23804, these test cases can be added to test/CodeGen/AArch64/aarch64-dynamic-stack-layout.ll

[t.p.northover]

lib/Target/AArch64/AArch64FrameLowering.cpp

356–358

Can't we just nick any temporary register we please (I think the original comment miscategorises X9) and make sure it's used?

scratchSPReg = AArch64::X9;
MF.getRegInfo().setPhysRegUsed(scratchSPReg); // Take that!

[eugenis]

lib/Target/AArch64/AArch64FrameLowering.cpp

356–358

The caller in PrologEpilogInserter.cpp does

TFI->processFunctionBeforeCalleeSavedScan(Fn, RS);

// Scan the function for modified callee saved registers and insert spill code
// for any callee saved registers that are modified.
calculateCalleeSavedRegisters(Fn);

...

if (!F->hasFnAttribute(Attribute::Naked))
  insertPrologEpilogCode(Fn);

So it looks like it is too late to add used regisgters at this point (in emitPrologue).

[eugenis]

Addressed reviewer's comments.

[t.p.northover]

lib/Target/AArch64/AArch64FrameLowering.cpp

354–358

But X9 isn't callee-saved (AAPCS says it's temporary), so I don't think we need to worry about stack layouts: it's basically dead at this point.

Unless there's some assertion that *everything* is fixed by now?

[eugenis]

lib/Target/AArch64/AArch64FrameLowering.cpp

354–358

Well, all tests pass.
Would you like me to switch to your variant?

[t.p.northover]

lib/Target/AArch64/AArch64FrameLowering.cpp

354–358

It seems like the simplest variant so far. I'd say so, unless Kristof objects.

[kristof.beyls]

lib/Target/AArch64/AArch64FrameLowering.cpp

354–358

I can't think of a reason not to do the simpler variant (always using X9 as the scratch register) - so no objections from my side.

[eugenis]

PTAL

[kristof.beyls]

I've only got two very minor suggestions for improvements.
With those improvements, LGTM!

lib/Target/AArch64/AArch64FrameLowering.cpp

354–358

I'd keep the comment

// Use the first callee-saved register as a scratch register

otherwise looks ok.

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

525–526

This comment isn't correct anymore in the newest version of the patch. I would just remove the comment.

[eugenis]

updated comments

[kristof.beyls]

LGTM

[eugenis]

Committed as r241943. Thanks for the review!

lib/Target/AArch64/AArch64FrameLowering.cpp

void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
   }
 
   // All of the remaining stack allocations are for locals.
   AFI->setLocalStackSize(NumBytes);
 
   // Allocate space for the rest of the frame.
 
   const unsigned Alignment = MFI->getMaxAlignment();
-  const bool NeedsRealignment = (Alignment > 16);
+  const bool NeedsRealignment = RegInfo->needsStackRealignment(MF);
   unsigned scratchSPReg = AArch64::SP;
-  if (NeedsRealignment) {
-    // Use the first callee-saved register as a scratch register
-    assert(MF.getRegInfo().isPhysRegUsed(AArch64::X9) &&
-           "No scratch register to align SP!");
+  if (NumBytes && NeedsRealignment) {
+    // Use the first callee-saved register as a scratch register.
     scratchSPReg = AArch64::X9;
+    MF.getRegInfo().setPhysRegUsed(scratchSPReg);
   }
 
   // If we're a leaf function, try using the red zone.
   if (NumBytes && !canUseRedZone(MF))
     // FIXME: in the case of dynamic re-alignment, NumBytes doesn't have
     // the correct value here, as NumBytes also includes padding bytes,
     // which shouldn't be counted here.
     emitFrameOffset(MBB, MBBI, DL, scratchSPReg, AArch64::SP, -NumBytes, TII,
                     MachineInstr::FrameSetup);
 
-  assert(!(NeedsRealignment && NumBytes==0) &&
-         "NumBytes should never be 0 when realignment is needed");
-
   if (NumBytes && NeedsRealignment) {
     const unsigned NrBitsToZero = countTrailingZeros(Alignment);
     assert(NrBitsToZero > 1);
     assert(scratchSPReg != AArch64::SP);
 
     // SUB X9, SP, NumBytes
     //   -- X9 is temporary register, so shouldn't contain any live data here,
     //   -- free to use. This is already produced by emitFrameOffset above.

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

 ; CHECK: ldr   w[[VLA:[0-9]+]], [x[[VLASPTMP]]]
 ;   Check epilogue:
 ;     Check that stack pointer get restored from frame pointer.
 ; CHECK: sub  sp, x29, #16
 ; CHECK: ldp  x29, x30, [sp, #16]
 ; CHECK: ldp  x20, x19, [sp], #32
 ; CHECK: ret
 
+
+define void @realign_conditional(i1 %b) {
+entry:
+  br i1 %b, label %bb0, label %bb1
+
+bb0:
+  %MyAlloca = alloca i8, i64 64, align 32
+  br label %bb1
+
+bb1:
+  ret void
+}
+
+; CHECK-LABEL: realign_conditional
+; No realignment in the prologue.
+; CHECK-NOT:  and
+; CHECK-NOT:  0xffffffffffffffe0
+; CHECK:  tbz  {{.*}} .[[LABEL:.*]]
+; Stack is realigned in a non-entry BB.
+; 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
+  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).
+; CHECK:  sub  x9, sp, #32            // =32
+; CHECK:  and  sp, x9, #0xffffffffffffffe0
+; CHECK:  mov   x19, sp
+; CHECK:  tbz  {{.*}} .[[LABEL:.*]]
+; Stack is realigned in a non-entry BB.
+; CHECK:  sub  [[REG:x[01-9]+]], sp, #64
+; CHECK:  and  sp, [[REG]], #0xffffffffffffffe0
+; CHECK:  .[[LABEL]]:
+; CHECK:  ret
+
 attributes #0 = { "less-precise-fpmad"="false" "no-frame-pointer-elim"="true" "no-frame-pointer-elim-non-leaf" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "stack-protector-buffer-size"="8" "unsafe-fp-math"="false" "use-soft-float"="false" }
 attributes #1 = { nounwind "less-precise-fpmad"="false" "no-frame-pointer-elim"="true" "no-frame-pointer-elim-non-leaf" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "stack-protector-buffer-size"="8" "unsafe-fp-math"="false" "use-soft-float"="false" }
 
 !1 = !{!2, !2, i64 0}
 !2 = !{!"int", !3, i64 0}
 !3 = !{!"omnipotent char", !4, i64 0}
 !4 = !{!"Simple C/C++ TBAA"}