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

[Sebastian Kreutzer via llvm-commits]

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

This addresses issues #141051.
XRay uses a special event kind for tail calls on some architectures. This feature is implemented on AArch64, but wasn't fully activated.
Tests in `llvm/test/CodeGen/AArch64/xray-tail-call-sled.ll` were incomplete and did not check for the emitted sled type.
This patch correctly enables emission of tail call sleds on AArch64 and fixes the tests to check the sled kind.

>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] [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]+]]: