[lld][WebAssembly] Support for the custom-page-sizes WebAssembly proposal (#128942)

This commit adds support for WebAssembly's custom-page-sizes proposal to
`wasm-ld`. An overview of the proposal can be found
[here](https://github.com/WebAssembly/custom-page-sizes/blob/main/proposals/custom-page-sizes/Overview.md).
In a sentence, it allows customizing a Wasm memory's page size, enabling
Wasm to target environments with less than 64KiB of memory (the default
Wasm page size) available for Wasm memories.

This commit contains the following:

* Adds a `--page-size=N` CLI flag to `wasm-ld` for configuring the
linked Wasm binary's linear memory's page size.

* When the page size is configured to a non-default value, then the
final Wasm binary will use the encodings defined in the
custom-page-sizes proposal to declare the linear memory's page size.

* Defines a `__wasm_first_page_end` symbol, whose address points to the
first page in the Wasm linear memory, a.k.a. is the Wasm memory's page
size. This allows writing code that is compatible with any page size,
and doesn't require re-compiling its object code. At the same time,
because it just lowers to a constant rather than a memory access or
something, it enables link-time optimization.

* Adds tests for these new features.

r? @sbc100 

cc @sunfishcode

lld/test/wasm/initial-heap.test

@@ -20,7 +20,7 @@ RUN: not wasm-ld %t.o -o %t5.wasm --stack-first -z stack-size=131072 --initial-h
 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
+NOT-PAGE-MULTIPLE: initial heap must be aligned to the page size (65536 bytes)
 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

lld/test/wasm/mutable-global-exports.s

@@ -101,6 +101,9 @@ _start:
 # CHECK-ALL-NEXT:      - Name:            __table_base
 # CHECK-ALL-NEXT:        Kind:            GLOBAL
 # CHECK-ALL-NEXT:        Index:           10
+# CHECK-ALL-NEXT:      - Name:            __wasm_first_page_end
+# CHECK-ALL-NEXT:        Kind:            GLOBAL
+# CHECK-ALL-NEXT:        Index:           11
 # CHECK-ALL-NEXT:  - Type:            CODE
 
 # CHECK-ALL:         Name:            target_features
@@ -110,4 +113,3 @@ _start:
 # CHECK-ALL-NEXT:      - Prefix:          USED
 # CHECK-ALL-NEXT:        Name:            mutable-globals
 # CHECK-ALL-NEXT: ...
-

lld/test/wasm/page-size.s

@@ -0,0 +1,43 @@
+# RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown %s -o %t.o
+
+        .globl  _start
+_start:
+        .functype _start () -> (i32)
+        i32.const __wasm_first_page_end
+        end_function
+
+# Add a symbol to smoke test that `__wasm_first_page_end` is absolute and not
+# relative to other data.
+        .section .data.foo,"",@
+foo:
+        .int32  0x11111111
+        .size   foo, 4
+
+# RUN: wasm-ld -no-gc-sections -o %t.custom.wasm %t.o --page-size=1
+# RUN: obj2yaml %t.custom.wasm | FileCheck %s --check-prefix=CHECK-CUSTOM
+
+# CHECK-CUSTOM:      - Type:            MEMORY
+# CHECK-CUSTOM-NEXT:   Memories:
+# CHECK-CUSTOM-NEXT:   - Flags:           [ HAS_PAGE_SIZE ]
+# CHECK-CUSTOM-NEXT:     Minimum:         0x10410
+# CHECK-CUSTOM-NEXT:     PageSize:        0x1
+
+# RUN: llvm-objdump --disassemble-symbols=_start %t.custom.wasm | FileCheck %s --check-prefix=CHECK-CUSTOM-DIS
+
+# CHECK-CUSTOM-DIS:      <_start>:
+# CHECK-CUSTOM-DIS:          i32.const 1
+# CHECK-CUSTOM-DIS-NEXT:     end
+
+# RUN: wasm-ld -no-gc-sections -o %t.default.wasm %t.o
+# RUN: obj2yaml %t.default.wasm | FileCheck %s --check-prefix=CHECK-DEFAULT
+
+# CHECK-DEFAULT:      - Type:            MEMORY
+# CHECK-DEFAULT-NEXT:   Memories:
+# CHECK-DEFAULT-NEXT:     Minimum:         0x2
+# CHECK-DEFAULT-NEXT: - Type:            GLOBAL
+
+# RUN: llvm-objdump --disassemble-symbols=_start %t.default.wasm | FileCheck %s --check-prefix=CHECK-DEFAULT-DIS
+
+# CHECK-DEFAULT-DIS:      <_start>:
+# CHECK-DEFAULT-DIS:          i32.const 65536
+# CHECK-DEFAULT-DIS-NEXT:     end

lld/test/wasm/shared-memory.yaml

@@ -56,7 +56,7 @@ Sections:
         Flags:           [  ]
 ...
 
-# SHARED-UNALIGNED: maximum memory must be 65536-byte aligned{{$}}
+# SHARED-UNALIGNED: maximum memory must be aligned to the page size (65536 bytes)
 
 # SHARED-NO-ATOMICS: 'atomics' feature must be used in order to use shared memory
 

lld/wasm/Config.h

@@ -94,6 +94,7 @@ struct Config {
   // runtime).
   uint64_t tableBase;
   uint64_t zStackSize;
+  uint64_t pageSize;
   unsigned ltoPartitions;
   unsigned ltoo;
   llvm::CodeGenOptLevel ltoCgo;

lld/wasm/Driver.cpp

@@ -643,7 +643,10 @@ static void readConfigs(opt::InputArgList &args) {
   ctx.arg.maxMemory = args::getInteger(args, OPT_max_memory, 0);
   ctx.arg.noGrowableMemory = args.hasArg(OPT_no_growable_memory);
   ctx.arg.zStackSize =
-      args::getZOptionValue(args, OPT_z, "stack-size", WasmPageSize);
+      args::getZOptionValue(args, OPT_z, "stack-size", WasmDefaultPageSize);
+  ctx.arg.pageSize = args::getInteger(args, OPT_page_size, WasmDefaultPageSize);
+  if (ctx.arg.pageSize != 1 && ctx.arg.pageSize != WasmDefaultPageSize)
+    error("--page_size=N must be either 1 or 65536");
 
   // -Bdynamic by default if -pie or -shared is specified.
   if (ctx.arg.pie || ctx.arg.shared)
@@ -1000,6 +1003,11 @@ static void createOptionalSymbols() {
     WasmSym::definedTableBase = symtab->addOptionalDataSymbol("__table_base");
   }
 
+  WasmSym::firstPageEnd =
+      symtab->addOptionalDataSymbol("__wasm_first_page_end");
+  if (WasmSym::firstPageEnd)
+    WasmSym::firstPageEnd->setVA(ctx.arg.pageSize);
+
   // For non-shared memory programs we still need to define __tls_base since we
   // allow object files built with TLS to be linked into single threaded
   // programs, and such object files can contain references to this symbol.

lld/wasm/Options.td

@@ -232,6 +232,9 @@ def import_table: FF<"import-table">,
 def initial_heap: JJ<"initial-heap=">,
   HelpText<"Initial size of the heap">;
 
+def page_size: JJ<"page-size=">,
+  HelpText<"The Wasm page size (Defaults to 65536)">;
+
 def initial_memory: JJ<"initial-memory=">,
   HelpText<"Initial size of the linear memory">;
 

lld/wasm/SymbolTable.cpp

@@ -792,7 +792,7 @@ Symbol *SymbolTable::addUndefinedTag(StringRef name,
 }
 
 TableSymbol *SymbolTable::createUndefinedIndirectFunctionTable(StringRef name) {
-  WasmLimits limits{0, 0, 0}; // Set by the writer.
+  WasmLimits limits{0, 0, 0, 0}; // Set by the writer.
   WasmTableType *type = make<WasmTableType>();
   type->ElemType = ValType::FUNCREF;
   type->Limits = limits;
@@ -807,7 +807,7 @@ TableSymbol *SymbolTable::createUndefinedIndirectFunctionTable(StringRef name) {
 
 TableSymbol *SymbolTable::createDefinedIndirectFunctionTable(StringRef name) {
   const uint32_t invalidIndex = -1;
-  WasmLimits limits{0, 0, 0}; // Set by the writer.
+  WasmLimits limits{0, 0, 0, 0}; // Set by the writer.
   WasmTableType type{ValType::FUNCREF, limits};
   WasmTable desc{invalidIndex, type, name};
   InputTable *table = make<InputTable>(desc, nullptr);

lld/wasm/Symbols.cpp

@@ -84,6 +84,7 @@ DefinedFunction *WasmSym::applyGlobalRelocs;
 DefinedFunction *WasmSym::applyTLSRelocs;
 DefinedFunction *WasmSym::applyGlobalTLSRelocs;
 DefinedFunction *WasmSym::initTLS;
+DefinedData *WasmSym::firstPageEnd;
 DefinedFunction *WasmSym::startFunction;
 DefinedData *WasmSym::dsoHandle;
 DefinedData *WasmSym::dataEnd;

lld/wasm/Symbols.h

@@ -577,6 +577,10 @@ struct WasmSym {
   static DefinedData *heapBase;
   static DefinedData *heapEnd;
 
+  // __wasm_first_page_end
+  // A symbol whose address is the end of the first page in memory (if any).
+  static DefinedData *firstPageEnd;
+
   // __wasm_init_memory_flag
   // Symbol whose contents are nonzero iff memory has already been initialized.
   static DefinedData *initMemoryFlag;

lld/wasm/SyntheticSections.cpp

@@ -369,10 +369,14 @@ void MemorySection::writeBody() {
     flags |= WASM_LIMITS_FLAG_IS_SHARED;
   if (ctx.arg.is64.value_or(false))
     flags |= WASM_LIMITS_FLAG_IS_64;
+  if (ctx.arg.pageSize != WasmDefaultPageSize)
+    flags |= WASM_LIMITS_FLAG_HAS_PAGE_SIZE;
   writeUleb128(os, flags, "memory limits flags");
   writeUleb128(os, numMemoryPages, "initial pages");
   if (hasMax)
     writeUleb128(os, maxMemoryPages, "max pages");
+  if (ctx.arg.pageSize != WasmDefaultPageSize)
+    writeUleb128(os, llvm::Log2_64(ctx.arg.pageSize), "page size");
 }
 
 void TagSection::writeBody() {

lld/wasm/Writer.cpp

@@ -320,6 +320,12 @@ static void setGlobalPtr(DefinedGlobal *g, uint64_t memoryPtr) {
   g->global->setPointerValue(memoryPtr);
 }
 
+static void checkPageAligned(StringRef name, uint64_t value) {
+  if (value != alignTo(value, ctx.arg.pageSize))
+    error(name + " must be aligned to the page size (" +
+          Twine(ctx.arg.pageSize) + " bytes)");
+}
+
 // Fix the memory layout of the output binary.  This assigns memory offsets
 // to each of the input data sections as well as the explicit stack region.
 // The default memory layout is as follows, from low to high.
@@ -445,8 +451,7 @@ void Writer::layoutMemory() {
   }
 
   if (ctx.arg.initialHeap != 0) {
-    if (ctx.arg.initialHeap != alignTo(ctx.arg.initialHeap, WasmPageSize))
-      error("initial heap must be " + Twine(WasmPageSize) + "-byte aligned");
+    checkPageAligned("initial heap", ctx.arg.initialHeap);
     uint64_t maxInitialHeap = maxMemorySetting - memoryPtr;
     if (ctx.arg.initialHeap > maxInitialHeap)
       error("initial heap too large, cannot be greater than " +
@@ -455,8 +460,7 @@ void Writer::layoutMemory() {
   }
 
   if (ctx.arg.initialMemory != 0) {
-    if (ctx.arg.initialMemory != alignTo(ctx.arg.initialMemory, WasmPageSize))
-      error("initial memory must be " + Twine(WasmPageSize) + "-byte aligned");
+    checkPageAligned("initial memory", ctx.arg.initialMemory);
     if (memoryPtr > ctx.arg.initialMemory)
       error("initial memory too small, " + Twine(memoryPtr) + " bytes needed");
     if (ctx.arg.initialMemory > maxMemorySetting)
@@ -465,9 +469,9 @@ void Writer::layoutMemory() {
     memoryPtr = ctx.arg.initialMemory;
   }
 
-  memoryPtr = alignTo(memoryPtr, WasmPageSize);
+  memoryPtr = alignTo(memoryPtr, ctx.arg.pageSize);
 
-  out.memorySec->numMemoryPages = memoryPtr / WasmPageSize;
+  out.memorySec->numMemoryPages = memoryPtr / ctx.arg.pageSize;
   log("mem: total pages = " + Twine(out.memorySec->numMemoryPages));
 
   if (WasmSym::heapEnd) {
@@ -480,8 +484,7 @@ void Writer::layoutMemory() {
 
   uint64_t maxMemory = 0;
   if (ctx.arg.maxMemory != 0) {
-    if (ctx.arg.maxMemory != alignTo(ctx.arg.maxMemory, WasmPageSize))
-      error("maximum memory must be " + Twine(WasmPageSize) + "-byte aligned");
+    checkPageAligned("maximum memory", ctx.arg.maxMemory);
     if (memoryPtr > ctx.arg.maxMemory)
       error("maximum memory too small, " + Twine(memoryPtr) + " bytes needed");
     if (ctx.arg.maxMemory > maxMemorySetting)
@@ -503,7 +506,7 @@ void Writer::layoutMemory() {
   }
 
   if (maxMemory != 0) {
-    out.memorySec->maxMemoryPages = maxMemory / WasmPageSize;
+    out.memorySec->maxMemoryPages = maxMemory / ctx.arg.pageSize;
     log("mem: max pages   = " + Twine(out.memorySec->maxMemoryPages));
   }
 }
@@ -932,7 +935,7 @@ static void finalizeIndirectFunctionTable() {
   }
 
   uint32_t tableSize = ctx.arg.tableBase + out.elemSec->numEntries();
-  WasmLimits limits = {0, tableSize, 0};
+  WasmLimits limits = {0, tableSize, 0, 0};
   if (WasmSym::indirectFunctionTable->isDefined() && !ctx.arg.growableTable) {
     limits.Flags |= WASM_LIMITS_FLAG_HAS_MAX;
     limits.Maximum = limits.Minimum;

lld/wasm/WriterUtils.cpp

@@ -69,6 +69,8 @@ static std::string toString(const llvm::wasm::WasmLimits &limits) {
   ret += "; min=" + std::to_string(limits.Minimum);
   if (limits.Flags & WASM_LIMITS_FLAG_HAS_MAX)
     ret += "; max=" + std::to_string(limits.Maximum);
+  if (limits.Flags & WASM_LIMITS_FLAG_HAS_PAGE_SIZE)
+    ret += "; pagesize=" + std::to_string(limits.PageSize);
   return ret;
 }
 
@@ -200,6 +202,8 @@ void writeLimits(raw_ostream &os, const WasmLimits &limits) {
   writeUleb128(os, limits.Minimum, "limits min");
   if (limits.Flags & WASM_LIMITS_FLAG_HAS_MAX)
     writeUleb128(os, limits.Maximum, "limits max");
+  if (limits.Flags & WASM_LIMITS_FLAG_HAS_PAGE_SIZE)
+    writeUleb128(os, llvm::Log2_64(limits.PageSize), "page size");
 }
 
 void writeGlobalType(raw_ostream &os, const WasmGlobalType &type) {

llvm/include/llvm/BinaryFormat/Wasm.h

@@ -28,8 +28,9 @@ const char WasmMagic[] = {'\0', 'a', 's', 'm'};
 const uint32_t WasmVersion = 0x1;
 // Wasm linking metadata version
 const uint32_t WasmMetadataVersion = 0x2;
-// Wasm uses a 64k page size
-const uint32_t WasmPageSize = 65536;
+// Wasm uses a 64k page size by default (but the custom-page-sizes proposal
+// allows changing it)
+const uint32_t WasmDefaultPageSize = 65536;
 
 enum : unsigned {
   WASM_SEC_CUSTOM = 0,     // Custom / User-defined section
@@ -157,6 +158,7 @@ enum : unsigned {
   WASM_LIMITS_FLAG_HAS_MAX = 0x1,
   WASM_LIMITS_FLAG_IS_SHARED = 0x2,
   WASM_LIMITS_FLAG_IS_64 = 0x4,
+  WASM_LIMITS_FLAG_HAS_PAGE_SIZE = 0x8,
 };
 
 enum : unsigned {
@@ -319,6 +321,7 @@ struct WasmLimits {
   uint8_t Flags;
   uint64_t Minimum;
   uint64_t Maximum;
+  uint32_t PageSize;
 };
 
 struct WasmTableType {
@@ -534,7 +537,10 @@ inline bool operator!=(const WasmGlobalType &LHS, const WasmGlobalType &RHS) {
 inline bool operator==(const WasmLimits &LHS, const WasmLimits &RHS) {
   return LHS.Flags == RHS.Flags && LHS.Minimum == RHS.Minimum &&
          (LHS.Flags & WASM_LIMITS_FLAG_HAS_MAX ? LHS.Maximum == RHS.Maximum
-                                               : true);
+                                               : true) &&
+         (LHS.Flags & WASM_LIMITS_FLAG_HAS_PAGE_SIZE
+              ? LHS.PageSize == RHS.PageSize
+              : true);
 }
 
 inline bool operator==(const WasmTableType &LHS, const WasmTableType &RHS) {

llvm/include/llvm/BinaryFormat/WasmTraits.h

@@ -64,10 +64,10 @@ template <> struct DenseMapInfo<wasm::WasmGlobalType, void> {
 // Traits for using WasmLimits in a DenseMap
 template <> struct DenseMapInfo<wasm::WasmLimits, void> {
   static wasm::WasmLimits getEmptyKey() {
-    return wasm::WasmLimits{0xff, 0xff, 0xff};
+    return wasm::WasmLimits{0xff, 0xff, 0xff, 0xff};
   }
   static wasm::WasmLimits getTombstoneKey() {
-    return wasm::WasmLimits{0xee, 0xee, 0xee};
+    return wasm::WasmLimits{0xee, 0xee, 0xee, 0xee};
   }
   static unsigned getHashValue(const wasm::WasmLimits &Limits) {
     unsigned Hash = hash_value(Limits.Flags);

llvm/include/llvm/MC/MCSymbolWasm.h

@@ -147,7 +147,7 @@ public:
                     uint8_t flags = wasm::WASM_LIMITS_FLAG_NONE) {
     // Declare a table with element type VT and no limits (min size 0, no max
     // size).
-    wasm::WasmLimits Limits = {flags, 0, 0};
+    wasm::WasmLimits Limits = {flags, 0, 0, 0};
     setTableType({VT, Limits});
   }
 };

llvm/include/llvm/ObjectYAML/WasmYAML.h

@@ -48,6 +48,7 @@ struct Limits {
   LimitFlags Flags;
   yaml::Hex32 Minimum;
   yaml::Hex32 Maximum;
+  yaml::Hex32 PageSize;
 };
 
 struct Table {

llvm/lib/MC/WasmObjectWriter.cpp

@@ -845,7 +845,8 @@ void WasmObjectWriter::writeImportSection(ArrayRef<wasm::WasmImport> Imports,
   if (Imports.empty())
     return;
 
-  uint64_t NumPages = (DataSize + wasm::WasmPageSize - 1) / wasm::WasmPageSize;
+  uint64_t NumPages =
+      (DataSize + wasm::WasmDefaultPageSize - 1) / wasm::WasmDefaultPageSize;
 
   SectionBookkeeping Section;
   startSection(Section, wasm::WASM_SEC_IMPORT);

llvm/lib/Object/WasmObjectFile.cpp

@@ -291,6 +291,12 @@ static wasm::WasmLimits readLimits(WasmObjectFile::ReadContext &Ctx) {
   Result.Minimum = readVaruint64(Ctx);
   if (Result.Flags & wasm::WASM_LIMITS_FLAG_HAS_MAX)
     Result.Maximum = readVaruint64(Ctx);
+  if (Result.Flags & wasm::WASM_LIMITS_FLAG_HAS_PAGE_SIZE) {
+    uint32_t PageSizeLog2 = readVaruint32(Ctx);
+    if (PageSizeLog2 >= 32)
+      report_fatal_error("log2(wasm page size) too large");
+    Result.PageSize = 1 << PageSizeLog2;
+  }
   return Result;
 }
 

llvm/lib/ObjectYAML/WasmYAML.cpp

@@ -373,6 +373,8 @@ void MappingTraits<WasmYAML::Limits>::mapping(IO &IO,
   IO.mapRequired("Minimum", Limits.Minimum);
   if (!IO.outputting() || Limits.Flags & wasm::WASM_LIMITS_FLAG_HAS_MAX)
     IO.mapOptional("Maximum", Limits.Maximum);
+  if (!IO.outputting() || Limits.Flags & wasm::WASM_LIMITS_FLAG_HAS_PAGE_SIZE)
+    IO.mapOptional("PageSize", Limits.PageSize);
 }
 
 void MappingTraits<WasmYAML::ElemSegment>::mapping(
@@ -553,6 +555,7 @@ void ScalarBitSetTraits<WasmYAML::LimitFlags>::bitset(
   BCase(HAS_MAX);
   BCase(IS_SHARED);
   BCase(IS_64);
+  BCase(HAS_PAGE_SIZE);
 #undef BCase
 }
 

llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmParser.cpp

@@ -204,7 +204,7 @@ struct WebAssemblyOperand : public MCParsedAsmOperand {
 
 // Perhaps this should go somewhere common.
 static wasm::WasmLimits defaultLimits() {
-  return {wasm::WASM_LIMITS_FLAG_NONE, 0, 0};
+  return {wasm::WASM_LIMITS_FLAG_NONE, 0, 0, 0};
 }
 
 static MCSymbolWasm *getOrCreateFunctionTableSymbol(MCContext &Ctx,

llvm/lib/Target/WebAssembly/WebAssemblyUtilities.cpp

@@ -135,7 +135,7 @@ MCSymbolWasm *WebAssembly::getOrCreateFuncrefCallTableSymbol(
     // modules define the table.
     Sym->setWeak(true);
 
-    wasm::WasmLimits Limits = {0, 1, 1};
+    wasm::WasmLimits Limits = {0, 1, 1, 0};
     wasm::WasmTableType TableType = {wasm::ValType::FUNCREF, Limits};
     Sym->setType(wasm::WASM_SYMBOL_TYPE_TABLE);
     Sym->setTableType(TableType);

llvm/tools/obj2yaml/wasm2yaml.cpp

@@ -37,6 +37,7 @@ static WasmYAML::Limits makeLimits(const wasm::WasmLimits &Limits) {
   L.Flags = Limits.Flags;
   L.Minimum = Limits.Minimum;
   L.Maximum = Limits.Maximum;
+  L.PageSize = Limits.PageSize;
   return L;
 }