[BOLT] Preserve Offset annotation in SCTC (#91693)

aaupov · web-flow · commit 6b9bca8faaa7 · 2024-05-10T13:20:51.000-07:00
Offset annotation is used in writing BAT tables.

Test Plan: updated sctc-bug4.test
diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
@@ -910,6 +910,11 @@ uint64_t SimplifyConditionalTailCalls::fixTailCalls(BinaryFunction &BF) {
       auto &CTCAnnotation =
           MIB->getOrCreateAnnotationAs<uint64_t>(*CondBranch, "CTCTakenCount");
       CTCAnnotation = CTCTakenFreq;
+      // Preserve Offset annotation, used in BAT.
+      // Instr is a direct tail call instruction that was created when CTCs are
+      // first expanded, and has the original CTC offset set.
+      if (std::optional<uint32_t> Offset = MIB->getOffset(*Instr))
+        MIB->setOffset(*CondBranch, *Offset);
 
       // Remove the unused successor which may be eliminated later
       // if there are no other users.
diff --git a/bolt/test/X86/sctc-bug4.test b/bolt/test/X86/sctc-bug4.test
@@ -1,25 +1,29 @@
-# Check that fallthrough blocks are handled properly.
+# Check that fallthrough blocks are handled properly and Offset annotation is
+# set for conditional tail calls.
 
 RUN: %clang %cflags %S/Inputs/sctc_bug4.s -o %t
-RUN: llvm-bolt %t -o %t.null \
+RUN: llvm-bolt %t -o %t.null --enable-bat \
 RUN:   -funcs=test_func -print-sctc -sequential-disassembly 2>&1 | FileCheck %s
 
 CHECK:      .Ltmp2 (3 instructions, align : 1)
 CHECK-NEXT:   CFI State : 0
+CHECK-NEXT:   Input offset: 0x24
 CHECK-NEXT:   Predecessors: .LFT1
 CHECK-NEXT:     00000024: 	cmpq	$0x20, %rsi
-CHECK-NEXT:     00000028: 	ja	dummy # TAILCALL {{.*}}# CTCTakenCount: 0
+CHECK-NEXT:     00000028: 	ja	dummy # TAILCALL # Offset: 53 # CTCTakenCount: 0
 CHECK-NEXT:     0000002a: 	jmp .Ltmp4
 CHECK-NEXT:   Successors: .Ltmp4
 CHECK-NEXT:   CFI State: 0
 
 CHECK:      .Ltmp1 (2 instructions, align : 1)
 CHECK-NEXT:   CFI State : 0
+CHECK-NEXT:   Input offset: 0x2c
 CHECK-NEXT:   Predecessors: .LFT0
 CHECK-NEXT:     0000002c: 	xorq	%r11, %rax
 CHECK-NEXT:     0000002f: 	retq
 CHECK-NEXT:   CFI State: 0
 
 CHECK:      .Ltmp4 (4 instructions, align : 1)
 CHECK-NEXT:  CFI State : 0
+CHECK-NEXT:  Input offset: 0x3a
 CHECK-NEXT:  Predecessors: .Ltmp2
