[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/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,15 @@ 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))));
+                     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 +188,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 +204,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 +215,9 @@ 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))),
+                            b.makeBinary(BinaryOp::ShlInt32,
+                                         b.makeMemorySize(memory_name),
+                                         b.makeConst(memory_page_size))),
                b.makeUnreachable()));
 
     body->list.push_back(b.makeBlock(
@@ -241,7 +244,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();
         }
       }
     }

src/passes/MemoryPacking.cpp

@@ -315,7 +315,7 @@ void MemoryPacking::calculateRanges(Module* module,
     // Check if we can rule out a trap by it being in bounds.
     if (auto* c = segment->offset->dynCast<Const>()) {
       auto* memory = module->getMemory(segment->memory);
-      auto memorySize = memory->initial * Memory::kPageSize;
+      auto memorySize = memory->initial << memory->pageSizeLog2;
       Index start = c->value.getUnsigned();
       Index size = segment->data.size();
       Index end;

src/passes/MultiMemoryLowering.cpp

@@ -68,9 +68,11 @@ struct MultiMemoryLowering : public Pass {
   // properties will be set
   Name module;
   Name base;
-  // The initial page size of the combined memory
+  // The page size of the combined memory
+  uint8_t pageSizeLog2;
+  // The initial page count of the combined memory
   Address totalInitialPages;
-  // The max page size of the combined memory
+  // The max page count of the combined memory
   Address totalMaxPages;
   // There is no offset for the first memory, so offsetGlobalNames will always
   // have a size that is one less than the count of memories at the time this
@@ -435,6 +437,7 @@ struct MultiMemoryLowering : public Pass {
                                           : Builder::MemoryInfo::Memory64;
     isShared = getFirstMemory().shared;
     isImported = getFirstMemory().imported();
+    pageSizeLog2 = Memory::kDefaultPageSizeLog2;
     for (auto& memory : wasm->memories) {
       // We are assuming that each memory is configured the same as the first
       // and assert if any of the memories does not match this configuration
@@ -446,18 +449,21 @@ struct MultiMemoryLowering : public Pass {
         Fatal() << "MultiMemoryLowering: only the first memory can be imported";
       }
 
+      // Calculating the page size of the combined memory.
+      // This corresponds to the smaller granularity among combined memories
+      pageSizeLog2 = std::min(pageSizeLog2,memory->pageSizeLog2);
+
       // Calculating the total initial and max page size for the combined memory
       // by totaling the initial and max page sizes for the memories in the
       // module
-      totalInitialPages = totalInitialPages + memory->initial;
+      totalInitialPages = totalInitialPages + (memory->initial << (memory->pageSizeLog2 - pageSizeLog2));
       if (memory->hasMax()) {
-        totalMaxPages = totalMaxPages + memory->max;
+        totalMaxPages = totalMaxPages + (memory->max << (memory->pageSizeLog2 - pageSizeLog2));
       }
     }
     // Ensuring valid initial and max page sizes that do not exceed the number
     // of pages addressable by the pointerType
-    Address maxSize =
-      pointerType == Type::i32 ? Memory::kMaxSize32 : Memory::kMaxSize64;
+    Address maxSize = pointerType == Type::i32 ? 1ull<<(32-pageSizeLog2) : 1ull<<(64-pageSizeLog2);
     if (totalMaxPages > maxSize || totalMaxPages == 0) {
       totalMaxPages = Memory::kUnlimitedSize;
     }
@@ -504,9 +510,9 @@ struct MultiMemoryLowering : public Pass {
         Name name = Names::getValidGlobalName(
           *wasm, memory->name.toString() + "_byte_offset");
         offsetGlobalNames.push_back(std::move(name));
-        addGlobal(name, offsetRunningTotal * Memory::kPageSize);
+        addGlobal(name, offsetRunningTotal << pageSizeLog2);
       }
-      offsetRunningTotal += memory->initial;
+      offsetRunningTotal += memory->initial << (memory->pageSizeLog2 - pageSizeLog2);
     }
   }
 
@@ -554,10 +560,10 @@ struct MultiMemoryLowering : public Pass {
       functionName, Signature(pointerType, pointerType), {});
     function->setLocalName(0, "page_delta");
     auto pageSizeConst = [&]() {
-      return builder.makeConst(Literal(Memory::kPageSize));
+      return builder.makeConst(Literal(wasm->memories[memIdx]->pageSizeLog2));
     };
     auto getOffsetDelta = [&]() {
-      return builder.makeBinary(Abstract::getBinary(pointerType, Abstract::Mul),
+      return builder.makeBinary(Abstract::getBinary(pointerType, Abstract::Shl),
                                 builder.makeLocalGet(0, pointerType),
                                 pageSizeConst());
     };
@@ -588,7 +594,8 @@ struct MultiMemoryLowering : public Pass {
         builder.makeBinary(
           EqInt32,
           builder.makeMemoryGrow(
-            builder.makeLocalGet(0, pointerType), combinedMemory, memoryInfo),
+            builder.makeBinary( Abstract::getBinary(pointerType, Abstract::Shl),
+            builder.makeLocalGet(0, pointerType), builder.makeConst(Literal(wasm->memories[memIdx]->pageSizeLog2 - pageSizeLog2))) , combinedMemory, memoryInfo),
           builder.makeConst(-1)),
         builder.makeReturn(builder.makeConst(-1))));
 
@@ -609,7 +616,7 @@ struct MultiMemoryLowering : public Pass {
           // size
           builder.makeBinary(
             Abstract::getBinary(pointerType, Abstract::Sub),
-            builder.makeBinary(Abstract::getBinary(pointerType, Abstract::Mul),
+            builder.makeBinary(Abstract::getBinary(pointerType, Abstract::Shl),
                                builder.makeLocalGet(sizeLocal, pointerType),
                                pageSizeConst()),
             getMoveSource(offsetGlobalName)),
@@ -646,11 +653,11 @@ struct MultiMemoryLowering : public Pass {
       functionName, Signature(Type::none, pointerType), {});
     Expression* functionBody;
     auto pageSizeConst = [&]() {
-      return builder.makeConst(Literal(Memory::kPageSize));
+      return builder.makeConst(Literal(pageSizeLog2));
     };
     auto getOffsetInPageUnits = [&](Name global) {
       return builder.makeBinary(
-        Abstract::getBinary(pointerType, Abstract::DivU),
+        Abstract::getBinary(pointerType, Abstract::ShrU),
         builder.makeGlobalGet(global, pointerType),
         pageSizeConst());
     };
@@ -697,6 +704,7 @@ struct MultiMemoryLowering : public Pass {
       memory->base = base;
       memory->module = module;
     }
+    memory->pageSizeLog2 = pageSizeLog2;
     wasm->addMemory(std::move(memory));
   }