[BOLT] Split functions: support fragments with multiple parents

Summary:
Gracefully handle binaries with split functions where two fragments are folded
into one, resulting in a fragment with two parent functions.

This behavior is expected in GCC8+ with -O2 optimization level, where both
function splitting and ICF are enabled by default.

On the BOLT side, the changes are:
- BinaryFunction: allow multiple parent fragments:
  - `ParentFragment` --> `ParentFragments`,
  - `setParentFragment` --> `addParentFragment`.
- BinaryContext:
  - `populateJumpTables`: mark fragments to be skipped later,
  - `registerFragment`: add a name heuristic check, return false if it failed,
  - `processInterproceduralReferences`: check if `registerFragment`
succeeded, otherwise issue a warning,
  - `skipMarkedFragments`: move out fragment traversal and skipping from
  `populateJumpTables` into a separate function.

This change fixes an issue where unrelated functions might be registered
as fragments:

```
BOLT-WARNING: interprocedural reference between unrelated fragments:
bad_gs/1(*2) and amd_decode_mce.cold.27/1(*2)
```

(Linux kernel binary)

(cherry picked from FBD32786688)

bolt/include/bolt/Core/BinaryContext.h

@@ -195,6 +195,9 @@ class BinaryContext {
   /// with size won't overflow
   uint32_t DuplicatedJumpTables{0x10000000};
 
+  /// Function fragments to skip.
+  std::vector<BinaryFunction *> FragmentsToSkip;
+
   /// The runtime library.
   std::unique_ptr<RuntimeLibrary> RtLibrary;
 
@@ -864,13 +867,19 @@ public:
 
   /// @}
 
-  /// Register \p TargetFunction as fragment of \p Function.
-  void registerFragment(BinaryFunction &TargetFunction,
+  /// Register \p TargetFunction as a fragment of \p Function if checks pass:
+  /// - if \p TargetFunction name matches \p Function name with a suffix:
+  ///   fragment_name == parent_name.cold(.\d+)?
+  /// True if the Function is registered, false if the check failed.
+  bool registerFragment(BinaryFunction &TargetFunction,
                         BinaryFunction &Function) const;
 
   /// Resolve inter-procedural dependencies from \p Function.
   void processInterproceduralReferences(BinaryFunction &Function);
 
+  /// Skip functions with all parent and child fragments transitively.
+  void skipMarkedFragments();
+
   /// Perform any necessary post processing on the symbol table after
   /// function disassembly is complete.  This processing fixes top
   /// level data holes and makes sure the symbol table is valid.

bolt/include/bolt/Core/BinaryFunction.h

@@ -354,7 +354,7 @@ private:
   std::string ColdCodeSectionName;
 
   /// Parent function fragment for split function fragments.
-  BinaryFunction *ParentFragment{nullptr};
+  SmallPtrSet<BinaryFunction *, 1> ParentFragments;
 
   /// Indicate if the function body was folded into another function.
   /// Used by ICF optimization.
@@ -633,15 +633,15 @@ private:
 
   /// If the function represents a secondary split function fragment, set its
   /// parent fragment to \p BF.
-  void setParentFragment(BinaryFunction &BF) {
+  void addParentFragment(BinaryFunction &BF) {
+    assert(this != &BF);
     assert(IsFragment && "function must be a fragment to have a parent");
-    assert((!ParentFragment || ParentFragment == &BF) &&
-           "cannot have more than one parent function");
-    ParentFragment = &BF;
+    ParentFragments.insert(&BF);
   }
 
   /// Register a child fragment for the main fragment of a split function.
   void addFragment(BinaryFunction &BF) {
+    assert(this != &BF);
     Fragments.insert(&BF);
   }
 
@@ -2000,20 +2000,9 @@ public:
     return IsFragment;
   }
 
-  /// Return parent function fragment if this function is a secondary (child)
-  /// fragment of another function.
-  BinaryFunction *getParentFragment() const {
-    return ParentFragment;
-  }
-
-  /// If the function is a nested child fragment of another function, return its
-  /// topmost parent fragment.
-  const BinaryFunction *getTopmostFragment() const {
-    const BinaryFunction *BF = this;
-    while (BF->getParentFragment())
-      BF = BF->getParentFragment();
-
-    return BF;
+  /// Returns if the given function is a parent fragment of this function.
+  bool isParentFragment(BinaryFunction *Parent) const {
+    return ParentFragments.count(Parent);
   }
 
   /// Set the profile data for the number of times the function was called.
@@ -2558,12 +2547,6 @@ public:
   FragmentInfo &cold() { return ColdFragment; }
 
   const FragmentInfo &cold() const { return ColdFragment; }
-
-  /// Mark child fragments as ignored.
-  void ignoreFragments() {
-    for (BinaryFunction *Fragment : Fragments)
-      Fragment->setIgnored();
-  }
 };
 
 inline raw_ostream &operator<<(raw_ostream &OS,

bolt/lib/Core/BinaryContext.cpp

@@ -29,8 +29,10 @@
 #include "llvm/MC/MCSymbol.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Regex.h"
+#include <algorithm>
 #include <functional>
 #include <iterator>
+#include <unordered_set>
 
 using namespace llvm;
 
@@ -483,6 +485,18 @@ BinaryContext::analyzeMemoryAt(uint64_t Address, BinaryFunction &BF) {
   return MemoryContentsType::UNKNOWN;
 }
 
+/// Check if <fragment restored name> == <parent restored name>.cold(.\d+)?
+bool isPotentialFragmentByName(BinaryFunction &Fragment,
+                               BinaryFunction &Parent) {
+  for (StringRef Name : Parent.getNames()) {
+    std::string NamePrefix = Regex::escape(NameResolver::restore(Name));
+    std::string NameRegex = Twine(NamePrefix, "\\.cold(\\.[0-9]+)?").str();
+    if (Fragment.hasRestoredNameRegex(NameRegex))
+      return true;
+  }
+  return false;
+}
+
 bool BinaryContext::analyzeJumpTable(const uint64_t Address,
                                      const JumpTable::JumpTableType Type,
                                      BinaryFunction &BF,
@@ -500,19 +514,6 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
       Offsets->emplace_back(Offset);
   };
 
-  auto isFragment = [](BinaryFunction &Fragment,
-                       BinaryFunction &Parent) -> bool {
-    // Check if <fragment restored name> == <parent restored name>.cold(.\d+)?
-    for (StringRef BFName : Parent.getNames()) {
-      std::string BFNamePrefix = Regex::escape(NameResolver::restore(BFName));
-      std::string BFNameRegex =
-          Twine(BFNamePrefix, "\\.cold(\\.[0-9]+)?").str();
-      if (Fragment.hasRestoredNameRegex(BFNameRegex))
-        return true;
-    }
-    return false;
-  };
-
   auto doesBelongToFunction = [&](const uint64_t Addr,
                                   BinaryFunction *TargetBF) -> bool {
     if (BF.containsAddress(Addr))
@@ -522,25 +523,12 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
       return false;
     // Case 1: check if BF is a fragment and TargetBF is its parent.
     if (BF.isFragment()) {
-      // BF is a fragment, but parent function is not registered.
-      // This means there's no direct jump between parent and fragment.
-      // Set parent link here in jump table analysis, based on function name
-      // matching heuristic.
-      if (!BF.getParentFragment() && isFragment(BF, *TargetBF))
-        registerFragment(BF, *TargetBF);
-      return BF.getParentFragment() == TargetBF;
+      // Parent function may or may not be already registered.
+      // Set parent link based on function name matching heuristic.
+      return registerFragment(BF, *TargetBF);
     }
     // Case 2: check if TargetBF is a fragment and BF is its parent.
-    if (TargetBF->isFragment()) {
-      // TargetBF is a fragment, but parent function is not registered.
-      // This means there's no direct jump between parent and fragment.
-      // Set parent link here in jump table analysis, based on function name
-      // matching heuristic.
-      if (!TargetBF->getParentFragment() && isFragment(*TargetBF, BF))
-        registerFragment(*TargetBF, BF);
-      return TargetBF->getParentFragment() == &BF;
-    }
-    return false;
+    return TargetBF->isFragment() && registerFragment(*TargetBF, BF);
   };
 
   ErrorOr<BinarySection &> Section = getSectionForAddress(Address);
@@ -608,10 +596,10 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
                    << TargetBF->getPrintName() << '\n';
             if (TargetBF->isFragment())
               dbgs() << "  ! is a fragment\n";
-            BinaryFunction *TargetParent = TargetBF->getParentFragment();
-            dbgs() << "  ! its parent is "
-                   << (TargetParent ? TargetParent->getPrintName() : "(none)")
-                   << '\n';
+            for (BinaryFunction *TargetParent : TargetBF->ParentFragments)
+              dbgs() << "  ! its parent is "
+                     << (TargetParent ? TargetParent->getPrintName() : "(none)")
+                     << '\n';
           }
         }
         if (Value == BF.getAddress())
@@ -638,7 +626,8 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
       BF.setHasSplitJumpTable(true);
       // Add invalid offset for proper identification of jump table size.
       addOffset(INVALID_OFFSET);
-      LLVM_DEBUG(dbgs() << "OK: address in split fragment\n");
+      LLVM_DEBUG(dbgs() << "OK: address in split fragment "
+                        << TargetBF->getPrintName() << '\n');
     }
   }
 
