[BOLT] Introduce getOutputBinaryFunctions(). NFCI

[maksfb]

To gain better control over the functions that go into the output file and their order, introduce BinaryContext::getOutputBinaryFunctions().

The new API returns a modifiable list of functions in output order.

This list is filled by a new PopulateOutputFunctions pass and includes emittable functions from the input file, plus functions added by BOLT (injected functions).

The new functionality allows to freely intermix input functions with injected ones in the output, which will be used in new PRs.

The new function replaces BinaryContext::getSortedFunctions(), but unlike its predecessor, it includes injected functions in the returned list.

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Maksim Panchenko (maksfb)

Changes

To gain better control over the functions that go into the output file and their order, introduce BinaryContext::getOutputBinaryFunctions().

The new API returns a modifiable list of functions in output order.

This list is filled by a new PopulateOutputFunctions pass and includes emittable functions from the input file, plus functions added by BOLT (injected functions).

The new functionality allows to freely intermix input functions with injected ones in the output, which will be used in new PRs.

The new function replaces BinaryContext::getSortedFunctions(), but unlike its predecessor, it includes injected functions in the returned list.

---

Full diff: https://github.com/llvm/llvm-project/pull/172174.diff

10 Files Affected:

• (modified) bolt/include/bolt/Core/BinaryContext.h (+6-3)
• (modified) bolt/include/bolt/Passes/BinaryPasses.h (+9)
• (modified) bolt/lib/Core/BinaryContext.cpp (+4-10)
• (modified) bolt/lib/Core/BinaryEmitter.cpp (+1-5)
• (modified) bolt/lib/Passes/BinaryPasses.cpp (+22)
• (modified) bolt/lib/Passes/LongJmp.cpp (+1-1)
• (modified) bolt/lib/Passes/SplitFunctions.cpp (+3-3)
• (modified) bolt/lib/Rewrite/BinaryPassManager.cpp (+3)
• (modified) bolt/lib/Rewrite/MachORewriteInstance.cpp (+1)
• (modified) bolt/lib/Rewrite/RewriteInstance.cpp (+1-1)

