[llvm] [X86] Do not elect to tail call if caller must preserve all registers (PR #112098)

[Antonio Frighetto via llvm-commits]

https://github.com/antoniofrighetto created https://github.com/llvm/llvm-project/pull/112098

A miscompilation issue has been addressed with improved checking.

Fixes: https://github.com/llvm/llvm-project/issues/97758.

>From c04361dd3db85d3f5855c780efeab7b1f1aa113e Mon Sep 17 00:00:00 2001
From: Antonio Frighetto <me at antoniofrighetto.com>
Date: Sat, 12 Oct 2024 17:10:21 +0200
Subject: [PATCH 1/2] [X86] Introduce test for PR112098 (NFC)

---
 .../test/CodeGen/X86/tailcall-caller-nocsr.ll | 32 +++++++++++++++++++
 1 file changed, 32 insertions(+)
 create mode 100644 llvm/test/CodeGen/X86/tailcall-caller-nocsr.ll

diff --git a/llvm/test/CodeGen/X86/tailcall-caller-nocsr.ll b/llvm/test/CodeGen/X86/tailcall-caller-nocsr.ll
new file mode 100644
index 00000000000000..5606fbb27032fb
--- /dev/null
+++ b/llvm/test/CodeGen/X86/tailcall-caller-nocsr.ll
@@ -0,0 +1,32 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -mtriple=x86_64-linux-gnu -mattr=-sse,-avx | FileCheck %s
+
+ at .str = private unnamed_addr constant [6 x i8] c"%d %d\00", align 1
+
+define void @caller(i32 %0, i32 %1) #0 {
+; CHECK-LABEL: caller:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    pushq %r11
+; CHECK-NEXT:    pushq %r10
+; CHECK-NEXT:    pushq %r9
+; CHECK-NEXT:    pushq %r8
+; CHECK-NEXT:    pushq %rdx
+; CHECK-NEXT:    pushq %rcx
+; CHECK-NEXT:    pushq %rax
+; CHECK-NEXT:    movl %edi, %esi
+; CHECK-NEXT:    movl $.L.str, %edi
+; CHECK-NEXT:    popq %rax
+; CHECK-NEXT:    popq %rcx
+; CHECK-NEXT:    popq %rdx
+; CHECK-NEXT:    popq %r8
+; CHECK-NEXT:    popq %r9
+; CHECK-NEXT:    popq %r10
+; CHECK-NEXT:    popq %r11
+; CHECK-NEXT:    jmp printf at PLT # TAILCALL
+  %3 = tail call i32 @printf(ptr @.str, i32 %0, i32 %1)
+  ret void
+}
+
+declare i32 @printf(ptr, ...) nounwind
+
+attributes #0 = { mustprogress nounwind "no_caller_saved_registers" }

>From ec6b31130f29370648c384e94310beeeabe588c1 Mon Sep 17 00:00:00 2001
From: Antonio Frighetto <me at antoniofrighetto.com>
Date: Sat, 12 Oct 2024 17:14:20 +0200
Subject: [PATCH 2/2] [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 at 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 at PLT # TAILCALL
+; CHECK-NEXT:    retq
   %3 = tail call i32 @printf(ptr @.str, i32 %0, i32 %1)
   ret void
 }