Optimize getSymbols

Author: MaskRay

lld/ELF/BPSectionOrderer.cpp

@@ -26,17 +26,17 @@ template <> struct lld::BPOrdererTraits<struct BPOrdererELF> {
 };
 namespace {
 struct BPOrdererELF : lld::BPOrderer<BPOrdererELF> {
+  DenseMap<const InputSectionBase *, Defined *> secToSym;
+
   static uint64_t getSize(const Section &sec) { return sec.getSize(); }
   static bool isCodeSection(const Section &sec) {
     return sec.flags & llvm::ELF::SHF_EXECINSTR;
   }
-  static SmallVector<Defined *, 0> getSymbols(const Section &sec) {
-    SmallVector<Defined *, 0> symbols;
-    for (auto *sym : sec.file->getSymbols())
-      if (auto *d = llvm::dyn_cast_or_null<Defined>(sym))
-        if (d->size > 0 && d->section == &sec)
-          symbols.emplace_back(d);
-    return symbols;
+  ArrayRef<Defined *> getSymbols(const Section &sec) {
+    auto it = secToSym.find(&sec);
+    if (it == secToSym.end())
+      return {};
+    return ArrayRef(it->second);
   }
 
   static void
@@ -69,28 +69,27 @@ DenseMap<const InputSectionBase *, int> elf::runBalancedPartitioning(
   // Collect candidate sections and associated symbols.
   SmallVector<InputSectionBase *> sections;
   DenseMap<CachedHashStringRef, DenseSet<unsigned>> rootSymbolToSectionIdxs;
-  DenseSet<const InputSectionBase *> seenSections;
+  BPOrdererELF orderer;
 
   auto addSection = [&](Symbol &sym) {
     auto *d = dyn_cast<Defined>(&sym);
-    if (!d || d->size == 0)
+    if (!d)
       return;
     auto *sec = dyn_cast_or_null<InputSectionBase>(d->section);
-    if (sec && seenSections.insert(sec).second) {
-      rootSymbolToSectionIdxs[CachedHashStringRef(getRootSymbol(sym.getName()))]
-          .insert(sections.size());
-      sections.emplace_back(sec);
-    }
+    if (!sec || sec->size == 0 || !orderer.secToSym.try_emplace(sec, d).second)
+      return;
+    rootSymbolToSectionIdxs[CachedHashStringRef(getRootSymbol(sym.getName()))]
+        .insert(sections.size());
+    sections.emplace_back(sec);
   };
 
   for (Symbol *sym : ctx.symtab->getSymbols())
     addSection(*sym);
   for (ELFFileBase *file : ctx.objectFiles)
     for (Symbol *sym : file->getLocalSymbols())
       addSection(*sym);
-
-  return BPOrdererELF::computeOrder(profilePath, forFunctionCompression,
-                                    forDataCompression,
-                                    compressionSortStartupFunctions, verbose,
-                                    sections, rootSymbolToSectionIdxs);
+  return orderer.computeOrder(profilePath, forFunctionCompression,
+                              forDataCompression,
+                              compressionSortStartupFunctions, verbose,
+                              sections, rootSymbolToSectionIdxs);
 }

lld/test/ELF/bp-section-orderer.s

@@ -21,6 +21,7 @@
 # RUN: ld.lld a.o --irpgo-profile=a.profdata --bp-startup-sort=function --verbose-bp-section-orderer --icf=all 2>&1 | FileCheck %s --check-prefix=STARTUP-FUNC-ORDER
 
 # STARTUP-FUNC-ORDER: Ordered 3 sections using balanced partitioning
+# STARTUP-FUNC-ORDER: Total area under the page fault curve: 3.
 
 # RUN: ld.lld -o - a.o --symbol-ordering-file a.orderfile --irpgo-profile=a.profdata --bp-startup-sort=function | llvm-nm --numeric-sort --format=just-symbols - | FileCheck %s --check-prefix=ORDERFILE
 # RUN: ld.lld -o - a.o --symbol-ordering-file a.orderfile --bp-compression-sort=both | llvm-nm --numeric-sort --format=just-symbols - | FileCheck %s --check-prefix=ORDERFILE