[COFF] Optimize range extension thunk insertion memory usage

[rnk]

Summary

This avoids allocating O(#relocs) of intermediate data for each section
when range extension thunks aren't needed for that section. This also
removes a std::vector from SectionChunk, which further reduces its size.

Instead, this change adds the range extension thunk symbols to the
object files that contain sections that need extension thunks. By adding
them to the symbol table of the parent object, that means they now have
a symbol table index. Then we can then modify the original relocation,
after copying it to read-write memory, to use the new symbol table
index.

This makes linking browser_tests.exe with no PDB 10.46% faster, moving
it from 11.364s to 10.288s averaged over five runs.

[mstorsjo]

This looks good from both a performance and memory usage perspective - I'll try to test run it on a largish executable that needs thunks. The code takes a little while to wrap the head around, but I think I understand it.

I tested it on a largish testcase that uses thunks, where linking takes around 24 seconds, where code layout was pretty consistently 2450 ms before this patch and 2500 ms after it. So to my eye, it's a very, very marginal slowdown for the actual thunk cases, but clearly outweighed by everything else.

lld/COFF/Writer.cpp

460 ↗

(On Diff #192515)

Does the initial value ~0U actually matter here? If Insertion.second is true (if this was the first insertion), the value will be unconditionally overwritten anyway, right? (I guess it's a good sentinel to safeguard against not updating it though, so it's probably ok as such, just trying to grasp all aspects of this.)

[mstorsjo]

I forgot to say... Thanks for doing this! This eases my conscience a little, for bloating things for others due to this feature :-)

[ruiu]

Nice!

lld/COFF/Writer.cpp

460 ↗

(On Diff #192515)

Yeah, it looks to me that this can be a set instead of a map, as it doesn't seem like it is using map values.

[mstorsjo]

lld/COFF/Writer.cpp

460 ↗

(On Diff #192515)

It does use them (that's ThunkSymbolIndex), but the initial value inserted into the map with the insert mrthod could be anything, as it is always overwritten if it actually was a new key.

[rnk]

Thanks! Looks like I never sent my reply comment. Oops :(

lld/COFF/Writer.cpp

460 ↗

(On Diff #192515)

Right, it's going to be overwritten after insertion, I was just trying to pick a sentinel that would be more likely to crash.

lld/trunk/COFF/Chunks.h

public:
   virtual size_t getSize() const = 0;
 
   // Write this chunk to a mmap'ed file, assuming Buf is pointing to
   // beginning of the file. Because this function may use RVA values
   // of other chunks for relocations, you need to set them properly
   // before calling this function.
   virtual void writeTo(uint8_t *Buf) const {}
 
-  // Called by the writer once before assigning addresses and writing
-  // the output.
-  virtual void readRelocTargets() {}
-
-  // Called if restarting thunk addition.
-  virtual void resetRelocTargets() {}
-
   // Called by the writer after an RVA is assigned, but before calling
   // getSize().
   virtual void finalizeContents() {}
 
   // The writer sets and uses the addresses.
   uint64_t getRVA() const { return RVA; }
   void setRVA(uint64_t V) { RVA = V; }
 
public:
   public:
     symbol_iterator() = default;
 
     Symbol *operator*() const { return File->getSymbol(I->SymbolTableIndex); }
   };
 
   SectionChunk(ObjFile *File, const coff_section *Header);
   static bool classof(const Chunk *C) { return C->kind() == SectionKind; }
-  void readRelocTargets() override;
-  void resetRelocTargets() override;
   size_t getSize() const override { return Header->SizeOfRawData; }
   ArrayRef<uint8_t> getContents() const;
   void writeTo(uint8_t *Buf) const override;
   bool hasData() const override;
   uint32_t getOutputCharacteristics() const override;
   StringRef getSectionName() const override { return SectionName; }
   void getBaserels(std::vector<Baserel> *Res) override;
   bool isCOMDAT() const;
public:
   // The COMDAT selection if this is a COMDAT chunk.
   llvm::COFF::COMDATType Selection = (llvm::COFF::COMDATType)0;
 
   ArrayRef<coff_relocation> Relocs;
 
   // Used by the garbage collector.
   bool Live;
 
-  // When inserting a thunk, we need to adjust a relocation to point to
-  // the thunk instead of the actual original target Symbol.
-  std::vector<Symbol *> RelocTargets;
-
 private:
   StringRef SectionName;
   std::vector<SectionChunk *> AssocChildren;
 
   // Used for ICF (Identical COMDAT Folding)
   void replace(SectionChunk *Other);
   uint32_t Class[2] = {0, 0};
 };

lld/trunk/COFF/Chunks.cpp

SectionChunk::SectionChunk(ObjFile *F, const coff_section *H)
 
   // If linker GC is disabled, every chunk starts out alive.  If linker GC is
   // enabled, treat non-comdat sections as roots. Generally optimized object
   // files will be built with -ffunction-sections or /Gy, so most things worth
   // stripping will be in a comdat.
   Live = !Config->DoGC || !isCOMDAT();
 }
 
-// Initialize the RelocTargets vector, to allow redirecting certain relocations
-// to a thunk instead of the actual symbol the relocation's symbol table index
-// indicates.
-void SectionChunk::readRelocTargets() {
-  assert(RelocTargets.empty());
-  RelocTargets.reserve(Relocs.size());
-  for (const coff_relocation &Rel : Relocs)
-    RelocTargets.push_back(File->getSymbol(Rel.SymbolTableIndex));
-}
-
-// Reset RelocTargets to their original targets before thunks were added.
-void SectionChunk::resetRelocTargets() {
-  for (size_t I = 0, E = Relocs.size(); I < E; ++I)
-    RelocTargets[I] = File->getSymbol(Relocs[I].SymbolTableIndex);
-}
-
 static void add16(uint8_t *P, int16_t V) { write16le(P, read16le(P) + V); }
 static void add32(uint8_t *P, int32_t V) { write32le(P, read32le(P) + V); }
 static void add64(uint8_t *P, int64_t V) { write64le(P, read64le(P) + V); }
 static void or16(uint8_t *P, uint16_t V) { write16le(P, read16le(P) | V); }
 static void or32(uint8_t *P, uint32_t V) { write32le(P, read32le(P) | V); }
 
 // Verify that given sections are appropriate targets for SECREL
 // relocations. This check is relaxed because unfortunately debug
for (size_t I = 0, E = Relocs.size(); I < E; I++) {
     // beginning of the following input section.
     if (Rel.VirtualAddress >= InputSize) {
       error("relocation points beyond the end of its parent section");
       continue;
     }
 
     uint8_t *Off = Buf + OutputSectionOff + Rel.VirtualAddress;
 
-    // Use the potentially remapped Symbol instead of the one that the
-    // relocation points to.
-    auto *Sym = dyn_cast_or_null<Defined>(RelocTargets[I]);
+    auto *Sym =
+        dyn_cast_or_null<Defined>(File->getSymbol(Rel.SymbolTableIndex));
 
     // Get the output section of the symbol for this relocation.  The output
     // section is needed to compute SECREL and SECTION relocations used in debug
     // info.
     Chunk *C = Sym ? Sym->getChunk() : nullptr;
     OutputSection *OS = C ? C->getOutputSection() : nullptr;
 
     // Skip the relocation if it refers to a discarded section, and diagnose it

 // fixed by the loader if load-time relocation is needed.
 // Only called when base relocation is enabled.
 void SectionChunk::getBaserels(std::vector<Baserel> *Res) {
   for (size_t I = 0, E = Relocs.size(); I < E; I++) {
     const coff_relocation &Rel = Relocs[I];
     uint8_t Ty = getBaserelType(Rel);
     if (Ty == IMAGE_REL_BASED_ABSOLUTE)
       continue;
-    // Use the potentially remapped Symbol instead of the one that the
-    // relocation points to.
-    Symbol *Target = RelocTargets[I];
+    Symbol *Target = File->getSymbol(Rel.SymbolTableIndex);
     if (!Target || isa<DefinedAbsolute>(Target))
       continue;
     Res->emplace_back(RVA + Rel.VirtualAddress, Ty);
   }
 }
 
 // MinGW specific.
 // Check whether a static relocation of type Type can be deferred and

lld/trunk/COFF/InputFiles.h

public:
   // underlying object file.
   Symbol *getSymbol(uint32_t SymbolIndex) {
     return Symbols[SymbolIndex];
   }
 
   // Returns the underlying COFF file.
   COFFObjectFile *getCOFFObj() { return COFFObj.get(); }
 
+  // Add a symbol for a range extension thunk. Return the new symbol table
+  // index. This index can be used to modify a relocation.
+  uint32_t addRangeThunkSymbol(Symbol *Thunk) {
+    Symbols.push_back(Thunk);
+    return Symbols.size() - 1;
+  }
+
   static std::vector<ObjFile *> Instances;
 
   // Flags in the absolute @feat.00 symbol if it is present. These usually
   // indicate if an object was compiled with certain security features enabled
   // like stack guard, safeseh, /guard:cf, or other things.
   uint32_t Feat00Flags = 0;
 
   // True if this object file is compatible with SEH.  COFF-specific and

lld/trunk/COFF/PDB.cpp

 }
 
 // Allocate memory for a .debug$S / .debug$F section and relocate it.
 static ArrayRef<uint8_t> relocateDebugChunk(BumpPtrAllocator &Alloc,
                                             SectionChunk &DebugChunk) {
   uint8_t *Buffer = Alloc.Allocate<uint8_t>(DebugChunk.getSize());
   assert(DebugChunk.OutputSectionOff == 0 &&
          "debug sections should not be in output sections");
-  DebugChunk.readRelocTargets();
   DebugChunk.writeTo(Buffer);
   return makeArrayRef(Buffer, DebugChunk.getSize());
 }
 
 static pdb::SectionContrib createSectionContrib(const Chunk *C, uint32_t Modi) {
   OutputSection *OS = C ? C->getOutputSection() : nullptr;
   pdb::SectionContrib SC;
   memset(&SC, 0, sizeof(SC));

lld/trunk/COFF/Writer.cpp

 private:
   void createSections();
   void createMiscChunks();
   void createImportTables();
   void appendImportThunks();
   void locateImportTables();
   void createExportTable();
   void mergeSections();
-  void readRelocTargets();
   void removeUnusedSections();
   void assignAddresses();
   void finalizeAddresses();
   void removeEmptySections();
   void createSymbolAndStringTable();
   void openFile(StringRef OutputPath);
   template <typename PEHeaderTy> void writeHeader();
   void createSEHTable();

 // the target go out of range.
 //
 // After adding thunks, we verify that all relocations are in range (with
 // no extra margin requirements). If this failed, we restart (throwing away
 // the previously created thunks) and retry with a wider margin.
 static bool createThunks(OutputSection *OS, int Margin) {
   bool AddressesChanged = false;
   DenseMap<uint64_t, Defined *> LastThunks;
+  DenseMap<std::pair<ObjFile *, Defined *>, uint32_t> ThunkSymtabIndices;
   size_t ThunksSize = 0;
   // Recheck Chunks.size() each iteration, since we can insert more
   // elements into it.
   for (size_t I = 0; I != OS->Chunks.size(); ++I) {
     SectionChunk *SC = dyn_cast_or_null<SectionChunk>(OS->Chunks[I]);
     if (!SC)
       continue;
     size_t ThunkInsertionSpot = I + 1;
 
     // Try to get a good enough estimate of where new thunks will be placed.
     // Offset this by the size of the new thunks added so far, to make the
     // estimate slightly better.
     size_t ThunkInsertionRVA = SC->getRVA() + SC->getSize() + ThunksSize;
-    for (size_t J = 0, E = SC->Relocs.size(); J < E; ++J) {
-      const coff_relocation &Rel = SC->Relocs[J];
-      Symbol *&RelocTarget = SC->RelocTargets[J];
+    ObjFile *File = SC->File;
+    std::vector<std::pair<uint32_t, uint32_t>> RelocReplacements;
+    ArrayRef<coff_relocation> OriginalRelocs =
+        File->getCOFFObj()->getRelocations(SC->Header);
+    for (size_t J = 0, E = OriginalRelocs.size(); J < E; ++J) {
+      const coff_relocation &Rel = OriginalRelocs[J];
+      Symbol *RelocTarget = File->getSymbol(Rel.SymbolTableIndex);
 
       // The estimate of the source address P should be pretty accurate,
       // but we don't know whether the target Symbol address should be
       // offset by ThunkSize or not (or by some of ThunksSize but not all of
       // it), giving us some uncertainty once we have added one thunk.
       uint64_t P = SC->getRVA() + Rel.VirtualAddress + ThunksSize;
 
       Defined *Sym = dyn_cast_or_null<Defined>(RelocTarget);
for (size_t J = 0, E = OriginalRelocs.size(); J < E; ++J) {
             ThunkInsertionRVA); // Estimate of where it will be located.
         ThunkChunk->setOutputSection(OS);
         OS->Chunks.insert(OS->Chunks.begin() + ThunkInsertionSpot, ThunkChunk);
         ThunkInsertionSpot++;
         ThunksSize += ThunkChunk->getSize();
         ThunkInsertionRVA += ThunkChunk->getSize();
         AddressesChanged = true;
       }
-      RelocTarget = Thunk;
+
+      // To redirect the relocation, add a symbol to the parent object file's
+      // symbol table, and replace the relocation symbol table index with the
+      // new index.
+      auto Insertion = ThunkSymtabIndices.insert({{File, Thunk}, ~0U});
+      uint32_t &ThunkSymbolIndex = Insertion.first->second;
+      if (Insertion.second)
+        ThunkSymbolIndex = File->addRangeThunkSymbol(Thunk);
+      RelocReplacements.push_back({J, ThunkSymbolIndex});
+    }
+
+    // Get a writable copy of this section's relocations so they can be
+    // modified. If the relocations point into the object file, allocate new
+    // memory. Otherwise, this must be previously allocated memory that can be
+    // modified in place.
+    MutableArrayRef<coff_relocation> NewRelocs;
+    if (OriginalRelocs.data() == SC->Relocs.data()) {
+      NewRelocs = makeMutableArrayRef(
+          BAlloc.Allocate<coff_relocation>(OriginalRelocs.size()),
+          OriginalRelocs.size());
+    } else {
+      NewRelocs = makeMutableArrayRef(
+          const_cast<coff_relocation *>(SC->Relocs.data()), SC->Relocs.size());
+    }
+
+    // Copy each relocation, but replace the symbol table indices which need
+    // thunks.
+    auto NextReplacement = RelocReplacements.begin();
+    auto EndReplacement = RelocReplacements.end();
+    for (size_t I = 0, E = OriginalRelocs.size(); I != E; ++I) {
+      NewRelocs[I] = OriginalRelocs[I];
+      if (NextReplacement != EndReplacement && NextReplacement->first == I) {
+        NewRelocs[I].SymbolTableIndex = NextReplacement->second;
+        ++NextReplacement;
+      }
     }
+
+    SC->Relocs = makeArrayRef(NewRelocs.data(), NewRelocs.size());
   }
   return AddressesChanged;
 }
 
 // Verify that all relocations are in range, with no extra margin requirements.
 static bool verifyRanges(const std::vector<Chunk *> Chunks) {
   for (Chunk *C : Chunks) {
     SectionChunk *SC = dyn_cast_or_null<SectionChunk>(C);
     if (!SC)
       continue;
 
     for (size_t J = 0, E = SC->Relocs.size(); J < E; ++J) {
       const coff_relocation &Rel = SC->Relocs[J];
-      Symbol *RelocTarget = SC->RelocTargets[J];
+      Symbol *RelocTarget = SC->File->getSymbol(Rel.SymbolTableIndex);
 
       Defined *Sym = dyn_cast_or_null<Defined>(RelocTarget);
       if (!Sym)
         continue;
 
       uint64_t P = SC->getRVA() + Rel.VirtualAddress;
       uint64_t S = Sym->getRVA();
 
while (true) {
 
     if (Pass >= 10)
       fatal("adding thunks hasn't converged after " + Twine(Pass) + " passes");
 
     if (Pass > 0) {
       // If the previous pass didn't work out, reset everything back to the
       // original conditions before retrying with a wider margin. This should
       // ideally never happen under real circumstances.
-      for (OutputSection *Sec : OutputSections) {
+      for (OutputSection *Sec : OutputSections)
         Sec->Chunks = Sec->OrigChunks;
-        for (Chunk *C : Sec->Chunks)
-          C->resetRelocTargets();
-      }
       Margin *= 2;
     }
 
     // Try adding thunks everywhere where it is needed, with a margin
     // to avoid things going out of range due to the added thunks.
     bool AddressesChanged = false;
     for (OutputSection *Sec : OutputSections)
       AddressesChanged |= createThunks(Sec, Margin);
void Writer::run() {
   ScopedTimer T1(CodeLayoutTimer);
 
   createImportTables();
   createSections();
   createMiscChunks();
   appendImportThunks();
   createExportTable();
   mergeSections();
-  readRelocTargets();
   removeUnusedSections();
   finalizeAddresses();
   removeEmptySections();
   setSectionPermissions();
   createSymbolAndStringTable();
 
   if (FileSize > UINT32_MAX)
     fatal("image size (" + Twine(FileSize) + ") " +
for (auto &P : Config->Merge) {
     if (!To) {
       From->Name = ToName;
       continue;
     }
     To->merge(From);
   }
 }
 
-// Visits all sections to initialize their relocation targets.
-void Writer::readRelocTargets() {
-  for (OutputSection *Sec : OutputSections)
-    parallelForEach(Sec->Chunks, [&](Chunk *C) { C->readRelocTargets(); });
-}
-
 // Visits all sections to assign incremental, non-overlapping RVAs and
 // file offsets.
 void Writer::assignAddresses() {
   SizeOfHeaders = DOSStubSize + sizeof(PEMagic) + sizeof(coff_file_header) +
                   sizeof(data_directory) * NumberOfDataDirectory +
                   sizeof(coff_section) * OutputSections.size();
   SizeOfHeaders +=
       Config->is64() ? sizeof(pe32plus_header) : sizeof(pe32_header);