[BOLT] Verify externally referenced blocks against jump table targets

[Amir]

Summary

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.

[Amir]

Verify the set of internal references against jump table targets.

[rafauler]

bolt/lib/Core/BinaryFunction.cpp

1897

What happens if we do have JumpTables and yet we have internal references that do not appear in any jump table?

1898

From reading this code I got the impression that UnknownControlFlow is only used in strict mode?

[rafauler]

I guess I reviewed an older version of the diff

[Amir]

HasUnknownControlFlow -> !Strict

[Amir]

Move checks past HasFixedIndirectBranch

[Amir]

clang-format

[Amir]

Update warning message to make it clear for end user what's going on

bolt/lib/Core/BinaryFunction.cpp

1923

@ayermolo: hope the warning makes sense now

[Amir]

Add extra check

[Amir]

Mark as non-simple

[Amir]

Added assembly test

[Amir]

Fix shrinkwrapping-restore-position test

[Amir]

Add runtime test checking if the moved code with external data references to it
works correctly. If there's a reference to offset which is not a registered BB
or label, ignore the function.

[Amir]

clang-format

[Amir]

Check for constant islands

[maksfb]

bolt/lib/Core/BinaryFunction.cpp

1911

Let's refactor the code into a function since it's quite large.

1916

1931

1933

1939

1940

[Amir]

Move checks into BF::validateExternalReferences

[Amir]

Address nits

[Amir]

rebase

[Amir]

Ping @maksfb

[maksfb]

bolt/include/bolt/Core/BinaryFunction.h

2122

It's reasonable to expect for validate...() function to return true if all checks pass. Please invert the return code.

bolt/lib/Core/BinaryFunction.cpp

1766

Is there a reason to keep Found in a variable?

1778

If you don't expect this warning to be firing frequently, let's print it at default verbosity level.

[Amir]

Address Maksim's comments

[maksfb]

LGTM, but please rename the method before the commit.

bolt/lib/Core/BinaryFunction.cpp

1911–1912

There's an apparent confusion with "internal" vs "external" reference in this context. Let's rename the method to validateExternallyReferencedOffsets() or validateReferencedOffsets() to avoid it.

[Amir]

Rename method to validateExternallyReferencedOffsets

[Amir]

clang-format

bolt/include/bolt/Core/BinaryFunction.h

   /// first we make our best guess about a branch without knowing the CFG,
   /// and later after we have the CFG for the function, we verify our earlier
   /// assumptions and also do our best at processing unknown indirect branches.
   ///
   /// Return true upon successful processing, or false if the control flow
   /// 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; }
 
   const IndirectCallSiteProfile &getAllCallSites() const {
     return AllCallSites;
   }
 
   /// Walks the list of basic blocks filling in missing information about

bolt/lib/Core/BinaryFunction.cpp

void BinaryFunction::postProcessJumpTables() {
   // Remove duplicates branches. We can get a bunch of them from jump tables.
   // Without doing jump table value profiling we don't have use for extra
   // (duplicate) branches.
   llvm::sort(TakenBranches);
   auto NewEnd = std::unique(TakenBranches.begin(), TakenBranches.end());
   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) {
     HasUnknownControlFlow = true;
     BB.removeAllSuccessors();
     for (uint64_t PossibleDestination : ExternallyReferencedOffsets)
       if (BinaryBasicBlock *SuccBB = getBasicBlockAtOffset(PossibleDestination))
         BB.addSuccessor(SuccBB);
bool BinaryFunction::postProcessIndirectBranches(
   }
 
   if (HasInternalLabelReference)
     return false;
 
   // If there's only one jump table, and one indirect jump, and no other
   // references, then we should be able to derive the jump table even if we
   // fail to match the pattern.
   if (HasUnknownControlFlow && NumIndirectJumps == 1 &&
       JumpTables.size() == 1 && LastIndirectJump) {
     BC.MIB->setJumpTable(*LastIndirectJump, LastJT, LastJTIndexReg, AllocId);
     HasUnknownControlFlow = false;
 
     LastIndirectJumpBB->updateJumpTableSuccessors();
   }
 
   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;
 
   return true;
 }
 
 void BinaryFunction::recomputeLandingPads() {
   updateBBIndices(0);
 
   for (BinaryBasicBlock *BB : BasicBlocks) {
     BB->LandingPads.clear();
     BB->Throwers.clear();
   }
 
   for (BinaryBasicBlock *BB : BasicBlocks) {
     std::unordered_set<const BinaryBasicBlock *> BBLandingPads;
     for (MCInst &Instr : *BB) {
       if (!BC.MIB->isInvoke(Instr))
         continue;
 
       const Optional<MCPlus::MCLandingPad> EHInfo = BC.MIB->getEHInfo(Instr);
       if (!EHInfo || !EHInfo->first)
         continue;
 
       BinaryBasicBlock *LPBlock = getBasicBlockForLabel(EHInfo->first);
       if (!BBLandingPads.count(LPBlock)) {
         BBLandingPads.insert(LPBlock);
         BB->LandingPads.emplace_back(LPBlock);
         LPBlock->Throwers.emplace_back(BB);
       }
     }
   }
 }

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

cold_path:
   ret
 hot_path:
   addq  $0x20, %rsp
   pop %r14
   pop %rbx
   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
 
 # CHECK-OBJDUMP:     <_start>:
 # CHECK-OBJDUMP:         leaq    (%rbp), %rsp
 # CHECK-OBJDUMP-NEXT:    popq    %rbp
 # CHECK-OBJDUMP-NEXT:    movq    -0x10(%rsp), %rbx
 # CHECK-OBJDUMP-NEXT:    movq    -0x18(%rsp), %r14
 # CHECK-OBJDUMP-NEXT:    retq

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

+# 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

+#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

+# 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