[BOLT] Process fragment siblings in lite mode, keep lite mode on

In lite mode, include split function fragments to the list of functions to
process even if a fragment has no samples. This is required to properly
detect and update split jump tables (jump tables that contain pointers to code
in the main and cold fragments).

Reviewed By: #bolt, maksfb

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

bolt/include/bolt/Core/BinaryContext.h

@@ -601,6 +601,9 @@ public:
   /// Indicates if the binary is stripped
   bool IsStripped{false};
 
+  /// Indicates if the binary contains split functions.
+  bool HasSplitFunctions{false};
+
   /// Is the binary always loaded at a fixed address. Shared objects and
   /// position-independent executables (PIEs) are examples of binaries that
   /// will have HasFixedLoadAddress set to false.

bolt/include/bolt/Rewrite/RewriteInstance.h

@@ -95,6 +95,9 @@ private:
   /// from meta data in the file.
   void discoverFileObjects();
 
+  /// Process fragments, locate parent functions.
+  void registerFragments();
+
   /// Read info from special sections. E.g. eh_frame and .gcc_except_table
   /// for exception and stack unwinding information.
   Error readSpecialSections();

bolt/lib/Core/BinaryContext.cpp

@@ -479,18 +479,6 @@ MemoryContentsType BinaryContext::analyzeMemoryAt(uint64_t Address,
   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, const uint64_t NextJTAddress,
@@ -513,14 +501,9 @@ bool BinaryContext::analyzeJumpTable(
     // Nothing to do if we failed to identify the containing function.
     if (!TargetBF)
       return false;
-    // Case 1: check if BF is a fragment and TargetBF is its parent.
-    if (BF.isFragment()) {
-      // 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.
-    return TargetBF->isFragment() && registerFragment(*TargetBF, BF);
+    // Check if BF is a fragment of TargetBF or vice versa.
+    return (BF.isFragment() && BF.isParentFragment(TargetBF)) ||
+           (TargetBF->isFragment() && TargetBF->isParentFragment(&BF));
   };
 
   ErrorOr<BinarySection &> Section = getSectionForAddress(Address);
@@ -1116,8 +1099,6 @@ void BinaryContext::generateSymbolHashes() {
 
 bool BinaryContext::registerFragment(BinaryFunction &TargetFunction,
                                      BinaryFunction &Function) const {
-  if (!isPotentialFragmentByName(TargetFunction, Function))
-    return false;
   assert(TargetFunction.isFragment() && "TargetFunction must be a fragment");
   if (TargetFunction.isParentFragment(&Function))
     return true;
@@ -1242,7 +1223,7 @@ void BinaryContext::processInterproceduralReferences() {
 
     if (TargetFunction) {
       if (TargetFunction->isFragment() &&
-          !registerFragment(*TargetFunction, Function)) {
+          !TargetFunction->isParentFragment(&Function)) {
         errs() << "BOLT-WARNING: interprocedural reference between unrelated "
                   "fragments: "
                << Function.getPrintName() << " and "

bolt/lib/Rewrite/RewriteInstance.cpp

@@ -52,6 +52,7 @@
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/LEB128.h"
 #include "llvm/Support/ManagedStatic.h"
+#include "llvm/Support/Regex.h"
 #include "llvm/Support/Timer.h"
 #include "llvm/Support/ToolOutputFile.h"
 #include "llvm/Support/raw_ostream.h"
@@ -928,6 +929,9 @@ void RewriteInstance::discoverFileObjects() {
   BinaryFunction *PreviousFunction = nullptr;
   unsigned AnonymousId = 0;
 
+  // Regex object for matching cold fragments.
+  Regex ColdFragment(".*\\.cold(\\.[0-9]+)?");
+
   const auto SortedSymbolsEnd = LastSymbol == SortedFileSymbols.end()
                                     ? LastSymbol
                                     : std::next(LastSymbol);
@@ -1186,6 +1190,21 @@ void RewriteInstance::discoverFileObjects() {
       if (!IsSimple)
         BF->setSimple(false);
     }
+
+    // Check if it's a cold function fragment.
+    if (ColdFragment.match(SymName)) {
+      static bool PrintedWarning = false;
+      if (!PrintedWarning) {
+        PrintedWarning = true;
+        errs() << "BOLT-WARNING: split function detected on input : "
+               << SymName;
+        if (BC->HasRelocations)
+          errs() << ". The support is limited in relocation mode\n";
+      }
+      BC->HasSplitFunctions = true;
+      BF->IsFragment = true;
+    }
+
     if (!AlternativeName.empty())
       BF->addAlternativeName(AlternativeName);
 
@@ -1266,6 +1285,56 @@ void RewriteInstance::discoverFileObjects() {
     // Read all relocations now that we have binary functions mapped.
     processRelocations();
   }
+  registerFragments();
+}
+
+void RewriteInstance::registerFragments() {
+  if (!BC->HasSplitFunctions)
+    return;
+
+  for (auto &BFI : BC->getBinaryFunctions()) {
+    BinaryFunction &Function = BFI.second;
+    if (!Function.isFragment())
+      continue;
+    unsigned ParentsFound = 0;
+    for (StringRef Name : Function.getNames()) {
+      StringRef BaseName, Suffix;
+      std::tie(BaseName, Suffix) = Name.split('/');
+      const size_t ColdSuffixPos = BaseName.find(".cold");
+      if (ColdSuffixPos == StringRef::npos)
+        continue;
+      // For cold function with local (foo.cold/1) symbol, prefer a parent with
+      // local symbol as well (foo/1) over global symbol (foo).
+      std::string ParentName = BaseName.substr(0, ColdSuffixPos).str();
+      const BinaryData *BD = BC->getBinaryDataByName(ParentName);
+      if (Suffix != "") {
+        ParentName.append(Twine("/", Suffix).str());
+        const BinaryData *BDLocal = BC->getBinaryDataByName(ParentName);
+        if (BDLocal || !BD)
+          BD = BDLocal;
+      }
+      if (!BD) {
+        if (opts::Verbosity >= 1)
+          outs() << "BOLT-INFO: parent function not found for " << Name << "\n";
+        continue;
+      }
+      const uint64_t Address = BD->getAddress();
+      BinaryFunction *BF = BC->getBinaryFunctionAtAddress(Address);
+      if (!BF) {
+        if (opts::Verbosity >= 1)
+          outs() << formatv("BOLT-INFO: parent function not found at {0:x}\n",
+                            Address);
+        continue;
+      }
+      BC->registerFragment(Function, *BF);
+      ++ParentsFound;
+    }
+    if (!ParentsFound) {
+      errs() << "BOLT-ERROR: parent function not found for " << Function
+             << '\n';
+      exit(1);
+    }
+  }
 }
 
 void RewriteInstance::createPLTBinaryFunction(uint64_t TargetAddress,
@@ -1456,26 +1525,6 @@ void RewriteInstance::adjustFunctionBoundaries() {
     if (std::next(BFI) != BFE)
       NextFunction = &std::next(BFI)->second;
 
-    // Check if it's a fragment of a function.
-    std::optional<StringRef> FragName =
-        Function.hasRestoredNameRegex(".*\\.cold(\\.[0-9]+)?");
-    if (FragName) {
-      static bool PrintedWarning = false;
-      if (!PrintedWarning) {
-        PrintedWarning = true;
-        errs() << "BOLT-WARNING: split function detected on input : "
-               << *FragName;
-        if (BC->HasRelocations)
-          errs() << ". The support is limited in relocation mode";
-        if (opts::Lite) {
-          opts::Lite = false;
-          errs() << "\nBOLT-WARNING: disabling lite mode (-lite) when split "
-                 << "functions are present\n";
-        }
-      }
-      Function.IsFragment = true;
-    }
-
     // Check if there's a symbol or a function with a larger address in the
     // same section. If there is - it determines the maximum size for the
     // current function. Otherwise, it is the size of a containing section
@@ -2785,8 +2834,18 @@ void RewriteInstance::selectFunctionsToProcess() {
   }
 
   uint64_t NumFunctionsToProcess = 0;
-  auto shouldProcess = [&](const BinaryFunction &Function) {
+  auto mustSkip = [&](const BinaryFunction &Function) {
     if (opts::MaxFunctions && NumFunctionsToProcess > opts::MaxFunctions)
+      return true;
+    for (std::string &Name : opts::SkipFunctionNames)
+      if (Function.hasNameRegex(Name))
+        return true;
+
+    return false;
+  };
+
+  auto shouldProcess = [&](const BinaryFunction &Function) {
+    if (mustSkip(Function))
       return false;
 
     // If the list is not empty, only process functions from the list.
@@ -2804,10 +2863,6 @@ void RewriteInstance::selectFunctionsToProcess() {
       return Match.has_value();
     }
 
-    for (std::string &Name : opts::SkipFunctionNames)
-      if (Function.hasNameRegex(Name))
-        return false;
-
     if (opts::Lite) {
       // Forcibly include functions specified in the -function-order file.
       if (opts::ReorderFunctions == ReorderFunctions::RT_USER) {
@@ -2843,9 +2898,15 @@ void RewriteInstance::selectFunctionsToProcess() {
       continue;
     }
 
+    // Decide what to do with fragments after parent functions are processed.
+    if (Function.isFragment())
+      continue;
+
     if (!shouldProcess(Function)) {
-      LLVM_DEBUG(dbgs() << "BOLT-INFO: skipping processing of function "
-                        << Function << " per user request\n");
+      if (opts::Verbosity >= 1) {
+        outs() << "BOLT-INFO: skipping processing " << Function
+               << " per user request\n";
+      }
       Function.setIgnored();
     } else {
       ++NumFunctionsToProcess;
@@ -2853,6 +2914,52 @@ void RewriteInstance::selectFunctionsToProcess() {
         outs() << "BOLT-INFO: processing ending on " << Function << '\n';
     }
   }
+
+  if (!BC->HasSplitFunctions)
+    return;
+
+  // Fragment overrides:
+  // - If the fragment must be skipped, then the parent must be skipped as well.
+  // Otherwise, fragment should follow the parent function:
+  // - if the parent is skipped, skip fragment,
+  // - if the parent is processed, process the fragment(s) as well.
+  for (auto &BFI : BC->getBinaryFunctions()) {
+    BinaryFunction &Function = BFI.second;
+    if (!Function.isFragment())
+      continue;
+    if (mustSkip(Function)) {
+      for (BinaryFunction *Parent : Function.ParentFragments) {
+        if (opts::Verbosity >= 1) {
+          outs() << "BOLT-INFO: skipping processing " << *Parent
+                 << " together with fragment function\n";
+        }
+        Parent->setIgnored();
+        --NumFunctionsToProcess;
+      }
+      Function.setIgnored();
+      continue;
+    }
+
+    bool IgnoredParent =
+        llvm::any_of(Function.ParentFragments, [&](BinaryFunction *Parent) {
+          return Parent->isIgnored();
+        });
+    if (IgnoredParent) {
+      if (opts::Verbosity >= 1) {
+        outs() << "BOLT-INFO: skipping processing " << Function
+               << " together with parent function\n";
+      }
+      Function.setIgnored();
+    } else {
+      ++NumFunctionsToProcess;
+      if (opts::Verbosity >= 1) {
+        outs() << "BOLT-INFO: processing " << Function
+               << " as a sibling of non-ignored function\n";
+      }
+      if (opts::MaxFunctions && NumFunctionsToProcess == opts::MaxFunctions)
+        outs() << "BOLT-INFO: processing ending on " << Function << '\n';
+    }
+  }
 }
 
 void RewriteInstance::readDebugInfo() {

bolt/test/CMakeLists.txt

@@ -52,6 +52,7 @@ list(APPEND BOLT_TEST_DEPS
   merge-fdata
   not
   perf2bolt
+  split-file
   yaml2obj
   )
 

bolt/test/X86/fragment-lite-reverse.s

@@ -0,0 +1,81 @@
+# Check that BOLT in lite mode processes fragments as expected.
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
+# RUN: link_fdata %s %t.o %t.fdata
+# RUN: %clang %cflags %t.o -o %t.exe -Wl,-q
+# RUN: llvm-bolt %t.exe -o %t.out --lite=1 --data %t.fdata -v=1 2>&1 \
+# RUN:   | FileCheck %s
+
+# CHECK: BOLT-INFO: skipping processing main.cold.1 together with parent function
+# CHECK: BOLT-INFO: skipping processing foo.cold.1/1 together with parent function
+# CHECK: BOLT-INFO: skipping processing bar.cold.1/1 together with parent function
+
+  .globl main
+  .type main, %function
+main:
+  .cfi_startproc
+  cmpl	$0x0, %eax
+  je	main.cold.1
+  retq
+  .cfi_endproc
+.size main, .-main
+
+  .globl foo
+  .type foo, %function
+foo:
+  .cfi_startproc
+  cmpl	$0x0, %eax
+  je	foo.cold.1
+  retq
+  .cfi_endproc
+.size foo, .-foo
+
+  .local bar
+  .type bar, %function
+bar:
+  .cfi_startproc
+  cmpl	$0x0, %eax
+  je	bar.cold.1
+  retq
+  .cfi_endproc
+.size bar, .-bar
+
+  .section .text.cold
+  .globl main.cold.1
+  .type main.cold.1, %function
+main.cold.1:
+# FDATA: 0 [unknown] 0 1 main.cold.1 0 1 0
+  .cfi_startproc
+  pushq	%rbp
+  movq	%rsp, %rbp
+  movl	$0x0, %eax
+  popq	%rbp
+  retq
+  .cfi_endproc
+.size main.cold.1, .-main.cold.1
+
+  .local foo.cold.1
+  .type foo.cold.1, %function
+foo.cold.1:
+# FDATA: 0 [unknown] 0 1 foo.cold.1/1 0 1 0
+  .cfi_startproc
+  pushq	%rbp
+  movq	%rsp, %rbp
+  movl	$0x0, %eax
+  popq	%rbp
+  retq
+  .cfi_endproc
+.size foo.cold.1, .-foo.cold.1
+
+  .local bar.cold.1
+  .type bar.cold.1, %function
+bar.cold.1:
+# FDATA: 0 [unknown] 0 1 bar.cold.1/1 0 1 0
+  .cfi_startproc
+  pushq	%rbp
+  movq	%rsp, %rbp
+  movl	$0x0, %eax
+  popq	%rbp
+  retq
+  .cfi_endproc
+.size bar.cold.1, .-bar.cold.1

bolt/test/X86/fragment-lite.s

@@ -0,0 +1,151 @@
+# Check that BOLT in lite mode processes fragments as expected.
+
+# RUN: split-file %s %t
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %t/main.s -o %t.o
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %t/baz.s -o %t.baz.o
+# RUN: link_fdata %s %t.o %t.main.fdata
+# RUN: link_fdata %s %t.baz.o %t.baz.fdata
+# RUN: merge-fdata %t.main.fdata %t.baz.fdata > %t.fdata
+# RUN: %clang %cflags %t.o %t.baz.o -o %t.exe -Wl,-q
+# RUN: llvm-bolt %t.exe -o %t.out --lite=1 --data %t.fdata -v=1 -print-cfg \
+# RUN:   2>&1 | FileCheck %s
+
+# CHECK: BOLT-INFO: processing main.cold.1 as a sibling of non-ignored function
+# CHECK: BOLT-INFO: processing foo.cold.1/1 as a sibling of non-ignored function
+# CHECK: BOLT-INFO: processing bar.cold.1/1 as a sibling of non-ignored function
+# CHECK: BOLT-INFO: processing baz.cold.1 as a sibling of non-ignored function
+# CHECK: BOLT-INFO: processing baz.cold.1/1 as a sibling of non-ignored function
+
+# CHECK: Binary Function "main.cold.1" after building cfg
+# CHECK: Parent : main
+
+# CHECK: Binary Function "foo.cold.1/1" after building cfg
+# CHECK: Parent : foo
+
+# CHECK: Binary Function "bar.cold.1/1" after building cfg
+# CHECK: Parent : bar/1
+
+# CHECK: Binary Function "baz.cold.1" after building cfg
+# CHECK: Parent : baz{{$}}
+
+# CHECK: Binary Function "baz.cold.1/1" after building cfg
+# CHECK: Parent : baz/1
+
+#--- main.s
+  .globl main
+  .type main, %function
+main:
+  .cfi_startproc
+# FDATA: 0 [unknown] 0 1 main 0 1 0
+  cmpl	$0x0, %eax
+  je	main.cold.1
+  retq
+  .cfi_endproc
+.size main, .-main
+
+  .globl foo
+  .type foo, %function
+foo:
+  .cfi_startproc
+# FDATA: 0 [unknown] 0 1 foo 0 1 0
+  cmpl	$0x0, %eax
+  je	foo.cold.1
+  retq
+  .cfi_endproc
+.size foo, .-foo
+
+  .local bar
+  .type bar, %function
+bar:
+  .cfi_startproc
+# FDATA: 0 [unknown] 0 1 bar/1 0 1 0
+  cmpl	$0x0, %eax
+  je	bar.cold.1
+  retq
+  .cfi_endproc
+.size bar, .-bar
+
+  .globl baz
+  .type baz, %function
+baz:
+  .cfi_startproc
+# FDATA: 0 [unknown] 0 1 baz 0 1 0
+  cmpl	$0x0, %eax
+  je	baz.cold.1
+  retq
+  .cfi_endproc
+.size baz, .-baz
+
+  .section .text.cold
+  .globl main.cold.1
+  .type main.cold.1, %function
+main.cold.1:
+  .cfi_startproc
+  pushq	%rbp
+  movq	%rsp, %rbp
+  movl	$0x0, %eax
+  popq	%rbp
+  retq
+  .cfi_endproc
+.size main.cold.1, .-main.cold.1
+
+  .local foo.cold.1
+  .type foo.cold.1, %function
+foo.cold.1:
+  .cfi_startproc
+  pushq	%rbp
+  movq	%rsp, %rbp
+  movl	$0x0, %eax
+  popq	%rbp
+  retq
+  .cfi_endproc
+.size foo.cold.1, .-foo.cold.1
+
+  .local bar.cold.1
+  .type bar.cold.1, %function
+bar.cold.1:
+  .cfi_startproc
+  pushq	%rbp
+  movq	%rsp, %rbp
+  movl	$0x0, %eax
+  popq	%rbp
+  retq
+  .cfi_endproc
+.size bar.cold.1, .-bar.cold.1
+
+  .globl baz.cold.1
+  .type baz.cold.1, %function
+baz.cold.1:
+  .cfi_startproc
+  pushq	%rbp
+  movq	%rsp, %rbp
+  movl	$0x0, %eax
+  popq	%rbp
+  retq
+  .cfi_endproc
+.size baz.cold.1, .-baz.cold.1
+
+#--- baz.s
+  .local baz
+  .type baz, %function
+baz:
+  .cfi_startproc
+# FDATA: 0 [unknown] 0 1 baz/1 0 1 0
+  cmpl	$0x0, %eax
+  je	baz.cold.1
+  retq
+  .cfi_endproc
+.size baz, .-baz
+
+  .section .text.cold
+  .local baz.cold.1
+  .type baz.cold.1, %function
+baz.cold.1:
+  .cfi_startproc
+  pushq	%rbp
+  movq	%rsp, %rbp
+  movl	$0x0, %eax
+  popq	%rbp
+  retq
+  .cfi_endproc
+.size baz.cold.1, .-baz.cold.1

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

@@ -11,8 +11,8 @@
 # RUN: llvm-bolt %t.exe -o %t.out -v=1 --print-only=main2.cold.1 --print-disasm 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-INFO: marking main2.cold.1(*2) as a fragment of main
 # CHECK: Binary Function "main2.cold.1(*2)" after disassembly
 # CHECK: End of Function "main2.cold.1(*2)"
 # CHECK-DAG: BOLT-WARNING: Ignoring main2

bolt/test/lit.cfg.py

@@ -93,6 +93,7 @@ tools = [
     ToolSubst('merge-fdata', unresolved='fatal'),
     ToolSubst('llvm-readobj', unresolved='fatal'),
     ToolSubst('llvm-dwp', unresolved='fatal'),
+    ToolSubst('split-file', unresolved='fatal'),
 ]
 llvm_config.add_tool_substitutions(tools, tool_dirs)