[llvm] b1a5ee1 - [ARM] Check all terms in emitPopInst when clearing Restored for LR. (#75527)

[via llvm-commits]

Author: Florian Hahn
Date: 2023-12-20T16:56:15+01:00
New Revision: b1a5ee1febd8a903cec3dfdad61d57900dc3823e

URL: https://github.com/llvm/llvm-project/commit/b1a5ee1febd8a903cec3dfdad61d57900dc3823e
DIFF: https://github.com/llvm/llvm-project/commit/b1a5ee1febd8a903cec3dfdad61d57900dc3823e.diff

LOG: [ARM] Check all terms in emitPopInst when clearing Restored for LR. (#75527)

emitPopInst checks a single function exit MBB. If other paths also exit
the function and any of there terminators uses LR implicitly, it is not
save to clear the Restored bit.

Check all terminators for the function before clearing Restored.

This fixes a mis-compile in outlined-fn-may-clobber-lr-in-caller.ll
where the machine-outliner previously introduced BLs that clobbered LR
which in turn is used by the tail call return.

Alternative to #73553

Added: 
    

Modified: 
    llvm/lib/Target/ARM/ARMFrameLowering.cpp
    llvm/lib/Target/ARM/ARMFrameLowering.h
    llvm/test/CodeGen/Thumb2/outlined-fn-may-clobber-lr-in-caller.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/ARM/ARMFrameLowering.cpp b/llvm/lib/Target/ARM/ARMFrameLowering.cpp
index a3a71a8ec09a45..10d9c7f275bebb 100644
--- a/llvm/lib/Target/ARM/ARMFrameLowering.cpp
+++ b/llvm/lib/Target/ARM/ARMFrameLowering.cpp
@@ -1645,9 +1645,6 @@ void ARMFrameLowering::emitPopInst(MachineBasicBlock &MBB,
         // Fold the return instruction into the LDM.
         DeleteRet = true;
         LdmOpc = AFI->isThumbFunction() ? ARM::t2LDMIA_RET : ARM::LDMIA_RET;
-        // We 'restore' LR into PC so it is not live out of the return block:
-        // Clear Restored bit.
-        Info.setRestored(false);
       }
 
       // If NoGap is true, pop consecutive registers and then leave the rest
@@ -2785,6 +2782,33 @@ void ARMFrameLowering::determineCalleeSaves(MachineFunction &MF,
   AFI->setLRIsSpilled(SavedRegs.test(ARM::LR));
 }
 
+void ARMFrameLowering::processFunctionBeforeFrameFinalized(
+    MachineFunction &MF, RegScavenger *RS) const {
+  TargetFrameLowering::processFunctionBeforeFrameFinalized(MF, RS);
+
+  MachineFrameInfo &MFI = MF.getFrameInfo();
+  if (!MFI.isCalleeSavedInfoValid())
+    return;
+
+  // Check if all terminators do not implicitly use LR. Then we can 'restore' LR
+  // into PC so it is not live out of the return block: Clear the Restored bit
+  // in that case.
+  for (CalleeSavedInfo &Info : MFI.getCalleeSavedInfo()) {
+    if (Info.getReg() != ARM::LR)
+      continue;
+    if (all_of(MF, [](const MachineBasicBlock &MBB) {
+          return all_of(MBB.terminators(), [](const MachineInstr &Term) {
+            return !Term.isReturn() || Term.getOpcode() == ARM::LDMIA_RET ||
+                   Term.getOpcode() == ARM::t2LDMIA_RET ||
+                   Term.getOpcode() == ARM::tPOP_RET;
+          });
+        })) {
+      Info.setRestored(false);
+      break;
+    }
+  }
+}
+
 void ARMFrameLowering::getCalleeSaves(const MachineFunction &MF,
                                       BitVector &SavedRegs) const {
   TargetFrameLowering::getCalleeSaves(MF, SavedRegs);

diff  --git a/llvm/lib/Target/ARM/ARMFrameLowering.h b/llvm/lib/Target/ARM/ARMFrameLowering.h
index 16f2ce6bea6f1a..8d2b8beb9a58fb 100644
--- a/llvm/lib/Target/ARM/ARMFrameLowering.h
+++ b/llvm/lib/Target/ARM/ARMFrameLowering.h
@@ -59,6 +59,9 @@ class ARMFrameLowering : public TargetFrameLowering {
   void determineCalleeSaves(MachineFunction &MF, BitVector &SavedRegs,
                             RegScavenger *RS) const override;
 
+  void processFunctionBeforeFrameFinalized(
+      MachineFunction &MF, RegScavenger *RS = nullptr) const override;
+
   void adjustForSegmentedStacks(MachineFunction &MF,
                                 MachineBasicBlock &MBB) const override;
 

diff  --git a/llvm/test/CodeGen/Thumb2/outlined-fn-may-clobber-lr-in-caller.ll b/llvm/test/CodeGen/Thumb2/outlined-fn-may-clobber-lr-in-caller.ll
index d81d008b44bed8..1dbb21f40a761c 100644
--- a/llvm/test/CodeGen/Thumb2/outlined-fn-may-clobber-lr-in-caller.ll
+++ b/llvm/test/CodeGen/Thumb2/outlined-fn-may-clobber-lr-in-caller.ll
@@ -11,8 +11,6 @@ target datalayout = "e-m:o-p:32:32-Fi8-i64:64-v128:64:128-a:0:32-n32-S64"
 
 ; Test case to make sure calling an outlined function does not clobber LR used
 ; by a tail call in caller.
-; FIXME: Currently bl OUTLINED_FUNCTION_0 clobbers LR, which in turn is used
-;        by the later call to memcpy to return to the caller.
 define void @test(ptr nocapture noundef writeonly %arg, i32 noundef %arg1, i8 noundef zeroext %arg2) unnamed_addr #0 {
 ; CHECK-LABEL: test:
 ; CHECK:       @ %bb.0: @ %bb
@@ -22,11 +20,19 @@ define void @test(ptr nocapture noundef writeonly %arg, i32 noundef %arg1, i8 no
 ; CHECK-NEXT:    cmp r1, #1
 ; CHECK-NEXT:    bne .LBB0_5
 ; CHECK-NEXT:  @ %bb.2: @ %bb4
-; CHECK-NEXT:    bl OUTLINED_FUNCTION_0
+; CHECK-NEXT:    movs r1, #1
+; CHECK-NEXT:    strb.w r1, [r0, #36]
+; CHECK-NEXT:    movs r1, #30
+; CHECK-NEXT:    strb.w r1, [r0, #34]
+; CHECK-NEXT:    add.w r1, r2, r2, lsl #3
 ; CHECK-NEXT:    ldr r2, .LCPI0_1
 ; CHECK-NEXT:    b .LBB0_4
 ; CHECK-NEXT:  .LBB0_3: @ %bb14
-; CHECK-NEXT:    bl OUTLINED_FUNCTION_0
+; CHECK-NEXT:    movs r1, #1
+; CHECK-NEXT:    strb.w r1, [r0, #36]
+; CHECK-NEXT:    movs r1, #30
+; CHECK-NEXT:    strb.w r1, [r0, #34]
+; CHECK-NEXT:    add.w r1, r2, r2, lsl #3
 ; CHECK-NEXT:    ldr r2, .LCPI0_0
 ; CHECK-NEXT:  .LBB0_4: @ %bb4
 ; CHECK-NEXT:    add.w r1, r2, r1, lsl #2