[BOLT] Verify externally referenced blocks against jump table targets

For functions with references to internal offsets from data, verify externally
referenced blocks against the set of jump table targets. Mark the function
as non-simple if there are any unclaimed data to code references.

Reviewed By: #bolt, maksfb

Differential Revision: https://reviews.llvm.org/D132495

bolt/include/bolt/Core/BinaryFunction.h

@@ -2115,6 +2115,12 @@ public:
   /// cannot be statically evaluated for any given indirect branch.
   bool postProcessIndirectBranches(MCPlusBuilder::AllocatorIdTy AllocId);
 
+  /// Validate that all data references to function offsets are claimed by
+  /// recognized jump tables. Register externally referenced blocks as entry
+  /// points. Returns true if there are no unclaimed externally referenced
+  /// offsets.
+  bool validateExternallyReferencedOffsets();
+
   /// Return all call site profile info for this function.
   IndirectCallSiteProfile &getAllCallSites() { return AllCallSites; }
 

bolt/lib/Core/BinaryFunction.cpp

@@ -1748,6 +1748,43 @@ void BinaryFunction::postProcessJumpTables() {
   TakenBranches.erase(NewEnd, TakenBranches.end());
 }
 
+bool BinaryFunction::validateExternallyReferencedOffsets() {
+  SmallPtrSet<MCSymbol *, 4> JTTargets;
+  for (const JumpTable *JT : llvm::make_second_range(JumpTables))
+    JTTargets.insert(JT->Entries.begin(), JT->Entries.end());
+
+  bool HasUnclaimedReference = false;
+  for (uint64_t Destination : ExternallyReferencedOffsets) {
+    // Ignore __builtin_unreachable().
+    if (Destination == getSize())
+      continue;
+    // Ignore constant islands
+    if (isInConstantIsland(Destination + getAddress()))
+      continue;
+
+    if (BinaryBasicBlock *BB = getBasicBlockAtOffset(Destination)) {
+      // Check if the externally referenced offset is a recognized jump table
+      // target.
+      if (JTTargets.contains(BB->getLabel()))
+        continue;
+
+      if (opts::Verbosity >= 1) {
+        errs() << "BOLT-WARNING: unclaimed data to code reference (possibly "
+               << "an unrecognized jump table entry) to " << BB->getName()
+               << " in " << *this << "\n";
+      }
+      auto L = BC.scopeLock();
+      addEntryPoint(*BB);
+    } else {
+      errs() << "BOLT-WARNING: unknown data to code reference to offset "
+             << Twine::utohexstr(Destination) << " in " << *this << "\n";
+      setIgnored();
+    }
+    HasUnclaimedReference = true;
+  }
+  return !HasUnclaimedReference;
+}
+
 bool BinaryFunction::postProcessIndirectBranches(
     MCPlusBuilder::AllocatorIdTy AllocId) {
   auto addUnknownControlFlow = [&](BinaryBasicBlock &BB) {
@@ -1868,6 +1905,14 @@ bool BinaryFunction::postProcessIndirectBranches(
   if (HasFixedIndirectBranch)
     return false;
 
+  // Validate that all data references to function offsets are claimed by
+  // recognized jump tables. Register externally referenced blocks as entry
+  // points.
+  if (!opts::StrictMode && hasInternalReference()) {
+    if (!validateExternallyReferencedOffsets())
+      return false;
+  }
+
   if (HasUnknownControlFlow && !BC.HasRelocations)
     return false;
 

bolt/test/X86/shrinkwrapping-restore-position.s

@@ -50,10 +50,11 @@ hot_path:
   pop %rbp
   ret
   .cfi_endproc
+end:
   .size _start, .-_start
 
   .data
-rel:  .quad cold_path
+rel:  .quad end
 
 # CHECK:   BOLT-INFO: Shrink wrapping moved 2 spills inserting load/stores and 0 spills inserting push/pops
 

bolt/test/X86/unclaimed-jt-entries.s

@@ -0,0 +1,66 @@
+# This test ensures that "unclaimed" jump table entries are accounted later
+# in postProcessIndirectBranches and the function is marked as non-simple.
+
+# The test is compiled from the following source using GCC 12.2 -O3:
+# https://godbolt.org/z/YcPG131s6
+# int func(long long Input) {
+#   switch(Input) {
+#   case 3: return 1;
+#   case 4: return 2;
+#   case 6: return 3;
+#   case 8: return 4;
+#   case 13: return 5;
+#   default: __builtin_unreachable();
+#   }
+# }
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
+# RUN: %clang %cflags -no-pie %t.o -o %t.exe -Wl,-q
+# RUN: llvm-bolt %t.exe -v=1 -o %t.out |& FileCheck %s
+
+# CHECK: BOLT-WARNING: unclaimed data to code reference (possibly an unrecognized jump table entry) to .Ltmp[[#]] in main
+# CHECK: BOLT-WARNING: unclaimed data to code reference (possibly an unrecognized jump table entry) to .Ltmp[[#]] in main
+# CHECK: BOLT-WARNING: unclaimed data to code reference (possibly an unrecognized jump table entry) to .Ltmp[[#]] in main
+# CHECK: BOLT-WARNING: unclaimed data to code reference (possibly an unrecognized jump table entry) to .Ltmp[[#]] in main
+# CHECK: BOLT-WARNING: unclaimed data to code reference (possibly an unrecognized jump table entry) to .Ltmp[[#]] in main
+# CHECK: BOLT-WARNING: failed to post-process indirect branches for main
+
+  .text
+  .globl main
+  .type main, %function
+  .size main, .Lend-main
+main:
+  jmp *L4-24(,%rdi,8)
+.L5:
+  movl $4, %eax
+  ret
+.L9:
+  movl $2, %eax
+  ret
+.L8:
+  movl $1, %eax
+  ret
+.L3:
+  movl $5, %eax
+  ret
+.L6:
+  movl $3, %eax
+  ret
+.Lend:
+
+.section .rodata
+  .globl L4
+L4:
+  .quad .L8
+  .quad .L9
+  .quad .L3
+  .quad .L6
+  .quad .L3
+  .quad .L5
+  .quad .L3
+  .quad .L3
+  .quad .L3
+  .quad .L3
+  .quad .L3

bolt/test/runtime/X86/Inputs/unclaimed-jt-entries.c

@@ -0,0 +1,10 @@
+#include <stdio.h>
+#include <stdlib.h>
+
+int func(long long Input);
+
+int main(int argc, char *argv[]) {
+  int arg = atoi(argv[1]);
+  printf("%d\n", func(arg));
+  return 0;
+}

bolt/test/runtime/X86/unclaimed-jt-entries.s

@@ -0,0 +1,78 @@
+# This test ensures that "unclaimed" jump table entries are accounted later
+# in postProcessIndirectBranches and the function is marked as non-simple.
+
+# The test is compiled from the following source using GCC 12.2 -O3:
+# https://godbolt.org/z/YcPG131s6
+# int func(long long Input) {
+#   switch(Input) {
+#   case 3: return 1;
+#   case 4: return 2;
+#   case 6: return 3;
+#   case 8: return 4;
+#   case 13: return 5;
+#   default: __builtin_unreachable();
+#   }
+# }
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
+# RUN: %clang %cflags %S/Inputs/unclaimed-jt-entries.c -no-pie %t.o -o %t.exe -Wl,-q
+# RUN: llvm-bolt %t.exe -v=1 -o %t.out --sequential-disassembly |& FileCheck %s
+
+# CHECK: BOLT-WARNING: unclaimed data to code reference (possibly an unrecognized jump table entry) to .Ltmp[[#]] in func
+# CHECK: BOLT-WARNING: unclaimed data to code reference (possibly an unrecognized jump table entry) to .Ltmp[[#]] in func
+# CHECK: BOLT-WARNING: unclaimed data to code reference (possibly an unrecognized jump table entry) to .Ltmp[[#]] in func
+# CHECK: BOLT-WARNING: unclaimed data to code reference (possibly an unrecognized jump table entry) to .Ltmp[[#]] in func
+# CHECK: BOLT-WARNING: unclaimed data to code reference (possibly an unrecognized jump table entry) to .Ltmp[[#]] in func
+# CHECK: BOLT-WARNING: failed to post-process indirect branches for func
+
+# Run the optimized binary
+# RUN: %t.out 3 | FileCheck %s --check-prefix=CHECK3
+# CHECK3: 1
+# RUN: %t.out 4 | FileCheck %s --check-prefix=CHECK4
+# CHECK4: 2
+# RUN: %t.out 6 | FileCheck %s --check-prefix=CHECK6
+# CHECK6: 3
+# RUN: %t.out 8 | FileCheck %s --check-prefix=CHECK8
+# CHECK8: 4
+# RUN: %t.out 13 | FileCheck %s --check-prefix=CHECK13
+# CHECK13: 5
+
+  .text
+  .globl func
+  .type func, %function
+  .size func, .Lend-func
+func:
+  jmp *L4-24(,%rdi,8)
+.L5:
+  movl $4, %eax
+  ret
+.L9:
+  movl $2, %eax
+  ret
+.L8:
+  movl $1, %eax
+  ret
+.L3:
+  movl $5, %eax
+  ret
+.L6:
+  movl $3, %eax
+  ret
+.Lend:
+
+.section .rodata
+  .globl L4
+L4:
+  .quad .L8
+  .quad .L9
+  .quad .L3
+  .quad .L6
+  .quad .L3
+  .quad .L5
+  .quad .L3
+  .quad .L3
+  .quad .L3
+  .quad .L3
+  .quad .L3