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

[Florian Hahn via llvm-commits]

https://github.com/fhahn created https://github.com/llvm/llvm-project/pull/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

>From f6896d14fb4cf03db8c94a402edbd22626d7a2dd Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Thu, 14 Dec 2023 19:46:56 +0000
Subject: [PATCH] [ARM] Check all terms in emitPopInst when clearing Restored
 for LR.

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.
---
 llvm/lib/Target/ARM/ARMFrameLowering.cpp       | 18 +++++++++++++++---
 .../test/CodeGen/Thumb2/mve-float16regloops.ll |  1 +
 .../test/CodeGen/Thumb2/mve-float32regloops.ll |  1 +
 .../outlined-fn-may-clobber-lr-in-caller.ll    | 12 ++++++++++--
 4 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Target/ARM/ARMFrameLowering.cpp b/llvm/lib/Target/ARM/ARMFrameLowering.cpp
index a3a71a8ec09a45..79b70cab82f145 100644
--- a/llvm/lib/Target/ARM/ARMFrameLowering.cpp
+++ b/llvm/lib/Target/ARM/ARMFrameLowering.cpp
@@ -1645,9 +1645,21 @@ 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);
+        // 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
+        // Restored bit.
+        if (all_of(MF, [MI](const MachineBasicBlock &MBB) {
+              return all_of(MBB.terminators(), [MI](const MachineInstr &Term) {
+                //  MI's terminator is to be re-written, don't check the old
+                //  opcode.
+                if (&*MI == &Term)
+                  return true;
+                return Term.getOpcode() == ARM::LDMIA_RET ||
+                       Term.getOpcode() == ARM::t2LDMIA_RET ||
+                       Term.getOpcode() == ARM::tPOP_RET;
+              });
+            }))
+          Info.setRestored(false);
       }
 
       // If NoGap is true, pop consecutive registers and then leave the rest
diff --git a/llvm/test/CodeGen/Thumb2/mve-float16regloops.ll b/llvm/test/CodeGen/Thumb2/mve-float16regloops.ll
index 1c95d28b5eed1b..a75f445097f28b 100644
--- a/llvm/test/CodeGen/Thumb2/mve-float16regloops.ll
+++ b/llvm/test/CodeGen/Thumb2/mve-float16regloops.ll
@@ -831,6 +831,7 @@ define void @arm_fir_f32_1_4_mve(ptr nocapture readonly %S, ptr nocapture readon
 ; CHECK-NEXT:    mov r0, r1
 ; CHECK-NEXT:  .LBB15_10: @ %while.end55
 ; CHECK-NEXT:    ands r1, r9, #3
+; CHECK-NEXT:    @ implicit-def: $lr
 ; CHECK-NEXT:    beq .LBB15_12
 ; CHECK-NEXT:  @ %bb.11: @ %if.then59
 ; CHECK-NEXT:    vldrw.u32 q0, [r0]
diff --git a/llvm/test/CodeGen/Thumb2/mve-float32regloops.ll b/llvm/test/CodeGen/Thumb2/mve-float32regloops.ll
index 808626d9a0aebe..c29653e6827263 100644
--- a/llvm/test/CodeGen/Thumb2/mve-float32regloops.ll
+++ b/llvm/test/CodeGen/Thumb2/mve-float32regloops.ll
@@ -822,6 +822,7 @@ define void @arm_fir_f32_1_4_mve(ptr nocapture readonly %S, ptr nocapture readon
 ; CHECK-NEXT:    mov r0, r1
 ; CHECK-NEXT:  .LBB15_10: @ %while.end55
 ; CHECK-NEXT:    ands r1, r10, #3
+; CHECK-NEXT:    @ implicit-def: $lr
 ; CHECK-NEXT:    beq .LBB15_12
 ; CHECK-NEXT:  @ %bb.11: @ %if.then59
 ; CHECK-NEXT:    vldrw.u32 q0, [r0]
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..34d93c985e7204 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
@@ -22,11 +22,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