diff --git a/bolt/include/bolt/Core/BinaryContext.h b/bolt/include/bolt/Core/BinaryContext.h
index ebfc72456f52d..31c90d2c502bd 100644
--- a/bolt/include/bolt/Core/BinaryContext.h
+++ b/bolt/include/bolt/Core/BinaryContext.h
@@ -231,6 +231,9 @@ class BinaryContext {
   /// Store all functions in the binary, sorted by original address.
   std::map<uint64_t, BinaryFunction> BinaryFunctions;
 
+  /// Functions to be considered for the output in a sorted order.
+  BinaryFunctionListType OutputFunctions;
+
   /// A mutex that is used to control parallel accesses to BinaryFunctions.
   mutable llvm::sys::RWMutex BinaryFunctionsMutex;
 
@@ -554,6 +557,9 @@ class BinaryContext {
     return BinaryFunctions;
   }
 
+  /// Return functions meant for the output in a sorted order.
+  BinaryFunctionListType &getOutputBinaryFunctions() { return OutputFunctions; }
+
   /// Create BOLT-injected function
   BinaryFunction *createInjectedBinaryFunction(const std::string &Name,
                                                bool IsSimple = true);
@@ -1387,9 +1393,6 @@ class BinaryContext {
   unsigned addDebugFilenameToUnit(const uint32_t DestCUID,
                                   const uint32_t SrcCUID, unsigned FileIndex);
 
-  /// Return functions in output layout order
-  BinaryFunctionListType getSortedFunctions();
-
   /// Do the best effort to calculate the size of the function by emitting
   /// its code, and relaxing branch instructions. By default, branch
   /// instructions are updated to match the layout. Pass \p FixBranches set to
diff --git a/bolt/include/bolt/Passes/BinaryPasses.h b/bolt/include/bolt/Passes/BinaryPasses.h
index ad8473c4aae02..b5971d315d764 100644
--- a/bolt/include/bolt/Passes/BinaryPasses.h
+++ b/bolt/include/bolt/Passes/BinaryPasses.h
@@ -195,6 +195,15 @@ class FixupBranches : public BinaryFunctionPass {
   Error runOnFunctions(BinaryContext &BC) override;
 };
 
+/// Initialize the output function list.
+class PopulateOutputFunctions : public BinaryFunctionPass {
+public:
+  explicit PopulateOutputFunctions() : BinaryFunctionPass(false) {}
+
+  const char *getName() const override { return "populate-output-functions"; }
+  Error runOnFunctions(BinaryContext &BC) override;
+};
+
 /// Fix the CFI state and exception handling information after all other
 /// passes have completed.
 class FinalizeFunctions : public BinaryFunctionPass {
diff --git a/bolt/lib/Core/BinaryContext.cpp b/bolt/lib/Core/BinaryContext.cpp
index a5ced5d14f2e6..f0541921c70a8 100644
--- a/bolt/lib/Core/BinaryContext.cpp
+++ b/bolt/lib/Core/BinaryContext.cpp
@@ -1715,16 +1715,6 @@ unsigned BinaryContext::addDebugFilenameToUnit(const uint32_t DestCUID,
                                DestCUID, DstUnit->getVersion()));
 }
 
-BinaryFunctionListType BinaryContext::getSortedFunctions() {
-  BinaryFunctionListType SortedFunctions(BinaryFunctions.size());
-  llvm::transform(llvm::make_second_range(BinaryFunctions),
-                  SortedFunctions.begin(),
-                  [](BinaryFunction &BF) { return &BF; });
-
-  llvm::stable_sort(SortedFunctions, compareBinaryFunctionByIndex);
-  return SortedFunctions;
-}
-
 BinaryFunctionListType BinaryContext::getAllBinaryFunctions() {
   BinaryFunctionListType AllFunctions;
   AllFunctions.reserve(BinaryFunctions.size() + InjectedBinaryFunctions.size());
@@ -2569,6 +2559,10 @@ BinaryContext::createInjectedBinaryFunction(const std::string &Name,
   BinaryFunction *BF = InjectedBinaryFunctions.back();
   setSymbolToFunctionMap(BF->getSymbol(), BF);
   BF->CurrentState = BinaryFunction::State::CFG;
+
+  if (!getOutputBinaryFunctions().empty())
+    getOutputBinaryFunctions().push_back(BF);
+
   return BF;
 }
 
diff --git a/bolt/lib/Core/BinaryEmitter.cpp b/bolt/lib/Core/BinaryEmitter.cpp
index 9f9607562a0a5..025516af82023 100644
--- a/bolt/lib/Core/BinaryEmitter.cpp
+++ b/bolt/lib/Core/BinaryEmitter.cpp
@@ -282,11 +282,7 @@ void BinaryEmitter::emitFunctions() {
   }
 
   // Emit functions in sorted order.
-  BinaryFunctionListType SortedFunctions = BC.getSortedFunctions();
-  emit(SortedFunctions);
-
-  // Emit functions added by BOLT.
-  emit(BC.getInjectedBinaryFunctions());
+  emit(BC.getOutputBinaryFunctions());
 
   // Mark the end of hot text.
   if (opts::HotText) {
diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index 999de6f63e7a5..35e30a39872d1 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -554,6 +554,28 @@ Error FixupBranches::runOnFunctions(BinaryContext &BC) {
   return Error::success();
 }
 
+Error PopulateOutputFunctions::runOnFunctions(BinaryContext &BC) {
+  BinaryFunctionListType &OutputFunctions = BC.getOutputBinaryFunctions();
+
+  assert(OutputFunctions.empty() && "Output function list already initialized");
+
+  OutputFunctions.reserve(BC.getBinaryFunctions().size() +
+                          BC.getInjectedBinaryFunctions().size());
+  llvm::transform(llvm::make_second_range(BC.getBinaryFunctions()),
+                  std::back_inserter(OutputFunctions),
+                  [](BinaryFunction &BF) { return &BF; });
+
+  llvm::erase_if(OutputFunctions,
+                 [&BC](BinaryFunction *BF) { return !BC.shouldEmit(*BF); });
+
+  llvm::stable_sort(OutputFunctions, compareBinaryFunctionByIndex);
+
+  llvm::copy(BC.getInjectedBinaryFunctions(),
+             std::back_inserter(OutputFunctions));
+
+  return Error::success();
+}
+
 Error FinalizeFunctions::runOnFunctions(BinaryContext &BC) {
   std::atomic<bool> HasFatal{false};
   ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index b96d054119bf9..a9f81a0480138 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -921,7 +921,7 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
   }
 
   BC.outs() << "BOLT-INFO: Starting stub-insertion pass\n";
-  BinaryFunctionListType Sorted = BC.getSortedFunctions();
+  BinaryFunctionListType Sorted = BC.getOutputBinaryFunctions();
   bool Modified;
   uint32_t Iterations = 0;
   do {
diff --git a/bolt/lib/Passes/SplitFunctions.cpp b/bolt/lib/Passes/SplitFunctions.cpp
index 66a373ad2de72..8a6f555c491a3 100644
--- a/bolt/lib/Passes/SplitFunctions.cpp
+++ b/bolt/lib/Passes/SplitFunctions.cpp
@@ -206,7 +206,7 @@ struct SplitCacheDirected final : public SplitStrategy {
   }
 
   void initializeAuxiliaryVariables() {
-    for (BinaryFunction *BF : BC.getSortedFunctions()) {
+    for (BinaryFunction *BF : BC.getOutputBinaryFunctions()) {
       if (!shouldConsiderForCallGraph(*BF))
         continue;
 
@@ -234,7 +234,7 @@ struct SplitCacheDirected final : public SplitStrategy {
   void buildCallGraph() {
     Callers.resize(TotalNumBlocks);
     Callees.resize(TotalNumBlocks);
-    for (const BinaryFunction *SrcFunction : BC.getSortedFunctions()) {
+    for (const BinaryFunction *SrcFunction : BC.getOutputBinaryFunctions()) {
       if (!shouldConsiderForCallGraph(*SrcFunction))
         continue;
 
@@ -337,7 +337,7 @@ struct SplitCacheDirected final : public SplitStrategy {
     const BinaryBasicBlock *ThisBB = &(ThisBF->front());
     const size_t ThisGI = GlobalIndices[ThisBB];
 
-    for (const BinaryFunction *DstBF : BC.getSortedFunctions()) {
+    for (const BinaryFunction *DstBF : BC.getOutputBinaryFunctions()) {
       if (!shouldConsiderForCallGraph(*DstBF))
         continue;
 
diff --git a/bolt/lib/Rewrite/BinaryPassManager.cpp b/bolt/lib/Rewrite/BinaryPassManager.cpp
index 85f23dceefe26..58d24e15cde01 100644
--- a/bolt/lib/Rewrite/BinaryPassManager.cpp
+++ b/bolt/lib/Rewrite/BinaryPassManager.cpp
@@ -484,6 +484,9 @@ Error BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
   Manager.registerPass(
       std::make_unique<ReorderFunctions>(PrintReorderedFunctions));
 
+  // Produce the list of functions for the output file in a sorted order.
+  Manager.registerPass(std::make_unique<PopulateOutputFunctions>());
+
   // This is the second run of the SplitFunctions pass required by certain
   // splitting strategies (e.g. cdsplit). Running the SplitFunctions pass again
   // after ReorderFunctions allows the finalized function order to be utilized
diff --git a/bolt/lib/Rewrite/MachORewriteInstance.cpp b/bolt/lib/Rewrite/MachORewriteInstance.cpp
index a1c2cef601e43..17f726ae13945 100644
--- a/bolt/lib/Rewrite/MachORewriteInstance.cpp
+++ b/bolt/lib/Rewrite/MachORewriteInstance.cpp
@@ -354,6 +354,7 @@ void MachORewriteInstance::runOptimizationPasses() {
       std::make_unique<ReorderBasicBlocks>(opts::PrintReordered));
   Manager.registerPass(
       std::make_unique<FixupBranches>(opts::PrintAfterBranchFixup));
+  Manager.registerPass(std::make_unique<PopulateOutputFunctions>());
   // This pass should always run last.*
   Manager.registerPass(
       std::make_unique<FinalizeFunctions>(opts::PrintFinalized));
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 2ef03d283366b..eed2c2093598d 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -3996,7 +3996,7 @@ void RewriteInstance::emitAndLink() {
 
   if (opts::PrintCacheMetrics) {
     BC->outs() << "BOLT-INFO: cache metrics after emitting functions:\n";
-    CacheMetrics::printAll(BC->outs(), BC->getSortedFunctions());
+    CacheMetrics::printAll(BC->outs(), BC->getAllBinaryFunctions());
   }
 }