[AVR] fix interrupt stack pointer restoration

[couchand]

Summary

This patch fixes a corruption of the stack pointer and several registers in any AVR interrupt with non-empty stack frame. Previously, the callee-saved registers were popped before restoring the stack pointer, causing the pointer math to use the wrong base value while also corrupting the caller's register. This change fixes the code to restore the stack pointer last before exiting the interrupt service routine.

https://bugs.llvm.org/show_bug.cgi?id=47253

[couchand]

Fixes formatting.

[dylanmckay]

This is a great patch, thanks @couchand!

LGTM - do you need someone to commit this for you?

[couchand]

Thanks for the code review @dylanmckay!

Yes, please commit it for me, thanks.

llvm/lib/Target/AVR/AVRFrameLowering.cpp

void AVRFrameLowering::emitPrologue(MachineFunction &MF,
   MI->getOperand(3).setIsDead();
 
   // Write back R29R28 to SP and temporarily disable interrupts.
   BuildMI(MBB, MBBI, DL, TII.get(AVR::SPWRITE), AVR::SP)
       .addReg(AVR::R29R28)
       .setMIFlag(MachineInstr::FrameSetup);
 }
 
+static void restoreStatusRegister(MachineFunction &MF, MachineBasicBlock &MBB) {
+  const AVRMachineFunctionInfo *AFI = MF.getInfo<AVRMachineFunctionInfo>();
+
+  MachineBasicBlock::iterator MBBI = MBB.getLastNonDebugInstr();
+
+  DebugLoc DL = MBBI->getDebugLoc();
+  const AVRSubtarget &STI = MF.getSubtarget<AVRSubtarget>();
+  const AVRInstrInfo &TII = *STI.getInstrInfo();
+
+  // Emit special epilogue code to restore R1, R0 and SREG in interrupt/signal
+  // handlers at the very end of the function, just before reti.
+  if (AFI->isInterruptOrSignalHandler()) {
+    BuildMI(MBB, MBBI, DL, TII.get(AVR::POPRd), AVR::R0);
+    BuildMI(MBB, MBBI, DL, TII.get(AVR::OUTARr))
+        .addImm(0x3f)
+        .addReg(AVR::R0, RegState::Kill);
+    BuildMI(MBB, MBBI, DL, TII.get(AVR::POPWRd), AVR::R1R0);
+  }
+}
+
 void AVRFrameLowering::emitEpilogue(MachineFunction &MF,
                                     MachineBasicBlock &MBB) const {
   const AVRMachineFunctionInfo *AFI = MF.getInfo<AVRMachineFunctionInfo>();
 
   // Early exit if the frame pointer is not needed in this function except for
   // signal/interrupt handlers where special code generation is required.
   if (!hasFP(MF) && !AFI->isInterruptOrSignalHandler()) {
     return;
   }
 
   MachineBasicBlock::iterator MBBI = MBB.getLastNonDebugInstr();
   assert(MBBI->getDesc().isReturn() &&
          "Can only insert epilog into returning blocks");
 
   DebugLoc DL = MBBI->getDebugLoc();
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   unsigned FrameSize = MFI.getStackSize() - AFI->getCalleeSavedFrameSize();
   const AVRSubtarget &STI = MF.getSubtarget<AVRSubtarget>();
   const AVRInstrInfo &TII = *STI.getInstrInfo();
 
-  // Emit special epilogue code to restore R1, R0 and SREG in interrupt/signal
-  // handlers at the very end of the function, just before reti.
-  if (AFI->isInterruptOrSignalHandler()) {
-    BuildMI(MBB, MBBI, DL, TII.get(AVR::POPRd), AVR::R0);
-    BuildMI(MBB, MBBI, DL, TII.get(AVR::OUTARr))
-        .addImm(0x3f)
-        .addReg(AVR::R0, RegState::Kill);
-    BuildMI(MBB, MBBI, DL, TII.get(AVR::POPWRd), AVR::R1R0);
-  }
-
   // Early exit if there is no need to restore the frame pointer.
   if (!FrameSize) {
+    restoreStatusRegister(MF, MBB);
     return;
   }
 
   // Skip the callee-saved pop instructions.
   while (MBBI != MBB.begin()) {
     MachineBasicBlock::iterator PI = std::prev(MBBI);
     int Opc = PI->getOpcode();
 
MachineInstr *MI = BuildMI(MBB, MBBI, DL, TII.get(Opcode), AVR::R29R28)
                          .addReg(AVR::R29R28, RegState::Kill)
                          .addImm(FrameSize);
   // The SREG implicit def is dead.
   MI->getOperand(3).setIsDead();
 
   // Write back R29R28 to SP and temporarily disable interrupts.
   BuildMI(MBB, MBBI, DL, TII.get(AVR::SPWRITE), AVR::SP)
       .addReg(AVR::R29R28, RegState::Kill);
+
+  restoreStatusRegister(MF, MBB);
 }
 
 // Return true if the specified function should have a dedicated frame
 // pointer register. This is true if the function meets any of the following
 // conditions:
 //  - a register has been spilled
 //  - has allocas
 //  - input arguments are passed using the stack

llvm/test/CodeGen/AVR/interrupts.ll

 ; CHECK: pop r0
 ; CHECK-NEXT: out 63, r0
 ; CHECK-NEXT: pop r1
 ; CHECK-NEXT: pop r0
 ; CHECK-NEXT: reti
   ret void
 }
 
+define avr_intrcc void @interrupt_alloca() {
+; CHECK-LABEL: interrupt_alloca:
+; CHECK: sei
+; CHECK-NEXT: push r0
+; CHECK-NEXT: push r1
+; CHECK-NEXT: in r0, 63
+; CHECK-NEXT: push r0
+; CHECK: clr r0
+; CHECK: push r28
+; CHECK-NEXT: push r29
+; CHECK-NEXT: in r28, 61
+; CHECK-NEXT: in r29, 62
+; CHECK-NEXT: sbiw r28, 1
+; CHECK-NEXT: in r0, 63
+; CHECK-NEXT: cli
+; CHECK-NEXT: out 62, r29
+; CHECK-NEXT: out 63, r0
+; CHECK-NEXT: out 61, r28
+; CHECK: adiw r28, 1
+; CHECK-NEXT: in r0, 63
+; CHECK-NEXT: cli
+; CHECK-NEXT: out 62, r29
+; CHECK-NEXT: out 63, r0
+; CHECK-NEXT: out 61, r28
+; CHECK-NEXT: pop r29
+; CHECK-NEXT: pop r28
+; CHECK: pop r0
+; CHECK-NEXT: out 63, r0
+; CHECK-NEXT: pop r1
+; CHECK-NEXT: pop r0
+; CHECK-NEXT: reti
+  alloca i8
+  ret void
+}
+
 attributes #0 = { "interrupt" }
 attributes #1 = { "signal" }