Winch: implement rmw add ops

tests/disas/winch/x64/atomic/rmw/i32_atomic_rmw16_addu.wat

@@ -0,0 +1,34 @@
+;;! target = "x86_64"
+;;! test = "winch"
+
+(module
+  (memory 1 1 shared)
+  (func (export "_start") (result i32)
+        (i32.atomic.rmw16.add_u (i32.const 0) (i32.const 42))))
+;; wasm[0]::function[0]:
+;;       pushq   %rbp
+;;       movq    %rsp, %rbp
+;;       movq    8(%rdi), %r11
+;;       movq    0x10(%r11), %r11
+;;       addq    $0x10, %r11
+;;       cmpq    %rsp, %r11
+;;       ja      0x5e
+;;   1c: movq    %rdi, %r14
+;;       subq    $0x10, %rsp
+;;       movq    %rdi, 8(%rsp)
+;;       movq    %rsi, (%rsp)
+;;       movl    $0x2a, %eax
+;;       movl    $0, %ecx
+;;       andw    $1, %cx
+;;       cmpw    $0, %cx
+;;       jne     0x60
+;;   44: movl    $0, %ecx
+;;       movq    0x58(%r14), %r11
+;;       movq    (%r11), %rdx
+;;       addq    %rcx, %rdx
+;;       lock xaddw %ax, (%rdx)
+;;       addq    $0x10, %rsp
+;;       popq    %rbp
+;;       retq
+;;   5e: ud2
+;;   60: ud2

tests/disas/winch/x64/atomic/rmw/i32_atomic_rmw8_addu.wat

@@ -0,0 +1,29 @@
+;;! target = "x86_64"
+;;! test = "winch"
+
+(module
+  (memory 1 1 shared)
+  (func (export "_start") (result i32)
+        (i32.atomic.rmw8.add_u (i32.const 0) (i32.const 42))))
+;; wasm[0]::function[0]:
+;;       pushq   %rbp
+;;       movq    %rsp, %rbp
+;;       movq    8(%rdi), %r11
+;;       movq    0x10(%r11), %r11
+;;       addq    $0x10, %r11
+;;       cmpq    %rsp, %r11
+;;       ja      0x4a
+;;   1c: movq    %rdi, %r14
+;;       subq    $0x10, %rsp
+;;       movq    %rdi, 8(%rsp)
+;;       movq    %rsi, (%rsp)
+;;       movl    $0x2a, %eax
+;;       movl    $0, %ecx
+;;       movq    0x58(%r14), %r11
+;;       movq    (%r11), %rdx
+;;       addq    %rcx, %rdx
+;;       lock xaddb %al, (%rdx)
+;;       addq    $0x10, %rsp
+;;       popq    %rbp
+;;       retq
+;;   4a: ud2

tests/disas/winch/x64/atomic/rmw/i32_atomic_rmw_add.wat

@@ -0,0 +1,34 @@
+;;! target = "x86_64"
+;;! test = "winch"
+
+(module
+  (memory 1 1 shared)
+  (func (export "_start") (result i32)
+        (i32.atomic.rmw.add (i32.const 0) (i32.const 42))))
+;; wasm[0]::function[0]:
+;;       pushq   %rbp
+;;       movq    %rsp, %rbp
+;;       movq    8(%rdi), %r11
+;;       movq    0x10(%r11), %r11
+;;       addq    $0x10, %r11
+;;       cmpq    %rsp, %r11
+;;       ja      0x5b
+;;   1c: movq    %rdi, %r14
+;;       subq    $0x10, %rsp
+;;       movq    %rdi, 8(%rsp)
+;;       movq    %rsi, (%rsp)
+;;       movl    $0x2a, %eax
+;;       movl    $0, %ecx
+;;       andl    $3, %ecx
+;;       cmpl    $0, %ecx
+;;       jne     0x5d
+;;   42: movl    $0, %ecx
+;;       movq    0x58(%r14), %r11
+;;       movq    (%r11), %rdx
+;;       addq    %rcx, %rdx
+;;       lock xaddl %eax, (%rdx)
+;;       addq    $0x10, %rsp
+;;       popq    %rbp
+;;       retq
+;;   5b: ud2
+;;   5d: ud2

tests/disas/winch/x64/atomic/rmw/i64_atomic_rmw16_addu.wat

