[lld][wasm] Split __wasm_apply_data_relocs when it exceeds the function size limit

[k1nder10]

Summary

__wasm_apply_data_relocs can grow drastically exceeding the function size limit.
This change introduces __wasm_apply_data_relocs_tail function which is supposed
to include instructions that caused the limit hit. __wasm_apply_data_relocs
calls __wasm_apply_data_relocs_tail at the end of the function body.
__wasm_apply_data_relocs_tail is empty if wasm module doesn't hit the function
size limit.

discussed here:
https://github.com/llvm/llvm-project/issues/55608
https://github.com/emscripten-core/emscripten/discussions/17374

[k1nder10]

It's my first time pushing the changes to LLVM.
Should I add some tests for this? If so, where could I find examples of them?

lld/wasm/Driver.cpp

767

I tried to move the creation of this function to Writer.cpp 1349 under 'else' statement, so that we don't have any changes in out .wasm files
unless we exceed the function size limit. However, I got some strange runtime errors.

So, for now, we will have this function created even If we don't hit the limit (it just will be empty). Is it okay?

[k1nder10]

lld/wasm/SplitApplyDataRelocsStrategy.cpp

41

This should be changed to work properly. Will upload a new patchset tomorrow.

const size_t pos = function_to_split_.rfind(instruction_opcode, config->maxFunctionSize - skip_instruction_count);

[k1nder10]

[lld][wasm] Split __wasm_apply_data_relocs when it exceeds the function size limit

wasm_apply_data_relocs can grow drastically exceeding the function size limit.
This change introduces wasm_apply_data_relocs_tail function which is supposed
to include instructions that caused the limit hit. wasm_apply_data_relocs
calls wasm_apply_data_relocs_tail at the end of the function body.
__wasm_apply_data_relocs_tail is empty if wasm module doesn't hit the function
size limit.

discussed here:
https://github.com/llvm/llvm-project/issues/55608
https://github.com/emscripten-core/emscripten/discussions/17374

[k1nder10]

[lld][wasm] Split __wasm_apply_data_relocs when it exceeds the function size limit

wasm_apply_data_relocs can grow drastically exceeding the function size limit.
This change introduces wasm_apply_data_relocs_tail function which is supposed
to include instructions that caused the limit hit. wasm_apply_data_relocs
calls wasm_apply_data_relocs_tail at the end of the function body.
__wasm_apply_data_relocs_tail is empty if wasm module doesn't hit the function
size limit.

discussed here:
https://github.com/llvm/llvm-project/issues/55608
https://github.com/emscripten-core/emscripten/discussions/17374

[sbc100]

I feel like a better approach might be to generate N different functions during the generation phase rather than trying to split the function after generation?

Also, I think we (also) want to investigate using a data driven approach for the relocations that don't involve symbol-specific globals. See the discussion on that here: https://github.com/emscripten-core/emscripten/issues/12682#issuecomment-1297755521

lld/test/wasm/shared.s

1 ↗

