[llvm] [BOLT] Register Linux kernel dynamic branch offsets (PR #90677)

[Maksim Panchenko via llvm-commits]

https://github.com/maksfb created https://github.com/llvm/llvm-project/pull/90677

To match profile data to code we need to know instruction offsets within a function. For this reason, we mark branches with an "Offset" annotation while disassembling the code. However, dynamic branches could be NOPs in the code. We need to explicitly add the "Offset" annotation in such cases.

>From 99050838a44bd92ac83ab47e9dd32d54a046f1de Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks at fb.com>
Date: Mon, 29 Apr 2024 17:54:15 -0700
Subject: [PATCH] [BOLT] Register Linux kernel dynamic branch offsets

To match profile data to code we need to know instruction offsets within
a function. For this reason, we mark branches with an "Offset"
annotation while disassembling the code. However, dynamic branches could
be NOPs in the code. We need to explicitly add the "Offset" annotation
in such cases.
---
 bolt/lib/Rewrite/LinuxKernelRewriter.cpp |  3 +++
 bolt/test/X86/linux-static-keys.s        | 29 +++++++++++++++++++-----
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/bolt/lib/Rewrite/LinuxKernelRewriter.cpp b/bolt/lib/Rewrite/LinuxKernelRewriter.cpp
index 3944bb742938a7..d95fcedddf7a9e 100644
--- a/bolt/lib/Rewrite/LinuxKernelRewriter.cpp
+++ b/bolt/lib/Rewrite/LinuxKernelRewriter.cpp
@@ -1695,6 +1695,9 @@ Error LinuxKernelRewriter::readStaticKeysJumpTable() {
     if (!BC.MIB->getSize(*Inst))
       BC.MIB->setSize(*Inst, Size);
 
+    if (!BC.MIB->getOffset(*Inst))
+      BC.MIB->setOffset(*Inst, JumpAddress - BF->getAddress());
+
     if (opts::LongJumpLabels)
       BC.MIB->setSize(*Inst, 5);
   }
diff --git a/bolt/test/X86/linux-static-keys.s b/bolt/test/X86/linux-static-keys.s
index 08454bf9763193..fb419e0f762755 100644
--- a/bolt/test/X86/linux-static-keys.s
+++ b/bolt/test/X86/linux-static-keys.s
@@ -3,6 +3,8 @@
 ## Check that BOLT correctly updates the Linux kernel static keys jump table.
 
 # RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
+# RUN: link_fdata %s %t.o %t.fdata
+# RUN: llvm-strip --strip-unneeded %t.o
 # RUN: %clang %cflags -nostdlib %t.o -o %t.exe \
 # RUN:   -Wl,--image-base=0xffffffff80000000,--no-dynamic-linker,--no-eh-frame-hdr
 
@@ -11,6 +13,12 @@
 # RUN: llvm-bolt %t.exe --print-normalized -o %t.out --keep-nops=0 \
 # RUN:   --bolt-info=0 |& FileCheck %s
 
+## Verify that profile is matched correctly.
+
+# RUN: llvm-bolt %t.exe --print-normalized -o %t.out --keep-nops=0 \
+# RUN:   --bolt-info=0 --data %t.fdata |& \
+# RUN:   FileCheck --check-prefix=CHECK-FDATA %s
+
 ## Verify the bindings again on the rewritten binary with nops removed.
 
 # RUN: llvm-bolt %t.out -o %t.out.1 --print-normalized |& FileCheck %s
@@ -25,15 +33,24 @@ _start:
 # CHECK: Binary Function "_start"
   nop
 .L0:
-  jmp .L1
+  jmp L1
 # CHECK:      jit
 # CHECK-SAME: # ID: 1 {{.*}} # Likely: 0 # InitValue: 1
   nop
-.L1:
+L1:
   .nops 5
+  jmp .L0
 # CHECK:      jit
 # CHECK-SAME: # ID: 2 {{.*}} # Likely: 1 # InitValue: 1
-.L2:
+
+## Check that a branch profile associated with a NOP is handled properly when
+## dynamic branch is created.
+
+# FDATA: 1 _start #L1# 1 _start #L2# 3 42
+# CHECK-FDATA: jit {{.*}} # ID: 2
+# CHECK-FDATA-NEXT: jmp
+# CHECK-FDATA-NEXT: Successors: {{.*}}  (mispreds: 3, count: 42)
+L2:
   nop
   .size _start, .-_start
 
@@ -51,11 +68,11 @@ foo:
 __start___jump_table:
 
   .long .L0 - . # Jump address
-  .long .L1 - . # Target address
+  .long L1 - . # Target address
   .quad 1       # Key address
 
-  .long .L1 - . # Jump address
-  .long .L2 - . # Target address
+  .long L1 - . # Jump address
+  .long L2 - . # Target address
   .quad 0       # Key address
 
   .globl __stop___jump_table