@@ -0,0 +1,34 @@
+;;! target = "x86_64"
+;;! test = "winch"
+
+(module
+  (memory 1 1 shared)
+  (func (export "_start") (result i64)
+        (i64.atomic.rmw16.add_u (i32.const 0) (i64.const 42))))
+;; wasm[0]::function[0]:
+;;       pushq   %rbp
+;;       movq    %rsp, %rbp
+;;       movq    8(%rdi), %r11
+;;       movq    0x10(%r11), %r11
+;;       addq    $0x10, %r11
+;;       cmpq    %rsp, %r11
+;;       ja      0x60
+;;   1c: movq    %rdi, %r14
+;;       subq    $0x10, %rsp
+;;       movq    %rdi, 8(%rsp)
+;;       movq    %rsi, (%rsp)
+;;       movq    $0x2a, %rax
+;;       movl    $0, %ecx
+;;       andw    $1, %cx
+;;       cmpw    $0, %cx
+;;       jne     0x62
+;;   46: movl    $0, %ecx
+;;       movq    0x58(%r14), %r11
+;;       movq    (%r11), %rdx
+;;       addq    %rcx, %rdx
+;;       lock xaddw %ax, (%rdx)
+;;       addq    $0x10, %rsp
+;;       popq    %rbp
+;;       retq
+;;   60: ud2
+;;   62: ud2

tests/disas/winch/x64/atomic/rmw/i64_atomic_rmw32_addu.wat

@@ -0,0 +1,29 @@
+;;! target = "x86_64"
+;;! test = "winch"
+
+(module
+  (memory 1 1 shared)
+  (func (export "_start") (result i64)
+        (i64.atomic.rmw8.add_u (i32.const 0) (i64.const 42))))
+;; wasm[0]::function[0]:
+;;       pushq   %rbp
+;;       movq    %rsp, %rbp
+;;       movq    8(%rdi), %r11
+;;       movq    0x10(%r11), %r11
+;;       addq    $0x10, %r11
+;;       cmpq    %rsp, %r11
+;;       ja      0x4c
+;;   1c: movq    %rdi, %r14
+;;       subq    $0x10, %rsp
+;;       movq    %rdi, 8(%rsp)
+;;       movq    %rsi, (%rsp)
+;;       movq    $0x2a, %rax
+;;       movl    $0, %ecx
+;;       movq    0x58(%r14), %r11
+;;       movq    (%r11), %rdx
+;;       addq    %rcx, %rdx
+;;       lock xaddb %al, (%rdx)
+;;       addq    $0x10, %rsp
+;;       popq    %rbp
+;;       retq
+;;   4c: ud2

tests/disas/winch/x64/atomic/rmw/i64_atomic_rmw8_addu.wat

@@ -0,0 +1,29 @@
+;;! target = "x86_64"
+;;! test = "winch"
+
+(module
+  (memory 1 1 shared)
+  (func (export "_start") (result i64)
+        (i64.atomic.rmw8.add_u (i32.const 0) (i64.const 42))))
+;; wasm[0]::function[0]:
+;;       pushq   %rbp
+;;       movq    %rsp, %rbp
+;;       movq    8(%rdi), %r11
+;;       movq    0x10(%r11), %r11
+;;       addq    $0x10, %r11
+;;       cmpq    %rsp, %r11
+;;       ja      0x4c
+;;   1c: movq    %rdi, %r14
+;;       subq    $0x10, %rsp
+;;       movq    %rdi, 8(%rsp)
+;;       movq    %rsi, (%rsp)
+;;       movq    $0x2a, %rax
+;;       movl    $0, %ecx
+;;       movq    0x58(%r14), %r11
+;;       movq    (%r11), %rdx
+;;       addq    %rcx, %rdx
+;;       lock xaddb %al, (%rdx)
+;;       addq    $0x10, %rsp
+;;       popq    %rbp
+;;       retq
+;;   4c: ud2

tests/disas/winch/x64/atomic/rmw/i64_atomic_rmw_add.wat

@@ -0,0 +1,34 @@
+;;! target = "x86_64"
+;;! test = "winch"
+
+(module
+  (memory 1 1 shared)
+  (func (export "_start") (result i64)
+        (i64.atomic.rmw.add (i32.const 0) (i64.const 42))))
+;; wasm[0]::function[0]:
+;;       pushq   %rbp
+;;       movq    %rsp, %rbp
+;;       movq    8(%rdi), %r11
+;;       movq    0x10(%r11), %r11
+;;       addq    $0x10, %r11
+;;       cmpq    %rsp, %r11
+;;       ja      0x60
+;;   1c: movq    %rdi, %r14
+;;       subq    $0x10, %rsp
+;;       movq    %rdi, 8(%rsp)
+;;       movq    %rsi, (%rsp)
+;;       movq    $0x2a, %rax
+;;       movl    $0, %ecx
+;;       andq    $7, %rcx
+;;       cmpq    $0, %rcx
+;;       jne     0x62
+;;   46: movl    $0, %ecx
+;;       movq    0x58(%r14), %r11
+;;       movq    (%r11), %rdx
+;;       addq    %rcx, %rdx
+;;       lock xaddq %rax, (%rdx)
+;;       addq    $0x10, %rsp
+;;       popq    %rbp
+;;       retq
+;;   60: ud2
+;;   62: ud2

winch/codegen/src/codegen/error.rs

@@ -32,6 +32,8 @@ pub(crate) enum CodeGenError {
     /// implies a compiler bug.
     #[error("Winch internal error: {0}")]
     Internal(InternalError),
+    #[error("Unsupported extend kind")]
+    UnsupportedExtendKind,
 }
 
 /// An internal error.
