[Wasm64] The binary format offset of load/store should be u64leb in wasm64 (#5038)

4 files changed, 57 insertions, 17 deletions

diff --git a/src/wasm-stack.h b/src/wasm-stack.h
index dc55c80ee..8f72c1384 100644
--- a/src/wasm-stack.h
+++ b/src/wasm-stack.h
@@ -124,7 +124,7 @@ public:
 private:
   void emitMemoryAccess(size_t alignment,
                         size_t bytes,
-                        uint32_t offset,
+                        uint64_t offset,
                         Name memory);
   int32_t getBreakIndex(Name name);
 
diff --git a/src/wasm/wasm-s-parser.cpp b/src/wasm/wasm-s-parser.cpp
index a7d63a7c6..d48e48585 100644
--- a/src/wasm/wasm-s-parser.cpp
+++ b/src/wasm/wasm-s-parser.cpp
@@ -1895,7 +1895,8 @@ static size_t parseMemAttributes(size_t i,
                                  Element& s,
                                  Address& offset,
                                  Address& align,
-                                 Address fallbackAlign) {
+                                 Address fallbackAlign,
+                                 bool memory64) {
   offset = 0;
   align = fallbackAlign;
   // Parse "align=X" and "offset=X" arguments, bailing out on anything else.
@@ -1926,7 +1927,7 @@ static size_t parseMemAttributes(size_t i,
       }
       align = value;
     } else if (str[0] == 'o') {
-      if (value > std::numeric_limits<uint32_t>::max()) {
+      if (!memory64 && value > std::numeric_limits<uint32_t>::max()) {
         throw ParseException("bad offset", s[i]->line, s[i]->col);
       }
       offset = value;
@@ -1984,7 +1985,8 @@ SExpressionWasmBuilder::makeLoad(Element& s, Type type, bool isAtomic) {
     memory = getMemoryNameAtIdx(0);
   }
   ret->memory = memory;
-  i = parseMemAttributes(i, s, ret->offset, ret->align, ret->bytes);
+  i = parseMemAttributes(
+    i, s, ret->offset, ret->align, ret->bytes, isMemory64(memory));
   ret->ptr = parseExpression(s[i]);
   ret->finalize();
   return ret;
@@ -2007,7 +2009,8 @@ SExpressionWasmBuilder::makeStore(Element& s, Type type, bool isAtomic) {
     memory = getMemoryNameAtIdx(0);
   }
   ret->memory = memory;
-  i = parseMemAttributes(i, s, ret->offset, ret->align, ret->bytes);
+  i = parseMemAttributes(
+    i, s, ret->offset, ret->align, ret->bytes, isMemory64(memory));
   ret->ptr = parseExpression(s[i]);
   ret->value = parseExpression(s[i + 1]);
   ret->finalize();
@@ -2062,7 +2065,8 @@ Expression* SExpressionWasmBuilder::makeAtomicRMW(Element& s,
   }
   ret->memory = memory;
   Address align;
-  i = parseMemAttributes(i, s, ret->offset, align, ret->bytes);
+  i = parseMemAttributes(
+    i, s, ret->offset, align, ret->bytes, isMemory64(memory));
   if (align != ret->bytes) {
     throw ParseException("Align of Atomic RMW must match size", s.line, s.col);
   }
@@ -2090,7 +2094,8 @@ Expression* SExpressionWasmBuilder::makeAtomicCmpxchg(Element& s,
     memory = getMemoryNameAtIdx(0);
   }
   ret->memory = memory;
-  i = parseMemAttributes(i, s, ret->offset, align, ret->bytes);
+  i = parseMemAttributes(
+    i, s, ret->offset, align, ret->bytes, isMemory64(memory));
   if (align != ret->bytes) {
     throw ParseException(
       "Align of Atomic Cmpxchg must match size", s.line, s.col);
@@ -2125,7 +2130,8 @@ Expression* SExpressionWasmBuilder::makeAtomicWait(Element& s, Type type) {
     memory = getMemoryNameAtIdx(0);
   }
   ret->memory = memory;
-  i = parseMemAttributes(i, s, ret->offset, align, expectedAlign);
+  i = parseMemAttributes(
+    i, s, ret->offset, align, expectedAlign, isMemory64(memory));
   if (align != expectedAlign) {
     throw ParseException(
       "Align of memory.atomic.wait must match size", s.line, s.col);
@@ -2151,7 +2157,7 @@ Expression* SExpressionWasmBuilder::makeAtomicNotify(Element& s) {
   }
   ret->memory = memory;
   Address align;
-  i = parseMemAttributes(i, s, ret->offset, align, 4);
+  i = parseMemAttributes(i, s, ret->offset, align, 4, isMemory64(memory));
   if (align != 4) {
     throw ParseException(
       "Align of memory.atomic.notify must be 4", s.line, s.col);
@@ -2270,7 +2276,8 @@ Expression* SExpressionWasmBuilder::makeSIMDLoad(Element& s, SIMDLoadOp op) {
     memory = getMemoryNameAtIdx(0);
   }
   ret->memory = memory;
-  i = parseMemAttributes(i, s, ret->offset, ret->align, defaultAlign);
+  i = parseMemAttributes(
+    i, s, ret->offset, ret->align, defaultAlign, isMemory64(memory));
   ret->ptr = parseExpression(s[i]);
   ret->finalize();
   return ret;
@@ -2317,7 +2324,8 @@ SExpressionWasmBuilder::makeSIMDLoadStoreLane(Element& s,
     memory = getMemoryNameAtIdx(0);
   }
   ret->memory = memory;
-  i = parseMemAttributes(i, s, ret->offset, ret->align, defaultAlign);
+  i = parseMemAttributes(
+    i, s, ret->offset, ret->align, defaultAlign, isMemory64(memory));
   ret->index = parseLaneIndex(s[i++], lanes);
   ret->ptr = parseExpression(s[i++]);
   ret->vec = parseExpression(s[i]);
diff --git a/src/wasm/wasm-stack.cpp b/src/wasm/wasm-stack.cpp
index fddb7e41c..d88138075 100644
--- a/src/wasm/wasm-stack.cpp
+++ b/src/wasm/wasm-stack.cpp
@@ -2487,7 +2487,7 @@ void BinaryInstWriter::setScratchLocals() {
 
 void BinaryInstWriter::emitMemoryAccess(size_t alignment,
                                         size_t bytes,
-                                        uint32_t offset,
+                                        uint64_t offset,
                                         Name memory) {
   uint32_t alignmentBits = Bits::log2(alignment ? alignment : bytes);
   uint32_t memoryIdx = parent.getMemoryIndex(memory);
@@ -2500,7 +2500,13 @@ void BinaryInstWriter::emitMemoryAccess(size_t alignment,
   if (memoryIdx > 0) {
     o << U32LEB(memoryIdx);
   }
-  o << U32LEB(offset);
+
+  bool memory64 = parent.getModule()->getMemory(memory)->is64();
+  if (memory64) {
+    o << U64LEB(offset);
+  } else {
+    o << U32LEB(offset);
+  }
 }
 
 int32_t BinaryInstWriter::getBreakIndex(Name name) { // -1 if not found
diff --git a/test/lit/memory64-limits.wast b/test/lit/memory64-limits.wast
index 29e51d29e..f2d395a13 100644
--- a/test/lit/memory64-limits.wast
+++ b/test/lit/memory64-limits.wast
@@ -1,9 +1,35 @@
 ;; RUN: wasm-opt %s -all --roundtrip -S -o - | filecheck %s
 
+;; CHECK:     (module
+;; CHECK-NEXT: (type $none_=>_i64 (func (result i64)))
+;; CHECK-NEXT: (type $none_=>_none (func))
+;; CHECK-NEXT: (memory $0 i64 1 4294967296)
+;; CHECK-NEXT: (func $load_i64 (result i64)
+;; CHECK-NEXT:  (i64.load offset=8589934592
+;; CHECK-NEXT:   (i64.const 4294967296)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )
+;; CHECK-NEXT: (func $store
+;; CHECK-NEXT:  (i64.store offset=8589934592
+;; CHECK-NEXT:   (i64.const 4294967296)
+;; CHECK-NEXT:   (i64.const 123)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )
+;; CHECK-NEXT:)
+
 (module
  (memory $0 i64 1 4294967296)
-)
 
-;; CHECK:      (module
-;; CHECK-NEXT:  (memory $0 i64 1 4294967296)
-;; CHECK-NEXT: )
+  (func $load_i64 (result i64)
+    (i64.load offset=8589934592
+      (i64.const 0x100000000)
+    )
+  )
+
+  (func $store (result)
+    (i64.store offset=8589934592
+      (i64.const 0x100000000)
+      (i64.const 123)
+    )
+  )
+)