Splits __wasm_apply_relocs into multiple functions

[danchurch]

Summary

Rather than having one potentially large wasm_apply_relocs function, wasm_apply_relocs is split into 10 smaller functions, all executed in order.

Safari on iOS struggles with excessively large functions, and for a larger .wasm module (~7mb observed), _wasm_apply_relocs balloons to >500kb, causing Safari iOS to fail with a "max stack size exceeded" error. By splitting __wasm_apply_relocs into multiple smaller methods, relocation functionality is preserved while being able to load/run on iOS.

I understand that this code may not be perfect, setting the initial function count to 10 is a hack, etc. I'm submitting this patch to spark discussion, expecting that the overall approach may need to be modified (or at-least made tune-able with linker flags).

[jyknight]

Copying what I said from Discord --

If I were a reviewer on this, I'd say that's too much a hack, and that it would be a much better idea to implement the "TODO(sbc): Encode the relocations in the data section and write a loop here to apply them."

That would allow this function to be both constant and tiny.

[sbc100]

I tend to agree that maybe there is a better solution to be found here.

Can I ask what you use case for linking with -fpic is? Presumably you are building with emscripten's SIDE_MODULE/MAIN_MODULE. Would you mind opening an llvm bug for this were we can discuss this further?

lld/wasm/Driver.cpp