@@ -185,4 +187,8 @@ impl CodeGenError {
     pub(crate) const fn unimplemented_masm_instruction() -> Self {
         Self::UnimplementedMasmInstruction
     }
+
+    pub(crate) const fn unsupported_extend_kind() -> Self {
+        Self::UnsupportedExtendKind
+    }
 }

winch/codegen/src/codegen/mod.rs

@@ -3,8 +3,8 @@ use crate::{
     codegen::BlockSig,
     isa::reg::{writable, Reg},
     masm::{
-        Imm, IntCmpKind, LoadKind, MacroAssembler, MemOpKind, OperandSize, RegImm, SPOffset,
-        ShiftKind, TrapCode,
+        ExtendKind, Imm, IntCmpKind, LoadKind, MacroAssembler, MemOpKind, OperandSize, RegImm,
+        RmwOp, SPOffset, ShiftKind, TrapCode, UNTRUSTED_FLAGS,
     },
     stack::TypedReg,
 };
@@ -1363,6 +1363,37 @@ where
 
         Ok(())
     }
+
+    pub(crate) fn emit_atomic_rmw(
+        &mut self,
+        arg: &MemArg,
+        op: RmwOp,
+        size: OperandSize,
+        extend: Option<ExtendKind>,
+    ) -> Result<()> {
+        // Only unsigned extends are supported for atomic operations.
+        match extend {
+            Some(kind) if kind.signed() => bail!(CodeGenError::unsupported_extend_kind()),
+            _ => (),
+        }
+
+        let operand = self.context.pop_to_reg(self.masm, None).unwrap();
+        if let Some(addr) = self.emit_compute_heap_address_align_checked(arg, size)? {
+            let src = self.masm.address_at_reg(addr, 0)?;
+            self.masm.atomic_rmw(
+                src,
+                writable!(operand.reg),
+                size,
+                op,
+                UNTRUSTED_FLAGS,
+                extend,
+            )?;
+            self.context.stack.push(operand.into());
+            self.context.free_reg(addr);
+        }
+
+        Ok(())
+    }
 }
 
 /// Returns the index of the [`ControlStackFrame`] for the given

winch/codegen/src/isa/aarch64/masm.rs

@@ -13,15 +13,15 @@ use crate::{
     },
     masm::{
         CalleeKind, DivKind, ExtendKind, FloatCmpKind, Imm as I, IntCmpKind, LoadKind,
-        MacroAssembler as Masm, MemOpKind, MulWideKind, OperandSize, RegImm, RemKind, RoundingMode,
-        SPOffset, ShiftKind, StackSlot, TrapCode, TruncKind,
+        MacroAssembler as Masm, MemOpKind, MulWideKind, OperandSize, RegImm, RemKind, RmwOp,
+        RoundingMode, SPOffset, ShiftKind, StackSlot, TrapCode, TruncKind,
     },
     stack::TypedReg,
 };
 use anyhow::{anyhow, bail, Result};
 use cranelift_codegen::{
     binemit::CodeOffset,
-    ir::{RelSourceLoc, SourceLoc},
+    ir::{MemFlags, RelSourceLoc, SourceLoc},
     isa::aarch64::inst::{Cond, VectorSize},
     settings, Final, MachBufferFinalized, MachLabel,
 };
@@ -893,6 +893,18 @@ impl Masm for MacroAssembler {
         let _ = (context, kind);
         Err(anyhow!(CodeGenError::unimplemented_masm_instruction()))
     }
+
+    fn atomic_rmw(
+        &mut self,
+        _addr: Self::Address,
+        _operand: WritableReg,
+        _size: OperandSize,
+        _op: RmwOp,
+        _flags: MemFlags,
+        _extend: Option<ExtendKind>,
+    ) -> Result<()> {
+        Err(anyhow!(CodeGenError::unimplemented_masm_instruction()))
+    }
 }
 
 impl MacroAssembler {

winch/codegen/src/isa/x64/asm.rs

@@ -1073,6 +1073,31 @@ impl Assembler {
         });
     }
 
+    pub fn lock_xadd(
+        &mut self,
+        addr: Address,
+        operand: Reg,
+        dst: WritableReg,
+        size: OperandSize,
+        flags: MemFlags,
+    ) {
+        assert!(addr.is_offset());
+        let mem = Self::to_synthetic_amode(
+            &addr,
+            &mut self.pool,
+            &mut self.constants,
+            &mut self.buffer,
+            flags,
+        );
+
+        self.emit(Inst::LockXadd {
+            size: size.into(),
+            operand: operand.into(),
+            mem,
+            dst_old: dst.map(Into::into),
+        });
+    }
+
     pub fn cmp_ir(&mut self, src1: Reg, imm: i32, size: OperandSize) {
         let imm = RegMemImm::imm(imm as u32);