[BOLT] Add rewriting support for Linux kernel __bug_table

bolt/lib/Rewrite/LinuxKernelRewriter.cpp

@@ -212,6 +212,11 @@ class LinuxKernelRewriter final : public MetadataRewriter {
   /// Size of bug_entry struct.
   static constexpr size_t BUG_TABLE_ENTRY_SIZE = 12;
 
+  /// List of bug entries per function.
+  using FunctionBugListType =
+      DenseMap<BinaryFunction *, SmallVector<uint32_t, 2>>;
+  FunctionBugListType FunctionBugList;
+
   /// .pci_fixup section.
   ErrorOr<BinarySection &> PCIFixupSection = std::errc::bad_address;
   static constexpr size_t PCI_FIXUP_ENTRY_SIZE = 16;
@@ -254,7 +259,9 @@ class LinuxKernelRewriter final : public MetadataRewriter {
   Error readParaInstructions();
   Error rewriteParaInstructions();
 
+  /// __bug_table section handling.
   Error readBugTable();
+  Error rewriteBugTable();
 
   /// Do no process functions containing instruction annotated with
   /// \p Annotation.
@@ -339,6 +346,9 @@ class LinuxKernelRewriter final : public MetadataRewriter {
     if (Error E = rewriteStaticKeysJumpTable())
       return E;
 
+    if (Error E = rewriteBugTable())
+      return E;
+
     return Error::success();
   }
 
@@ -1164,15 +1174,17 @@ Error LinuxKernelRewriter::rewriteParaInstructions() {
 }
 
 /// Process __bug_table section.
-/// This section contains information useful for kernel debugging.
+/// This section contains information useful for kernel debugging, mostly
+/// utilized by WARN()/WARN_ON() macros and deprecated BUG()/BUG_ON().
+///
 /// Each entry in the section is a struct bug_entry that contains a pointer to
 /// the ud2 instruction corresponding to the bug, corresponding file name (both
 /// pointers use PC relative offset addressing), line number, and flags.
 /// The definition of the struct bug_entry can be found in
-/// `include/asm-generic/bug.h`
-///
-/// NB: find_bug() uses linear search to match an address to an entry in the bug
-///     table. Hence there is no need to sort entries when rewriting the table.
+/// `include/asm-generic/bug.h`. The first entry in the struct is an instruction
+/// address encoded as a PC-relative offset. In theory, it could be an absolute
+/// address if CONFIG_GENERIC_BUG_RELATIVE_POINTERS is not set, but in practice
+/// the kernel code relies on it being a relative offset on x86-64.
 Error LinuxKernelRewriter::readBugTable() {
   BugTableSection = BC.getUniqueSectionByName("__bug_table");
   if (!BugTableSection)
@@ -1215,6 +1227,8 @@ Error LinuxKernelRewriter::readBugTable() {
                                  " referenced by bug table entry %d",
                                  InstAddress, EntryID);
       BC.MIB->addAnnotation(*Inst, "BugEntry", EntryID);
+
+      FunctionBugList[BF].push_back(EntryID);
     }
   }
 
@@ -1223,6 +1237,52 @@ Error LinuxKernelRewriter::readBugTable() {
   return Error::success();
 }
 
+/// find_bug() uses linear search to match an address to an entry in the bug
+/// table. Hence, there is no need to sort entries when rewriting the table.
+/// When we need to erase an entry, we set its instruction address to zero.
+Error LinuxKernelRewriter::rewriteBugTable() {
+  if (!BugTableSection)
+    return Error::success();
+
+  for (BinaryFunction &BF : llvm::make_second_range(BC.getBinaryFunctions())) {
+    if (!BC.shouldEmit(BF))
+      continue;
+
+    if (!FunctionBugList.count(&BF))
+      continue;
+
+    // Bugs that will be emitted for this function.
+    DenseSet<uint32_t> EmittedIDs;
+    for (BinaryBasicBlock &BB : BF) {
+      for (MCInst &Inst : BB) {
+        if (!BC.MIB->hasAnnotation(Inst, "BugEntry"))
+          continue;
+        const uint32_t ID = BC.MIB->getAnnotationAs<uint32_t>(Inst, "BugEntry");
+        EmittedIDs.insert(ID);
+
+        // Create a relocation entry for this bug entry.
+        MCSymbol *Label =
+            BC.MIB->getOrCreateInstLabel(Inst, "__BUG_", BC.Ctx.get());
+        const uint64_t EntryOffset = (ID - 1) * BUG_TABLE_ENTRY_SIZE;
+        BugTableSection->addRelocation(EntryOffset, Label, ELF::R_X86_64_PC32,
+                                       /*Addend*/ 0);
+      }
+    }
+
+    // Clear bug entries that were not emitted for this function, e.g. as a
+    // result of DCE, but setting their instruction address to zero.
+    for (const uint32_t ID : FunctionBugList[&BF]) {
+      if (!EmittedIDs.count(ID)) {
+        const uint64_t EntryOffset = (ID - 1) * BUG_TABLE_ENTRY_SIZE;
+        BugTableSection->addRelocation(EntryOffset, nullptr, ELF::R_X86_64_PC32,
+                                       /*Addend*/ 0);
+      }
+    }
+  }
+
+  return Error::success();
+}
+
 /// The kernel can replace certain instruction sequences depending on hardware
 /// it is running on and features specified during boot time. The information
 /// about alternative instruction sequences is stored in .altinstructions

bolt/test/X86/linux-bug-table.s

@@ -1,14 +1,21 @@
 # REQUIRES: system-linux
 
-## Check that BOLT correctly parses the Linux kernel __bug_table section.
+## Check that BOLT correctly parses and updates the Linux kernel __bug_table
+## section.
 
 # RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
 # RUN: %clang %cflags -nostdlib %t.o -o %t.exe \
 # RUN:   -Wl,--image-base=0xffffffff80000000,--no-dynamic-linker,--no-eh-frame-hdr,--no-pie
 
 ## Verify bug entry bindings to instructions.
 
-# RUN: llvm-bolt %t.exe --print-normalized -o %t.out | FileCheck %s
+# RUN: llvm-bolt %t.exe --print-normalized --print-only=_start -o %t.out \
+# RUN:   --eliminate-unreachable=1 --bolt-info=0 | FileCheck %s
+
+## Verify bug entry bindings again after unreachable code elimination.
+
+# RUN: llvm-bolt %t.out -o %t.out.1 --print-only=_start --print-normalized \
+# RUN:   |& FileCheck --check-prefix=CHECK-REOPT %s
 
 # CHECK:      BOLT-INFO: Linux kernel binary detected
 # CHECK:      BOLT-INFO: parsed 2 bug table entries
@@ -17,17 +24,21 @@
   .globl _start
   .type _start, %function
 _start:
-# CHECK: Binary Function "_start"
-  nop
+  jmp .L1
 .L0:
   ud2
 # CHECK:      ud2
 # CHECK-SAME: BugEntry: 1
-  nop
 .L1:
   ud2
 # CHECK:      ud2
 # CHECK-SAME: BugEntry: 2
+
+## Only the second entry should remain after the first pass.
+
+# CHECK-REOPT: ud2
+# CHECK-REOPT-SAME: BugEntry: 2
+
   ret
   .size _start, .-_start