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

[jobnoorman]

Summary

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.

Note that I fully understand that adding complexity to improve the
performance of an assert might not be acceptable (since the checks
shouldn't exist in release builds). This patch has another advantage,
however: since external symbol names really have to be unique, it could
make sense to make this clear at the type level.

[rafauler]

I'll have to maintain this diff downstream when BOLT switches to JITLink. If this can be maintained by upstream JITLink devs, I would appreciate it. Thanks.

[rafauler]

Ping. Can't we just remove this assertion? Is it worth making the runtime much slower to check that we don't have a duplicate external symbol?

[jobnoorman]

Ping

[maksfb]

@lhames, this looks pretty harmless change to me. Please let us know if you disagree.

[maksfb]

@jobnoorman: should be fine to land it.

[jobnoorman]

Hi @lhames, I'm planning to land this at the beginning of next week. Please let me know if you have any issues with this.

llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h

 private:
   Block *First = nullptr;
   Block *Last = nullptr;
 };
 
 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>
   Addressable &createAddressable(ArgTs &&... Args) {
     Addressable *A =
         reinterpret_cast<Addressable *>(Allocator.Allocate<Addressable>());
     new (A) Addressable(std::forward<ArgTs>(Args)...);
     return *A;
getSectionSymbols(Section &S) {
     return S.symbols();
   }
 
   static iterator_range<Section::const_symbol_iterator>
   getSectionConstSymbols(const Section &S) {
     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; }
   };
 
   struct GetSectionMapEntryConstValue {
     const Section &operator()(const SectionMap::value_type &KV) const {
       return *KV.second;
     }
   };
 
 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>;
   using const_section_iterator =
       mapped_iterator<SectionMap::const_iterator, GetSectionMapEntryConstValue>;
 
   template <typename OuterItrT, typename InnerItrT, typename T,
             iterator_range<InnerItrT> getInnerRange(
public:
   /// size is not known, you should substitute '0'.
   /// The IsWeaklyReferenced argument determines whether the symbol must be
   /// present during lookup: Externals that are strongly referenced must be
   /// found or an error will be emitted. Externals that are weakly referenced
   /// are permitted to be undefined, in which case they are assigned an address
   /// 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;
   }
 
   /// Add an absolute symbol.
   Symbol &addAbsoluteSymbol(StringRef Name, orc::ExecutorAddr Address,
                             orc::ExecutorAddrDiff Size, Linkage L, Scope S,
                             bool IsLive) {
     assert((S == Scope::Local || llvm::count_if(AbsoluteSymbols,
public:
 
   iterator_range<const_block_iterator> blocks() const {
     auto Secs = sections();
     return make_range(const_block_iterator(Secs.begin(), Secs.end()),
                       const_block_iterator(Secs.end(), Secs.end()));
   }
 
   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());
   }
 
   iterator_range<defined_symbol_iterator> defined_symbols() {
     auto Secs = sections();
     return make_range(defined_symbol_iterator(Secs.begin(), Secs.end()),
                       defined_symbol_iterator(Secs.end(), Secs.end()));
   }
if (Sym.isAbsolute()) {
       A.setAbsolute(false);
       A.setAddress(orc::ExecutorAddr());
     } else {
       assert(Sym.isDefined() && "Sym is not a defined symbol");
       Section &Sec = Sym.getBlock().getSection();
       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
   /// be absolute).
   ///
   /// The symbol's size, linkage, and callability, and liveness will be left
   /// unchanged, and its offset will be reset to 0.
   ///
   /// If the symbol was external then its scope will be set to local, otherwise
   /// it will be left unchanged.
   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);
       Sym.setScope(Scope::Local);
     } else {
       assert(Sym.isDefined() && "Sym is not a defined symbol");
       Section &Sec = Sym.getBlock().getSection();
       Sec.removeSymbol(Sym);
       Sym.makeAbsolute(createAddressable(Address));
     }
     AbsoluteSymbols.insert(&Sym);
   }
 
   /// Turn an absolute or external symbol into a defined one by attaching it to
   /// a block. Symbol must not already be defined.
   void makeDefined(Symbol &Sym, Block &Content, orc::ExecutorAddrDiff Offset,
                    orc::ExecutorAddrDiff Size, Linkage L, Scope S,
                    bool IsLive) {
     assert(!Sym.isDefined() && "Sym is already a defined symbol");
     if (Sym.isAbsolute()) {
       assert(AbsoluteSymbols.count(&Sym) &&
              "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);
     Sym.setOffset(Offset);
     Sym.setSize(Size);
     Sym.setLinkage(L);
     Sym.setScope(S);
     Sym.setLive(IsLive);
void mergeSections(Section &DstSection, Section &SrcSection,
     if (!PreserveSrcSection)
       removeSection(SrcSection);
   }
 
   /// Removes an external symbol. Also removes the underlying Addressable.
   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);
     destroyAddressable(Base);
   }
 
   /// Remove an absolute symbol. Also removes the underlying Addressable.
   void removeAbsoluteSymbol(Symbol &Sym) {
     assert(!Sym.isDefined() && Sym.isAbsolute() &&
            "Sym is not an absolute symbol");
     assert(AbsoluteSymbols.count(&Sym) &&
            "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);
     destroyAddressable(Base);
   }
 
   /// Removes defined symbols. Does not remove the underlying block.
   void removeDefinedSymbol(Symbol &Sym) {
private:
 
   std::string Name;
   Triple TT;
   SubtargetFeatures Features;
   unsigned PointerSize;
   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;
 };
 
 inline MutableArrayRef<char> Block::getMutableContent(LinkGraph &G) {
   if (!ContentMutable)
     setMutableContent(G.allocateContent({Data, Size}));
   return MutableArrayRef<char>(const_cast<char *>(Data), Size);
 }