[llvm] fe0e804 - [JITLink][NFC] Store external symbols in a StringMap

[Job Noorman via llvm-commits]

Author: Job Noorman
Date: 2023-08-29T09:21:02+02:00
New Revision: fe0e804a4ef4fdf0ae7b2654ad4ac8d1cfafec60

URL: https://github.com/llvm/llvm-project/commit/fe0e804a4ef4fdf0ae7b2654ad4ac8d1cfafec60
DIFF: https://github.com/llvm/llvm-project/commit/fe0e804a4ef4fdf0ae7b2654ad4ac8d1cfafec60.diff

LOG: [JITLink][NFC] Store external symbols in a StringMap

External symbols used to be stored in a `DenseSet`. An `assert` in
`addExternalSymbol` ensures that names of external symbols are unique.
However, for objects containing a huge number of external symbols, this
`assert` can be a performance bottleneck.

This patch proposes to store external symbols in a `StringMap` instead
making it significantly cheaper to check if a certain symbol name
already exists.

This issue came up while porting BOLT to JITLink (D147544): linking a
large binary using the JITLink port turned out to be about 4x slower
than the current version of BOLT that uses RuntimeDyld. This slowdown
was caused entirely by the `assert` in `addExternalSymbol`. Using this
patch, the JITLink port is slightly faster than RuntimeDyld.

Reviewed By: maksfb

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

Added: 
    

Modified: 
    llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h b/llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h
