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

maksfb · web-flow · commit 59ab29213def · 2024-05-01T21:56:55.000-07:00
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.
diff --git a/bolt/lib/Rewrite/LinuxKernelRewriter.cpp b/bolt/lib/Rewrite/LinuxKernelRewriter.cpp
@@ -1722,6 +1722,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
@@ -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

Original file line number	Diff line number	Diff line change
`@@ -1722,6 +1722,9 @@ Error LinuxKernelRewriter::readStaticKeysJumpTable() {`
`1722`	`1722`	`if (!BC.MIB->getSize(*Inst))`
`1723`	`1723`	`BC.MIB->setSize(*Inst, Size);`
`1724`	`1724`
	`1725`	`+ if (!BC.MIB->getOffset(*Inst))`
	`1726`	`+ BC.MIB->setOffset(*Inst, JumpAddress - BF->getAddress());`
	`1727`	`+`
`1725`	`1728`	`if (opts::LongJumpLabels)`
`1726`	`1729`	`BC.MIB->setSize(*Inst, 5);`
`1727`	`1730`	`}`