feat: add atomic instructions

Signed-off-by: Henry <mail@henrygressmann.de>

Author: explodingcamera

crates/parser/src/lib.rs

@@ -174,7 +174,6 @@ impl Parser {
         #[cfg(feature = "validate")]
         {
             let features = WasmFeatures::WASM3
-                .difference(WasmFeatures::THREADS)
                 .union(WasmFeatures::CUSTOM_PAGE_SIZES)
                 .union(WasmFeatures::WIDE_ARITHMETIC)
                 .union(WasmFeatures::COMPACT_IMPORTS);

crates/parser/src/macros.rs

@@ -203,6 +203,7 @@ pub(crate) mod visit {
         (@@function_references $($rest:tt)* ) => {};
         (@@gc $($rest:tt)* ) => {};
         (@@exceptions $($rest:tt)* ) => {};
+        (@@threads $($rest:tt)* ) => {};
 
         (@@$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*)) => {
             fn $visit(&mut self $($(,_: $argty)*)?) -> Self::Output {

crates/parser/src/visit.rs

@@ -11,9 +11,9 @@ use alloc::{
     vec::Vec,
 };
 use tinywasm_types::{
-    BinOp, BinOp128, CmpOp, ExceptionHandler, Global, Import, ImportKind, Instruction, MemoryType, Operand64,
-    Operand64Idx, Operand128, Operand128Idx, StorageType, TableDefinition, TagType, TypeSection, ValueCounts,
-    ValueLane, WasmFunctionData,
+    AtomicArg, AtomicOp, AtomicWidth, BinOp, BinOp128, CmpOp, ExceptionHandler, Global, Import, ImportKind,
+    Instruction, MemoryType, Operand64, Operand64Idx, Operand128, Operand128Idx, StorageType, TableDefinition, TagType,
+    TypeSection, ValueCounts, ValueLane, WasmFunctionData,
 };
 use wasmparser::{FunctionBody, OperatorsReader, OperatorsReaderAllocations, VisitSimdOperator};
 
@@ -497,6 +497,35 @@ pub(crate) fn process_operators_and_validate(
     Ok((instructions, builder.data, builder.uses_local_memory, validator.into_allocations(), reader.into_allocations()))
 }
 
+macro_rules! atomic_visitors {
+    ($($op:ident, $inputs:tt => $outputs:tt: $(($visit:ident, $width:literal)),+;)*) => {
+        $(atomic_visitors!(@group $op, $inputs => $outputs: $(($visit, $width)),+);)*
+    };
+    (@group $op:ident, $inputs:tt => $outputs:tt: $(($visit:ident, $width:literal)),+) => {
+        $(atomic_visitors!(@one $op, $inputs => $outputs: $visit, $width);)+
+    };
+    (@one $op:ident, [$($input:ident),*] => [$($output:ident),*]: $visit:ident, $width:literal) => {
+        fn $visit(&mut self, memarg: wasmparser::MemArg) -> Self::Output {
+            if memarg.align != ($width as u32).trailing_zeros() as u8 {
+                return Err(crate::ParseError::Other("invalid atomic alignment".into()));
+            }
+            let address = self.metadata.memory_size(memarg.memory)?;
+            self.mark_memory(memarg.memory);
+            let memory = self.push128(Operand128::<tinywasm_types::MemoryOperand>::new(memarg.offset, memarg.memory))?;
+            self.emit(
+                &[$(atomic_visitors!(@size $input, address)),*],
+                &[$(atomic_visitors!(@size $output, address)),*],
+                Instruction::Atomic(AtomicArg::new(memory, AtomicWidth::from_bytes($width), atomic_visitors!(@is64 [$($input),*] [$($output),*]), AtomicOp::$op)),
+            )
+        }
+    };
+    (@size Addr, $address:ident) => { $address };
+    (@size $lane:ident, $address:ident) => { ValueLane::$lane };
+    (@is64 [Addr, S64 $(, S64)*] $outputs:tt) => { true };
+    (@is64 [Addr] [S64]) => { true };
+    (@is64 $inputs:tt $outputs:tt) => { false };
+}
+
 impl<'a> wasmparser::VisitOperator<'a> for FunctionBuilder<'_> {
     type Output = Result<()>;
 
@@ -506,6 +535,41 @@ impl<'a> wasmparser::VisitOperator<'a> for FunctionBuilder<'_> {
 
     wasmparser::for_each_visit_operator!(impl_visit_operator);
 
+    fn visit_atomic_fence(&mut self) -> Self::Output {
+        self.emit(&[], &[], Instruction::AtomicFence)
+    }
+
+    fn visit_memory_atomic_notify(&mut self, _: wasmparser::MemArg) -> Self::Output {
+        Err(crate::ParseError::UnsupportedOperator("memory.atomic.notify".into()))
+    }
+    fn visit_memory_atomic_wait32(&mut self, _: wasmparser::MemArg) -> Self::Output {
+        Err(crate::ParseError::UnsupportedOperator("memory.atomic.wait32".into()))
+    }
+    fn visit_memory_atomic_wait64(&mut self, _: wasmparser::MemArg) -> Self::Output {
+        Err(crate::ParseError::UnsupportedOperator("memory.atomic.wait64".into()))
+    }
+
+    atomic_visitors! {
+        Load, [Addr] => [S32]: (visit_i32_atomic_load, 4), (visit_i32_atomic_load8_u, 1), (visit_i32_atomic_load16_u, 2);
+        Load, [Addr] => [S64]: (visit_i64_atomic_load, 8), (visit_i64_atomic_load8_u, 1), (visit_i64_atomic_load16_u, 2), (visit_i64_atomic_load32_u, 4);
+        Store, [Addr, S32] => []: (visit_i32_atomic_store, 4), (visit_i32_atomic_store8, 1), (visit_i32_atomic_store16, 2);
+        Store, [Addr, S64] => []: (visit_i64_atomic_store, 8), (visit_i64_atomic_store8, 1), (visit_i64_atomic_store16, 2), (visit_i64_atomic_store32, 4);
+        Add, [Addr, S32] => [S32]: (visit_i32_atomic_rmw_add, 4), (visit_i32_atomic_rmw8_add_u, 1), (visit_i32_atomic_rmw16_add_u, 2);
+        Add, [Addr, S64] => [S64]: (visit_i64_atomic_rmw_add, 8), (visit_i64_atomic_rmw8_add_u, 1), (visit_i64_atomic_rmw16_add_u, 2), (visit_i64_atomic_rmw32_add_u, 4);
+        Sub, [Addr, S32] => [S32]: (visit_i32_atomic_rmw_sub, 4), (visit_i32_atomic_rmw8_sub_u, 1), (visit_i32_atomic_rmw16_sub_u, 2);
+        Sub, [Addr, S64] => [S64]: (visit_i64_atomic_rmw_sub, 8), (visit_i64_atomic_rmw8_sub_u, 1), (visit_i64_atomic_rmw16_sub_u, 2), (visit_i64_atomic_rmw32_sub_u, 4);
+        And, [Addr, S32] => [S32]: (visit_i32_atomic_rmw_and, 4), (visit_i32_atomic_rmw8_and_u, 1), (visit_i32_atomic_rmw16_and_u, 2);
+        And, [Addr, S64] => [S64]: (visit_i64_atomic_rmw_and, 8), (visit_i64_atomic_rmw8_and_u, 1), (visit_i64_atomic_rmw16_and_u, 2), (visit_i64_atomic_rmw32_and_u, 4);
+        Or, [Addr, S32] => [S32]: (visit_i32_atomic_rmw_or, 4), (visit_i32_atomic_rmw8_or_u, 1), (visit_i32_atomic_rmw16_or_u, 2);
+        Or, [Addr, S64] => [S64]: (visit_i64_atomic_rmw_or, 8), (visit_i64_atomic_rmw8_or_u, 1), (visit_i64_atomic_rmw16_or_u, 2), (visit_i64_atomic_rmw32_or_u, 4);
+        Xor, [Addr, S32] => [S32]: (visit_i32_atomic_rmw_xor, 4), (visit_i32_atomic_rmw8_xor_u, 1), (visit_i32_atomic_rmw16_xor_u, 2);
+        Xor, [Addr, S64] => [S64]: (visit_i64_atomic_rmw_xor, 8), (visit_i64_atomic_rmw8_xor_u, 1), (visit_i64_atomic_rmw16_xor_u, 2), (visit_i64_atomic_rmw32_xor_u, 4);
+        Xchg, [Addr, S32] => [S32]: (visit_i32_atomic_rmw_xchg, 4), (visit_i32_atomic_rmw8_xchg_u, 1), (visit_i32_atomic_rmw16_xchg_u, 2);
+        Xchg, [Addr, S64] => [S64]: (visit_i64_atomic_rmw_xchg, 8), (visit_i64_atomic_rmw8_xchg_u, 1), (visit_i64_atomic_rmw16_xchg_u, 2), (visit_i64_atomic_rmw32_xchg_u, 4);
+        Cmpxchg, [Addr, S32, S32] => [S32]: (visit_i32_atomic_rmw_cmpxchg, 4), (visit_i32_atomic_rmw8_cmpxchg_u, 1), (visit_i32_atomic_rmw16_cmpxchg_u, 2);
+        Cmpxchg, [Addr, S64, S64] => [S64]: (visit_i64_atomic_rmw_cmpxchg, 8), (visit_i64_atomic_rmw8_cmpxchg_u, 1), (visit_i64_atomic_rmw16_cmpxchg_u, 2), (visit_i64_atomic_rmw32_cmpxchg_u, 4);
+    }
+
     lowering_ops! {
         memory [Addr] => [S32] {
             visit_i32_load => I32Load [load(Instruction::LoadLocal32)],

crates/tinywasm/src/error.rs

@@ -131,6 +131,9 @@ pub enum Trap {
         max: usize,
     },
 
+    /// An atomic memory access was not naturally aligned.
+    UnalignedAtomic,
+
     /// An out-of-bounds table access occurred
     TableOutOfBounds {
         /// The offset of the access
@@ -216,6 +219,7 @@ impl Trap {
         match self {
             Self::Unreachable => "unreachable",
             Self::MemoryOutOfBounds { .. } => "out of bounds memory access",
+            Self::UnalignedAtomic => "unaligned atomic",
             Self::TableOutOfBounds { .. } => "out of bounds table access",
             Self::ArrayOutOfBounds => "out of bounds array access",
             Self::DivisionByZero => "integer divide by zero",
@@ -324,6 +328,7 @@ impl Display for Trap {
             Self::MemoryOutOfBounds { offset, len, max } => {
                 write!(f, "out of bounds memory access: offset={offset}, len={len}, max={max}")
             }
+            Self::UnalignedAtomic => write!(f, "unaligned atomic"),
             Self::TableOutOfBounds { offset, len, max } => {
                 write!(f, "out of bounds table access: offset={offset}, len={len}, max={max}")
             }

crates/tinywasm/src/interpreter/executor/instructions.rs

@@ -323,6 +323,8 @@ macro_rules! instruction_handlers {
             ArrayInitElem(idx) => executor.exec_array_init_elem(*idx)?,
             MemorySize(addr) => executor.exec_memory_size(*addr)?,
             MemoryGrow(addr) => executor.exec_memory_grow(*addr)?,
+            Atomic(arg) => executor.exec_atomic(*arg)?,
+            AtomicFence => crate::shared::fence(crate::shared::Ordering::SeqCst),
 
             // Bulk memory operations
             MemoryCopy(idx) => executor.exec_memory_copy(*idx)?,

crates/tinywasm/src/interpreter/executor/mod.rs

@@ -1430,6 +1430,74 @@ impl<'store> Executor<'store> {
         TARGET::stack_push(&mut self.store.value_stack, cast(value))
     }
 
+    fn exec_atomic(&mut self, arg: AtomicArg) -> Result<(), Trap> {
+        match arg.width() {
+            AtomicWidth::Bytes1 => self.exec_atomic_width::<1>(arg),
+            AtomicWidth::Bytes2 => self.exec_atomic_width::<2>(arg),
+            AtomicWidth::Bytes4 => self.exec_atomic_width::<4>(arg),
+            AtomicWidth::Bytes8 => self.exec_atomic_width::<8>(arg),
+        }
+    }
+
+    fn exec_atomic_width<const N: usize>(&mut self, arg: AtomicArg) -> Result<(), Trap> {
+        let op = arg.op();
+        let value = if op == AtomicOp::Load {
+            0
+        } else if arg.is_64() {
+            u64::stack_pop(&mut self.store.value_stack)
+        } else {
+            u32::stack_pop(&mut self.store.value_stack) as u64
+        };
+        let expected = if op == AtomicOp::Cmpxchg {
+            if arg.is_64() {
+                u64::stack_pop(&mut self.store.value_stack)
+            } else {
+                u32::stack_pop(&mut self.store.value_stack) as u64
+            }
+        } else {
+            0
+        };
+        let memory = arg.memory.resolve(&self.func.data);
+        let mem = self.store.state.get_mem_mut(self.mem_addr(memory.memory()));
+        let base = self.store.value_stack.pop_memory_operand(mem.kind.arch())?;
+        let addr = cold_err!(mem.effective_addr::<N>(base, memory.offset()))?;
+        if addr % N != 0 {
+            return cold!(Err(Trap::UnalignedAtomic));
+        }
+        if op == AtomicOp::Store {
+            let mut bytes = [0u8; N];
+            bytes.copy_from_slice(&value.to_le_bytes()[..N]);
+            return cold_err!(mem.inner.write_fixed::<N>(addr, &bytes));
+        }
+        let bytes = cold_err!(mem.inner.read_fixed::<N>(addr))?;
+        let mut padded = [0u8; 8];
+        padded[..N].copy_from_slice(&bytes);
+        let old = u64::from_le_bytes(padded);
+        let mask = if N == 8 { u64::MAX } else { (1u64 << (N * 8)) - 1 };
+        let next = match op {
+            AtomicOp::Load => None,
+            AtomicOp::Store => unreachable!(),
+            AtomicOp::Xchg => Some(value),
+            AtomicOp::Add => Some(old.wrapping_add(value)),
+            AtomicOp::Sub => Some(old.wrapping_sub(value)),
+            AtomicOp::And => Some(old & value),
+            AtomicOp::Or => Some(old | value),
+            AtomicOp::Xor => Some(old ^ value),
+            AtomicOp::Cmpxchg => (old == (expected & mask)).then_some(value),
+        };
+        if let Some(next) = next {
+            let mut bytes = [0u8; N];
+            bytes.copy_from_slice(&next.to_le_bytes()[..N]);
+            cold_err!(mem.inner.write_fixed::<N>(addr, &bytes))?;
+        }
+        if arg.is_64() {
+            u64::stack_push(&mut self.store.value_stack, old)?;
+        } else {
+            u32::stack_push(&mut self.store.value_stack, old as u32)?;
+        }
+        Ok(())
+    }
+
     #[inline(always)]
     fn exec_mem_store_lane<U: MemValue<N> + Copy, const N: usize>(&mut self, arg: MemoryLaneArg) -> Result<(), Trap> {
         let bytes = Value128::stack_pop(&mut self.store.value_stack).to_mem_bytes();

crates/tinywasm/src/shared.rs

@@ -4,6 +4,6 @@ pub(crate) use alloc::rc::Rc as StoreShared;
 pub(crate) use tinywasm_types::Shared as StoreShared;
 
 #[cfg(not(feature = "portable-atomic"))]
-pub(crate) use core::sync::atomic::{AtomicU32, Ordering};
+pub(crate) use core::sync::atomic::{AtomicU32, Ordering, fence};
 #[cfg(feature = "portable-atomic")]
-pub(crate) use portable_atomic::{AtomicU32, Ordering};
+pub(crate) use portable_atomic::{AtomicU32, Ordering, fence};

crates/tinywasm/tests/atomics.rs

@@ -0,0 +1,121 @@
+use tinywasm::{Error, ModuleInstance, Store, Trap};
+
+type TestResult = Result<(), Box<dyn core::error::Error>>;
+
+#[test]
+fn atomic_operations_on_ordinary_memory() -> TestResult {
+    let wasm = wat::parse_str(
+        r#"
+        (module
+          (memory (export "memory") 1)
+          (func (export "run") (result i64)
+            i32.const 12
+            i32.const 255
+            i32.atomic.store8
+            i32.const 12
+            i32.const 2
+            i32.atomic.rmw8.add_u
+            drop
+            i32.const 12
+            i32.const 1
+            i32.const 7
+            i32.atomic.rmw8.cmpxchg_u
+            drop
+            i32.const 16
+            i64.const -1
+            i64.atomic.store32
+            i32.const 16
+            i64.atomic.load32_u
+            atomic.fence))
+    "#,
+    )?;
+    let module = tinywasm::parse_bytes(&wasm)?;
+    let mut store = Store::default();
+    let instance = ModuleInstance::instantiate(&mut store, &module, None)?;
+    assert_eq!(instance.func::<(), i64>(&store, "run")?.call(&mut store, ())?, 0xffff_ffff);
+    let memory = instance.memory("memory")?;
+    assert_eq!(memory.read_vec(&store, 12, 1)?, [7]);
+    Ok(())
+}
+
+#[test]
+fn atomic_alignment_and_bounds() -> TestResult {
+    let wasm = wat::parse_str(
+        r#"
+        (module (memory 1)
+          (func (export "load") (param i32) (result i32)
+            local.get 0
+            i32.atomic.load offset=2))
+    "#,
+    )?;
+    let module = tinywasm::parse_bytes(&wasm)?;
+    let mut store = Store::default();
+    let instance = ModuleInstance::instantiate(&mut store, &module, None)?;
+    let load = instance.func::<i32, i32>(&store, "load")?;
+    assert!(matches!(load.call(&mut store, 0), Err(Error::Trap(Trap::UnalignedAtomic))));
+    assert!(matches!(load.call(&mut store, 65534), Err(Error::Trap(Trap::MemoryOutOfBounds { .. }))));
+
+    let invalid_alignment = wat::parse_str("(module (memory 1) (func i32.const 0 i32.atomic.load align=1 drop))")?;
+    assert!(tinywasm::parse_bytes(&invalid_alignment).is_err());
+    Ok(())
+}
+
+#[test]
+fn atomic_rmw_operations_wrap_at_byte_width() -> TestResult {
+    for (op, expected, updated) in [
+        ("add", 0xff, 1),
+        ("sub", 0xff, 0xfd),
+        ("and", 0xff, 2),
+        ("or", 0xff, 0xff),
+        ("xor", 0xff, 0xfd),
+        ("xchg", 0xff, 2),
+    ] {
+        let wat = format!(
+            r#"
+            (module (memory 1)
+              (func (export "run") (result i32)
+                i32.const 0
+                i32.const 255
+                i32.atomic.store8
+                i32.const 0
+                i32.const 2
+                i32.atomic.rmw8.{op}_u)
+              (func (export "read") (result i32)
+                i32.const 0
+                i32.atomic.load8_u))
+        "#
+        );
+        let module = tinywasm::parse_bytes(&wat::parse_str(&wat)?)?;
+        let mut store = Store::default();
+        let instance = ModuleInstance::instantiate(&mut store, &module, None)?;
+        assert_eq!(instance.func::<(), i32>(&store, "run")?.call(&mut store, ())?, expected);
+        assert_eq!(instance.func::<(), i32>(&store, "read")?.call(&mut store, ())?, updated);
+    }
+    Ok(())
+}
+
+#[test]
+fn narrow_cmpxchg_truncates_expected_and_preserves_old_value() -> TestResult {
+    let wasm = wat::parse_str(
+        r#"
+        (module
+          (memory (export "memory") 1)
+          (func (export "exchange") (param i64 i64) (result i64)
+            i32.const 4
+            local.get 0
+            local.get 1
+            i64.atomic.rmw32.cmpxchg_u))
+    "#,
+    )?;
+    let module = tinywasm::parse_bytes(&wasm)?;
+    let mut store = Store::default();
+    let instance = ModuleInstance::instantiate(&mut store, &module, None)?;
+    let memory = instance.memory("memory")?;
+    let exchange = instance.func::<(i64, i64), i64>(&store, "exchange")?;
+    memory.copy_from_slice(&mut store, 4, &0x89ab_cdef_u32.to_le_bytes())?;
+
+    assert_eq!(exchange.call(&mut store, (0x1234_89ab_cdef, 7))?, 0x89ab_cdef);
+    assert_eq!(exchange.call(&mut store, (0x89ab_cdef, 9))?, 7);
+    assert_eq!(memory.read_vec(&store, 4, 4)?, 7u32.to_le_bytes());
+    Ok(())
+}

crates/tinywasm/tests/generated/wasm-threads.csv

@@ -1,3 +1,3 @@
 0.9.0,357,262,[{"name":"atomic.wast","passed":48,"failed":249},{"name":"exports.wast","passed":82,"failed":6},{"name":"imports.wast","passed":147,"failed":5},{"name":"memory.wast","passed":80,"failed":2}]
 0.10.0,357,262,[{"name":"atomic.wast","passed":48,"failed":249},{"name":"exports.wast","passed":82,"failed":6},{"name":"imports.wast","passed":147,"failed":5},{"name":"memory.wast","passed":80,"failed":2}]
-0.11.0-pre.0,357,262,[{"name":"atomic.wast","passed":48,"failed":249},{"name":"exports.wast","passed":82,"failed":6},{"name":"imports.wast","passed":147,"failed":5},{"name":"memory.wast","passed":80,"failed":2}]
+0.11.0,609,10,[{"name":"atomic.wast","passed":292,"failed":5},{"name":"exports.wast","passed":88,"failed":0},{"name":"imports.wast","passed":147,"failed":5},{"name":"memory.wast","passed":82,"failed":0}]