@@ -649,7 +638,6 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
 }
 
 void BinaryContext::populateJumpTables() {
-  std::vector<BinaryFunction *> FuncsToSkip;
   LLVM_DEBUG(dbgs() << "DataPCRelocations: " << DataPCRelocations.size()
                     << '\n');
   for (auto JTI = JumpTables.begin(), JTE = JumpTables.end(); JTI != JTE;
@@ -699,7 +687,7 @@ void BinaryContext::populateJumpTables() {
 
     // Mark to skip the function and all its fragments.
     if (BF.hasSplitJumpTable())
-      FuncsToSkip.push_back(&BF);
+      FragmentsToSkip.push_back(&BF);
   }
 
   if (opts::StrictMode && DataPCRelocations.size()) {
@@ -712,16 +700,36 @@ void BinaryContext::populateJumpTables() {
     assert(0 && "unclaimed PC-relative relocations left in data\n");
   }
   clearList(DataPCRelocations);
+}
+
+void BinaryContext::skipMarkedFragments() {
+  // Unique functions in the vector.
+  std::unordered_set<BinaryFunction *> UniqueFunctions(FragmentsToSkip.begin(),
+                                                       FragmentsToSkip.end());
+  // Copy the functions back to FragmentsToSkip.
+  FragmentsToSkip.assign(UniqueFunctions.begin(), UniqueFunctions.end());
+  auto addToWorklist = [&](BinaryFunction *Function) -> void {
+    if (UniqueFunctions.count(Function))
+      return;
+    FragmentsToSkip.push_back(Function);
+    UniqueFunctions.insert(Function);
+  };
   // Functions containing split jump tables need to be skipped with all
-  // fragments.
-  for (BinaryFunction *BF : FuncsToSkip) {
-    BinaryFunction *ParentBF =
-        const_cast<BinaryFunction *>(BF->getTopmostFragment());
-    LLVM_DEBUG(dbgs() << "Skipping " << ParentBF->getPrintName()
-                      << " family\n");
-    ParentBF->setIgnored();
-    ParentBF->ignoreFragments();
+  // fragments (transitively).
+  for (size_t I = 0; I != FragmentsToSkip.size(); I++) {
+    BinaryFunction *BF = FragmentsToSkip[I];
+    assert(UniqueFunctions.count(BF) &&
+           "internal error in traversing function fragments");
+    if (opts::Verbosity >= 1)
+      errs() << "BOLT-WARNING: Ignoring " << BF->getPrintName() << '\n';
+    BF->setIgnored();
+    std::for_each(BF->Fragments.begin(), BF->Fragments.end(), addToWorklist);
+    std::for_each(BF->ParentFragments.begin(), BF->ParentFragments.end(),
+                  addToWorklist);
   }
+  errs() << "BOLT-WARNING: Ignored " << FragmentsToSkip.size() << " functions "
+         << "due to cold fragments.\n";
+  FragmentsToSkip.clear();
 }
 
 MCSymbol *BinaryContext::getOrCreateGlobalSymbol(uint64_t Address,
@@ -1093,16 +1101,14 @@ void BinaryContext::generateSymbolHashes() {
   }
 }
 
-void BinaryContext::registerFragment(BinaryFunction &TargetFunction,
+bool BinaryContext::registerFragment(BinaryFunction &TargetFunction,
                                      BinaryFunction &Function) const {
-  // Only a parent function (or a sibling) can reach its fragment.
-  assert(!Function.IsFragment &&
-         "only one cold fragment is supported at this time");
-  if (BinaryFunction *TargetParent = TargetFunction.getParentFragment()) {
-    assert(TargetParent == &Function && "mismatching parent function");
-    return;
-  }
-  TargetFunction.setParentFragment(Function);
+  if (!isPotentialFragmentByName(TargetFunction, Function))
+    return false;
+  assert(TargetFunction.isFragment() && "TargetFunction must be a fragment");
+  if (TargetFunction.isParentFragment(&Function))
+    return true;
+  TargetFunction.addParentFragment(Function);
   Function.addFragment(TargetFunction);
   if (!HasRelocations) {
     TargetFunction.setSimple(false);
@@ -1112,6 +1118,7 @@ void BinaryContext::registerFragment(BinaryFunction &TargetFunction,
     outs() << "BOLT-INFO: marking " << TargetFunction
            << " as a fragment of " << Function << '\n';
   }
+  return true;
 }
 
 void BinaryContext::processInterproceduralReferences(BinaryFunction &Function) {
@@ -1125,8 +1132,13 @@ void BinaryContext::processInterproceduralReferences(BinaryFunction &Function) {
       continue;
 
     if (TargetFunction) {
-      if (TargetFunction->IsFragment)
-        registerFragment(*TargetFunction, Function);
+      if (TargetFunction->IsFragment &&
+          !registerFragment(*TargetFunction, Function)) {
+        errs() << "BOLT-WARNING: interprocedural reference between unrelated "
+                  "fragments: "
+               << Function.getPrintName() << " and "
+               << TargetFunction->getPrintName() << '\n';
+      }
       if (uint64_t Offset = Address - TargetFunction->getAddress())
         TargetFunction->addEntryPointAtOffset(Offset);
 

bolt/lib/Core/BinaryFunction.cpp

@@ -442,7 +442,7 @@ void BinaryFunction::print(raw_ostream &OS, std::string Annotation,
   if (isFolded()) {
     OS << "\n  FoldedInto  : " << *getFoldedIntoFunction();
   }
-  if (ParentFragment) {
+  for (BinaryFunction *ParentFragment : ParentFragments) {
     OS << "\n  Parent      : " << *ParentFragment;
   }
   if (!Fragments.empty()) {

bolt/lib/Rewrite/RewriteInstance.cpp

@@ -2789,6 +2789,7 @@ void RewriteInstance::disassembleFunctions() {
   }
 
   BC->populateJumpTables();
+  BC->skipMarkedFragments();
 
   for (auto &BFI : BC->getBinaryFunctions()) {
     BinaryFunction &Function = BFI.second;

bolt/test/X86/split-func-icf.s

@@ -0,0 +1,100 @@
+# This reproduces an issue where two cold fragments are folded into one, so the
+# fragment has two parents.
+# The fragment is only reachable through a jump table, so all functions must be
+# ignored.
+
+# 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 %t.exe -o %t.out -lite=0 -v=1 2>&1 | FileCheck %s
+
+# CHECK-NOT: unclaimed PC-relative relocations left in data
+# CHECK-DAG: BOLT-INFO: marking main2.cold.1(*2) as a fragment of main
+# CHECK-DAG: BOLT-INFO: marking main2.cold.1(*2) as a fragment of main2
+# CHECK-DAG: BOLT-WARNING: Ignoring main2
+# CHECK-DAG: BOLT-WARNING: Ignoring main
+# CHECK-DAG: BOLT-WARNING: Ignoring main2.cold.1(*2)
+# CHECK: BOLT-WARNING: Ignored 3 functions due to cold fragments.
+  .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
+
+  .globl main2
+  .type main2, %function
+  .p2align 2
+main2:
+LBB20:
+  andl $0xb, %ebx
+  cmpb $0x1, %cl
+  # exit through ret
+  ja LBB23
+
+# jump table dispatch, jumping to label indexed by val in %ecx
+LBB21:
+  leaq JUMP_TABLE2(%rip), %r8
+  movzbl %cl, %ecx
+  movslq (%r8,%rcx,4), %rax
+  addq %rax, %r8
+  jmpq *%r8
+
+LBB22:
+  xorq %rax, %rax
+LBB23:
+  addq $0x8, %rsp
+  ret
+.size main2, .-main2
+
+# cold fragment is only reachable through jump table
+  .globl main2.cold.1
+  .type main2.cold.1, %function
+main2.cold.1:
+  .globl main.cold.1
+  .type main.cold.1, %function
+  .p2align 2
+main.cold.1:
+  # load bearing nop: pad LBB4 so that it can't be treated
+  # as __builtin_unreachable by analyzeJumpTable
+  nop
+LBB4:
+  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 LBB4-JUMP_TABLE1
+  .long LBB3-JUMP_TABLE1
+
+  .globl JUMP_TABLE2
+JUMP_TABLE2:
+  .long LBB22-JUMP_TABLE2
+  .long LBB23-JUMP_TABLE2
+  .long LBB4-JUMP_TABLE2
+  .long LBB23-JUMP_TABLE2