[llvm] [xray] Do not emit tail exit hooks for conditional tail calls (PR #89364)

[Ricky Zhou via llvm-commits]

https://github.com/rickyz updated https://github.com/llvm/llvm-project/pull/89364

>From 0e2970b1c4ad8287656860e692133fc3c1227c75 Mon Sep 17 00:00:00 2001
From: Ricky Zhou <ricky at rzhou.org>
Date: Fri, 19 Apr 2024 02:44:07 -0700
Subject: [PATCH] [xray] Do not emit tail exit hooks for conditional tail calls

xray instruments tail call function exits by inserting a nop sled before
the tail call. When tracing is enabled, the nop sled is replaced with a
call to `__xray_FunctionTailExit()`. This currently does not work for
conditional tail calls, as the instrumentation assumes that the tail
call will be unconditional. This causes two issues:
 - `__xray_FunctionTailExit()` is inappropately called even when the
   tail call is not taken.
 - `__xray_FunctionTailExit()`'s prologue/epilogue adjusts the stack
   pointer with add/sub instructions. This clobbers condition flags,
   which can flip the condition used for the tail call, leading to
   incorrect program behavior.

Avoid this misbehavior by disabling instrumentation of conditional tail
calls. A more comprehensive fix could be to either rewrite conditional
tail calls to non-conditional ones, or to create specialized function
tail exit instrumentation functions for each of the possible types of
conditional tail calls.
---
 llvm/lib/CodeGen/XRayInstrumentation.cpp      |  4 +--
 .../CodeGen/X86/xray-conditional-tail-call.ll | 28 +++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/X86/xray-conditional-tail-call.ll

diff --git a/llvm/lib/CodeGen/XRayInstrumentation.cpp b/llvm/lib/CodeGen/XRayInstrumentation.cpp
index d40725838c943c..e5791ec12d3e5f 100644
--- a/llvm/lib/CodeGen/XRayInstrumentation.cpp
+++ b/llvm/lib/CodeGen/XRayInstrumentation.cpp
@@ -100,7 +100,7 @@ void XRayInstrumentation::replaceRetWithPatchableRet(
         //   PATCHABLE_RET <Opcode>, <Operand>...
         Opc = TargetOpcode::PATCHABLE_RET;
       }
-      if (TII->isTailCall(T) && op.HandleTailcall) {
+      if (TII->isTailCall(T) && T.isBarrier() && op.HandleTailcall) {
         // Treat the tail call as a return instruction, which has a
         // different-looking sled than the normal return case.
         Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
@@ -131,7 +131,7 @@ void XRayInstrumentation::prependRetWithPatchableExit(
           (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
         Opc = TargetOpcode::PATCHABLE_FUNCTION_EXIT;
       }
-      if (TII->isTailCall(T) && op.HandleTailcall) {
+      if (TII->isTailCall(T) && T.isBarrier() && op.HandleTailcall) {
         Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
       }
       if (Opc != 0) {
diff --git a/llvm/test/CodeGen/X86/xray-conditional-tail-call.ll b/llvm/test/CodeGen/X86/xray-conditional-tail-call.ll
new file mode 100644
index 00000000000000..270e4da3ce50bd
--- /dev/null
+++ b/llvm/test/CodeGen/X86/xray-conditional-tail-call.ll
@@ -0,0 +1,28 @@
+; RUN: llc -mtriple=x86_64 < %s | FileCheck %s
+
+declare void @tail_call_target()
+
+define void @conditional_tail_call(i32 %cond) "function-instrument"="xray-always" nounwind {
+  ; CHECK-LABEL: conditional_tail_call:
+  ; CHECK-NEXT:  .Lfunc_begin0:
+  ; CHECK-NEXT:  # %bb.0:
+  ; CHECK-NEXT:    .p2align 1, 0x90
+  ; CHECK-NEXT:  .Lxray_sled_0:
+  ; CHECK-NEXT:    .ascii "\353\t"
+  ; CHECK-NEXT:    nopw 512(%rax,%rax)
+  ; CHECK-NEXT:    testl %edi, %edi
+  ; CHECK-NEXT:    jne tail_call_target at PLT # TAILCALL
+  ; CHECK-NEXT:  # %bb.1:
+  ; CHECK-NEXT:    .p2align 1, 0x90
+  ; CHECK-NEXT:  .Lxray_sled_1:
+  ; CHECK-NEXT:    retq
+  ; CHECK-NEXT:    nopw %cs:512(%rax,%rax)
+  ; CHECK-NEXT:  .Lfunc_end0:
+  %cmp = icmp ne i32 %cond, 0
+  br i1 %cmp, label %docall, label %ret
+docall:
+  tail call void @tail_call_target()
+  ret void
+ret:
+  ret void
+}