[BOLT] Preserve Offset annotation in SCTC

[aaupov]

Offset annotation is used in writing BAT tables.

Test Plan: updated sctc-bug4.test

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Amir Ayupov (aaupov)

Changes

Offset annotation is used in writing BAT tables.

Test Plan: updated sctc-bug4.test

---

Full diff: https://github.com/llvm/llvm-project/pull/91693.diff

2 Files Affected:

• (modified) bolt/lib/Passes/BinaryPasses.cpp (+3)
• (modified) bolt/test/X86/sctc-bug4.test (+3)

diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index c0ba73108f577..40cb300a73aaf 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -910,6 +910,9 @@ uint64_t SimplifyConditionalTailCalls::fixTailCalls(BinaryFunction &BF) {
       auto &CTCAnnotation =
           MIB->getOrCreateAnnotationAs<uint64_t>(*CondBranch, "CTCTakenCount");
       CTCAnnotation = CTCTakenFreq;
+      // Preserve Offset annotation, used for BAT
+      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
index 00f5ee429b635..8f60e9359a36c 100644
--- a/bolt/test/X86/sctc-bug4.test
+++ b/bolt/test/X86/sctc-bug4.test
@@ -3,12 +3,15 @@
 RUN: %clang %cflags %S/Inputs/sctc_bug4.s -o %t
 RUN: llvm-bolt %t -o %t.null \
 RUN:   -funcs=test_func -print-sctc -sequential-disassembly 2>&1 | FileCheck %s
+RUN: llvm-bolt %t -o %t.null --enable-bat -funcs=test_func -print-sctc \
+RUN:   -sequential-disassembly 2>&1 | FileCheck %s --check-prefix=CHECK-BAT
 
 CHECK:      .Ltmp2 (3 instructions, align : 1)
 CHECK-NEXT:   CFI State : 0
 CHECK-NEXT:   Predecessors: .LFT1
 CHECK-NEXT:     00000024: 	cmpq	$0x20, %rsi
 CHECK-NEXT:     00000028: 	ja	dummy # TAILCALL {{.*}}# CTCTakenCount: 0
+CHECK-BAT:      00000028: 	ja	dummy # TAILCALL # Offset: 53
 CHECK-NEXT:     0000002a: 	jmp .Ltmp4
 CHECK-NEXT:   Successors: .Ltmp4
 CHECK-NEXT:   CFI State: 0