[AVR] Always expand STDSPQRr & STDWSPQRr

Currently, STDSPQRr and STDWSPQRr are expanded only during
AVRFrameLowering - this means that if any of those instructions happen
to appear _outside_ of the typical FrameSetup / FrameDestroy
context, they wouldn't get substituted, eventually leading to a crash:

```
LLVM ERROR: Not supported instr: <MCInst XXX <MCOperand Reg:1>
<MCOperand Imm:15> <MCOperand Reg:53>>
```

This commit fixes this issue by moving expansion of those two opcodes
into AVRExpandPseudo.

This bug was originally discovered due to the Rust compiler_builtins
library. Its 0.1.37 release contained a 128-bit software
division/remainder routine that exercised this buggy branch in the code.

Reviewed By: benshi001

Differential Revision: https://reviews.llvm.org/D123528

Author: Patryk27

llvm/lib/Target/AVR/AVRExpandPseudoInsts.cpp

@@ -1280,6 +1280,42 @@ bool AVRExpandPseudo::expand<AVR::STDWPtrQRr>(Block &MBB, BlockIt MBBI) {
   return true;
 }
 
+template <>
+bool AVRExpandPseudo::expand<AVR::STDSPQRr>(Block &MBB, BlockIt MBBI) {
+  MachineInstr &MI = *MBBI;
+  const MachineFunction &MF = *MBB.getParent();
+  const AVRSubtarget &STI = MF.getSubtarget<AVRSubtarget>();
+
+  assert(MI.getOperand(0).getReg() == AVR::SP &&
+         "SP is expected as base pointer");
+
+  assert(STI.getFrameLowering()->hasReservedCallFrame(MF) &&
+         "unexpected STDSPQRr pseudo instruction");
+
+  MI.setDesc(TII->get(AVR::STDPtrQRr));
+  MI.getOperand(0).setReg(AVR::R29R28);
+
+  return true;
+}
+
+template <>
+bool AVRExpandPseudo::expand<AVR::STDWSPQRr>(Block &MBB, BlockIt MBBI) {
+  MachineInstr &MI = *MBBI;
+  const MachineFunction &MF = *MBB.getParent();
+  const AVRSubtarget &STI = MF.getSubtarget<AVRSubtarget>();
+
+  assert(MI.getOperand(0).getReg() == AVR::SP &&
+         "SP is expected as base pointer");
+
+  assert(STI.getFrameLowering()->hasReservedCallFrame(MF) &&
+         "unexpected STDWSPQRr pseudo instruction");
+
+  MI.setDesc(TII->get(AVR::STDWPtrQRr));
+  MI.getOperand(0).setReg(AVR::R29R28);
+
+  return true;
+}
+
 template <>
 bool AVRExpandPseudo::expand<AVR::INWRdA>(Block &MBB, BlockIt MBBI) {
   MachineInstr &MI = *MBBI;
@@ -2365,6 +2401,8 @@ bool AVRExpandPseudo::expandMI(Block &MBB, BlockIt MBBI) {
     EXPAND(AVR::STWPtrPiRr);
     EXPAND(AVR::STWPtrPdRr);
     EXPAND(AVR::STDWPtrQRr);
+    EXPAND(AVR::STDSPQRr);
+    EXPAND(AVR::STDWSPQRr);
     EXPAND(AVR::INWRdA);
     EXPAND(AVR::OUTWARr);
     EXPAND(AVR::PUSHWRr);

llvm/lib/Target/AVR/AVRFrameLowering.cpp

@@ -201,8 +201,8 @@ void AVRFrameLowering::emitEpilogue(MachineFunction &MF,
 
     // Restore the frame pointer by doing FP += <size>.
     MachineInstr *MI = BuildMI(MBB, MBBI, DL, TII.get(Opcode), AVR::R29R28)
-                          .addReg(AVR::R29R28, RegState::Kill)
-                          .addImm(FrameSize);
+                           .addReg(AVR::R29R28, RegState::Kill)
+                           .addImm(FrameSize);
     // The SREG implicit def is dead.
     MI->getOperand(3).setIsDead();
   }
@@ -299,7 +299,7 @@ bool AVRFrameLowering::restoreCalleeSavedRegisters(
 /// real instructions.
 static void fixStackStores(MachineBasicBlock &MBB,
                            MachineBasicBlock::iterator StartMI,
-                           const TargetInstrInfo &TII, Register FP) {
+                           const TargetInstrInfo &TII) {
   // Iterate through the BB until we hit a call instruction or we reach the end.
   for (MachineInstr &MI :
        llvm::make_early_inc_range(llvm::make_range(StartMI, MBB.end()))) {
@@ -313,15 +313,15 @@ static void fixStackStores(MachineBasicBlock &MBB,
       continue;
 
     assert(MI.getOperand(0).getReg() == AVR::SP &&
-           "Invalid register, should be SP!");
+           "SP is expected as base pointer");
 
     // Replace this instruction with a regular store. Use Y as the base
     // pointer since it is guaranteed to contain a copy of SP.
     unsigned STOpc =
         (Opcode == AVR::STDWSPQRr) ? AVR::STDWPtrQRr : AVR::STDPtrQRr;
 
     MI.setDesc(TII.get(STOpc));
-    MI.getOperand(0).setReg(FP);
+    MI.getOperand(0).setReg(AVR::R31R30);
   }
 }
 
@@ -331,69 +331,66 @@ MachineBasicBlock::iterator AVRFrameLowering::eliminateCallFramePseudoInstr(
   const AVRSubtarget &STI = MF.getSubtarget<AVRSubtarget>();
   const AVRInstrInfo &TII = *STI.getInstrInfo();
 
-  // There is nothing to insert when the call frame memory is allocated during
-  // function entry. Delete the call frame pseudo and replace all pseudo stores
-  // with real store instructions.
   if (hasReservedCallFrame(MF)) {
-    fixStackStores(MBB, MI, TII, AVR::R29R28);
     return MBB.erase(MI);
   }
 
   DebugLoc DL = MI->getDebugLoc();
   unsigned int Opcode = MI->getOpcode();
   int Amount = TII.getFrameSize(*MI);
 
-  // ADJCALLSTACKUP and ADJCALLSTACKDOWN are converted to adiw/subi
-  // instructions to read and write the stack pointer in I/O space.
-  if (Amount != 0) {
-    assert(getStackAlign() == Align(1) && "Unsupported stack alignment");
-
-    if (Opcode == TII.getCallFrameSetupOpcode()) {
-      // Update the stack pointer.
-      // In many cases this can be done far more efficiently by pushing the
-      // relevant values directly to the stack. However, doing that correctly
-      // (in the right order, possibly skipping some empty space for undef
-      // values, etc) is tricky and thus left to be optimized in the future.
-      BuildMI(MBB, MI, DL, TII.get(AVR::SPREAD), AVR::R31R30).addReg(AVR::SP);
-
-      MachineInstr *New =
-          BuildMI(MBB, MI, DL, TII.get(AVR::SUBIWRdK), AVR::R31R30)
-              .addReg(AVR::R31R30, RegState::Kill)
-              .addImm(Amount);
-      New->getOperand(3).setIsDead();
-
-      BuildMI(MBB, MI, DL, TII.get(AVR::SPWRITE), AVR::SP).addReg(AVR::R31R30);
-
-      // Make sure the remaining stack stores are converted to real store
-      // instructions.
-      fixStackStores(MBB, MI, TII, AVR::R31R30);
-    } else {
-      assert(Opcode == TII.getCallFrameDestroyOpcode());
-
-      // Note that small stack changes could be implemented more efficiently
-      // with a few pop instructions instead of the 8-9 instructions now
-      // required.
-
-      // Select the best opcode to adjust SP based on the offset size.
-      unsigned addOpcode;
-      if (isUInt<6>(Amount)) {
-        addOpcode = AVR::ADIWRdK;
-      } else {
-        addOpcode = AVR::SUBIWRdK;
-        Amount = -Amount;
-      }
+  if (Amount == 0) {
+    return MBB.erase(MI);
+  }
+
+  assert(getStackAlign() == Align(1) && "Unsupported stack alignment");
+
+  if (Opcode == TII.getCallFrameSetupOpcode()) {
+    // Update the stack pointer.
+    // In many cases this can be done far more efficiently by pushing the
+    // relevant values directly to the stack. However, doing that correctly
+    // (in the right order, possibly skipping some empty space for undef
+    // values, etc) is tricky and thus left to be optimized in the future.
+    BuildMI(MBB, MI, DL, TII.get(AVR::SPREAD), AVR::R31R30).addReg(AVR::SP);
+
+    MachineInstr *New =
+        BuildMI(MBB, MI, DL, TII.get(AVR::SUBIWRdK), AVR::R31R30)
+            .addReg(AVR::R31R30, RegState::Kill)
+            .addImm(Amount);
+    New->getOperand(3).setIsDead();
 
-      // Build the instruction sequence.
-      BuildMI(MBB, MI, DL, TII.get(AVR::SPREAD), AVR::R31R30).addReg(AVR::SP);
+    BuildMI(MBB, MI, DL, TII.get(AVR::SPWRITE), AVR::SP).addReg(AVR::R31R30);
 
-      MachineInstr *New = BuildMI(MBB, MI, DL, TII.get(addOpcode), AVR::R31R30)
-                              .addReg(AVR::R31R30, RegState::Kill)
-                              .addImm(Amount);
-      New->getOperand(3).setIsDead();
+    // Make sure the remaining stack stores are converted to real store
+    // instructions.
+    fixStackStores(MBB, MI, TII);
+  } else {
+    assert(Opcode == TII.getCallFrameDestroyOpcode());
 
-      BuildMI(MBB, MI, DL, TII.get(AVR::SPWRITE), AVR::SP)
-          .addReg(AVR::R31R30, RegState::Kill);
+    // Note that small stack changes could be implemented more efficiently
+    // with a few pop instructions instead of the 8-9 instructions now
+    // required.
+
+    // Select the best opcode to adjust SP based on the offset size.
+    unsigned AddOpcode;
+
+    if (isUInt<6>(Amount)) {
+      AddOpcode = AVR::ADIWRdK;
+    } else {
+      AddOpcode = AVR::SUBIWRdK;
+      Amount = -Amount;
     }
+
+    // Build the instruction sequence.
+    BuildMI(MBB, MI, DL, TII.get(AVR::SPREAD), AVR::R31R30).addReg(AVR::SP);
+
+    MachineInstr *New = BuildMI(MBB, MI, DL, TII.get(AddOpcode), AVR::R31R30)
+                            .addReg(AVR::R31R30, RegState::Kill)
+                            .addImm(Amount);
+    New->getOperand(3).setIsDead();
+
+    BuildMI(MBB, MI, DL, TII.get(AVR::SPWRITE), AVR::SP)
+        .addReg(AVR::R31R30, RegState::Kill);
   }
 
   return MBB.erase(MI);
@@ -420,7 +417,7 @@ struct AVRFrameAnalyzer : public MachineFunctionPass {
 
   bool runOnMachineFunction(MachineFunction &MF) override {
     const MachineFrameInfo &MFI = MF.getFrameInfo();
-    AVRMachineFunctionInfo *FuncInfo = MF.getInfo<AVRMachineFunctionInfo>();
+    AVRMachineFunctionInfo *AFI = MF.getInfo<AVRMachineFunctionInfo>();
 
     // If there are no fixed frame indexes during this stage it means there
     // are allocas present in the function.
@@ -431,7 +428,7 @@ struct AVRFrameAnalyzer : public MachineFunctionPass {
       for (unsigned i = 0, e = MFI.getObjectIndexEnd(); i != e; ++i) {
         // Variable sized objects have size 0.
         if (MFI.getObjectSize(i)) {
-          FuncInfo->setHasAllocas(true);
+          AFI->setHasAllocas(true);
           break;
         }
       }
@@ -460,7 +457,7 @@ struct AVRFrameAnalyzer : public MachineFunctionPass {
           }
 
           if (MFI.isFixedObjectIndex(MO.getIndex())) {
-            FuncInfo->setHasStackArgs(true);
+            AFI->setHasStackArgs(true);
             return false;
           }
         }

llvm/test/CodeGen/AVR/stdwstk.ll

@@ -0,0 +1,15 @@
+; RUN: llc < %s -march=avr -mcpu=atmega328 -O1 | FileCheck %s
+; CHECK-NOT: stdwstk
+
+; Checks that we expand STDWSPQRr always - even if it appears outside of the
+; FrameSetup/FrameDestroy context.
+
+declare { } @foo(i128, i128) addrspace(1)
+
+define i128 @bar(i128 %a, i128 %b) addrspace(1) {
+  %b_neg = icmp slt i128 %b, 0
+  %divisor = select i1 %b_neg, i128 0, i128 %b
+  %result = tail call fastcc addrspace(1) { } @foo(i128 undef, i128 %divisor)
+
+  ret i128 0
+}