From ec6b31130f29370648c384e94310beeeabe588c1 Mon Sep 17 00:00:00 2001
From: Antonio Frighetto <me@antoniofrighetto.com>
Date: Sat, 12 Oct 2024 17:14:20 +0200
Subject: [PATCH] [X86] Do not elect to tail call if caller must preserve all
 registers

A miscompilation issue has been addressed with improved checking.

Fixes: https://github.com/llvm/llvm-project/issues/97758.
---
 llvm/lib/Target/X86/X86ISelLoweringCall.cpp    | 7 +++++++
 llvm/test/CodeGen/X86/tailcall-caller-nocsr.ll | 4 +++-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/X86/X86ISelLoweringCall.cpp b/llvm/lib/Target/X86/X86ISelLoweringCall.cpp
index 8561658379f7e4..eb166ff7002df5 100644
--- a/llvm/lib/Target/X86/X86ISelLoweringCall.cpp
+++ b/llvm/lib/Target/X86/X86ISelLoweringCall.cpp
@@ -2856,6 +2856,13 @@ bool X86TargetLowering::IsEligibleForTailCallOptimization(
       return false;
   }
 
+  // The stack frame of the caller cannot be replaced by the tail-callee one's
+  // if the function is required to preserve all the registers. Conservatively
+  // prevent tail optimization even if hypothetically all the registers are used
+  // for passing formal parameters or returning values.
+  if (CLI.CB->getFunction()->hasFnAttribute("no_caller_saved_registers"))
+    return false;
+
   unsigned StackArgsSize = CCInfo.getStackSize();
 
   // If the callee takes no arguments then go on to check the results of the
diff --git a/llvm/test/CodeGen/X86/tailcall-caller-nocsr.ll b/llvm/test/CodeGen/X86/tailcall-caller-nocsr.ll
index 5606fbb27032fb..0385017a1ced73 100644
--- a/llvm/test/CodeGen/X86/tailcall-caller-nocsr.ll
+++ b/llvm/test/CodeGen/X86/tailcall-caller-nocsr.ll
@@ -13,8 +13,10 @@ define void @caller(i32 %0, i32 %1) #0 {
 ; CHECK-NEXT:    pushq %rdx
 ; CHECK-NEXT:    pushq %rcx
 ; CHECK-NEXT:    pushq %rax
+; CHECK-NEXT:    movl %esi, %edx
 ; CHECK-NEXT:    movl %edi, %esi
 ; CHECK-NEXT:    movl $.L.str, %edi
+; CHECK-NEXT:    callq printf@PLT
 ; CHECK-NEXT:    popq %rax
 ; CHECK-NEXT:    popq %rcx
 ; CHECK-NEXT:    popq %rdx
@@ -22,7 +24,7 @@ define void @caller(i32 %0, i32 %1) #0 {
 ; CHECK-NEXT:    popq %r9
 ; CHECK-NEXT:    popq %r10
 ; CHECK-NEXT:    popq %r11
-; CHECK-NEXT:    jmp printf@PLT # TAILCALL
+; CHECK-NEXT:    retq
   %3 = tail call i32 @printf(ptr @.str, i32 %0, i32 %1)
   ret void
 }