[Custom Page Sizes] Support Custom Page Sizes proposal (#6873)

Author: BilelGho

scripts/test/shared.py

@@ -405,6 +405,8 @@ def get_tests(test_dir, extensions=[], recursive=False):
 
     # Unlinkable module accepted
     'linking.wast',
+    'memory_max.wast',
+    'memory_max_i64.wast',
 
     # Invalid module accepted
     'unreached-invalid.wast',

src/binaryen-c.cpp

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

src/binaryen-c.h

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

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

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

src/parser/context-decls.cpp

@@ -143,6 +143,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;
 };
 
@@ -353,7 +354,9 @@ template<typename Ctx> struct TypeParserCtx {
   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, std::optional<uint8_t>) {
+    return Ok{};
+  }
 
   HeapType getBlockTypeFromResult(const std::vector<Type> results) {
     assert(results.size() == 1);
@@ -1052,13 +1055,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, std::optional<uint8_t> pageSizeLog2) {
+    uint8_t _pageSizeLog2 = pageSizeLog2.value_or(16);
+    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,
+                      std::optional<uint8_t> pageSize) {
+    uint8_t pageSizeLog2 = pageSize.value_or(16);
+    return {addressType, limits, pageSizeLog2, shared};
   }
 
   Result<TypeUseT>
@@ -1400,8 +1408,10 @@ 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, std::optional<uint8_t>) {
+    return Ok{};
+  }
+  MemTypeT makeMemType(Type, LimitsT, bool, std::optional<uint8_t>) { return Ok{}; }
 
   Result<> addFunc(Name name,
                    const std::vector<Name>&,

src/parser/parsers.h

@@ -818,7 +818,36 @@ 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> Result<std::optional<uint8_t>> mempagesize(Ctx& ctx) {
+  if (!ctx.in.takeSExprStart("pagesize"sv)) {
+    return std::nullopt; // 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");
+  }
+
+  // return the log2 of the page size, which is the number of trailing zeros
+  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 std::make_optional<uint8_t>(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) {
@@ -840,7 +869,9 @@ Result<typename Ctx::MemTypeT> memtypeContinued(Ctx& ctx, Type addressType) {
   if (ctx.in.takeKeyword("shared"sv)) {
     shared = true;
   }
-  return ctx.makeMemType(addressType, *limits, shared);
+  auto pageSize = mempagesize(ctx);
+  CHECK_ERR(pageSize);
+  return ctx.makeMemType(addressType, *limits, shared, *pageSize);
 }
 
 // memorder ::= '' | 'seqcst' | 'acqrel'
@@ -3434,6 +3465,8 @@ template<typename Ctx> MaybeResult<> memory(Ctx& ctx) {
 
   std::optional<typename Ctx::MemTypeT> mtype;
   std::optional<typename Ctx::DataStringT> data;
+  auto mempageSize = mempagesize(ctx);
+  CHECK_ERR(mempageSize);
   if (ctx.in.takeSExprStart("data"sv)) {
     if (import) {
       return ctx.in.err("imported memories cannot have inline data");
@@ -3443,9 +3476,13 @@ 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);
+    mtype = ctx.makeMemType(
+      addressType, ctx.getLimitsFromData(*datastr, *mempageSize), false, *mempageSize);
     data = *datastr;
+  } else if ((*mempageSize).has_value()) {
+    // 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,16 @@ 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 memory_name = module->memories.front()->name;
+    Address::address32_t memory_page_size = 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))));
+                     memory_page_size == 0 ? static_cast<Expression*>(b.makeMemorySize(memory_name)):
+                    static_cast<Expression*>(b.makeBinary(BinaryOp::ShlInt32,
+                                  b.makeMemorySize(memory_name),
+                                  b.makeConst(memory_page_size)))));
     // if dst + size > memsize or src + size > memsize, then trap.
     body->list.push_back(b.makeIf(
       b.makeBinary(BinaryOp::OrInt32,
@@ -187,9 +189,9 @@ struct LLVMMemoryCopyFillLowering
                                                b.makeLocalGet(src, Type::i32),
                                                b.makeLocalGet(i, Type::i32)),
                                   Type::i32,
-                                  memory),
+                                  memory_name),
                        Type::i32,
-                       memory),
+                       memory_name),
            // i += step
            b.makeLocalSet(i,
                           b.makeBinary(BinaryOp::AddInt32,
@@ -203,7 +205,9 @@ struct LLVMMemoryCopyFillLowering
   void createMemoryFillFunc(Module* module) {
     Builder b(*module);
     Index dst = 0, val = 1, size = 2;
-    Name memory = module->memories.front()->name;
+    Name memory_name = module->memories.front()->name;
+    Address::address32_t memory_page_size =
+      module->memories.front()->pageSizeLog2;
     Block* body = b.makeBlock();
 
     // if dst + size > memsize in bytes, then trap.
@@ -212,9 +216,10 @@ struct LLVMMemoryCopyFillLowering
                             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))),
+                            memory_page_size == 0 ? static_cast<Expression*>(b.makeMemorySize(memory_name)):
+                            static_cast<Expression*>(b.makeBinary(BinaryOp::ShlInt32,
+                                         b.makeMemorySize(memory_name),
+                                         b.makeConst(memory_page_size)))),
                b.makeUnreachable()));
 
     body->list.push_back(b.makeBlock(
@@ -241,7 +246,7 @@ struct LLVMMemoryCopyFillLowering
                                     b.makeLocalGet(size, Type::i32)),
                        b.makeLocalGet(val, Type::i32),
                        Type::i32,
-                       memory),
+                       memory_name),
            b.makeBreak("copy", nullptr)}))));
     module->getFunction(memFillFuncName)->body = body;
   }

src/passes/Memory64Lowering.cpp

@@ -293,8 +293,8 @@ struct Memory64Lowering : public WalkerPass<PostWalker<Memory64Lowering>> {
     for (auto& memory : module->memories) {
       if (memory->is64()) {
         memory->addressType = Type::i32;
-        if (memory->hasMax() && memory->max > Memory::kMaxSize32) {
-          memory->max = Memory::kMaxSize32;
+        if (memory->hasMax() && memory->max > memory->maxSize32()) {
+          memory->max = memory->maxSize32();
         }
       }
     }