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

[Amir Ayupov via llvm-commits]

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

>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 1/3] [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

>From 0bc277df6d538daf65b0a17810a9f3929d6958d7 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 2/3] Address comments

---
 bolt/lib/Passes/BinaryPasses.cpp | 4 +++-
 bolt/test/X86/sctc-bug4.test     | 8 +++++---
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index 40cb300a73aaf..df6dbcddeed56 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -910,7 +910,9 @@ uint64_t SimplifyConditionalTailCalls::fixTailCalls(BinaryFunction &BF) {
       auto &CTCAnnotation =
           MIB->getOrCreateAnnotationAs<uint64_t>(*CondBranch, "CTCTakenCount");
       CTCAnnotation = CTCTakenFreq;
-      // Preserve Offset annotation, used for BAT
+      // 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);
 
diff --git a/bolt/test/X86/sctc-bug4.test b/bolt/test/X86/sctc-bug4.test
index 8f60e9359a36c..912637bffcf07 100644
--- a/bolt/test/X86/sctc-bug4.test
+++ b/bolt/test/X86/sctc-bug4.test
@@ -1,23 +1,24 @@
 # Check that fallthrough blocks are handled properly.
 
 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
 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:   Input offset: 0x24
 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:     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
@@ -25,4 +26,5 @@ 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

>From 8fdd3a022737bd8407afd033ac04f9d520579d8f Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov at fb.com>
Date: Fri, 10 May 2024 12:07:22 -0700
Subject: [PATCH 3/3] Update test description

---
 bolt/test/X86/sctc-bug4.test | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/bolt/test/X86/sctc-bug4.test b/bolt/test/X86/sctc-bug4.test
index 912637bffcf07..890b442140bfe 100644
--- a/bolt/test/X86/sctc-bug4.test
+++ b/bolt/test/X86/sctc-bug4.test
@@ -1,4 +1,5 @@
-# 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 --enable-bat \