[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

4 ↗

(On Diff #437416)

Please updated the message

[Amir]

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

11 ↗

(On Diff #437416)

[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 meory used by jump table offsets
+  /// 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.

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 ||
-            (JT->Parent->isFragment() &&
-             JT->Parent->isParentFragment(&Function)) ||
-            (Function.isFragment() && Function.isParentFragment(JT->Parent))) &&
+    assert((JT->Parent == &Function || isFragmentOf(JT->Parent, &Function) ||
+            isFragmentOf(&Function, JT->Parent)) &&
            "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 ((JT->Parent->isFragment() && JT->Parent->isParentFragment(&Function)) ||
-        (Function.isFragment() && Function.isParentFragment(JT->Parent))) {
+    if (isFragmentOf(JT->Parent, &Function) ||
+        isFragmentOf(&Function, JT->Parent)) {
       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();
   }
 

bolt/lib/Core/BinaryFunction.cpp

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;