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

[Amir Ayupov via llvm-commits]

https://github.com/aaupov created https://github.com/llvm/llvm-project/pull/91693

Offset annotation is used in writing BAT tables.

Test Plan: updated sctc-bug4.test


>From 2ec259b50889d25e4b0c63ababc0e9af40c081b9 Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov at fb.com>
Date: Thu, 9 May 2024 20:03:15 -0700
Subject: [PATCH] [BOLT] Preserve Offset annotation in SCTC

Offset annotation is used in writing BAT tables.

Test Plan: updated sctc-bug4.test
---
 bolt/lib/Passes/BinaryPasses.cpp | 3 +++
 bolt/test/X86/sctc-bug4.test     | 3 +++
 2 files changed, 6 insertions(+)

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