[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

40 ↗

(On Diff #483201)

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

8 ↗

(On Diff #483898)

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

42

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/test/wasm/data-segments.ll

 ; Also test in combination with PIC/pie
 ; RUN: wasm-ld --experimental-pic -pie -no-gc-sections --no-entry --shared-memory --max-memory=131072 %t.atomics.bulk-mem.pic.o -o %t.pic.wasm
 ; RUN: obj2yaml %t.pic.wasm | FileCheck %s --check-prefixes PASSIVE-PIC,PASSIVE32-PIC
-; RUN: llvm-objdump --disassemble-symbols=__wasm_call_ctors,__wasm_apply_data_relocs,__wasm_init_memory --no-show-raw-insn --no-leading-addr %t.pic.wasm | FileCheck %s --check-prefixes DIS,PIC-DIS -DPTR=i32
+; RUN: llvm-objdump --disassemble-symbols=__wasm_call_ctors,__wasm_apply_data_relocs,__wasm_apply_data_relocs_tail,__wasm_init_memory --no-show-raw-insn --no-leading-addr %t.pic.wasm | FileCheck %s --check-prefixes DIS,PIC-DIS -DPTR=i32
 
 ; Also test in combination with PIC/pie + wasm64
 ; RUN: wasm-ld -mwasm64 --experimental-pic -pie -no-gc-sections --no-entry --shared-memory --max-memory=131072 %t.atomics.bulk-mem.pic-mem64.o -o %t.pic-mem64.wasm
 ; RUN: obj2yaml %t.pic-mem64.wasm | FileCheck %s --check-prefixes PASSIVE-PIC,PASSIVE64-PIC
-; RUN: llvm-objdump --disassemble-symbols=__wasm_call_ctors,__wasm_apply_data_relocs,__wasm_init_memory --no-show-raw-insn --no-leading-addr %t.pic-mem64.wasm | FileCheck %s --check-prefixes DIS,PIC-DIS -DPTR=i64
+; RUN: llvm-objdump --disassemble-symbols=__wasm_call_ctors,__wasm_apply_data_relocs,__wasm_apply_data_relocs_tail,__wasm_init_memory --no-show-raw-insn --no-leading-addr %t.pic-mem64.wasm | FileCheck %s --check-prefixes DIS,PIC-DIS -DPTR=i64
 
 @a = hidden global [6 x i8] c"hello\00", align 1
 @b = hidden global [8 x i8] c"goodbye\00", align 1

 ; PASSIVE-NEXT:        Name:            __wasm_init_memory
 
 ;      PASSIVE-PIC:  - Type:            START
-; PASSIVE-PIC-NEXT:    StartFunction:   3
+; PASSIVE-PIC-NEXT:    StartFunction:   4
 ; PASSIVE-PIC-NEXT:  - Type:            DATACOUNT
 ; PASSIVE-PIC-NEXT:    Count:           3
 ; PASSIVE-PIC-NEXT:  - Type:            CODE

 ; PASSIVE-PIC-NEXT:        Locals:          []
 ; PASSIVE-PIC-NEXT:        Body:            0B
 ; PASSIVE-PIC-NEXT:      - Index:           3
+; PASSIVE-PIC-NEXT:        Locals:          []
+; PASSIVE-PIC-NEXT:        Body:            0B
+; PASSIVE-PIC-NEXT:      - Index:           4
 ; PASSIVE-PIC-NEXT:        Locals:
 ; PASSIVE32-PIC-NEXT:          - Type:            I32
 ; PASSIVE64-PIC-NEXT:          - Type:            I64

 ; PASSIVE-PIC-NEXT:      - Index:           2
 ; PASSIVE-PIC-NEXT:        Name:            __wasm_apply_data_relocs
 ; PASSIVE-PIC-NEXT:      - Index:           3
+; PASSIVE-PIC-NEXT:        Name:            __wasm_apply_data_relocs_tail
+; PASSIVE-PIC-NEXT:      - Index:           4
 ; PASSIVE-PIC-NEXT:        Name:            __wasm_init_memory
 
 ; no data relocations.

 ; DIS-EMPTY:
 ; DIS-NEXT:        end
 
-; In PIC mode __wasm_apply_data_relocs is export seperatly to __wasm_call_ctors
+; In PIC mode __wasm_apply_data_relocs and __wasm_apply_data_relocs_tail
+; are exported seperately to __wasm_call_ctors
 ; PIC-DIS:     <__wasm_apply_data_relocs>:
 ; PIC-DIS-EMPTY:
 
+; PIC-DIS:     <__wasm_apply_data_relocs_tail>:
+; PIC-DIS-EMPTY:
+
 ; DIS-LABEL:       <__wasm_init_memory>:
 
 ; PIC-DIS:           .local [[PTR]]

lld/test/wasm/pie.ll

 ; RUN: llc -relocation-model=pic -mattr=+mutable-globals -filetype=obj %s -o %t.o
 ; RUN: wasm-ld --no-gc-sections --experimental-pic -pie -o %t.wasm %t.o
 ; RUN: obj2yaml %t.wasm | FileCheck %s
-; RUN: llvm-objdump --disassemble-symbols=__wasm_call_ctors,__wasm_apply_data_relocs --no-show-raw-insn --no-leading-addr %t.wasm | FileCheck %s --check-prefixes DISASSEM
+; RUN: llvm-objdump --disassemble-symbols=__wasm_call_ctors,__wasm_apply_data_relocs,__wasm_apply_data_relocs_tail --no-show-raw-insn --no-leading-addr %t.wasm | FileCheck %s --check-prefixes DISASSEM
 
 target triple = "wasm32-unknown-emscripten"
 

 ; CHECK-NEXT:         GlobalMutable:   false
 
 ; CHECK:        - Type:            START
-; CHECK-NEXT:     StartFunction:   3
+; CHECK-NEXT:     StartFunction:   4
 
 ; CHECK:        - Type:            CUSTOM
 ; CHECK-NEXT:     Name:            name

 ; CHECK-NEXT:       - Index:           2
 ; CHECK-NEXT:         Name:            __wasm_apply_data_relocs
 ; CHECK-NEXT:       - Index:           3
-; CHECK-NEXT:         Name:            __wasm_apply_global_relocs
+; CHECK-NEXT:         Name:            __wasm_apply_data_relocs_tail
 ; CHECK-NEXT:       - Index:           4
-; CHECK-NEXT:         Name:            foo
+; CHECK-NEXT:         Name:            __wasm_apply_global_relocs
 ; CHECK-NEXT:       - Index:           5
-; CHECK-NEXT:         Name:            get_data_address
+; CHECK-NEXT:         Name:            foo
 ; CHECK-NEXT:       - Index:           6
+; CHECK-NEXT:         Name:            get_data_address
+; CHECK-NEXT:       - Index:           7
 ; CHECK-NEXT:         Name:            _start
 ; CHECK-NEXT:     GlobalNames:
 

 ; DISASSEM-LABEL:  <__wasm_apply_data_relocs>:
 ; DISASSEM:        end
 
+; DISASSEM-LABEL:  <__wasm_apply_data_relocs_tail>:
+; DISASSEM:        end
+
 ; Run the same test with extended-const support.  When this is available
 ; we don't need __wasm_apply_global_relocs and instead rely on the add
 ; instruction in the InitExpr.  We also, therefore, do not need these globals

 ; EXTENDED-CONST-NEXT:        Name:            __wasm_call_ctors
 ; EXTENDED-CONST-NEXT:      - Index:           2
 ; EXTENDED-CONST-NEXT:        Name:            __wasm_apply_data_relocs
+; EXTENDED-CONST-NEXT:      - Index:           3
+; EXTENDED-CONST-NEXT:        Name:            __wasm_apply_data_relocs_tail
 
-; Run the same test with threading support.  In this mode
-; we expect __wasm_init_memory and __wasm_apply_data_relocs
+; Run the same test with threading support.  In this mode we expect
+; __wasm_init_memory, __wasm_apply_data_relocs and __wasm_apply_data_relocs
 ; to be generated along with __wasm_start as the start
 ; function.
 

 ; RUN: llvm-objdump --disassemble-symbols=__wasm_start --no-show-raw-insn --no-leading-addr %t.shmem.wasm | FileCheck %s --check-prefix DISASSEM-SHMEM
 
 ; SHMEM:         - Type:            START
-; SHMEM-NEXT:      StartFunction:   6
+; SHMEM-NEXT:      StartFunction:   7
 
 ; DISASSEM-SHMEM-LABEL:  <__wasm_start>:
 ; DISASSEM-SHMEM-EMPTY:
+; DISASSEM-SHMEM-NEXT:   call 6
 ; DISASSEM-SHMEM-NEXT:   call 5
-; DISASSEM-SHMEM-NEXT:   call 4
 ; DISASSEM-SHMEM-NEXT:   end
 
 ; SHMEM:         FunctionNames:

 ; SHMEM-NEXT:      - Index:           3
 ; SHMEM-NEXT:        Name:            __wasm_apply_data_relocs
 ; SHMEM-NEXT:      - Index:           4
-; SHMEM-NEXT:        Name:            __wasm_init_memory
+; SHMEM-NEXT:        Name:            __wasm_apply_data_relocs_tail
 ; SHMEM-NEXT:      - Index:           5
-; SHMEM-NEXT:        Name:            __wasm_apply_global_relocs
+; SHMEM-NEXT:        Name:            __wasm_init_memory
 ; SHMEM-NEXT:      - Index:           6
-; SHMEM-NEXT:        Name:            __wasm_start
+; SHMEM-NEXT:        Name:            __wasm_apply_global_relocs
 ; SHMEM-NEXT:      - Index:           7
-; SHMEM-NEXT:        Name:            foo
+; SHMEM-NEXT:        Name:            __wasm_start
 ; SHMEM-NEXT:      - Index:           8
-; SHMEM-NEXT:        Name:            get_data_address
+; SHMEM-NEXT:        Name:            foo
 ; SHMEM-NEXT:      - Index:           9
+; SHMEM-NEXT:        Name:            get_data_address
+; SHMEM-NEXT:      - Index:           10
 ; SHMEM-NEXT:        Name:            _start

lld/wasm/ApplyDataRelocsSplitter.h

+#ifndef LLD_WASM_APPLY_DATA_RELOCS_SPLITTER_H
+#define LLD_WASM_APPLY_DATA_RELOCS_SPLITTER_H
+
+#include <string>
+#include <utility>
+
+namespace lld::wasm {
+
+class ApplyDataRelocsSplitter {
+public:
+  using SplittedFunctions = std::pair<std::string, std::string>;
+
+  ApplyDataRelocsSplitter(const std::string &functionBody);
+
+  // 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)
+  // If the function doesn't exceed the limit returns the pair of the given
+  // function and the empty one.
+  SplittedFunctions split();
+
+private:
+  size_t findPositionToCutInstructionsFrom(unsigned instructionOpcode) const;
+  void appendApplyDataRelocsTailCall(std::string &functionBody,
+                                     size_t pos) const;
+  void appendCutOffInstructions(std::string &functionBody, size_t pos) const;
+
+  std::string createEmptyFunctionBody() const;
+
+  const std::string &function_to_split_;
+};
+
+} // namespace lld::wasm
+
+#endif

lld/wasm/ApplyDataRelocsSplitter.cpp

+#include "ApplyDataRelocsSplitter.h"
+#include "Config.h"
+#include "Symbols.h"
+#include "WriterUtils.h"
+#include "llvm/BinaryFormat/Wasm.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace lld::wasm {
+
+ApplyDataRelocsSplitter::ApplyDataRelocsSplitter(
+    const std::string &functionBody)
+    : function_to_split_(functionBody) {}
+
+ApplyDataRelocsSplitter::SplittedFunctions ApplyDataRelocsSplitter::split() {
+  // don't split the function if it doesn't exceed the limit
+  if (function_to_split_.size() < config->maxFunctionSize)
+    return {function_to_split_, createEmptyFunctionBody()};
+
+  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);
+  // don't split the function if we didn't find an appropriate instruction
+  if (pos == std::string::npos)
+    return {function_to_split_, createEmptyFunctionBody()};
+
+  appendApplyDataRelocsTailCall(functions.first, pos);
+  appendCutOffInstructions(functions.second, pos);
+
+  return functions;
+}
+
+size_t ApplyDataRelocsSplitter::findPositionToCutInstructionsFrom(
+    unsigned instructionOpcode) const {
+  // Skip some instructions to leave the place for CALL instruction and other
+  constexpr size_t skip_instruction_count = 10;
+  const size_t instruction_opcode_pos = function_to_split_.rfind(
+      instructionOpcode, config->maxFunctionSize - skip_instruction_count);
+
+  // includes the following instructions: current instruction + align + offset
+  constexpr size_t offset = 3;
+
+  return instruction_opcode_pos + offset;
+}
+
+void ApplyDataRelocsSplitter::appendApplyDataRelocsTailCall(
+    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 ApplyDataRelocsSplitter::appendCutOffInstructions(
+    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));
+}
+
+std::string ApplyDataRelocsSplitter::createEmptyFunctionBody() const {
+  std::string bodyContent;
+  llvm::raw_string_ostream os(bodyContent);
+
+  writeUleb128(os, 0, "num locals");
+  writeU8(os, llvm::wasm::WASM_OPCODE_END, "END");
+
+  return bodyContent;
+}
+
+} // namespace lld::wasm

