Add parsing/printing for relaxed atomics in RMW

Author: stevenfontanella

scripts/test/generate-atomic-spec-test.py

@@ -0,0 +1,212 @@
+import itertools
+from dataclasses import astuple, dataclass
+
+# Workaround for python <3.10, escape characters can't appear in f-strings.
+# Although we require 3.10 in some places, the formatter complains without this.
+newline = "\n"
+
+backslash = '\\'
+
+
+@dataclass
+class instruction_test:
+    op: str
+    arg: str
+    should_drop: bool
+    bin: bytes
+
+
+ALL_OPS = [
+    instruction_test("i32.atomic.load", "(i32.const 51)", True, b"\x41\x33\xfe\x10%(align)s%(memidx)s%(ordering)s\x00\x1a"),
+    instruction_test("i32.atomic.store", "(i32.const 51) (i32.const 51)", False, b"\x41\x33\x41\x33\xfe\x17%(align)s%(memidx)s%(ordering)s\x00"),
+    instruction_test("i32.atomic.rmw.add", "(i32.const 51) (i32.const 51)", True, b"\x41\x33\x41\x33\xfe\x1e%(align)s%(memidx)s%(ordering)s\x00\x1a"),
+]
+
+
+def indent(s):
+    return "\n".join(f"  {line}" if line else "" for line in s.split("\n"))
+
+
+# skips None for convenience
+def instruction(*args):
+    return f"({' '.join(arg for arg in args if arg is not None)})"
+
+
+def atomic_instruction(op, memid, ordering, /, *args, drop):
+    if drop:
+        return f"(drop {instruction(op, memid, ordering, *args)})"
+    return instruction(op, memid, ordering, *args)
+
+
+def func(memid, ordering):
+    return f'''(func ${ordering if ordering is not None else "no_ordering"}{"_with_memid" if memid is not None else "_without_memid"}
+{indent(newline.join(atomic_instruction(op, memid, ordering, arg, drop=should_drop) for op, arg, should_drop, _ in map(astuple, ALL_OPS)))}
+)'''
+
+
+def module(*statements):
+    return f'''(module
+{newline.join(map(indent, statements))}
+)'''
+
+
+def module_binary(bin):
+    return f'''(module binary "{''.join(f'{backslash}{byte:02x}' for byte in bin)}")'''
+
+
+def assert_invalid(module, reason):
+    return f'''(assert_invalid {module} "{reason}")'''
+
+
+def text_test():
+    # Declare two memories so we have control over whether the memory idx is printed
+    # A memory idx of 0 is allowed to be omitted.
+    return module(
+        "(memory 1 1 shared)",
+        "(memory 1 1 shared)",
+        "",
+        "\n\n".join([f'{func(memid, ordering)}' for memid in [None, "1"] for ordering in [None, "acqrel", "seqcst"]]))
+
+
+def to_unsigned_leb(num):
+    ret = bytearray()
+
+    if num == 0:
+        ret = bytearray()
+        ret.append(0)
+        return ret
+    ret = bytearray()
+    while num > 0:
+        rem = num >> 7
+        ret.append((num & 0x7F) | (bool(rem) << 7))
+
+        num = rem
+    return ret
+
+
+def bin_to_str(bin):
+    return ''.join(f'{backslash}{byte:02x}' for byte in bin)
+
+
+@dataclass
+class statement:
+    bin: bytes
+    text: str
+
+
+@dataclass
+class function:
+    body: [statement]
+    memidx: bytes
+    ordering: bytes
+
+
+def normalize_spaces(s):
+    return " ".join(s.split())
+
+
+def binary_line(bin):
+    return f'"{bin_to_str(bin)}"\n'
+
+
+def binary_test_example():
+    return r'''(module binary
+  "\00asm\01\00\00\00" ;; header + version
+  "\01\05\01\60\00\01\7f\03\02\01\00\05\05\01\03\17\80\02" ;; other sections
+  "\0a\0c\01" ;; code section
+    "\0a\00" ;; func size + decl count
+    "\41\33" ;; i32.const 51
+    "\fe\10" ;; i32.atomic.load
+    "\62" ;; 2 | (1<<5) | (1<<6):  Alignment of 2 (32-bit load), with bit 5 set indicating that the next byte is a memory ordering
+    "\00" ;; memory index
+    "\01" ;; acqrel ordering
+    "\00" ;; offset
+    "\0b" ;; end
+)'''
+
+
+def binary_tests():
+    funcs: [function] = []
+    for (memidx_bytes, memidx), (ordering_bytes, ordering) in itertools.product([(b'', None), (b'\x01', "1")], [(b'', None), (b'\x00', "seqcst"), (b'\x01', "acqrel")]):
+        func = function([], memidx, ordering)
+        for test_case in ALL_OPS:
+            align = 2 | (bool(memidx_bytes) << 6) | (bool(ordering_bytes) << 5)
+            s = statement(
+                bin=test_case.bin % {b'align': int.to_bytes(align), b'ordering': ordering_bytes if "rmw" not in test_case.op else int.to_bytes(ordering_bytes[0] | ordering_bytes[0] << 4) if ordering_bytes else b'', b'memidx': memidx_bytes},
+                text=atomic_instruction(test_case.op, memidx, ordering, test_case.arg, drop=test_case.should_drop))
+            func.body.append(s)
+
+        # Functions end with 0x0b.
+        func.body[-1].bin += b'\x0b'
+        funcs.append(func)
+
+    str_builder = []
+
+    for func in funcs:
+        bin_size = sum(len(statement.bin) for statement in func.body)
+        # body size plus 1 byte for the number of locals (0)
+        func_bytes = to_unsigned_leb(bin_size + 1)
+        # number of locals, none in our case
+        func_bytes.append(0x00)
+        str_builder.append(f'"{bin_to_str(func_bytes)}" ;; func\n')
+        for stmt in func.body:
+            str_builder.append(f'"{bin_to_str(stmt.bin)}" ;; {stmt.text}\n')
+
+    section_size = (
+        # function body size
+        sum(len(statement.bin) for func in funcs for statement in func.body) +
+        # function count byte
+        1 +
+        # num locals per function (always 0)
+        len(funcs) +
+        # each function declares its size, add bytes for the LEB encoding of each function's size
+        sum(len(to_unsigned_leb(sum(len(statement.bin) for statement in func.body))) for func in funcs))
+
+    code_section = bytearray(b"\x0a") + to_unsigned_leb(section_size) + to_unsigned_leb(len(funcs))
+
+    '''(module
+        (memory 1 1 shared)
+        (memory 1 1 shared)
+       )
+    '''
+    module = b"\x00\x61\x73\x6d\x01\x00\x00\x00\x01\x04\01\x60\x00\x00\x03\x07\06\x00\x00\x00\x00\x00\x00\x05\07\x02\x03\x01\x01\x03\x01\x01"
+    with open("asdf.wasm", "wb") as f:
+        f.write(module + code_section)
+
+    str_builder = [binary_line(module), f'"{bin_to_str(code_section)}" ;; code section\n'] + str_builder
+
+    return f"(module binary\n{indent(''.join(str_builder))})"
+
+
+def failing_test(instruction, arg, /, memidx, drop):
+    """Module assertion that sets a memory ordering immediate for a non-atomic instruction"""
+
+    func = f"(func ${''.join(filter(str.isalnum, instruction))} {atomic_instruction(instruction, memidx, 'acqrel', arg, drop=drop)})"
+    return assert_invalid(module("(memory 1 1 shared)", "", func), f"Can't set memory ordering for non-atomic {instruction}")
+
+
+def drop_atomic(instruction):
+    first, atomic, last = instruction.partition(".atomic")
+    return first + last
+
+
+def failing_tests():
+    inst = ALL_OPS[0]
+    op = drop_atomic(inst.op)
+
+    return failing_test(op, inst.arg, memidx=None, drop=inst.should_drop)
+
+
+def main():
+    print(text_test())
+    print()
+    print(binary_test_example())
+    print()
+    print(binary_tests())
+    print()
+    print(failing_tests())
+    print()
+
+
+if __name__ == "__main__":
+    main()