static void createSyntheticSymbols() {
   static llvm::wasm::WasmGlobalType globalTypeI32 = {WASM_TYPE_I32, false};
   static llvm::wasm::WasmGlobalType mutableGlobalTypeI32 = {WASM_TYPE_I32,
                                                             true};
   WasmSym::callCtors = symtab->addSyntheticFunction(
       "__wasm_call_ctors", WASM_SYMBOL_VISIBILITY_HIDDEN,
       make<SyntheticFunction>(nullSignature, "__wasm_call_ctors"));
 
   if (config->isPic) {
+    for (auto i = 0; i < WasmSym::applyRelocFuncCount; i++) {
       // For PIC code we create a synthetic function __wasm_apply_relocs which
       // is called from __wasm_call_ctors before the user-level constructors.
-    WasmSym::applyRelocs = symtab->addSyntheticFunction(
-        "__wasm_apply_relocs", WASM_SYMBOL_VISIBILITY_HIDDEN,
-        make<SyntheticFunction>(nullSignature, "__wasm_apply_relocs"));
+      char buff[1024];
+      sprintf(buff, "__wasm__apply_relocs_%d", i);
+      WasmSym::applyRelocFuncs[i] = symtab->addSyntheticFunction(
+          buff, WASM_SYMBOL_VISIBILITY_HIDDEN,
+          make<SyntheticFunction>(nullSignature, buff));
+    }
   }
 
 
   if (config->isPic) {
     WasmSym::stackPointer =
         createUndefinedGlobal("__stack_pointer", &mutableGlobalTypeI32);
     // For PIC code, we import two global variables (__memory_base and
     // __table_base) from the environment and use these as the offset at

lld/wasm/InputChunks.h

 // each global variable.
 class InputSegment : public InputChunk {
 public:
   InputSegment(const WasmSegment &seg, ObjFile *f)
       : InputChunk(f, InputChunk::DataSegment), segment(seg) {}
 
   static bool classof(const InputChunk *c) { return c->kind() == DataSegment; }
 
-  void generateRelocationCode(raw_ostream &os) const;
+  void generateRelocationCode(raw_ostream &os, int relocFuncIdx,
+                              int relocFuncCount) const;
 
   uint32_t getAlignment() const { return segment.Data.Alignment; }
   StringRef getName() const override { return segment.Data.Name; }
   StringRef getDebugName() const override { return StringRef(); }
   uint32_t getComdat() const override { return segment.Data.Comdat; }
   uint32_t getInputSectionOffset() const override {
     return segment.SectionOffset;
   }

lld/wasm/InputChunks.cpp

void InputFunction::writeTo(uint8_t *buf) const {
   LLVM_DEBUG(dbgs() << "  write final chunk: " << chunkSize << "\n");
   memcpy(buf, lastRelocEnd, chunkSize);
   LLVM_DEBUG(dbgs() << "  total: " << (buf + chunkSize - orig) << "\n");
 }
 
 // Generate code to apply relocations to the data section at runtime.
 // This is only called when generating shared libaries (PIC) where address are
 // not known at static link time.
-void InputSegment::generateRelocationCode(raw_ostream &os) const {
+void InputSegment::generateRelocationCode(raw_ostream &os, int relocFuncIdx,
+                                          int relocFuncCount) const {
   LLVM_DEBUG(dbgs() << "generating runtime relocations: " << getName()
-                    << " count=" << relocations.size() << "\n");
+                    << " count=" << relocations.size() << " chunk "
+                    << relocFuncidx << "/" << relocFuncCount << "\n");
+  const auto relocsPerFunc =
+      static_cast<int>(relocations.size() / relocFuncCount);
+  const auto startIdx = relocsPerFunc * relocFuncIdx;
+  const auto endIdx = relocFuncIdx + 1 < relocFuncCount
+                          ? relocsPerFunc * (relocFuncIdx + 1)
+                          : static_cast<int>(relocations.size());
 
   // TODO(sbc): Encode the relocations in the data section and write a loop
   // here to apply them.
   uint32_t segmentVA = outputSeg->startVA + outputSegmentOffset;
-  for (const WasmRelocation &rel : relocations) {
+  for (auto relIdx = startIdx; relIdx < endIdx; relIdx++) {
+    const WasmRelocation &rel = relocations[relIdx];
     uint32_t offset = rel.Offset - getInputSectionOffset();
     uint32_t outputOffset = segmentVA + offset;
 
     LLVM_DEBUG(dbgs() << "gen reloc: type=" << relocTypeToString(rel.Type)
                       << " addend=" << rel.Addend << " index=" << rel.Index
                       << " output offset=" << outputOffset << "\n");
 
     // Get __memory_base

lld/wasm/MarkLive.cpp

void MarkLive::enqueue(Symbol *sym) {
   if (InputChunk *chunk = sym->getChunk())
     queue.push_back(chunk);
 
   // The ctor functions are all referenced by the synthetic callCtors
   // function.  However, this function does not contain relocations so we
   // have to manually mark the ctors as live if callCtors itself is live.
   if (sym == WasmSym::callCtors) {
     if (config->isPic)
-      enqueue(WasmSym::applyRelocs);
+      for (auto i = 0; i < WasmSym::applyRelocFuncCount; i++)
+        enqueue(WasmSym::applyRelocFuncs[i]);
     for (const ObjFile *obj : symtab->objectFiles) {
       const WasmLinkingData &l = obj->getWasmObj()->linkingData();
       for (const WasmInitFunc &f : l.InitFunctions) {
         auto* initSym = obj->getFunctionSymbol(f.Symbol);
         if (!initSym->isDiscarded())
           enqueue(initSym);
       }
     }

lld/wasm/Symbols.h

struct WasmSym {
   static DefinedFunction *initMemory;
 
   // __wasm_call_ctors
   // Function that directly calls all ctors in priority order.
   static DefinedFunction *callCtors;
 
   // __wasm_apply_relocs
   // Function that applies relocations to data segment post-instantiation.
-  static DefinedFunction *applyRelocs;
+  static DefinedFunction *applyRelocFuncs[10];
+  static int applyRelocFuncCount;
 
   // __wasm_init_tls
   // Function that allocates thread-local storage and initializes it.
   static DefinedFunction *initTLS;
 
   // __dso_handle
   // Symbol used in calls to __cxa_atexit to determine current DLL
   static DefinedData *dsoHandle;

lld/wasm/Symbols.cpp

case wasm::Symbol::OutputSectionKind:
     return "OutputSectionKind";
   }
   llvm_unreachable("invalid symbol kind");
 }
 
 namespace wasm {
 DefinedFunction *WasmSym::callCtors;
 DefinedFunction *WasmSym::initMemory;
-DefinedFunction *WasmSym::applyRelocs;
+DefinedFunction *WasmSym::applyRelocFuncs[10];
+int WasmSym::applyRelocFuncCount = 10;
 DefinedFunction *WasmSym::initTLS;
 DefinedData *WasmSym::dsoHandle;
 DefinedData *WasmSym::dataEnd;
 DefinedData *WasmSym::globalBase;
 DefinedData *WasmSym::heapBase;
 DefinedData *WasmSym::initMemoryFlag;
 GlobalSymbol *WasmSym::stackPointer;
 GlobalSymbol *WasmSym::tlsBase;

lld/wasm/Writer.cpp

 }
 
 // For -shared (PIC) output, we create create a synthetic function which will
 // apply any relocations to the data segments on startup.  This function is
 // called __wasm_apply_relocs and is added at the beginning of __wasm_call_ctors
 // before any of the constructors run.
 void Writer::createApplyRelocationsFunction() {
   LLVM_DEBUG(dbgs() << "createApplyRelocationsFunction\n");
+  for (auto i = 0; i < WasmSym::applyRelocFuncCount; i++) {
     // First write the body's contents to a string.
     std::string bodyContent;
     {
       raw_string_ostream os(bodyContent);
       writeUleb128(os, 0, "num locals");
       for (const OutputSegment *seg : segments)
         for (const InputSegment *inSeg : seg->inputSegments)
-        inSeg->generateRelocationCode(os);
+          inSeg->generateRelocationCode(os, i, WasmSym::applyRelocFuncCount);
       writeU8(os, WASM_OPCODE_END, "END");
     }
 
-  createFunction(WasmSym::applyRelocs, bodyContent);
+    createFunction(WasmSym::applyRelocFuncs[i], bodyContent);
+  }
 }
 
 // Create synthetic "__wasm_call_ctors" function based on ctor functions
 // in input object.
 void Writer::createCallCtorsFunction() {
   if (!WasmSym::callCtors->isLive())
     return;
 
   // First write the body's contents to a string.
   std::string bodyContent;
   {
     raw_string_ostream os(bodyContent);
     writeUleb128(os, 0, "num locals");
 
     if (config->isPic) {
+      for (auto i = 0; i < WasmSym::applyRelocFuncCount; i++) {
         writeU8(os, WASM_OPCODE_CALL, "CALL");
-      writeUleb128(os, WasmSym::applyRelocs->getFunctionIndex(),
+        writeUleb128(os, WasmSym::applyRelocFuncs[i]->getFunctionIndex(),
                      "function index");
       }
+    }
 
     // Call constructors
     for (const WasmInitEntry &f : initFunctions) {
       writeU8(os, WASM_OPCODE_CALL, "CALL");
       writeUleb128(os, f.sym->getFunctionIndex(), "function index");
     }
     writeU8(os, WASM_OPCODE_END, "END");
   }