[BOLT][Instrumentation] Preserve red zone for functions with tail calls only

Allow a function with tail calls only to clobber its red zone.

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

Reviewed By: #bolt, yota9

Differential Revision: https://reviews.llvm.org/D145202

2 files changed, 58 insertions, 6 deletions

diff --git a/bolt/lib/Passes/Instrumentation.cpp b/bolt/lib/Passes/Instrumentation.cpp
index a0350b14d2b5..c6e1bf9e6b16 100644
--- a/bolt/lib/Passes/Instrumentation.cpp
+++ b/bolt/lib/Passes/Instrumentation.cpp
@@ -357,12 +357,13 @@ void Instrumentation::instrumentFunction(BinaryFunction &Function,
   // instructions to protect the red zone
   bool IsLeafFunction = true;
   DenseSet<const BinaryBasicBlock *> InvokeBlocks;
-  for (auto BBI = Function.begin(), BBE = Function.end(); BBI != BBE; ++BBI) {
-    for (auto I = BBI->begin(), E = BBI->end(); I != E; ++I) {
-      if (BC.MIB->isCall(*I)) {
-        if (BC.MIB->isInvoke(*I))
-          InvokeBlocks.insert(&*BBI);
-        IsLeafFunction = false;
+  for (const BinaryBasicBlock &BB : Function) {
+    for (const MCInst &Inst : BB) {
+      if (BC.MIB->isCall(Inst)) {
+        if (BC.MIB->isInvoke(Inst))
+          InvokeBlocks.insert(&BB);
+        if (!BC.MIB->isTailCall(Inst))
+          IsLeafFunction = false;
       }
     }
   }
diff --git a/bolt/test/runtime/X86/instrumentation-tail-call.s b/bolt/test/runtime/X86/instrumentation-tail-call.s
new file mode 100644
index 000000000000..792d084e3f3d
--- /dev/null
+++ b/bolt/test/runtime/X86/instrumentation-tail-call.s
@@ -0,0 +1,51 @@
+# This reproduces a bug with instrumentation when trying to instrument
+# a function with only tail calls. Such functions can clobber red zone,
+# see https://github.com/llvm/llvm-project/issues/61114.
+
+# REQUIRES: system-linux,bolt-runtime
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
+# RUN: %clang %cflags -no-pie %t.o -o %t.exe -Wl,-q
+
+# RUN: llvm-bolt %t.exe --instrument --instrumentation-file=%t.fdata \
+# RUN:   -o %t.instrumented
+# RUN: %t.instrumented arg1 arg2
+# RUN: llvm-objdump %t.instrumented --disassemble-symbols=main | FileCheck %s
+
+# CHECK: leaq 0x80(%rsp), %rsp
+
+  .text
+  .globl  main
+  .type main, %function
+  .p2align  4
+main:
+  pushq %rbp
+  movq  %rsp, %rbp
+  mov   %rax,-0x10(%rsp)
+  leaq targetFunc, %rax
+  pushq %rax                  # We save the target function address in the stack
+  subq  $0x18, %rsp           # Set up a dummy stack frame
+  cmpl  $0x2, %edi
+  jb    .LBBerror             # Add control flow so we don't have a trivial case
+.LBB2:
+  addq $0x20, %rsp
+  movq %rbp, %rsp
+  pop %rbp
+  mov -0x10(%rsp),%rax
+  jmp targetFunc
+
+.LBBerror:
+  addq $0x20, %rsp
+  movq %rbp, %rsp
+  pop %rbp
+  movq $1, %rax               # Finish with an error if we go this path
+  retq
+  .size main, .-main
+
+  .globl targetFunc
+  .type targetFunc, %function
+  .p2align  4
+targetFunc:
+  xorq %rax, %rax
+  retq
+  .size targetFunc, .-targetFunc