Custom Page Sizes proposal support (#8307)

Proposal:

https://github.com/WebAssembly/custom-page-sizes/blob/main/proposals/custom-page-sizes/Overview.md

Fixes: #6873

Author: HerrCai0907

scripts/test/shared.py

@@ -405,7 +405,6 @@ def get_tests(test_dir, extensions=[], recursive=False):
     'threads/atomic.wast',
 ]
 SPEC_TESTSUITE_PROPOSALS_TO_SKIP = [
-    'custom-page-sizes',
     'wide-arithmetic',
 ]
 

src/binaryen-c.cpp

@@ -502,6 +502,9 @@ BinaryenFeatures BinaryenFeatureCallIndirectOverlong(void) {
 BinaryenFeatures BinaryenFeatureRelaxedAtomics(void) {
   return static_cast<BinaryenFeatures>(FeatureSet::RelaxedAtomics);
 }
+BinaryenFeatures BinaryenFeatureCustomPageSizes(void) {
+  return static_cast<BinaryenFeatures>(FeatureSet::CustomPageSizes);
+}
 BinaryenFeatures BinaryenFeatureAll(void) {
   return static_cast<BinaryenFeatures>(FeatureSet::All);
 }

src/binaryen-c.h

@@ -244,6 +244,7 @@ BINARYEN_API BinaryenFeatures BinaryenFeatureFP16(void);
 BINARYEN_API BinaryenFeatures BinaryenFeatureBulkMemoryOpt(void);
 BINARYEN_API BinaryenFeatures BinaryenFeatureCallIndirectOverlong(void);
 BINARYEN_API BinaryenFeatures BinaryenFeatureRelaxedAtomics(void);
+BINARYEN_API BinaryenFeatures BinaryenFeatureCustomPageSizes(void);
 BINARYEN_API BinaryenFeatures BinaryenFeatureAll(void);
 
 // Modules

src/ir/memory-utils.cpp

@@ -22,7 +22,8 @@ namespace wasm::MemoryUtils {
 
 bool isSubType(const Memory& a, const Memory& b) {
   return a.shared == b.shared && a.addressType == b.addressType &&
-         a.initial >= b.initial && a.max <= b.max;
+         a.initial >= b.initial && a.max <= b.max &&
+         a.pageSizeLog2 == b.pageSizeLog2;
 }
 
 bool flatten(Module& wasm) {

src/ir/module-utils.cpp

@@ -173,6 +173,7 @@ Memory* copyMemory(const Memory* memory, Module& out) {
   ret->hasExplicitName = memory->hasExplicitName;
   ret->initial = memory->initial;
   ret->max = memory->max;
+  ret->pageSizeLog2 = memory->pageSizeLog2;
   ret->shared = memory->shared;
   ret->addressType = memory->addressType;
   ret->module = memory->module;

src/js/binaryen.js-post.js

@@ -193,6 +193,7 @@ function initializeConstants() {
     'BulkMemoryOpt',
     'CallIndirectOverlong',
     'RelaxedAtomics',
+    'CustomPageSizes',
     'All'
   ].forEach(name => {
     Module['Features'][name] = Module['_BinaryenFeature' + name]();

src/parser/context-decls.cpp

@@ -144,6 +144,7 @@ Result<Memory*> ParseDeclsCtx::addMemoryDecl(Index pos,
   m->initial = type.limits.initial;
   m->max = type.limits.max ? *type.limits.max : Memory::kUnlimitedSize;
   m->shared = type.shared;
+  m->pageSizeLog2 = type.pageSizeLog2;
   if (name) {
     // TODO: if the existing memory is not explicitly named, fix its name
     // and continue.

src/parser/contexts.h

@@ -49,6 +49,7 @@ struct Limits {
 struct MemType {
   Type addressType;
   Limits limits;
+  uint8_t pageSizeLog2;
   bool shared;
 };
 
@@ -174,8 +175,6 @@ struct NullTypeParserCtx {
   DataStringT makeDataString() { return Ok{}; }
   void appendDataString(DataStringT&, std::string_view) {}
 
-  MemTypeT makeMemType(Type, LimitsT, bool) { return Ok{}; }
-
   BlockTypeT getBlockTypeFromResult(size_t results) { return Ok{}; }
 
   Result<> getBlockTypeFromTypeUse(Index, TypeUseT) { return Ok{}; }
@@ -353,9 +352,8 @@ template<typename Ctx> struct TypeParserCtx {
   void appendDataString(DataStringT&, std::string_view) {}
 
   Result<LimitsT> makeLimits(uint64_t, std::optional<uint64_t>) { return Ok{}; }
-  LimitsT getLimitsFromData(DataStringT) { return Ok{}; }
 
-  MemTypeT makeMemType(Type, LimitsT, bool) { return Ok{}; }
+  MemTypeT makeMemType(Type, LimitsT, bool, uint8_t) { return Ok{}; }
 
   HeapType getBlockTypeFromResult(const std::vector<Type> results) {
     assert(results.size() == 1);
@@ -1086,13 +1084,18 @@ struct ParseDeclsCtx : NullTypeParserCtx, NullInstrParserCtx {
     data.insert(data.end(), str.begin(), str.end());
   }
 
-  Limits getLimitsFromData(const std::vector<char>& data) {
-    uint64_t size = (data.size() + Memory::kPageSize - 1) / Memory::kPageSize;
+  Limits getLimitsFromData(const std::vector<char>& data,
+                           uint8_t pageSizeLog2) {
+    uint64_t size =
+      (data.size() + (1 << pageSizeLog2) - 1) / (1 << pageSizeLog2);
     return {size, size};
   }
 
-  MemType makeMemType(Type addressType, Limits limits, bool shared) {
-    return {addressType, limits, shared};
+  MemType makeMemType(Type addressType,
+                      Limits limits,
+                      bool shared,
+                      uint8_t pageSizeLog2) {
+    return {addressType, limits, pageSizeLog2, shared};
   }
 
   Result<TypeUseT>
@@ -1474,8 +1477,8 @@ struct ParseModuleTypesCtx : TypeParserCtx<ParseModuleTypesCtx>,
 
   Type makeTableType(Type addressType, LimitsT, Type type) { return type; }
 
-  LimitsT getLimitsFromData(DataStringT) { return Ok{}; }
-  MemTypeT makeMemType(Type, LimitsT, bool) { return Ok{}; }
+  LimitsT getLimitsFromData(DataStringT, uint8_t) { return Ok{}; }
+  MemTypeT makeMemType(Type, LimitsT, bool, uint8_t) { return Ok{}; }
 
   Result<> addFunc(Name name,
                    const std::vector<Name>&,

src/parser/parsers.h

@@ -20,6 +20,7 @@
 #include "common.h"
 #include "contexts.h"
 #include "lexer.h"
+#include "wasm.h"
 #include "wat-parser-internal.h"
 
 namespace wasm::WATParser {
@@ -829,7 +830,34 @@ template<typename Ctx> Result<typename Ctx::LimitsT> limits64(Ctx& ctx) {
   return ctx.makeLimits(uint64_t(*n), m);
 }
 
-// memtype ::= (limits32 | 'i32' limits32 | 'i64' limit64) shared?
+// mempagesize? ::= ('(' 'pagesize' u64 ')') ?
+template<typename Ctx> MaybeResult<uint8_t> mempagesize(Ctx& ctx) {
+  if (!ctx.in.takeSExprStart("pagesize"sv)) {
+    return {}; // No pagesize specified
+  }
+  auto pageSize = ctx.in.takeU64();
+  if (!pageSize) {
+    return ctx.in.err("expected page size");
+  }
+
+  if (!Bits::isPowerOf2(*pageSize)) {
+    return ctx.in.err("page size must be a power of two");
+  }
+
+  if (!ctx.in.takeRParen()) {
+    return ctx.in.err("expected end of mempagesize");
+  }
+
+  uint8_t pageSizeLog2 = (uint8_t)Bits::ceilLog2(*pageSize);
+
+  if (pageSizeLog2 != 0 && pageSizeLog2 != Memory::kDefaultPageSizeLog2) {
+    return ctx.in.err("memory page size can only be 1 or 64 KiB");
+  }
+
+  return pageSizeLog2;
+}
+
+// memtype ::= (limits32 | 'i32' limits32 | 'i64' limit64) shared? mempagesize?
 //  note: the index type 'i32' or 'i64' is already parsed to simplify parsing of
 //  memory abbreviations.
 template<typename Ctx> Result<typename Ctx::MemTypeT> memtype(Ctx& ctx) {
@@ -851,7 +879,11 @@ Result<typename Ctx::MemTypeT> memtypeContinued(Ctx& ctx, Type addressType) {
   if (ctx.in.takeKeyword("shared"sv)) {
     shared = true;
   }
-  return ctx.makeMemType(addressType, *limits, shared);
+  MaybeResult<uint8_t> mempageSize = mempagesize(ctx);
+  CHECK_ERR(mempageSize);
+  const uint8_t pageSizeLog2 =
+    mempageSize ? *mempageSize : Memory::kDefaultPageSizeLog2;
+  return ctx.makeMemType(addressType, *limits, shared, pageSizeLog2);
 }
 
 // memorder ::= 'seqcst' | 'acqrel'
@@ -3584,6 +3616,8 @@ template<typename Ctx> MaybeResult<> memory(Ctx& ctx) {
 
   std::optional<typename Ctx::MemTypeT> mtype;
   std::optional<typename Ctx::DataStringT> data;
+  MaybeResult<uint8_t> mempageSize = mempagesize(ctx);
+  CHECK_ERR(mempageSize);
   if (ctx.in.takeSExprStart("data"sv)) {
     if (import) {
       return ctx.in.err("imported memories cannot have inline data");
@@ -3593,9 +3627,17 @@ template<typename Ctx> MaybeResult<> memory(Ctx& ctx) {
     if (!ctx.in.takeRParen()) {
       return ctx.in.err("expected end of inline data");
     }
-    mtype =
-      ctx.makeMemType(addressType, ctx.getLimitsFromData(*datastr), false);
+    const uint8_t pageSizeLog2 =
+      mempageSize.getPtr() ? *mempageSize : Memory::kDefaultPageSizeLog2;
+    mtype = ctx.makeMemType(addressType,
+                            ctx.getLimitsFromData(*datastr, pageSizeLog2),
+                            false,
+                            pageSizeLog2);
     data = *datastr;
+  } else if (mempageSize) {
+    // If we have a memory page size not within a memtype expression, we expect
+    // a memory abbreviation.
+    return ctx.in.err("expected data segment in memory abbreviation");
   } else {
     auto type = memtypeContinued(ctx, addressType);
     CHECK_ERR(type);

src/passes/LLVMMemoryCopyFillLowering.cpp

@@ -117,14 +117,17 @@ struct LLVMMemoryCopyFillLowering
   void createMemoryCopyFunc(Module* module) {
     Builder b(*module);
     Index dst = 0, src = 1, size = 2, start = 3, end = 4, step = 5, i = 6;
-    Name memory = module->memories.front()->name;
+    Name memoryName = module->memories.front()->name;
+    Address::address32_t pageSizeLog2 = module->memories.front()->pageSizeLog2;
     Block* body = b.makeBlock();
     // end = memory size in bytes
-    body->list.push_back(
-      b.makeLocalSet(end,
-                     b.makeBinary(BinaryOp::MulInt32,
-                                  b.makeMemorySize(memory),
-                                  b.makeConst(Memory::kPageSize))));
+    body->list.push_back(b.makeLocalSet(
+      end,
+      pageSizeLog2 == 0
+        ? static_cast<Expression*>(b.makeMemorySize(memoryName))
+        : static_cast<Expression*>(b.makeBinary(BinaryOp::ShlInt32,
+                                                b.makeMemorySize(memoryName),
+                                                b.makeConst(pageSizeLog2)))));
     // if dst + size > memsize or src + size > memsize, then trap.
     body->list.push_back(b.makeIf(
       b.makeBinary(BinaryOp::OrInt32,
@@ -187,9 +190,9 @@ struct LLVMMemoryCopyFillLowering
                                                b.makeLocalGet(src, Type::i32),
                                                b.makeLocalGet(i, Type::i32)),
                                   Type::i32,
-                                  memory),
+                                  memoryName),
                        Type::i32,
-                       memory),
+                       memoryName),
            // i += step
            b.makeLocalSet(i,
                           b.makeBinary(BinaryOp::AddInt32,
@@ -203,19 +206,23 @@ struct LLVMMemoryCopyFillLowering
   void createMemoryFillFunc(Module* module) {
     Builder b(*module);
     Index dst = 0, val = 1, size = 2;
-    Name memory = module->memories.front()->name;
+    Name memoryName = module->memories.front()->name;
+    Address::address32_t pageSizeLog2 = module->memories.front()->pageSizeLog2;
     Block* body = b.makeBlock();
 
     // if dst + size > memsize in bytes, then trap.
-    body->list.push_back(
-      b.makeIf(b.makeBinary(BinaryOp::GtUInt32,
-                            b.makeBinary(BinaryOp::AddInt32,
-                                         b.makeLocalGet(dst, Type::i32),
-                                         b.makeLocalGet(size, Type::i32)),
-                            b.makeBinary(BinaryOp::MulInt32,
-                                         b.makeMemorySize(memory),
-                                         b.makeConst(Memory::kPageSize))),
-               b.makeUnreachable()));
+    body->list.push_back(b.makeIf(
+      b.makeBinary(
+        BinaryOp::GtUInt32,
+        b.makeBinary(BinaryOp::AddInt32,
+                     b.makeLocalGet(dst, Type::i32),
+                     b.makeLocalGet(size, Type::i32)),
+        pageSizeLog2 == 0
+          ? static_cast<Expression*>(b.makeMemorySize(memoryName))
+          : static_cast<Expression*>(b.makeBinary(BinaryOp::ShlInt32,
+                                                  b.makeMemorySize(memoryName),
+                                                  b.makeConst(pageSizeLog2)))),
+      b.makeUnreachable()));
 
     body->list.push_back(b.makeBlock(
       "out",
@@ -241,7 +248,7 @@ struct LLVMMemoryCopyFillLowering
                                     b.makeLocalGet(size, Type::i32)),
                        b.makeLocalGet(val, Type::i32),
                        Type::i32,
-                       memory),
+                       memoryName),
            b.makeBreak("copy", nullptr)}))));
     module->getFunction(memFillFuncName)->body = body;
   }