[BOLT] Register Linux kernel dynamic branch offsets

[maksfb]

To match profile data to code we need to know branch instruction offsets within a function. For this reason, we mark branches with the "Offset" annotation while disassembling the code. However, dynamic branches in the Linux kernel could be NOPs in disassembled code, and we ignore them while adding annotations. We need to explicitly add the "Offset" annotation while creating dynamic branches.

Note that without this change, getInstructionAtOffset() would still return a branch instruction if the offset matched the last instruction in a basic block (and the profile data was matched correctly). However, the function failed for cases when the searched instruction was followed by an unconditional jump. "Offset" annotation solves this case.

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Maksim Panchenko (maksfb)

Changes

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.

---

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

2 Files Affected:

• (modified) bolt/lib/Rewrite/LinuxKernelRewriter.cpp (+3)
• (modified) bolt/test/X86/linux-static-keys.s (+23-6)

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

[dcci]

SG, we discussed this at length some time ago (and I also remember the explanation at EuroLLVM 2023)

[maksfb]

SG, we discussed this at length some time ago (and I also remember the explanation at EuroLLVM 2023)

It's a bit more nuanced. Updated the summary with details.