[BOLT] Fix instrumenting conditional tail calls

[Amir]

Summary

We identify instructions to be instrumented based on Offset annotation.

BOLT "expands" conditional tail calls into a conditional jump to a basic block
with unconditional tail call. Move Offset annotation from former CTC to the tail
call.

For expanded CTC we keep Offset attached to the original instruction which is
converted into a regular conditional jump, while leaving the newly created tail
call without an Offset annotation. This leads to attempting the instrumentation
of the conditional jump which points to the basic block with an inherited input
offset thus creating an invalid edge description. At the same time, the newly
created tail call is skipped entirely which means we're not creating a call
description for it.

If we instead reassign Offset annotation from the conditional jump to the tail
call we fix both issues. The conditional jump will be skipped not creating an
invalid edge description, while tail call will be handled properly (unformly
with regular calls).

[maksfb]

Thanks for the quick fix. Can we use instruction annotations instead of IsCondTailCall basic block property for the conditional tail call identification?

[Amir]

This comment was removed by Amir.

[Amir]

Replace BB->isCondTailCall property with BB->getNumPseudos() == 1 and
BB->getLastNonPseudo is a tail call.

[Amir]

Identify cond tail call blocks by their offset (being equal to parent blocks' input offset)

[maksfb]

bolt/lib/Passes/Instrumentation.cpp

103–104 ↗

(On Diff #545015)

Can we merge IsInvoke and IsCondTailCall into a single argument, e.g. NeedsInstrumentation?

231–232 ↗

(On Diff #545015)

Same args can me merged here too.

396–397 ↗

(On Diff #545015)

While this approach works, a more explicit way to identify CTCs will be to mark TC instructions created from CTCs with a proper annotation and check for this annotation during instrumentation.

[Amir]

Move Offset annotation from former CTC to the tail call.

We identify instructions to be instrumented based on Offset annotation.

For expanded CTC we keep Offset attached to the original instruction which is
converted into a regular conditional jump, while leaving the newly created tail
call without an Offset annotation. This leads to attempting the instrumentation
of the conditional jump which points to the basic block with an inherited input
offset thus creating an invalid edge description. At the same time, the newly
created tail call is skipped entirely which means we're not creating a call
description for it.

If we instead reassign Offset annotation from the conditional jump to the tail
call we fix both issues. The conditional jump will be skipped not creating an
invalid edge description, while tail call will be handled properly (unformly
with regular calls).

[maksfb]

Nice find!

bolt/lib/Core/BinaryFunction.cpp

for (auto BBI = begin(); BBI != end(); ++BBI) {
     // Add basic block to the list that will be added to the end.
     NewBlocks.emplace_back(std::move(TailCallBB));
 
     // Swap edges as the TailCallBB corresponds to the taken branch.
     BB.swapConditionalSuccessors();
 
     // This branch is no longer a conditional tail call.
     BC.MIB->unsetConditionalTailCall(*CTCInstr);
+
+    // Move offset from CTCInstr to TailCallInstr.
+    if (std::optional<uint32_t> Offset = BC.MIB->getOffset(*CTCInstr)) {
+      BC.MIB->setOffset(TailCallInstr, *Offset);
+      BC.MIB->clearOffset(*CTCInstr);
+    }
   }
 
   insertBasicBlocks(std::prev(end()), std::move(NewBlocks),
                     /* UpdateLayout */ true,
                     /* UpdateCFIState */ false);
 }
 
 uint64_t BinaryFunction::getFunctionScore() const {

bolt/test/runtime/X86/instrumentation-tail-call.s

 # 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
 
+# RUN: FileCheck %s --input-file %t.fdata --check-prefix=CHECK-FDATA
+# CHECK-FDATA: 1 main {{.*}} 1 targetFunc 0 0 1
+
   .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
+  test %rsp, %rsp
+  jne 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