Add new ELF linker options for order section layout

Author: Colibrow

lld/ELF/BPSectionOrderer.cpp

@@ -0,0 +1,68 @@
+//===- BPSectionOrderer.cpp------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "BPSectionOrderer.h"
+#include "llvm/Support/xxhash.h"
+
+using namespace llvm;
+using namespace lld::elf;
+
+void BPSectionELF::getSectionHashes(
+    llvm::SmallVectorImpl<uint64_t> &hashes,
+    const llvm::DenseMap<const void *, uint64_t> &sectionToIdx) const {
+  constexpr unsigned windowSize = 4;
+
+  size_t size = isec->content().size();
+  for (size_t i = 0; i != size; ++i) {
+    auto window = isec->content().drop_front(i).take_front(windowSize);
+    hashes.push_back(xxHash64(window));
+  }
+
+  llvm::sort(hashes);
+  hashes.erase(std::unique(hashes.begin(), hashes.end()), hashes.end());
+}
+
+llvm::DenseMap<const lld::elf::InputSectionBase *, int>
+lld::elf::runBalancedPartitioning(Ctx &ctx, llvm::StringRef profilePath,
+                                  bool forFunctionCompression,
+                                  bool forDataCompression,
+                                  bool compressionSortStartupFunctions,
+                                  bool verbose) {
+  // Collect sections from symbols and wrap as BPSectionELF instances.
+  // Deduplicates sections referenced by multiple symbols.
+  SmallVector<std::unique_ptr<BPSectionBase>> sections;
+  DenseSet<const InputSectionBase *> seenSections;
+
+  for (Symbol *sym : ctx.symtab->getSymbols())
+    if (sym->getSize() > 0)
+      if (auto *d = dyn_cast<Defined>(sym))
+        if (auto *sec = dyn_cast_or_null<InputSectionBase>(d->section))
+          if (seenSections.insert(sec).second)
+            sections.emplace_back(std::make_unique<BPSectionELF>(sec));
+
+  for (ELFFileBase *file : ctx.objectFiles)
+    for (Symbol *sym : file->getLocalSymbols())
+      if (sym->getSize() > 0)
+        if (auto *d = dyn_cast<Defined>(sym))
+          if (auto *sec = dyn_cast_or_null<InputSectionBase>(d->section))
+            if (seenSections.insert(sec).second)
+              sections.emplace_back(std::make_unique<BPSectionELF>(sec));
+
+  auto reorderedSections = BPSectionBase::reorderSectionsByBalancedPartitioning(
+      profilePath, forFunctionCompression, forDataCompression,
+      compressionSortStartupFunctions, verbose, sections);
+
+  DenseMap<const InputSectionBase *, int> result;
+  for (const auto [sec, priority] : reorderedSections) {
+    auto *elfSection = cast<BPSectionELF>(sec);
+    result.try_emplace(
+        static_cast<const InputSectionBase *>(elfSection->getSection()),
+        static_cast<int>(priority));
+  }
+  return result;
+}

lld/ELF/BPSectionOrderer.h