(On Diff #483898)

Why are alll these test files showing up in the diff like every line of the file has changed? Did the line endings change maybe?

lld/wasm/Driver.cpp

767

I'd rather only emit the new function if needed.

Also, it seems like we should splitting into N possible sub-functions, not just 1, otherwise we just push the problem down the road a little.

769

We don't need either of these flags for the sub-function(s) I think

lld/wasm/SplitApplyDataRelocsStrategy.h

9

Creating a class hierarchy seems unnecessary since there is only a single sub-class. I'm not sure its worth adding new headers and source files either

[k1nder10]

In D140111#4004579, @sbc100 wrote:

I feel like a better approach might be to generate N different functions during the generation phase rather than trying to split the function after generation?

Also, I think we (also) want to investigate using a data driven approach for the relocations that don't involve symbol-specific globals. See the discussion on that here: https://github.com/emscripten-core/emscripten/issues/12682#issuecomment-1297755521

Do we know the number of relocations (the size of __wasm_apply_data_relocs) during the generation phase?

Speaking of the discussion you provided, I also think it is a good idea. However, I don't have much time to work on it.

[k1nder10]

Fix the line ending in tests. Fix tests. Removed unnecessary classes.

[k1nder10]

lld/wasm/ApplyDataRelocsSplitter.cpp

41 ↗

(On Diff #484242)

As I understand, you build a wasm linker with -fno-exceptions right?
How should I rewrite such code?
Adding an assertion would work in a debug build but what about the release one?

[k1nder10]

lld/wasm/Driver.cpp

767

Speaking of splitting this into N functions, it looks unnecessary as it must be enough to have 2 functions actually.
We have a more than million lines of code and a plethora of class hierarchies and we exceed the limit only by 200k instructions. So, there's a big stock (>7 million) of instructions to hit the limit again.

[k1nder10]

Got rid of exceptions

[k1nder10]

@sbc100
So, what do you think about this change? Can we merge it? Or should I try to split this up into N functions?
BTW, Is it okay that I have to modify the data-segments.ll & pie.ll files?

[sbc100]

I still have some concerns here. Firstly, I think it should be possible to do this without generating the tail function when its not needed. Can't applyDataRelocsTail simply be null in this case?

Secondly I'm not sure about and a new source file and header, or the name of the those new files, but thats mostly aesthetic.

[sbc100]

Also, I think it might be cleaner to simply generate the second function while the main function is being generated, rather then generate the long/unsplit function completely and then split it up.

Something like this phsudocode:

for (r in relocations) {
  if (current_size_is_over_limit()) {
    generate_tail_call_to_new_function():
    switch_current_output_function_to_new_function();
  }
  generate_code_for_relocation(r);
}

Then no splitting would be needed.

lld/wasm/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS Options.td)
 tablegen(LLVM Options.inc -gen-opt-parser-defs)
 add_public_tablegen_target(WasmOptionsTableGen)
 
 add_lld_library(lldWasm
   Driver.cpp
   InputChunks.cpp
   InputFiles.cpp
   LTO.cpp
   MapFile.cpp
   MarkLive.cpp
   OutputSections.cpp
   OutputSegment.cpp
   Relocations.cpp
+  SplitApplyDataRelocsStrategy.cpp
   SymbolTable.cpp
   Symbols.cpp
   SyntheticSections.cpp
   Writer.cpp
   WriterUtils.cpp
 
   LINK_COMPONENTS
   ${LLVM_TARGETS_TO_BUILD}

lld/wasm/Config.h

struct Configuration {
   uint64_t zStackSize;
   unsigned ltoPartitions;
   unsigned ltoo;
   unsigned optimize;
   llvm::StringRef thinLTOJobs;
   bool ltoDebugPassManager;
   UnresolvedPolicy unresolvedSymbols;
 
+  // https://github.com/v8/v8/blob/master/src/wasm/wasm-limits.h#L47
+  uint64_t maxFunctionSize = 7'654'321;
+
   llvm::StringRef entry;
   llvm::StringRef mapFile;
   llvm::StringRef outputFile;
   llvm::StringRef thinLTOCacheDir;
 
   llvm::StringSet<> allowUndefinedSymbols;
   llvm::StringSet<> exportedSymbols;
   std::vector<llvm::StringRef> requiredExports;

lld/wasm/Driver.cpp

if (config->isPic ||
       config->unresolvedSymbols == UnresolvedPolicy::ImportDynamic) {
     // For PIC code, or when dynamically importing addresses, we create
     // synthetic functions that apply relocations.  These get called from
     // __wasm_call_ctors before the user-level constructors.
     WasmSym::applyDataRelocs = symtab->addSyntheticFunction(
         "__wasm_apply_data_relocs",
         WASM_SYMBOL_VISIBILITY_DEFAULT | WASM_SYMBOL_EXPORTED,
         make<SyntheticFunction>(nullSignature, "__wasm_apply_data_relocs"));
+
+    WasmSym::applyDataRelocsTail = symtab->addSyntheticFunction(
+        "__wasm_apply_data_relocs_tail",
+        WASM_SYMBOL_VISIBILITY_DEFAULT | WASM_SYMBOL_EXPORTED,
+        make<SyntheticFunction>(nullSignature,
+                                "__wasm_apply_data_relocs_tail"));
   }
 }
 
 static void createOptionalSymbols() {
   if (config->relocatable)
     return;
 
   WasmSym::dsoHandle = symtab->addOptionalDataSymbol("__dso_handle");

lld/wasm/FunctionSplitter.h

+#ifndef LLD_WASM_FUNCTION_SPLITTER_H
+#define LLD_WASM_FUNCTION_SPLITTER_H
+
+#include "SplitStrategy.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+
+namespace lld::wasm {
+
+class FunctionSplitter {
+public:
+  FunctionSplitter(std::unique_ptr<SplitStrategy> strategy)
+      : split_strategy_(std::move(strategy)) {}
+
+  SplitStrategy::SplittedFunctions split() { return split_strategy_->split(); }
+
+private:
+  std::unique_ptr<SplitStrategy> split_strategy_;
+};
+
+} // namespace lld::wasm
+
+#endif

lld/wasm/SplitApplyDataRelocsStrategy.h

+#ifndef LLD_WASM_SPLIT_APPLY_DATA_RELOCS_STRATEGY_H
+#define LLD_WASM_SPLIT_APPLY_DATA_RELOCS_STRATEGY_H
+
+#include "SplitStrategy.h"
+
+namespace lld::wasm {
+
+class SplitApplyDataRelocsStrategy final : public SplitStrategy {
+public:
+  SplitApplyDataRelocsStrategy(const std::string &functionBody);
+  virtual ~SplitApplyDataRelocsStrategy() = default;
+
+  // Splits the function up into 2 functions. Cuts off the instructions
+  // that exceed the function size limit and put them into the second function.
+  // The first function calls the second one at the end.
+  // The last instruction that is gonna be left in the first function
+  // is WASM_OPCODE_I32_STORE/WASM_OPCODE_I64_STORE (+ align & offset)
+  SplittedFunctions split() override;
+
+private:
+  size_t findPositionToCutInstructionsFrom(unsigned instruction_opcode) const;
+  void addApplyDataRelocsTailCall(std::string &functionBody, size_t pos) const;
+  void addCutOffInstructions(std::string &functionBody, size_t pos) const;
+
+  const std::string &function_to_split_;
+};
+
+} // namespace lld::wasm
+
+#endif

lld/wasm/SplitApplyDataRelocsStrategy.cpp

+#include "SplitApplyDataRelocsStrategy.h"
+#include "Config.h"
+#include "Symbols.h"
+#include "WriterUtils.h"
+#include "llvm/BinaryFormat/Wasm.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace lld::wasm {
+
+SplitApplyDataRelocsStrategy::SplitApplyDataRelocsStrategy(
+    const std::string &functionBody)
+    : function_to_split_(functionBody) {}
+
+SplitApplyDataRelocsStrategy::SplittedFunctions
+SplitApplyDataRelocsStrategy::split() {
+  if (function_to_split_.size() < config->maxFunctionSize) {
+    throw "The function is intented to only split functions that exceed"
+          "the function size limit";
+  }
+
+  using namespace llvm::wasm;
+
+  SplittedFunctions functions;
+  const bool is64 = config->is64.value_or(false);
+  const unsigned store_opcode =
+      is64 ? WASM_OPCODE_I64_STORE : WASM_OPCODE_I32_STORE;
+
+  const size_t pos = findPositionToCutInstructionsFrom(store_opcode);
+  addApplyDataRelocsTailCall(functions.first, pos);
+  addCutOffInstructions(functions.second, pos);
+
+  return functions;
+}
+
+size_t SplitApplyDataRelocsStrategy::findPositionToCutInstructionsFrom(
+    unsigned instruction_opcode) const {
+  // Skip some instructions to leave the place for CALL instruction and others
+  constexpr size_t skip_instruction_count = 10;
+  const size_t pos = function_to_split_.rfind(
+      instruction_opcode, config->maxFunctionSize - skip_instruction_count);
+  if (pos == std::string::npos)
+    throw "No instruction opcode found";
+
+  // includes the following instructions: store + align + offset
+  constexpr size_t offset = 3;
+
+  return pos + offset;
+}
+
+void SplitApplyDataRelocsStrategy::addApplyDataRelocsTailCall(
+    std::string &functionBody, size_t pos) const {
+  llvm::raw_string_ostream first_os(functionBody);
+  functionBody = function_to_split_.substr(0, pos);
+
+  writeU8(first_os, llvm::wasm::WASM_OPCODE_CALL, "CALL");
+  writeUleb128(first_os, WasmSym::applyDataRelocsTail->getFunctionIndex(),
+               "function index");
+  writeU8(first_os, llvm::wasm::WASM_OPCODE_END, "END");
+}
+
+void SplitApplyDataRelocsStrategy::addCutOffInstructions(
+    std::string &functionBody, size_t pos) const {
+  llvm::raw_string_ostream second_os(functionBody);
+  writeUleb128(second_os, 0, "num locals");
+  functionBody.append(function_to_split_.substr(pos));
+}
+
+} // namespace lld::wasm

lld/wasm/SplitStrategy.h

+#ifndef LLD_WASM_SPLIT_STRATEGY_H
+#define LLD_WASM_SPLIT_STRATEGY_H
+
+#include <string>
+#include <utility>
+
+namespace lld::wasm {
+
+class SplitStrategy {
+public:
+  using SplittedFunctions = std::pair<std::string, std::string>;
+
+  virtual ~SplitStrategy() = default;
+  virtual SplittedFunctions split() = 0;
+};
+
+} // namespace lld::wasm
+
+#endif

lld/wasm/Symbols.h

struct WasmSym {
   // __wasm_call_dtors
   // Function that calls the libc/etc. cleanup function.
   static DefinedFunction *callDtors;
 
   // __wasm_apply_data_relocs
   // Function that applies relocations to data segment post-instantiation.
   static DefinedFunction *applyDataRelocs;
 
+  // __wasm_apply_data_relocs_tail
+  // __wasm_apply_data_relocs can grow drastically exceeding the function
+  // size limit. In that case, we split it up into 2 functions, so that
+  // the extra number of instructions go to applyDataRelocsTail.
+  // This function is called only from applyDataRelocs and only if
+  // the latter exceeds the function size limit. The rest of the time this
+  // function is empty.
+  static DefinedFunction *applyDataRelocsTail;
+
   // __wasm_apply_global_relocs
   // Function that applies relocations to wasm globals post-instantiation.
   // Unlike __wasm_apply_data_relocs this needs to run on every thread.
   static DefinedFunction *applyGlobalRelocs;
 
   // __wasm_apply_global_tls_relocs
   // Like applyGlobalRelocs but for globals that hold TLS addresses.  These
   // must be delayed until __wasm_init_tls.

lld/wasm/Symbols.cpp

std::string toString(wasm::Symbol::Kind kind) {
   llvm_unreachable("invalid symbol kind");
 }
 
 namespace wasm {
 DefinedFunction *WasmSym::callCtors;
 DefinedFunction *WasmSym::callDtors;
 DefinedFunction *WasmSym::initMemory;
 DefinedFunction *WasmSym::applyDataRelocs;
+DefinedFunction *WasmSym::applyDataRelocsTail;
 DefinedFunction *WasmSym::applyGlobalRelocs;
 DefinedFunction *WasmSym::applyGlobalTLSRelocs;
 DefinedFunction *WasmSym::initTLS;
 DefinedFunction *WasmSym::startFunction;
 DefinedData *WasmSym::dsoHandle;
 DefinedData *WasmSym::dataEnd;
 DefinedData *WasmSym::globalBase;
 DefinedData *WasmSym::heapBase;

lld/wasm/Writer.cpp

 //===- Writer.cpp ---------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Writer.h"
 #include "Config.h"
+#include "FunctionSplitter.h"
 #include "InputChunks.h"
 #include "InputElement.h"
 #include "MapFile.h"
 #include "OutputSections.h"
 #include "OutputSegment.h"
 #include "Relocations.h"
+#include "SplitApplyDataRelocsStrategy.h"
 #include "SymbolTable.h"
 #include "SyntheticSections.h"
 #include "WriterUtils.h"
 #include "lld/Common/CommonLinkerContext.h"
 #include "lld/Common/Strings.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
private:
   void createApplyDataRelocationsFunction();
   void createApplyGlobalRelocationsFunction();
   void createApplyGlobalTLSRelocationsFunction();
   void createCallCtorsFunction();
   void createInitTLSFunction();
   void createCommandExportWrappers();
   void createCommandExportWrapper(uint32_t functionIndex, DefinedFunction *f);
 
+  std::string createEmptyFunctionBody() const;
+
   void assignIndexes();
   void populateSymtab();
   void populateProducers();
   void populateTargetFeatures();
   // populateTargetFeatures happens early on so some checks are delayed
   // until imports and exports are finalized.  There are run unstead
   // in checkImportExportTargetFeatures
   void checkImportExportTargetFeatures();
std::string bodyContent;
     writeUleb128(os, 0, "num locals");
     for (const OutputSegment *seg : segments)
       for (const InputChunk *inSeg : seg->inputSegments)
         inSeg->generateRelocationCode(os);
 
     writeU8(os, WASM_OPCODE_END, "END");
   }
 
+  if (bodyContent.size() < config->maxFunctionSize) {
     createFunction(WasmSym::applyDataRelocs, bodyContent);
+    createFunction(WasmSym::applyDataRelocsTail, createEmptyFunctionBody());
+  } else {
+    FunctionSplitter splitter(
+        std::make_unique<SplitApplyDataRelocsStrategy>(bodyContent));
+    auto &&[firstFunctionBody, secondFunctionBody] = splitter.split();
+    createFunction(WasmSym::applyDataRelocs, firstFunctionBody);
+    createFunction(WasmSym::applyDataRelocsTail, secondFunctionBody);
+  }
+}
+
+std::string Writer::createEmptyFunctionBody() const {
+  std::string bodyContent;
+  raw_string_ostream os(bodyContent);
+
+  writeUleb128(os, 0, "num locals");
+  writeU8(os, WASM_OPCODE_END, "END");
+
+  return bodyContent;
 }
 
 // Similar to createApplyDataRelocationsFunction but generates relocation code
 // for WebAssembly globals. Because these globals are not shared between threads
 // these relocation need to run on every thread.
 void Writer::createApplyGlobalRelocationsFunction() {
   // First write the body's contents to a string.
   std::string bodyContent;