index 2f015c1d21f721..48268aa9f6899b 100644
--- a/llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h
+++ b/llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h
@@ -849,7 +849,8 @@ class SectionRange {
 class LinkGraph {
 private:
   using SectionMap = DenseMap<StringRef, std::unique_ptr<Section>>;
-  using ExternalSymbolSet = DenseSet<Symbol *>;
+  using ExternalSymbolMap = StringMap<Symbol *>;
+  using AbsoluteSymbolSet = DenseSet<Symbol *>;
   using BlockSet = DenseSet<Block *>;
 
   template <typename... ArgTs>
@@ -901,6 +902,12 @@ class LinkGraph {
     return S.symbols();
   }
 
+  struct GetExternalSymbolMapEntryValue {
+    Symbol *operator()(ExternalSymbolMap::value_type &KV) const {
+      return KV.second;
+    }
+  };
+
   struct GetSectionMapEntryValue {
     Section &operator()(SectionMap::value_type &KV) const { return *KV.second; }
   };
@@ -912,7 +919,10 @@ class LinkGraph {
   };
 
 public:
-  using external_symbol_iterator = ExternalSymbolSet::iterator;
+  using external_symbol_iterator =
+      mapped_iterator<ExternalSymbolMap::iterator,
+                      GetExternalSymbolMapEntryValue>;
+  using absolute_symbol_iterator = AbsoluteSymbolSet::iterator;
 
   using section_iterator =
       mapped_iterator<SectionMap::iterator, GetSectionMapEntryValue>;
@@ -1192,15 +1202,11 @@ class LinkGraph {
   /// of 0.
   Symbol &addExternalSymbol(StringRef Name, orc::ExecutorAddrDiff Size,
                             bool IsWeaklyReferenced) {
-    assert(llvm::count_if(ExternalSymbols,
-                          [&](const Symbol *Sym) {
-                            return Sym->getName() == Name;
-                          }) == 0 &&
-           "Duplicate external symbol");
+    assert(!ExternalSymbols.contains(Name) && "Duplicate external symbol");
     auto &Sym = Symbol::constructExternal(
         Allocator, createAddressable(orc::ExecutorAddr(), false), Name, Size,
         Linkage::Strong, IsWeaklyReferenced);
-    ExternalSymbols.insert(&Sym);
+    ExternalSymbols.insert({Sym.getName(), &Sym});
     return Sym;
   }
 
@@ -1281,10 +1287,14 @@ class LinkGraph {
   }
 
   iterator_range<external_symbol_iterator> external_symbols() {
-    return make_range(ExternalSymbols.begin(), ExternalSymbols.end());
+    return make_range(
+        external_symbol_iterator(ExternalSymbols.begin(),
+                                 GetExternalSymbolMapEntryValue()),
+        external_symbol_iterator(ExternalSymbols.end(),
+                                 GetExternalSymbolMapEntryValue()));
   }
 
-  iterator_range<external_symbol_iterator> absolute_symbols() {
+  iterator_range<absolute_symbol_iterator> absolute_symbols() {
     return make_range(AbsoluteSymbols.begin(), AbsoluteSymbols.end());
   }
 
@@ -1320,7 +1330,7 @@ class LinkGraph {
       Sec.removeSymbol(Sym);
       Sym.makeExternal(createAddressable(orc::ExecutorAddr(), false));
     }
-    ExternalSymbols.insert(&Sym);
+    ExternalSymbols.insert({Sym.getName(), &Sym});
   }
 
   /// Make the given symbol an absolute with the given address (must not already
@@ -1334,10 +1344,10 @@ class LinkGraph {
   void makeAbsolute(Symbol &Sym, orc::ExecutorAddr Address) {
     assert(!Sym.isAbsolute() && "Symbol is already absolute");
     if (Sym.isExternal()) {
-      assert(ExternalSymbols.count(&Sym) &&
+      assert(ExternalSymbols.contains(Sym.getName()) &&
              "Sym is not in the absolute symbols set");
       assert(Sym.getOffset() == 0 && "External is not at offset 0");
-      ExternalSymbols.erase(&Sym);
+      ExternalSymbols.erase(Sym.getName());
       auto &A = Sym.getAddressable();
       A.setAbsolute(true);
       A.setAddress(Address);
@@ -1362,9 +1372,9 @@ class LinkGraph {
              "Symbol is not in the absolutes set");
       AbsoluteSymbols.erase(&Sym);
     } else {
-      assert(ExternalSymbols.count(&Sym) &&
+      assert(ExternalSymbols.contains(Sym.getName()) &&
              "Symbol is not in the externals set");
-      ExternalSymbols.erase(&Sym);
+      ExternalSymbols.erase(Sym.getName());
     }
     Addressable &OldBase = *Sym.Base;
     Sym.setBlock(Content);
@@ -1449,10 +1459,11 @@ class LinkGraph {
   void removeExternalSymbol(Symbol &Sym) {
     assert(!Sym.isDefined() && !Sym.isAbsolute() &&
            "Sym is not an external symbol");
-    assert(ExternalSymbols.count(&Sym) && "Symbol is not in the externals set");
-    ExternalSymbols.erase(&Sym);
+    assert(ExternalSymbols.contains(Sym.getName()) &&
+           "Symbol is not in the externals set");
+    ExternalSymbols.erase(Sym.getName());
     Addressable &Base = *Sym.Base;
-    assert(llvm::none_of(ExternalSymbols,
+    assert(llvm::none_of(external_symbols(),
                          [&](Symbol *AS) { return AS->Base == &Base; }) &&
            "Base addressable still in use");
     destroySymbol(Sym);
@@ -1467,7 +1478,7 @@ class LinkGraph {
            "Symbol is not in the absolute symbols set");
     AbsoluteSymbols.erase(&Sym);
     Addressable &Base = *Sym.Base;
-    assert(llvm::none_of(ExternalSymbols,
+    assert(llvm::none_of(external_symbols(),
                          [&](Symbol *AS) { return AS->Base == &Base; }) &&
            "Base addressable still in use");
     destroySymbol(Sym);
@@ -1524,8 +1535,8 @@ class LinkGraph {
   support::endianness Endianness;
   GetEdgeKindNameFunction GetEdgeKindName = nullptr;
   DenseMap<StringRef, std::unique_ptr<Section>> Sections;
-  ExternalSymbolSet ExternalSymbols;
-  ExternalSymbolSet AbsoluteSymbols;
+  ExternalSymbolMap ExternalSymbols;
+  AbsoluteSymbolSet AbsoluteSymbols;
   orc::shared::AllocActions AAs;
 };