[BOLT] Allow function fragments to point to one jump table

[nhuhuan]

Summary

Resolve a crash related to split functions

Due to split function optimization, a function can be divided to two

fragments, and both fragments can access same jump table. This
violates 
the assumption that a jump table can only have one parent
function, 
which causes a crash during instrumentation.

We want to support the case: different functions cannot access same
jump tables, but different fragments of same function can!

As all fragments are from same function, we point JT::Parent to one
specific fragment. Right now it is the first disassembled fragment, but
we can point it to the function's main fragment later.

Functions are disassembled sequentially. Previously, at the end of
processing a function, JT::OffsetEntries is cleared, so other fragment
can no longer reuse JT::OffsetEntries. To extend the support for split
function, we only clear JT::OffsetEntries after all functions are
disassembled.

Let say A.hot and A.cold access JT of three targets {X, Y, Z}, where
X and Y are in A.hot, and Z is in A.cold. Suppose that A.hot is
disassembled first, JT::OffsetEntries = {X',Y',INVALID_OFFSET}. When
A.cold is disassembled, it cannot reuse JT::OffsetEntries above due to
different fragment start. A simple solution:
A.hot = {X',Y',INVALID_OFFSET}
A.cold = {INVALID_OFFSET, INVALID_OFFSET, INVALID_OFFSET}

We update the assertion to allow different fragments of same function
to get the same JumpTable object.

