[lld][WebAssembly] Add an --initial-heap option (#75594)

It is beneficial to preallocate a certain number of pages in the linear
memory (i. e. use the "minimum" field of WASM memories) so that fewer
"memory.grow"s are needed at startup.

So far, the way to do that has been to pass the "--initial-memory"
option to the linker. It works, but has the very significant downside of
requiring the user to know the size of static data beforehand, as it
must not exceed the number of bytes passed-in as "--initial-memory".

The new "--initial-heap" option avoids this downside by simply appending
the specified number of pages to static data (and stack), regardless of
how large they already are.

Ref: emscripten-core/emscripten#20888.

Author: SingleAccretion

lld/docs/WebAssembly.rst

@@ -123,9 +123,13 @@ WebAssembly-specific options:
    is not possible for undefined data symbols.  Undefined data symbols will
    still be reported as normal (in accordance with ``--unresolved-symbols``).
 
+.. option:: --initial-heap=<value>
+
+  Initial size of the heap. Default: zero.
+
 .. option:: --initial-memory=<value>
 
-  Initial size of the linear memory. Default: static data size.
+  Initial size of the linear memory. Default: the sum of stack, static data and heap sizes.
 
 .. option:: --max-memory=<value>
 

lld/test/wasm/initial-heap.test

@@ -0,0 +1,27 @@
+RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown %p/Inputs/start.s -o %t.o
+
+; The initial heap size will be added to the stack size
+RUN: wasm-ld %t.o -o %t1.wasm --stack-first -z stack-size=65536 --initial-heap=131072
+RUN: obj2yaml %t1.wasm | FileCheck %s --check-prefixes=CHECK,CHECK-2P
+
+; Also test that we can parse and process a large size correctly
+RUN: wasm-ld %t.o -o %t2.wasm --stack-first -z stack-size=65536 --initial-heap=4294901760
+RUN: obj2yaml %t2.wasm | FileCheck %s --check-prefixes=CHECK,CHECK-4G
+
+CHECK:      - Type:            MEMORY
+CHECK-NEXT:   Memories:
+CHECK-2P-NEXT:    Minimum:         0x3
+CHECK-4G-NEXT:    Minimum:         0x10000
+
+; Test various error cases.
+RUN: not wasm-ld %t.o -o %t3.wasm --initial-heap=131073 2>&1 | FileCheck %s --check-prefix NOT-PAGE-MULTIPLE
+RUN: not wasm-ld %t.o -o %t4.wasm --stack-first -z stack-size=65536 --initial-heap=4295032832 2>&1 | FileCheck %s --check-prefix TOO-LARGE-BY-ITSELF
+RUN: not wasm-ld %t.o -o %t5.wasm --stack-first -z stack-size=131072 --initial-heap=4294901760 2>&1 | FileCheck %s --check-prefix TOO-LARGE-WITH-STACK
+RUN: not wasm-ld %t.o -o %t6.wasm --stack-first -z stack-size=65536 --initial-heap=131072 --initial-memory=131072 2>&1 | FileCheck %s --check-prefix INITIAL-MEMORY-TOO-SMALL
+RUN: not wasm-ld %t.o -o %t7.wasm --stack-first -z stack-size=65536 --initial-heap=131072 --max-memory=131072 2>&1 | FileCheck %s --check-prefix MAX-MEMORY-TOO-SMALL
+
+NOT-PAGE-MULTIPLE: initial heap must be 65536-byte aligned
+TOO-LARGE-BY-ITSELF: initial heap too large, cannot be greater than 4294901760
+TOO-LARGE-WITH-STACK: initial heap too large, cannot be greater than 4294836224
+INITIAL-MEMORY-TOO-SMALL: initial memory too small, 196608 bytes needed
+MAX-MEMORY-TOO-SMALL: maximum memory too small, 196608 bytes needed

lld/wasm/Config.h

@@ -68,6 +68,7 @@ struct Configuration {
   bool isStatic = false;
   bool trace;
   uint64_t globalBase;
+  uint64_t initialHeap;
   uint64_t initialMemory;
   uint64_t maxMemory;
   uint64_t zStackSize;

lld/wasm/Driver.cpp

@@ -502,9 +502,10 @@ static void readConfigs(opt::InputArgList &args) {
   errorHandler().verbose = args.hasArg(OPT_verbose);
   LLVM_DEBUG(errorHandler().verbose = true);
 
-  config->initialMemory = args::getInteger(args, OPT_initial_memory, 0);
-  config->globalBase = args::getInteger(args, OPT_global_base, 0);
   config->tableBase = args::getInteger(args, OPT_table_base, 0);
+  config->globalBase = args::getInteger(args, OPT_global_base, 0);
+  config->initialHeap = args::getInteger(args, OPT_initial_heap, 0);
+  config->initialMemory = args::getInteger(args, OPT_initial_memory, 0);
   config->maxMemory = args::getInteger(args, OPT_max_memory, 0);
   config->zStackSize =
       args::getZOptionValue(args, OPT_z, "stack-size", WasmPageSize);

lld/wasm/Options.td

@@ -215,6 +215,9 @@ defm soname: Eq<"soname", "Set the module name in the generated name section">;
 def import_table: FF<"import-table">,
   HelpText<"Import function table from the environment">;
 
+def initial_heap: JJ<"initial-heap=">,
+  HelpText<"Initial size of the heap">;
+
 def initial_memory: JJ<"initial-memory=">,
   HelpText<"Initial size of the linear memory">;
 

lld/wasm/Writer.cpp

@@ -439,6 +439,16 @@ void Writer::layoutMemory() {
     maxMemorySetting = 1ULL << 34;
   }
 
+  if (config->initialHeap != 0) {
+    if (config->initialHeap != alignTo(config->initialHeap, WasmPageSize))
+      error("initial heap must be " + Twine(WasmPageSize) + "-byte aligned");
+    uint64_t maxInitialHeap = maxMemorySetting - memoryPtr;
+    if (config->initialHeap > maxInitialHeap)
+      error("initial heap too large, cannot be greater than " +
+            Twine(maxInitialHeap));
+    memoryPtr += config->initialHeap;
+  }
+
   if (config->initialMemory != 0) {
     if (config->initialMemory != alignTo(config->initialMemory, WasmPageSize))
       error("initial memory must be " + Twine(WasmPageSize) + "-byte aligned");