release/20.x: [lld][WebAssembly] Support for the custom-page-sizes WebAssembly proposal (#128942)

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

@@ -642,7 +642,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)
@@ -999,6 +1002,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

@@ -230,6 +230,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

@@ -360,10 +360,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 {
@@ -317,6 +319,7 @@ struct WasmLimits {
   uint8_t Flags;
   uint64_t Minimum;
   uint64_t Maximum;
+  uint32_t PageSize;
 };
 
 struct WasmTableType {
@@ -532,7 +535,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 @@ class MCSymbolWasm : public MCSymbol {
                     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

@@ -372,6 +372,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(
@@ -552,6 +554,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;
 }