Potential improvements:
A.hot = {X',Y',INVALID_OFFSET}
A.cold = {INVALID_OFFSET, INVALID_OFFSET, Z'}
The main issue is A.hot and A.cold have separate CFGs, thus jump table
targets are still constrained within fragment bounds.

Future improvements:
A.hot = {X, Y, Z}
A.cold = {X, Y, Z}

[nhuhuan]

Include all the changes

[Amir]

Please retitle the diff to state what is changed - e.g. "allow function fragments point to one jump table".

I like the summary, it explains the issue in a clear manner. Some nits:

• no need to explicitly mention python,
• I would rephrase 'next goal' and 'stretch goal' as potential or future improvements.
• please run a spell checker.

bolt/test/X86/split-func-jump-table-fragment-bidirection.s

5

Please updated the message

[Amir]

bolt/test/X86/split-func-jump-table-fragment-bidirection.s

12

[nhuhuan]

Update manual test

[Amir]

LGTM. Ready to land after fixing formatting issues: https://buildkite.com/llvm-project/premerge-checks/builds/98297#01816db9-c455-4e2c-8799-2bb1f9ec608d/3412-3413

[nhuhuan]

clang-formatted changes

bolt/include/bolt/Core/BinaryContext.h

public:
   std::pair<uint64_t, const MCSymbol *>
   duplicateJumpTable(BinaryFunction &Function, JumpTable *JT,
                      const MCSymbol *OldLabel);
 
   /// Generate a unique name for jump table at a given \p Address belonging
   /// to function \p BF.
   std::string generateJumpTableName(const BinaryFunction &BF, uint64_t Address);
 
+  /// Free memory used by jump table offsets
+  void clearJumpTableOffsets() {
+    for (auto &JTI : JumpTables) {
+      JumpTable &JT = *JTI.second;
+      JumpTable::OffsetsType Temp;
+      Temp.swap(JT.OffsetEntries);
+    }
+  }
   /// Return true if the array of bytes represents a valid code padding.
   bool hasValidCodePadding(const BinaryFunction &BF);
 
   /// Verify padding area between functions, and adjust max function size
   /// accordingly.
   void adjustCodePadding();
 
   /// Regular page size.

bolt/lib/Core/BinaryContext.cpp

registerNameAtAddress(Name, Address, SymbolSize ? SymbolSize : Size,
                         Alignment);
   setSymbolToFunctionMap(BF->getSymbol(), BF);
   return BF;
 }
 
 const MCSymbol *
 BinaryContext::getOrCreateJumpTable(BinaryFunction &Function, uint64_t Address,
                                     JumpTable::JumpTableType Type) {
+  auto isFragmentOf = [](BinaryFunction *Fragment, BinaryFunction *Parent) {
+    return (Fragment->isFragment() && Fragment->isParentFragment(Parent));
+  };
+
   if (JumpTable *JT = getJumpTableContainingAddress(Address)) {
     assert(JT->Type == Type && "jump table types have to match");
-    assert(JT->Parent == &Function &&
+    bool HasMultipleParents = isFragmentOf(JT->Parent, &Function) ||
+                              isFragmentOf(&Function, JT->Parent);
+    assert((JT->Parent == &Function || HasMultipleParents) &&
            "cannot re-use jump table of a different function");
     assert(Address == JT->getAddress() && "unexpected non-empty jump table");
 
+    // Flush OffsetEntries with INVALID_OFFSET if multiple parents
+    // Duplicate the entry for the parent function for easy access
+    if (HasMultipleParents) {
+      if (opts::Verbosity > 2) {
+        outs() << "BOLT-WARNING: Multiple fragments access same jump table: "
+               << JT->Parent->getPrintName() << "; " << Function.getPrintName()
+               << "\n";
+      }
+      constexpr uint64_t INVALID_OFFSET = std::numeric_limits<uint64_t>::max();
+      for (unsigned I = 0; I < JT->OffsetEntries.size(); ++I)
+        JT->OffsetEntries[I] = INVALID_OFFSET;
+      Function.JumpTables.emplace(Address, JT);
+    }
     return JT->getFirstLabel();
   }
 
   // Re-use the existing symbol if possible.
   MCSymbol *JTLabel = nullptr;
   if (BinaryData *Object = getBinaryDataAtAddress(Address)) {
     if (!isInternalSymbolName(Object->getSymbol()->getName()))
       JTLabel = Object->getSymbol();

bolt/lib/Core/BinaryFunction.cpp

void BinaryFunction::postProcessJumpTables() {
   for (auto &JTI : JumpTables) {
     JumpTable &JT = *JTI.second;
     if (JT.Type == JumpTable::JTT_PIC && opts::JumpTables == JTS_BASIC) {
       opts::JumpTables = JTS_MOVE;
       outs() << "BOLT-INFO: forcing -jump-tables=move as PIC jump table was "
                 "detected in function "
              << *this << '\n';
     }
+    if (JT.Entries.empty()) {
       for (unsigned I = 0; I < JT.OffsetEntries.size(); ++I) {
         MCSymbol *Label =
             getOrCreateLocalLabel(getAddress() + JT.OffsetEntries[I],
                                   /*CreatePastEnd*/ true);
         JT.Entries.push_back(Label);
       }
+    }
 
     const uint64_t BDSize =
         BC.getBinaryDataAtAddress(JT.getAddress())->getSize();
     if (!BDSize) {
       BC.setBinaryDataSize(JT.getAddress(), JT.getSize());
     } else {
       assert(BDSize >= JT.getSize() &&
              "jump table cannot be larger than the containing object");
while (EntryOffset < JT->getSize()) {
 
       // A label at the next entry means the end of this jump table.
       if (JT->Labels.count(EntryOffset))
         break;
     }
   }
   clearList(JTSites);
 
-  // Free memory used by jump table offsets.
-  for (auto &JTI : JumpTables) {
-    JumpTable &JT = *JTI.second;
-    clearList(JT.OffsetEntries);
-  }
-
   // Conservatively populate all possible destinations for unknown indirect
   // branches.
   if (opts::StrictMode && hasInternalReference()) {
     for (uint64_t Offset : UnknownIndirectBranchOffsets) {
       for (uint64_t PossibleDestination : ExternallyReferencedOffsets) {
         // Ignore __builtin_unreachable().
         if (PossibleDestination == getSize())
           continue;

bolt/lib/Rewrite/RewriteInstance.cpp

for (auto &BFI : BC->getBinaryFunctions()) {
     }
 
     if (opts::PrintAll || opts::PrintDisasm)
       Function.print(outs(), "after disassembly", true);
 
     BC->processInterproceduralReferences(Function);
   }
 
+  BC->clearJumpTableOffsets();
   BC->populateJumpTables();
   BC->skipMarkedFragments();
 
   for (auto &BFI : BC->getBinaryFunctions()) {
     BinaryFunction &Function = BFI.second;
 
     if (!shouldDisassemble(Function))
       continue;

bolt/test/X86/split-func-jump-table-fragment-bidirection.s

+# This reproduces an issue where two fragments of same function access same
+# jump table, which means at least one fragment visits the other, i.e., one
+# of them has split jump table. As a result, all of them will be marked as
+# non-simple function.
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
+# RUN: llvm-strip --strip-unneeded %t.o
+# RUN: %clang %cflags %t.o -o %t.exe -Wl,-q
+# RUN: llvm-bolt -v=3 %t.exe -o %t.out 2>&1 | FileCheck %s
+
+# CHECK: BOLT-WARNING: Multiple fragments access same jump table: main; main.cold.1
+
+  .text
+  .globl main
+  .type main, %function
+  .p2align 2
+main:
+LBB0:
+  andl $0xf, %ecx
+  cmpb $0x4, %cl
+  # exit through ret
+  ja LBB3
+
+# jump table dispatch, jumping to label indexed by val in %ecx
+LBB1:
+  leaq JUMP_TABLE1(%rip), %r8
+  movzbl %cl, %ecx
+  movslq (%r8,%rcx,4), %rax
+  addq %rax, %r8
+  jmpq *%r8
+
+LBB2:
+  xorq %rax, %rax
+LBB3:
+  addq $0x8, %rsp
+  ret
+.size main, .-main
+
+# cold fragment is only reachable
+  .globl main.cold.1
+  .type main.cold.1, %function
+  .p2align 2
+main.cold.1:
+  # load bearing nop: pad LBB8 so that it can't be treated
+  # as __builtin_unreachable by analyzeJumpTable
+  nop
+LBB4:
+  andl $0xb, %ebx
+  cmpb $0x1, %cl
+  # exit through ret
+  ja LBB7
+
+# jump table dispatch, jumping to label indexed by val in %ecx
+LBB5:
+  leaq JUMP_TABLE1(%rip), %r8
+  movzbl %cl, %ecx
+  movslq (%r8,%rcx,4), %rax
+  addq %rax, %r8
+  jmpq *%r8
+
+LBB6:
+  xorq %rax, %rax
+LBB7:
+  addq $0x8, %rsp
+  ret
+LBB8:
+  callq abort
+.size main.cold.1, .-main.cold.1
+
+  .rodata
+# jmp table, entries must be R_X86_64_PC32 relocs
+  .globl JUMP_TABLE1
+JUMP_TABLE1:
+  .long LBB2-JUMP_TABLE1
+  .long LBB3-JUMP_TABLE1
+  .long LBB8-JUMP_TABLE1
+  .long LBB6-JUMP_TABLE1
+  .long LBB7-JUMP_TABLE1