@@ -0,0 +1,109 @@
+//===- BPSectionOrderer.h -------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// This file uses Balanced Partitioning to order sections to improve startup
+/// time and compressed size.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLD_ELF_BPSECTION_ORDERER_H
+#define LLD_ELF_BPSECTION_ORDERER_H
+
+#include "InputFiles.h"
+#include "InputSection.h"
+#include "SymbolTable.h"
+#include "lld/Common/BPSectionOrdererBase.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/BinaryFormat/ELF.h"
+
+namespace lld::elf {
+
+class InputSection;
+
+class BPSymbolELF : public BPSymbol {
+  const Symbol *sym;
+
+public:
+  explicit BPSymbolELF(const Symbol *s) : sym(s) {}
+
+  llvm::StringRef getName() const override { return sym->getName(); }
+
+  const Defined *asDefined() const { return llvm::dyn_cast<Defined>(sym); }
+
+  std::optional<uint64_t> getValue() const override {
+    if (auto *d = asDefined())
+      return d->value;
+    return {};
+  }
+
+  std::optional<uint64_t> getSize() const override {
+    if (auto *d = asDefined())
+      return d->size;
+    return {};
+  }
+
+  InputSectionBase *getInputSection() const {
+    if (auto *d = asDefined())
+      return llvm::dyn_cast_or_null<InputSectionBase>(d->section);
+    return nullptr;
+  }
+
+  const Symbol *getSymbol() const { return sym; }
+};
+
+class BPSectionELF : public BPSectionBase {
+  const InputSectionBase *isec;
+
+public:
+  explicit BPSectionELF(const InputSectionBase *sec) : isec(sec) {}
+
+  const void *getSection() const override { return isec; }
+
+  uint64_t getSize() const override { return isec->getSize(); }
+
+  bool isCodeSection() const override {
+    return isec->flags & llvm::ELF::SHF_EXECINSTR;
+  }
+
+  SmallVector<std::unique_ptr<BPSymbol>> getSymbols() const override {
+    SmallVector<std::unique_ptr<BPSymbol>> symbols;
+    for (Symbol *sym : isec->file->getSymbols())
+      if (auto *d = dyn_cast<Defined>(sym))
+        if (d->size > 0 && d->section == isec)
+          symbols.emplace_back(std::make_unique<BPSymbolELF>(sym));
+
+    return symbols;
+  }
+
+  std::optional<StringRef>
+  getResolvedLinkageName(llvm::StringRef name) const override {
+    return {};
+  }
+
+  void getSectionHashes(llvm::SmallVectorImpl<uint64_t> &hashes,
+                        const llvm::DenseMap<const void *, uint64_t>
+                            &sectionToIdx) const override;
+
+  static bool classof(const BPSectionBase *s) { return true; }
+};
+
+/// Run Balanced Partitioning to find the optimal function and data order to
+/// improve startup time and compressed size.
+///
+/// It is important that -ffunction-sections and -fdata-sections are used to
+/// ensure functions and data are in their own sections and thus can be
+/// reordered.
+llvm::DenseMap<const InputSectionBase *, int>
+runBalancedPartitioning(Ctx &ctx, llvm::StringRef profilePath,
+                        bool forFunctionCompression, bool forDataCompression,
+                        bool compressionSortStartupFunctions, bool verbose);
+} // namespace lld::elf
+
+#endif

lld/ELF/CMakeLists.txt

