[llvm] [XRay] Fix tail call sleds for AArch64 (PR #141403)

[Sebastian Kreutzer via llvm-commits]

https://github.com/sebastiankreutzer updated https://github.com/llvm/llvm-project/pull/141403

>From cc60644270dac8f33748a51df94c60df0aee3772 Mon Sep 17 00:00:00 2001
From: Sebastian Kreutzer <SebastianKreutzer at gmx.net>
Date: Fri, 23 May 2025 16:08:04 +0200
Subject: [PATCH 1/2] [XRay] Fix tail call sleds for AArch64

---
 llvm/lib/CodeGen/XRayInstrumentation.cpp         |  5 +++--
 llvm/lib/Target/AArch64/AArch64InstrInfo.cpp     |  1 +
 llvm/test/CodeGen/AArch64/xray-tail-call-sled.ll | 16 ++++++++++++++++
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/XRayInstrumentation.cpp b/llvm/lib/CodeGen/XRayInstrumentation.cpp
index dbdb81d1e6b33..cd00e362e16db 100644
--- a/llvm/lib/CodeGen/XRayInstrumentation.cpp
+++ b/llvm/lib/CodeGen/XRayInstrumentation.cpp
@@ -296,8 +296,9 @@ bool XRayInstrumentation::run(MachineFunction &MF) {
     case Triple::ArchType::riscv64: {
       // For the architectures which don't have a single return instruction
       InstrumentationOptions op;
-      // RISC-V supports patching tail calls.
-      op.HandleTailcall = MF.getTarget().getTargetTriple().isRISCV();
+      // AArch64 and RISC-V support patching tail calls.
+      op.HandleTailcall = MF.getTarget().getTargetTriple().isAArch64() ||
+                          MF.getTarget().getTargetTriple().isRISCV() ;
       op.HandleAllReturns = true;
       prependRetWithPatchableExit(MF, TII, op);
       break;
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index 1a13adc300d2b..21e6fd28d3b71 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -163,6 +163,7 @@ unsigned AArch64InstrInfo::getInstSizeInBytes(const MachineInstr &MI) const {
         F.getFnAttributeAsParsedInteger("patchable-function-entry", 9) * 4;
     break;
   case TargetOpcode::PATCHABLE_FUNCTION_EXIT:
+  case TargetOpcode::PATCHABLE_TAIL_CALL:
   case TargetOpcode::PATCHABLE_TYPED_EVENT_CALL:
     // An XRay sled can be 4 bytes of alignment plus a 32-byte block.
     NumBytes = 36;
diff --git a/llvm/test/CodeGen/AArch64/xray-tail-call-sled.ll b/llvm/test/CodeGen/AArch64/xray-tail-call-sled.ll
index 484dca11104ad..a62257e5c9b40 100644
--- a/llvm/test/CodeGen/AArch64/xray-tail-call-sled.ll
+++ b/llvm/test/CodeGen/AArch64/xray-tail-call-sled.ll
@@ -21,8 +21,12 @@ define i32 @callee() nounwind noinline uwtable "function-instrument"="xray-alway
 ; CHECK-LINUX-LABEL: .Lxray_sleds_start0:
 ; CHECK-LINUX-NEXT:  [[TMP:.Ltmp[0-9]+]]:
 ; CHECK-LINUX:         .xword .Lxray_sled_0-[[TMP]]
+; CHECK-LINUX:         .xword
+; CHECK-LINUX-NEXT:    .byte   0x00
 ; CHECK-LINUX:       [[TMP:.Ltmp[0-9]+]]:
 ; CHECK-LINUX-NEXT:    .xword .Lxray_sled_1-[[TMP]]
+; CHECK-LINUX:         .xword
+; CHECK-LINUX-NEXT:    .byte   0x01
 ; CHECK-LINUX-LABEL: Lxray_sleds_end0:
 ; CHECK-LINUX-LABEL: .section xray_fn_idx,"ao", at progbits,callee{{$}}
 ; CHECK-LINUX:         .xword .Lxray_sleds_start0
@@ -32,8 +36,12 @@ define i32 @callee() nounwind noinline uwtable "function-instrument"="xray-alway
 ; CHECK-MACOS-LABEL: lxray_sleds_start0:
 ; CHECK-MACOS-NEXT:  [[TMP:Ltmp[0-9]+]]:
 ; CHECK-MACOS:         .quad Lxray_sled_0-[[TMP]]
+; CHECK-MACOS-NEXT:    .quad
+; CHECK-MACOS-NEXT:    .byte   0x00
 ; CHECK-MACOS:       [[TMP:Ltmp[0-9]+]]:
 ; CHECK-MACOS-NEXT:    .quad Lxray_sled_1-[[TMP]]
+; CHECK-MACOS-NEXT:    .quad
+; CHECK-MACOS-NEXT:    .byte   0x01
 ; CHECK-MACOS-LABEL: Lxray_sleds_end0:
 ; CHECK-MACOS-LABEL: .section __DATA,xray_fn_idx,regular,live_support{{$}}
 ; CHECK-MACOS:       [[IDX:lxray_fn_idx[0-9]+]]:
@@ -60,7 +68,11 @@ define i32 @caller() nounwind noinline uwtable "function-instrument"="xray-alway
 ; CHECK-LINUX-LABEL: .section xray_instr_map,"ao", at progbits,caller{{$}}
 ; CHECK-LINUX-LABEL: Lxray_sleds_start1:
 ; CHECK-LINUX:         .xword .Lxray_sled_2
+; CHECK-LINUX-NEXT:    .xword
+; CHECK-LINUX-NEXT:    .byte   0x00
 ; CHECK-LINUX:         .xword .Lxray_sled_3
+; CHECK-LINUX-NEXT:    .xword
+; CHECK-LINUX-NEXT:    .byte   0x02
 ; CHECK-LINUX-LABEL: Lxray_sleds_end1:
 ; CHECK-LINUX-LABEL: .section xray_fn_idx,"ao", at progbits,caller{{$}}
 ; CHECK-LINUX:       [[IDX:\.Lxray_fn_idx[0-9]+]]:
@@ -70,7 +82,11 @@ define i32 @caller() nounwind noinline uwtable "function-instrument"="xray-alway
 ; CHECK-MACOS-LABEL: .section __DATA,xray_instr_map,regular,live_support{{$}}
 ; CHECK-MACOS-LABEL: lxray_sleds_start1:
 ; CHECK-MACOS:         .quad Lxray_sled_2
+; CHECK-MACOS-NEXT:    .quad
+; CHECK-MACOS-NEXT:    .byte   0x00
 ; CHECK-MACOS:         .quad Lxray_sled_3
+; CHECK-MACOS-NEXT:    .quad
+; CHECK-MACOS-NEXT:    .byte   0x02
 ; CHECK-MACOS-LABEL: Lxray_sleds_end1:
 ; CHECK-MACOS-LABEL: .section __DATA,xray_fn_idx,regular,live_support{{$}}
 ; CHECK-MACOS:       [[IDX:lxray_fn_idx[0-9]+]]:

>From 5d15b98c8da1dcac01b13f53f88e9b6410022ce3 Mon Sep 17 00:00:00 2001
From: Sebastian Kreutzer <SebastianKreutzer at gmx.net>
Date: Sun, 25 May 2025 11:40:51 +0200
Subject: [PATCH 2/2] fixup! [XRay] Fix tail call sleds for AArch64

---
 llvm/lib/CodeGen/XRayInstrumentation.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/XRayInstrumentation.cpp b/llvm/lib/CodeGen/XRayInstrumentation.cpp
index cd00e362e16db..d2f40581ecebf 100644
--- a/llvm/lib/CodeGen/XRayInstrumentation.cpp
+++ b/llvm/lib/CodeGen/XRayInstrumentation.cpp
@@ -298,7 +298,7 @@ bool XRayInstrumentation::run(MachineFunction &MF) {
       InstrumentationOptions op;
       // AArch64 and RISC-V support patching tail calls.
       op.HandleTailcall = MF.getTarget().getTargetTriple().isAArch64() ||
-                          MF.getTarget().getTargetTriple().isRISCV() ;
+                          MF.getTarget().getTargetTriple().isRISCV();
       op.HandleAllReturns = true;
       prependRetWithPatchableExit(MF, TII, op);
       break;