Fix Thumb1 epilogue generation

[tyomitch]

Summary

This had been broken for a very long time, but nobody noticed until
D14357 enabled shrink-wrapping by default.

[tyomitch]

lib/Target/ARM/Thumb1FrameLowering.cpp

413

This must have been a leftover from some ARM lowering code: hasV4TOps() and hasV5TOps() have no relevance to Thumb1 epilogue generation, but this code used to prevent the "pop special fix-up" when it was necessary.

434

tPOP_RET is a return instruction, too; but it doesn't need a conversion to tPOP_RET, so we can return early.

462

If MBBI is the first instruction in the BB, then the last iteration of the while decrements InstUpToMBBI past the first instruction, crashing with an assertion failure.

test/CodeGen/Thumb/thumb-shrink-wrapping.ll

42

pop lr is illegal in Thumb1; in fact, llc -filetype=obj emitted pop r7 instead of this instruction, obviously unbalancing the stack.

[jroelofs]

lib/Target/ARM/Thumb1FrameLowering.cpp

413

IIRC, this check was to make sure that we didn't emit a pop {r[0-9], pc} on v4t, and expect it to change the thumb bit accordingly (because it can't).

534

Should assert that findRegisterDefOperandIdx didn't return -1.

537

Would this be better written as:

Pop = AddDefaultPred(BuildMI(MBB, MBBI, dl, TII.get(ARM::tPOP)))
          .addReg(PopReg, RegState::Define);

test/CodeGen/Thumb/thumb-shrink-wrapping.ll

43

Is the next instruction a bx lr?

If so, it'd be better to emit:

blx r1

on v5t+, and:

mov lr, r1
bx lr

on v4t.

[tyomitch]

lib/Target/ARM/Thumb1FrameLowering.cpp

413

If so, it seems it didn't have any regression tests -- so I'm not sure if my patch may lead to generation of the pop pc on v4T, or if it may not.

534

RemoveOperand asserts this the first thing it does.

537

No, because Pop may correspond either to the MachineInstrBuilder from line 519, or to the pre-existing tPOP from line 515.

test/CodeGen/Thumb/thumb-shrink-wrapping.ll

43

Is the next instruction a bx lr?

Yes it is; but it's also accessible via a by-pass.

Perhaps it would be beneficial to transform a

  pop {r7, r1}
  mov lr, r1
label:
  bx lr

into

  pop {r7, pc}
label:
  bx lr

--but that requires a more sophisticated analysis, reaching outside the BB with the epilogue, which emitPopSpecialFixUp() doesn't (and didn't) attempt to do.

[jroelofs]

lib/Target/ARM/Thumb1FrameLowering.cpp

413

test/CodeGen/ARM/thumb1_return_sequence.ll should have been testing it...

Looking at the svn blame, I think r242714 broke it. I've replied on that commit's email.

537

Oh, I see.

test/CodeGen/Thumb/thumb-shrink-wrapping.ll

43

Would you mind leaving a FIXME in emitPopSpecialFixUp explaining that we have to emit it that way for v4t, but for v5t+ it could be optimized to the shorter sequence?

[tyomitch]

Would you mind leaving a FIXME in emitPopSpecialFixUp

I can do better, and actually implement that missing optimization there :-)

[tyomitch]

This ^ does not look correct for v4t, which doesn't support changing the thumb
bit via a pop into pc.

That's right, it isn't -- same as thumb1_return_sequence.ll is incorrect for v4T.

I have a fix for v4T locally, but perhaps we can commit this fix (for v5T, broken since last week) now, and the fix for v4T (broken since July) in a separate, follow-up commit?
The two fixes are more or less independent.

[jroelofs]

In D14986#297333, @tyomitch wrote:

This ^ does not look correct for v4t, which doesn't support changing the thumb
bit via a pop into pc.

That's right, it isn't -- same as thumb1_return_sequence.ll is incorrect for v4T.

I have a fix for v4T locally, but perhaps we can commit this fix (for v5T, broken since last week) now, and the fix for v4T (broken since July) in a separate, follow-up commit?
The two fixes are more or less independent.