src/parser/parsers.h

@@ -1815,10 +1815,20 @@ Result<> makeAtomicRMW(Ctx& ctx,
                        uint8_t bytes) {
   auto mem = maybeMemidx(ctx);
   CHECK_ERR(mem);
+
+  auto maybeOrder = maybeMemOrder(ctx);
+  CHECK_ERR(maybeOrder);
+
   auto arg = memarg(ctx, bytes);
   CHECK_ERR(arg);
-  return ctx.makeAtomicRMW(
-    pos, annotations, op, type, bytes, mem.getPtr(), *arg, MemoryOrder::SeqCst);
+  return ctx.makeAtomicRMW(pos,
+                           annotations,
+                           op,
+                           type,
+                           bytes,
+                           mem.getPtr(),
+                           *arg,
+                           maybeOrder ? *maybeOrder : MemoryOrder::SeqCst);
 }
 
 template<typename Ctx>

src/passes/Print.cpp

@@ -661,6 +661,7 @@ struct PrintExpressionContents
     }
     restoreNormalColor(o);
     printMemoryName(curr->memory, o, wasm);
+    printMemoryOrder(curr->order);
     if (curr->offset) {
       o << " offset=" << curr->offset;
     }

src/wasm/wasm-binary.cpp

@@ -3705,37 +3705,37 @@ Result<> WasmBinaryReader::readInst() {
   case BinaryConsts::I32AtomicRMW##op: {                                       \
     auto [mem, align, offset, memoryOrder] = getRMWMemarg();                   \
     return builder.makeAtomicRMW(                                              \
-      RMW##op, 4, offset, Type::i32, mem, MemoryOrder::SeqCst);                \
+      RMW##op, 4, offset, Type::i32, mem, memoryOrder);                        \
   }                                                                            \
   case BinaryConsts::I32AtomicRMW##op##8U: {                                   \
     auto [mem, align, offset, memoryOrder] = getRMWMemarg();                   \
     return builder.makeAtomicRMW(                                              \
-      RMW##op, 1, offset, Type::i32, mem, MemoryOrder::SeqCst);                \
+      RMW##op, 1, offset, Type::i32, mem, memoryOrder);                        \
   }                                                                            \
   case BinaryConsts::I32AtomicRMW##op##16U: {                                  \
     auto [mem, align, offset, memoryOrder] = getRMWMemarg();                   \
     return builder.makeAtomicRMW(                                              \
-      RMW##op, 2, offset, Type::i32, mem, MemoryOrder::SeqCst);                \
+      RMW##op, 2, offset, Type::i32, mem, memoryOrder);                        \
   }                                                                            \
   case BinaryConsts::I64AtomicRMW##op: {                                       \
     auto [mem, align, offset, memoryOrder] = getRMWMemarg();                   \
     return builder.makeAtomicRMW(                                              \
-      RMW##op, 8, offset, Type::i64, mem, MemoryOrder::SeqCst);                \
+      RMW##op, 8, offset, Type::i64, mem, memoryOrder);                        \
   }                                                                            \
   case BinaryConsts::I64AtomicRMW##op##8U: {                                   \
     auto [mem, align, offset, memoryOrder] = getRMWMemarg();                   \
     return builder.makeAtomicRMW(                                              \
-      RMW##op, 1, offset, Type::i64, mem, MemoryOrder::SeqCst);                \
+      RMW##op, 1, offset, Type::i64, mem, memoryOrder);                        \
   }                                                                            \
   case BinaryConsts::I64AtomicRMW##op##16U: {                                  \
     auto [mem, align, offset, memoryOrder] = getRMWMemarg();                   \
     return builder.makeAtomicRMW(                                              \
-      RMW##op, 2, offset, Type::i64, mem, MemoryOrder::SeqCst);                \
+      RMW##op, 2, offset, Type::i64, mem, memoryOrder);                        \
   }                                                                            \
   case BinaryConsts::I64AtomicRMW##op##32U: {                                  \
     auto [mem, align, offset, memoryOrder] = getRMWMemarg();                   \
     return builder.makeAtomicRMW(                                              \
-      RMW##op, 4, offset, Type::i64, mem, MemoryOrder::SeqCst);                \
+      RMW##op, 4, offset, Type::i64, mem, memoryOrder);                        \
   }
 
           RMW(Add);

src/wasm/wasm-ir-builder.cpp

@@ -1075,7 +1075,7 @@ Result<> IRBuilder::visitEnd() {
     tryy->name = scope.label;
     tryy->finalize(tryy->type);
     push(maybeWrapForLabel(tryy));
-  } else if (Try* tryy;
+  } else if (Try * tryy;
              (tryy = scope.getCatch()) || (tryy = scope.getCatchAll())) {
     auto index = scope.getIndex();
     setCatchBody(tryy, *expr, index);

src/wasm/wasm-stack.cpp

@@ -540,7 +540,7 @@ void BinaryInstWriter::visitAtomicRMW(AtomicRMW* curr) {
                    curr->bytes,
                    curr->offset,
                    curr->memory,
-                   MemoryOrder::SeqCst,
+                   curr->order,
                    /*isRMW=*/true);
 }