[PowerPC] Restore stack ptr from frame ptr with setjmp

[nemanjai]

Summary

If a function happens to:

• call setjmp
• do a 16-byte stack allocation
• call a function that sets up a stack frame and longjmp's back

The stack pointer that is restores by setjmp will no longer point to a valid back chain. According to the ABI, stack accesses in such a function are to be frame pointer based - so it is an error (quite obviously) to restore the stack from the back chain.
We already restore the stack from the frame pointer when there are calls to fast_cc functions. We just need to also do that when there are calls to setjmp. This patch simply does that.

This was pointed out by the Julia team.

[vchuravy]

Looks good from the Julia side, this does indeed fix the upstream bug. Thanks @nemanjai!

[vchuravy]

This fixed a whole set of bugs for us but it also seems to have caused a regression.
https://github.com/JuliaLang/julia/issues/38859

The details are in the upstream issue with a pure llvm IR reproducer. I could reproduce on ToT and reverting this PR made me unable to reproduce it.
The reproducer is a bit flaky and needs to be run several times.

llvm/lib/Target/PowerPC/PPCFrameLowering.cpp

bool PPCFrameLowering::needsFP(const MachineFunction &MF) const {
   const MachineFrameInfo &MFI = MF.getFrameInfo();
 
   // Naked functions have no stack frame pushed, so we don't have a frame
   // pointer.
   if (MF.getFunction().hasFnAttribute(Attribute::Naked))
     return false;
 
   return MF.getTarget().Options.DisableFramePointerElim(MF) ||
          MFI.hasVarSizedObjects() || MFI.hasStackMap() || MFI.hasPatchPoint() ||
+         MF.exposesReturnsTwice() ||
          (MF.getTarget().Options.GuaranteedTailCallOpt &&
           MF.getInfo<PPCFunctionInfo>()->hasFastCall());
 }
 
 void PPCFrameLowering::replaceFPWithRealFP(MachineFunction &MF) const {
   bool is31 = needsFP(MF);
   unsigned FPReg  = is31 ? PPC::R31 : PPC::R1;
   unsigned FP8Reg = is31 ? PPC::X31 : PPC::X1;
 
   const PPCRegisterInfo *RegInfo = Subtarget.getRegisterInfo();
bool PPCFrameLowering::stackUpdateCanBeMoved(MachineFunction &MF) const {
   MachineFrameInfo &MFI = MF.getFrameInfo();
   unsigned FrameSize = MFI.getStackSize();
   if (!FrameSize || FrameSize > Subtarget.getRedZoneSize())
     return false;
 
   // Frame pointers and base pointers complicate matters so don't do anything
   // if we have them. For example having a frame pointer will sometimes require
   // a copy of r1 into r31 and that makes keeping track of updates to r1 more
-  // difficult.
-  if (hasFP(MF) || RegInfo->hasBasePointer(MF))
+  // difficult. Similar situation exists with setjmp.
+  if (hasFP(MF) || RegInfo->hasBasePointer(MF) || MF.exposesReturnsTwice())
     return false;
 
   // Calls to fast_cc functions use different rules for passing parameters on
   // the stack from the ABI and using PIC base in the function imposes
   // similar restrictions to using the base pointer. It is not generally safe
   // to move the stack pointer update in these situations.
   if (FI->hasFastCall() || FI->usesPICBase())
     return false;
void PPCFrameLowering::emitEpilogue(MachineFunction &MF,
   if (FrameSize) {
     // In the prologue, the loaded (or persistent) stack pointer value is
     // offset by the STDU/STDUX/STWU/STWUX instruction. For targets with red
     // zone add this offset back now.
 
     // If this function contained a fastcc call and GuaranteedTailCallOpt is
     // enabled (=> hasFastCall()==true) the fastcc call might contain a tail
     // call which invalidates the stack pointer value in SP(0). So we use the
-    // value of R31 in this case.
-    if (FI->hasFastCall()) {
+    // value of R31 in this case. Similar situation exists with setjmp.
+    if (FI->hasFastCall() || MF.exposesReturnsTwice()) {
       assert(HasFP && "Expecting a valid frame pointer.");
       if (!HasRedZone)
         RBReg = FPReg;
       if (!isLargeFrame) {
         BuildMI(MBB, MBBI, dl, AddImmInst, RBReg)
           .addReg(FPReg).addImm(FrameSize);
       } else {
         BuildMI(MBB, MBBI, dl, LoadImmShiftedInst, ScratchReg)

llvm/test/CodeGen/PowerPC/stack-restore-with-setjmp.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=powerpc64le-- -verify-machineinstrs | FileCheck %s
+; RUN: llc < %s -mtriple=powerpc64-- -verify-machineinstrs | FileCheck %s \
+; RUN:   --check-prefix=BE
+%struct.__jmp_buf_tag = type { [64 x i64], i32, %struct.__sigset_t, [8 x i8] }
+%struct.__sigset_t = type { [16 x i64] }
+
+@.str = private unnamed_addr constant [33 x i8] c"Successfully returned from main\0A\00", align 1
+
+; Function Attrs: nounwind
+define dso_local signext i32 @main(i32 signext %argc, i8** nocapture readnone %argv) local_unnamed_addr #0 {
+; CHECK-LABEL: main:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    mfocrf 12, 32
+; CHECK-NEXT:    mflr 0
+; CHECK-NEXT:    std 31, -8(1)
+; CHECK-NEXT:    std 0, 16(1)
+; CHECK-NEXT:    stw 12, 8(1)
+; CHECK-NEXT:    stdu 1, -784(1)
+; CHECK-NEXT:    # kill: def $r3 killed $r3 killed $x3
+; CHECK-NEXT:    cmpwi 2, 3, 2
+; CHECK-NEXT:    mr 31, 1
+; CHECK-NEXT:    li 3, 0
+; CHECK-NEXT:    blt 2, .LBB0_3
+; CHECK-NEXT:  # %bb.1: # %if.end
+; CHECK-NEXT:    addi 3, 31, 112
+; CHECK-NEXT:    bl _setjmp
+; CHECK-NEXT:    nop
+; CHECK-NEXT:    crmove 20, 10
+; CHECK-NEXT:    # kill: def $r3 killed $r3 killed $x3
+; CHECK-NEXT:    cmpwi 3, 0
+; CHECK-NEXT:    crorc 20, 10, 2
+; CHECK-NEXT:    crmove 21, 2
+; CHECK-NEXT:    bc 4, 20, .LBB0_4
+; CHECK-NEXT:  # %bb.2: # %if.end5
+; CHECK-NEXT:    addis 3, 2, .L.str@toc@ha
+; CHECK-NEXT:    addi 3, 3, .L.str@toc@l
+; CHECK-NEXT:    bl printf
+; CHECK-NEXT:    nop
+; CHECK-NEXT:    # kill: def $r3 killed $r3 killed $x3
+; CHECK-NEXT:  .LBB0_3: # %return
+; CHECK-NEXT:    extsw 3, 3
+; CHECK-NEXT:    addi 1, 31, 784
+; CHECK-NEXT:    ld 0, 16(1)
+; CHECK-NEXT:    lwz 12, 8(1)
+; CHECK-NEXT:    ld 31, -8(1)
+; CHECK-NEXT:    mtocrf 32, 12
+; CHECK-NEXT:    mtlr 0
+; CHECK-NEXT:    blr
+; CHECK-NEXT:  .LBB0_4: # %if.then3
+; CHECK-NEXT:    ld 4, 0(1)
+; CHECK-NEXT:    stdu 4, -16(1)
+; CHECK-NEXT:    addi 3, 1, 96
+; CHECK-NEXT:    li 4, -1
+; CHECK-NEXT:    stb 4, 0(3)
+; CHECK-NEXT:    addi 4, 31, 112
+; CHECK-NEXT:    bl test
+; CHECK-NEXT:    nop
+;
+; BE-LABEL: main:
+; BE:       # %bb.0: # %entry
+; BE-NEXT:    mflr 0
+; BE-NEXT:    std 31, -8(1)
+; BE-NEXT:    std 0, 16(1)
+; BE-NEXT:    mfcr 12
+; BE-NEXT:    stw 12, 8(1)
+; BE-NEXT:    stdu 1, -800(1)
+; BE-NEXT:    li 4, 0
+; BE-NEXT:    # kill: def $r3 killed $r3 killed $x3
+; BE-NEXT:    cmpwi 2, 3, 2
+; BE-NEXT:    mr 3, 4
+; BE-NEXT:    mr 31, 1
+; BE-NEXT:    blt 2, .LBB0_3
+; BE-NEXT:  # %bb.1: # %if.end
+; BE-NEXT:    addi 3, 31, 128
+; BE-NEXT:    bl _setjmp
+; BE-NEXT:    nop
+; BE-NEXT:    crmove 20, 10
+; BE-NEXT:    # kill: def $r3 killed $r3 killed $x3
+; BE-NEXT:    cmpwi 3, 0
+; BE-NEXT:    crorc 20, 10, 2
+; BE-NEXT:    crmove 21, 2
+; BE-NEXT:    bc 4, 20, .LBB0_4
+; BE-NEXT:  # %bb.2: # %if.end5
+; BE-NEXT:    addis 3, 2, .L.str@toc@ha
+; BE-NEXT:    addi 3, 3, .L.str@toc@l
+; BE-NEXT:    bl printf
+; BE-NEXT:    nop
+; BE-NEXT:    # kill: def $r3 killed $r3 killed $x3
+; BE-NEXT:  .LBB0_3: # %return
+; BE-NEXT:    extsw 3, 3
+; BE-NEXT:    addi 1, 31, 800
+; BE-NEXT:    ld 0, 16(1)
+; BE-NEXT:    lwz 12, 8(1)
+; BE-NEXT:    ld 31, -8(1)
+; BE-NEXT:    mtlr 0
+; BE-NEXT:    mtcrf 32, 12 # cr2
+; BE-NEXT:    blr
+; BE-NEXT:  .LBB0_4: # %if.then3
+; BE-NEXT:    ld 4, 0(1)
+; BE-NEXT:    stdu 4, -16(1)
+; BE-NEXT:    addi 3, 1, 112
+; BE-NEXT:    li 4, -1
+; BE-NEXT:    stb 4, 0(3)
+; BE-NEXT:    addi 4, 31, 128
+; BE-NEXT:    bl test
+; BE-NEXT:    nop
+entry:
+  %env_buffer = alloca [1 x %struct.__jmp_buf_tag], align 16
+  %cmp = icmp slt i32 %argc, 2
+  br i1 %cmp, label %return, label %if.end
+
+if.end:                                           ; preds = %entry
+  %0 = bitcast [1 x %struct.__jmp_buf_tag]* %env_buffer to i8*
+  call void @llvm.lifetime.start.p0i8(i64 656, i8* nonnull %0) #5
+  %arraydecay = getelementptr inbounds [1 x %struct.__jmp_buf_tag], [1 x %struct.__jmp_buf_tag]* %env_buffer, i64 0, i64 0
+  %call = call signext i32 @_setjmp(%struct.__jmp_buf_tag* nonnull %arraydecay) #6
+  %cmp1 = icmp ne i32 %argc, 2
+  %cmp2 = icmp eq i32 %call, 0
+  %or.cond = and i1 %cmp1, %cmp2
+  br i1 %or.cond, label %if.then3, label %if.end5
+
+if.then3:                                         ; preds = %if.end
+  %1 = alloca [8 x i8], align 16
+  %.sub = getelementptr inbounds [8 x i8], [8 x i8]* %1, i64 0, i64 0
+  store i8 -1, i8* %.sub, align 16
+  call void @test(i8* nonnull %.sub, %struct.__jmp_buf_tag* nonnull %arraydecay) #7
+  unreachable
+
+if.end5:                                          ; preds = %if.end
+  %call6 = call signext i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([33 x i8], [33 x i8]* @.str, i64 0, i64 0))
+  call void @llvm.lifetime.end.p0i8(i64 656, i8* nonnull %0) #5
+  br label %return
+
+return:                                           ; preds = %entry, %if.end5
+  %retval.0 = phi i32 [ %call6, %if.end5 ], [ 0, %entry ]
+  ret i32 %retval.0
+}
+
+; Function Attrs: argmemonly nofree nosync nounwind willreturn
+declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture)
+
+; Function Attrs: nounwind returns_twice
+declare signext i32 @_setjmp(%struct.__jmp_buf_tag*) local_unnamed_addr
+
+; Function Attrs: noreturn
+declare void @test(i8*, %struct.__jmp_buf_tag*) local_unnamed_addr
+
+; Function Attrs: nofree nounwind
+declare noundef signext i32 @printf(i8* nocapture noundef readonly, ...) local_unnamed_addr
+
+; Function Attrs: argmemonly nofree nosync nounwind willreturn
+declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture)
+
+attributes #0 = { nounwind }
+attributes #6 = { nounwind returns_twice }