[llvm] [ARM/Thumb] Mark TAILJMP*/TCRETURN* as using LR. (PR #73492)

[Florian Hahn via llvm-commits]

https://github.com/fhahn created https://github.com/llvm/llvm-project/pull/73492

TCRETURN/TAILJMP may use LR in the tail-called function to return to the original caller. At the moment, they are not marked as using LR, which can cause mis-compiles, e.g. if the MachineOutliner introduces new calls to outlined functions using BL (see outlined-fn-may-clobber-lr-in-caller.ll for an example).

I might be missing a better way to address this issue, but without marking LR as used, I am not sure how liveness of LR can be correctly tracked by the MachineOutliner to decide whether LR needs to be preserved across outlined calls.

The current version is causing some MachineVerifier failures in some cases, due to LR not being marked as live-in properly directly after ISel (see below for verifier error message). I am not really sure why this is happening.

 | *** Bad machine code: Using an undefined physical register ***
 | - function:    func
 | - basic block: %bb.2  (0x12280bc10)
 | - instruction: TCRETURNri %1:tcgpr, 0, <regmask $lr $d8 $d9 $d10 $d11 $d12 $d13 $d14 $d15 $q4 $q5 $q6 $q7 $r4 $r5 $r6 $r7 $r8 $r10 $r11 $s16 $s17 $s18 $s19 $s20 $s21 $s22 $s23 $s24 $s25 $s26 $s27 $s28 and 33 more...>, implicit $sp, implicit $lr, implicit $r0
 | - operand 4:   implicit $lr
 | LLVM ERROR: Found 1 machine code errors.

>From 80eacc16e7d2888807c7e0aff529431bdf9911d1 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Mon, 27 Nov 2023 09:24:39 +0000
Subject: [PATCH] [ARM/Thumb] Mark TAILJMP*/TCRETURN* as using LR.

TCRETURN/TAILJMP may use LR in the tail-called function to return to the
original caller. At the moment, they are not marked as using LR, which
can cause mis-compiles, e.g. if the MachineOutliner introduces new calls
to outlined functions using BL (see outlined-fn-may-clobber-lr-in-caller.ll
for an example).

I might be missing a better way to address this issue, but without
marking LR as used, I am not sure how liveness of LR can be correctly
tracked by the MachineOutliner to decide whether LR needs to be
preserved across outlined calls.

The current version is causing some MachineVerifier failures in some
cases, due to LR not being marked as live-in properly directly after
ISel (see below for verifier error message). I am not really sure why
this is happening.

 | *** Bad machine code: Using an undefined physical register ***
 | - function:    func
 | - basic block: %bb.2  (0x12280bc10)
 | - instruction: TCRETURNri %1:tcgpr, 0, <regmask $lr $d8 $d9 $d10 $d11 $d12 $d13 $d14 $d15 $q4 $q5 $q6 $q7 $r4 $r5 $r6 $r7 $r8 $r10 $r11 $s16 $s17 $s18 $s19 $s20 $s21 $s22 $s23 $s24 $s25 $s26 $s27 $s28 and 33 more...>, implicit $sp, implicit $lr, implicit $r0
 | - operand 4:   implicit $lr
 | LLVM ERROR: Found 1 machine code errors.
---
 llvm/lib/Target/ARM/ARMBaseInstrInfo.cpp                     | 4 ++++
 llvm/lib/Target/ARM/ARMInstrInfo.td                          | 2 +-
 llvm/lib/Target/ARM/ARMInstrThumb2.td                        | 2 +-
 .../CodeGen/Thumb2/outlined-fn-may-clobber-lr-in-caller.ll   | 5 +++--
 4 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Target/ARM/ARMBaseInstrInfo.cpp b/llvm/lib/Target/ARM/ARMBaseInstrInfo.cpp
index c09879fd9c2beb3..d3f896435862cca 100644
--- a/llvm/lib/Target/ARM/ARMBaseInstrInfo.cpp
+++ b/llvm/lib/Target/ARM/ARMBaseInstrInfo.cpp
@@ -6587,6 +6587,10 @@ void ARMBaseInstrInfo::buildOutlinedFrame(
   // For thunk outlining, rewrite the last instruction from a call to a
   // tail-call.
   if (OF.FrameConstructionID == MachineOutlinerThunk) {
+    // LR is used by the tail call.
+    if (!MBB.isLiveIn(ARM::LR))
+      MBB.addLiveIn(ARM::LR);
+
     MachineInstr *Call = &*--MBB.instr_end();
     bool isThumb = Subtarget.isThumb();
     unsigned FuncOp = isThumb ? 2 : 0;
diff --git a/llvm/lib/Target/ARM/ARMInstrInfo.td b/llvm/lib/Target/ARM/ARMInstrInfo.td
index 812b5730875d5d4..3e0df3afccb1028 100644
--- a/llvm/lib/Target/ARM/ARMInstrInfo.td
+++ b/llvm/lib/Target/ARM/ARMInstrInfo.td
@@ -2668,7 +2668,7 @@ def BXJ : ABI<0b0001, (outs), (ins GPR:$func), NoItinerary, "bxj", "\t$func",
 
 // Tail calls.
 
-let isCall = 1, isTerminator = 1, isReturn = 1, isBarrier = 1, Uses = [SP] in {
+let isCall = 1, isTerminator = 1, isReturn = 1, isBarrier = 1, Uses = [SP,LR] in {
   def TCRETURNdi : PseudoInst<(outs), (ins i32imm:$dst, i32imm:$SPDiff), IIC_Br, []>,
                    Sched<[WriteBr]>;
 
diff --git a/llvm/lib/Target/ARM/ARMInstrThumb2.td b/llvm/lib/Target/ARM/ARMInstrThumb2.td
index acd46e8093aa78d..0b1ae71b5bcb48b 100644
--- a/llvm/lib/Target/ARM/ARMInstrThumb2.td
+++ b/llvm/lib/Target/ARM/ARMInstrThumb2.td
@@ -4035,7 +4035,7 @@ def t2Bcc : T2I<(outs), (ins brtarget:$target), IIC_Br,
 // Windows SEH unwinding also needs a strict t2 branch for tail calls.
 let isCall = 1, isTerminator = 1, isReturn = 1, isBarrier = 1 in {
   // IOS version.
-  let Uses = [SP] in
+  let Uses = [SP,LR] in
   def tTAILJMPd: tPseudoExpand<(outs),
                    (ins thumb_br_target:$dst, pred:$p),
                    4, IIC_Br, [],
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 d81d008b44bed89..a9e59cecfc2417b 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
@@ -16,6 +16,8 @@ target datalayout = "e-m:o-p:32:32-Fi8-i64:64-v128:64:128-a:0:32-n32-S64"
 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
+; CHECK-NEXT:    .save {r7, lr}
+; CHECK-NEXT:    push {r7, lr}
 ; CHECK-NEXT:    cmp r1, #2
 ; CHECK-NEXT:    beq .LBB0_3
 ; CHECK-NEXT:  @ %bb.1: @ %bb
@@ -32,10 +34,9 @@ define void @test(ptr nocapture noundef writeonly %arg, i32 noundef %arg1, i8 no
 ; CHECK-NEXT:    add.w r1, r2, r1, lsl #2
 ; CHECK-NEXT:    adds r0, #4
 ; CHECK-NEXT:    movs r2, #30
+; CHECK-NEXT:    pop.w {r7, lr}
 ; CHECK-NEXT:    b __aeabi_memcpy
 ; CHECK-NEXT:  .LBB0_5: @ %bb24
-; CHECK-NEXT:    .save {r7, lr}
-; CHECK-NEXT:    push {r7, lr}
 ; CHECK-NEXT:    bl wombat
 ; CHECK-NEXT:    @APP
 ; CHECK-NEXT:    @NO_APP