lld/wasm/CMakeLists.txt

 add_public_tablegen_target(WasmOptionsTableGen)
 
 add_lld_library(lldWasm
+  ApplyDataRelocsSplitter.cpp
   Driver.cpp
   InputChunks.cpp
   InputFiles.cpp

lld/wasm/Config.h

   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;

lld/wasm/Driver.cpp

         "__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_HIDDEN,
+        make<SyntheticFunction>(nullSignature,
+                                "__wasm_apply_data_relocs_tail"));
   }
 }
 

lld/wasm/Symbols.h

   // 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.

lld/wasm/Symbols.cpp

 DefinedFunction *WasmSym::callDtors;
 DefinedFunction *WasmSym::initMemory;
 DefinedFunction *WasmSym::applyDataRelocs;
+DefinedFunction *WasmSym::applyDataRelocsTail;
 DefinedFunction *WasmSym::applyGlobalRelocs;
 DefinedFunction *WasmSym::applyGlobalTLSRelocs;
 DefinedFunction *WasmSym::initTLS;

lld/wasm/Writer.cpp

 //===----------------------------------------------------------------------===//
 
 #include "Writer.h"
+#include "ApplyDataRelocsSplitter.h"
 #include "Config.h"
 #include "InputChunks.h"
 #include "InputElement.h"

     writeU8(os, WASM_OPCODE_END, "END");
   }
 
-  createFunction(WasmSym::applyDataRelocs, bodyContent);
+  ApplyDataRelocsSplitter splitter(bodyContent);
+  auto &&[firstFunctionBody, secondFunctionBody] = splitter.split();
+  createFunction(WasmSym::applyDataRelocs, firstFunctionBody);
+  createFunction(WasmSym::applyDataRelocsTail, secondFunctionBody);
 }
 
 // Similar to createApplyDataRelocationsFunction but generates relocation code