Sounds reasonable to me.

[qcolombet]

lib/Target/ARM/Thumb1FrameLowering.cpp

411

Add a comment saying that LR cannot be encoded with Thumb1, i.e., it requires a special fix-up.

test/CodeGen/Thumb/pop-special-fixup.ll

3

Is there a chance this test could be merged with the shrink-wrapping one since it explicitly checks with shrink-wrapping enabled?

If that is the case, change the name of the function from test to popSpecialFixup.

8

Please use opt -instnamer to get rid of the %num variables.

[qcolombet]

lib/Target/ARM/Thumb1FrameLowering.cpp

456

Do you actually need that loop instead of the copyImplicitOps call?
If so, add a comment explaining why.

458

Period.

[tyomitch]

test/CodeGen/Thumb/pop-special-fixup.ll

3

I will merge the two test files in the following patch, which will address v4T specifically.

Currently, they are separate because thumb-shrink-wrapping.ll only tests v4T, and this one only tests v6M.

[tyomitch]

Quentin, thank you for the helpful suggestions!

I've incorporated them into the new revision of my patch.

[qcolombet]

LGTM too modulo the use of instnamer on the test case to get rid of the %[0-9]+ variables.

Cheers,
Q.

lib/Target/ARM/Thumb1FrameLowering.cpp

 }
 
 bool Thumb1FrameLowering::needPopSpecialFixUp(const MachineFunction &MF) const {
   ARMFunctionInfo *AFI =
       const_cast<MachineFunction *>(&MF)->getInfo<ARMFunctionInfo>();
   if (AFI->getArgRegsSaveSize())
     return true;
 
-  bool IsV4PopReturn = false;
   for (const CalleeSavedInfo &CSI : MF.getFrameInfo()->getCalleeSavedInfo())
     if (CSI.getReg() == ARM::LR)
-      IsV4PopReturn = true;
-  return IsV4PopReturn && STI.hasV4TOps() && !STI.hasV5TOps();
+      return true;
+
+  return false;
 }
 
 bool Thumb1FrameLowering::emitPopSpecialFixUp(MachineBasicBlock &MBB,
                                               bool DoIt) const {
   MachineFunction &MF = *MBB.getParent();
   ARMFunctionInfo *AFI = MF.getInfo<ARMFunctionInfo>();
   unsigned ArgRegsSaveSize = AFI->getArgRegsSaveSize();
   const TargetInstrInfo &TII = *STI.getInstrInfo();
   const ThumbRegisterInfo *RegInfo =
       static_cast<const ThumbRegisterInfo *>(STI.getRegisterInfo());
 
-  // If MBBI is a return instruction, we may be able to directly restore
+  // If MBBI is a return instruction, or is a tPOP followed by a return
+  // instruction in the successor BB, we may be able to directly restore
   // LR in the PC.
   // This is possible if we do not need to emit any SP update.
   // Otherwise, we need a temporary register to pop the value
   // and copy that value into LR.
   auto MBBI = MBB.getFirstTerminator();
-  if (!ArgRegsSaveSize && MBBI != MBB.end() &&
-      MBBI->getOpcode() == ARM::tBX_RET) {
-    if (!DoIt)
+  bool CanRestoreDirectly = !ArgRegsSaveSize;
+  if (CanRestoreDirectly) {
+    if (MBBI != MBB.end())
+      CanRestoreDirectly = (MBBI->getOpcode() == ARM::tBX_RET ||
+                            MBBI->getOpcode() == ARM::tPOP_RET);
+    else {
+      assert(MBB.back().getOpcode() == ARM::tPOP);
+      assert(MBB.succ_size() == 1);
+      if ((*MBB.succ_begin())->begin()->getOpcode() == ARM::tBX_RET)
+        MBBI--; // replace the final tPOP with a tPOP_RET
+      else
+        CanRestoreDirectly = false;
+    }
+  }
+
+  if (CanRestoreDirectly) {
+    if (!DoIt || MBBI->getOpcode() == ARM::tPOP_RET)
       return true;
     MachineInstrBuilder MIB =
         AddDefaultPred(
-            BuildMI(MBB, MBBI, MBBI->getDebugLoc(), TII.get(ARM::tPOP_RET)))
-            .addReg(ARM::PC, RegState::Define);
-    MIB.copyImplicitOps(&*MBBI);
-    // erase the old tBX_RET instruction
+            BuildMI(MBB, MBBI, MBBI->getDebugLoc(), TII.get(ARM::tPOP_RET)));
+    for (auto MO: MBBI->operands())
+      if (MO.isReg() && (MO.isImplicit() || MO.isDef()) &&
+          MO.getReg() != ARM::LR)
+        MIB.addOperand(MO);
+    MIB.addReg(ARM::PC, RegState::Define);
+    // erase the old instruction (tBX_RET or tPOP)
     MBB.erase(MBBI);
     return true;
   }
 
   // Look for a temporary register to use.
   // First, compute the liveness information.
   LivePhysRegs UsedRegs(STI.getRegisterInfo());
   UsedRegs.addLiveOuts(&MBB, /*AddPristines*/ true);
   // The semantic of pristines changed recently and now,
   // the callee-saved registers that are touched in the function
   // are not part of the pristines set anymore.
   // Add those callee-saved now.
   const TargetRegisterInfo *TRI = STI.getRegisterInfo();
   const MCPhysReg *CSRegs = TRI->getCalleeSavedRegs(&MF);
   for (unsigned i = 0; CSRegs[i]; ++i)
     UsedRegs.addReg(CSRegs[i]);
 
   DebugLoc dl = DebugLoc();
   if (MBBI != MBB.end()) {
     dl = MBBI->getDebugLoc();
     auto InstUpToMBBI = MBB.end();
-    // The post-decrement is on purpose here.
+    while (InstUpToMBBI != MBBI)
+      // The pre-decrement is on purpose here.
       // We want to have the liveness right before MBBI.
-    while (InstUpToMBBI-- != MBBI)
-      UsedRegs.stepBackward(*InstUpToMBBI);
+      UsedRegs.stepBackward(*--InstUpToMBBI);
   }
 
   // Look for a register that can be directly use in the POP.
   unsigned PopReg = 0;
   // And some temporary register, just in case.
   unsigned TemporaryReg = 0;
   BitVector PopFriendly =
       TRI->getAllocatableSet(MF, TRI->getRegClass(ARM::tGPRRegClassID));
if (TemporaryReg) {
     assert(!PopReg && "Unnecessary MOV is about to be inserted");
     PopReg = PopFriendly.find_first();
     AddDefaultPred(BuildMI(MBB, MBBI, dl, TII.get(ARM::tMOVr))
                        .addReg(TemporaryReg, RegState::Define)
                        .addReg(PopReg, RegState::Kill));
   }
 
   assert(PopReg && "Do not know how to get LR");
-  AddDefaultPred(BuildMI(MBB, MBBI, dl, TII.get(ARM::tPOP)))
-      .addReg(PopReg, RegState::Define);
+  MachineInstr *Pop;
+  if (MBBI == MBB.end()) {
+    Pop = &MBB.back();
+    assert(Pop->getOpcode() == ARM::tPOP);
+    Pop->RemoveOperand(Pop->findRegisterDefOperandIdx(ARM::LR));
+  } else
+    Pop = AddDefaultPred(BuildMI(MBB, MBBI, dl, TII.get(ARM::tPOP)));
+  Pop->addOperand(MachineOperand::CreateReg(PopReg, true));
 
   emitSPUpdate(MBB, MBBI, TII, dl, *RegInfo, ArgRegsSaveSize);
 
   if (!TemporaryReg && MBBI != MBB.end() && MBBI->getOpcode() == ARM::tBX_RET) {
     MachineInstrBuilder MIB = BuildMI(MBB, MBBI, dl, TII.get(ARM::tBX))
                                   .addReg(PopReg, RegState::Kill);
     AddDefaultPred(MIB);
     MIB.copyImplicitOps(&*MBBI);

test/CodeGen/Thumb/pop-special-fixup.ll

+; RUN: llc %s -enable-shrink-wrap=true -o - | FileCheck %s
+
+target triple = "thumbv6m-none-none-eabi"
+
+@retval = global i32 0, align 4
+
+define i32 @test(i32 %i, i32 %argc, i8** nocapture readonly %argv) {
+  %1 = icmp sgt i32 %argc, %i
+  br i1 %1, label %2, label %19
+
+  %3 = getelementptr inbounds i8*, i8** %argv, i32 %i
+  %4 = load i8*, i8** %3, align 4
+  %5 = load i8, i8* %4, align 1
+  %6 = icmp eq i8 %5, 45
+  %7 = getelementptr inbounds i8, i8* %4, i32 1
+  %. = select i1 %6, i8* %7, i8* %4
+  %.1 = select i1 %6, i32 -1, i32 1
+  %8 = load i8, i8* %., align 1
+  %.off2 = add i8 %8, -48
+  %9 = icmp ult i8 %.off2, 10
+  %.pre = load i32, i32* @retval, align 4
+  br i1 %9, label %.lr.ph.preheader, label %.critedge
+
+.lr.ph.preheader:                                 ; preds = %2
+  br label %.lr.ph
+
+.lr.ph:                                           ; preds = %.lr.ph.preheader, %.lr.ph
+  %10 = phi i32 [ %14, %.lr.ph ], [ %.pre, %.lr.ph.preheader ]
+  %11 = phi i8 [ %15, %.lr.ph ], [ %8, %.lr.ph.preheader ]
+  %valstring.03 = phi i8* [ %13, %.lr.ph ], [ %., %.lr.ph.preheader ]
+  %12 = zext i8 %11 to i32
+  %13 = getelementptr inbounds i8, i8* %valstring.03, i32 1
+  %14 = add nsw i32 %10, %12
+  store i32 %14, i32* @retval, align 4
+  %15 = load i8, i8* %13, align 1
+  %.off = add i8 %15, -48
+  %16 = icmp ult i8 %.off, 10
+  br i1 %16, label %.lr.ph, label %.critedge.loopexit
+
+.critedge.loopexit:                               ; preds = %.lr.ph
+  %.lcssa = phi i32 [ %14, %.lr.ph ]
+  br label %.critedge
+
+.critedge:                                        ; preds = %.critedge.loopexit, %2
+  %17 = phi i32 [ %.pre, %2 ], [ %.lcssa, %.critedge.loopexit ]
+  %18 = mul nsw i32 %17, %.1
+  store i32 %18, i32* @retval, align 4
+  br label %19
+
+; <label>:19                                      ; preds = %.critedge, %0
+  ret i32 0
+}
+
+; CHECK: push {r4, r5, r7, lr}
+; CHECK: pop {r4, r5, r7, r0}
+; CHECK: mov lr, r0
+; CHECK: movs r0, #0
+; CHECK: bx  lr
+

test/CodeGen/Thumb/thumb-shrink-wrapping.ll

 ; Set the alloca address in the second argument.
 ; Set the first argument to zero.
 ; CHECK: movs r0, #0
 ; CHECK-NEXT: add r1, sp, #4
 ; CHECK-NEXT: bl
 ;
 ; With shrink-wrapping, epilogue is just after the call.
 ; ENABLE-NEXT: add sp, #8
-; ENABLE-NEXT: pop {r7, lr}
+; ENABLE-NEXT: pop {r7, pc}
 ;
 ; CHECK: [[EXIT_LABEL]]:
 ;
 ; Without shrink-wrapping, epilogue is in the exit block.
 ; Epilogue code. (What we pop does not matter.)
 ; DISABLE: add sp, #8
 ; DISABLE-NEXT: pop {r7, pc}
 ;
 ; ENABLE-NEXT: bx lr