@@ -37,6 +37,7 @@ add_lld_library(lldELF
   Arch/X86.cpp
   Arch/X86_64.cpp
   ARMErrataFix.cpp
+  BPSectionOrderer.cpp
   CallGraphSort.cpp
   DWARF.cpp
   Driver.cpp

lld/ELF/Config.h

@@ -264,6 +264,12 @@ struct Config {
   bool armBe8 = false;
   BsymbolicKind bsymbolic = BsymbolicKind::None;
   CGProfileSortKind callGraphProfileSort;
+  llvm::StringRef irpgoProfilePath;
+  bool bpStartupFunctionSort = false;
+  bool bpCompressionSortStartupFunctions = false;
+  bool bpFunctionOrderForCompression = false;
+  bool bpDataOrderForCompression = false;
+  bool bpVerboseSectionOrderer = false;
   bool checkSections;
   bool checkDynamicRelocs;
   std::optional<llvm::DebugCompressionType> compressDebugSections;

lld/ELF/Driver.cpp

@@ -1259,6 +1259,55 @@ static void readConfigs(Ctx &ctx, opt::InputArgList &args) {
       ctx.arg.bsymbolic = BsymbolicKind::All;
   }
   ctx.arg.callGraphProfileSort = getCGProfileSortKind(ctx, args);
+  ctx.arg.irpgoProfilePath = args.getLastArgValue(OPT_irpgo_profile);
+  ctx.arg.bpCompressionSortStartupFunctions =
+      args.hasFlag(OPT_bp_compression_sort_startup_functions,
+                   OPT_no_bp_compression_sort_startup_functions, false);
+  if (auto *arg = args.getLastArg(OPT_bp_startup_sort)) {
+    StringRef startupSortStr = arg->getValue();
+    if (startupSortStr == "function") {
+      ctx.arg.bpStartupFunctionSort = true;
+    } else if (startupSortStr != "none") {
+      ErrAlways(ctx) << "unknown value '" + startupSortStr + "' for " +
+                            arg->getSpelling();
+    }
+    if (startupSortStr != "none")
+      if (args.hasArg(OPT_call_graph_ordering_file))
+        ErrAlways(ctx) << "--bp-startup-sort=function is incompatible with "
+                          "--call-graph-ordering-file";
+  }
+  if (ctx.arg.irpgoProfilePath.empty()) {
+    if (ctx.arg.bpStartupFunctionSort)
+      ErrAlways(ctx) << "--bp-startup-sort=function must be used with "
+                        "--irpgo-profile";
+    if (ctx.arg.bpCompressionSortStartupFunctions)
+      ErrAlways(ctx)
+          << "--bp-compression-sort-startup-functions must be used with "
+             "--irpgo-profile";
+  }
+
+  if (auto *arg = args.getLastArg(OPT_bp_compression_sort)) {
+    StringRef compressionSortStr = arg->getValue();
+    if (compressionSortStr == "function") {
+      ctx.arg.bpFunctionOrderForCompression = true;
+    } else if (compressionSortStr == "data") {
+      ctx.arg.bpDataOrderForCompression = true;
+    } else if (compressionSortStr == "both") {
+      ctx.arg.bpFunctionOrderForCompression = true;
+      ctx.arg.bpDataOrderForCompression = true;
+    } else if (compressionSortStr != "none") {
+      ErrAlways(ctx) << "unknown value '" + compressionSortStr + "' for " +
+                            arg->getSpelling();
+    }
+    if (ctx.arg.bpDataOrderForCompression ||
+        ctx.arg.bpFunctionOrderForCompression) {
+      if (args.getLastArg(OPT_call_graph_ordering_file) != nullptr) {
+        ErrAlways(ctx) << "--bp-compression-sort is incompatible with "
+                          "--call-graph-ordering-file";
+      }
+    }
+  }
+  ctx.arg.bpVerboseSectionOrderer = args.hasArg(OPT_verbose_bp_section_orderer);
   ctx.arg.checkSections =
       args.hasFlag(OPT_check_sections, OPT_no_check_sections, true);
   ctx.arg.chroot = args.getLastArgValue(OPT_chroot);

lld/ELF/Options.td

@@ -141,6 +141,24 @@ def call_graph_profile_sort: JJ<"call-graph-profile-sort=">,
 def : FF<"no-call-graph-profile-sort">, Alias<call_graph_profile_sort>, AliasArgs<["none"]>,
   Flags<[HelpHidden]>;
 
+defm irpgo_profile: Eq<"irpgo-profile",
+    "Read the IRPGO profile for use with -bp-startup-sort and other profile-guided optimizations">;
+
+def bp_startup_sort: JJ<"bp-startup-sort=">,
+    MetaVarName<"[none,function]">,
+    HelpText<"Order sections based on profile data to improve startup time">;
+
+defm bp_compression_sort_startup_functions: BB<"bp-compression-sort-startup-functions",
+    "Order startup functions by balanced partition to improve compressed size in addition to startup time",
+    "Do not order startup function for compression">;
+    
+def bp_compression_sort: JJ<"bp-compression-sort=">,
+    MetaVarName<"[none,function,data,both]">,
+    HelpText<"Order sections by balanced partition to improve compressed size">;
+
+def verbose_bp_section_orderer: FF<"verbose-bp-section-orderer">,
+    HelpText<"Print information on how many sections were ordered by balanced partitioning and a measure of the expected number of page faults">;
+    
 // --chroot doesn't have a help text because it is an internal option.
 def chroot: Separate<["--"], "chroot">;
 

lld/ELF/Writer.cpp

@@ -9,6 +9,7 @@
 #include "Writer.h"
 #include "AArch64ErrataFix.h"
 #include "ARMErrataFix.h"
+#include "BPSectionOrderer.h"
 #include "CallGraphSort.h"
 #include "Config.h"
 #include "InputFiles.h"
@@ -1082,6 +1083,16 @@ static void maybeShuffle(Ctx &ctx,
 // that don't appear in the order file.
 static DenseMap<const InputSectionBase *, int> buildSectionOrder(Ctx &ctx) {
   DenseMap<const InputSectionBase *, int> sectionOrder;
+  if (ctx.arg.bpStartupFunctionSort || ctx.arg.bpFunctionOrderForCompression ||
+      ctx.arg.bpDataOrderForCompression) {
+    TimeTraceScope timeScope("Balanced Partitioning Section Orderer");
+    sectionOrder = runBalancedPartitioning(
+        ctx, ctx.arg.bpStartupFunctionSort ? ctx.arg.irpgoProfilePath : "",
+        ctx.arg.bpFunctionOrderForCompression,
+        ctx.arg.bpDataOrderForCompression,
+        ctx.arg.bpCompressionSortStartupFunctions,
+        ctx.arg.bpVerboseSectionOrderer);
+  }
   if (!ctx.arg.callGraphProfile.empty())
     sectionOrder = computeCallGraphProfileOrder(ctx);