From f9d6efe2b53ce6466ee4e95fe162348f087886ad Mon Sep 17 00:00:00 2001
From: adhoc <postma.marin@protonmail.com>
Date: Sun, 12 Jan 2025 12:29:53 +0100
Subject: [PATCH 01/17] add MemOpKind to wasm_store

---
 winch/codegen/src/codegen/mod.rs      |  2 +-
 winch/codegen/src/isa/aarch64/masm.rs | 11 ++++++++---
 winch/codegen/src/isa/x64/masm.rs     |  9 +++++++--
 winch/codegen/src/masm.rs             |  2 +-
 4 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/winch/codegen/src/codegen/mod.rs b/winch/codegen/src/codegen/mod.rs
index 43470f12e979..2373a44e11bf 100644
--- a/winch/codegen/src/codegen/mod.rs
+++ b/winch/codegen/src/codegen/mod.rs
@@ -873,7 +873,7 @@ where
         let addr = self.emit_compute_heap_address(&arg, size)?;
         if let Some(addr) = addr {
             self.masm
-                .wasm_store(src.reg.into(), self.masm.address_at_reg(addr, 0)?, size)?;
+                .wasm_store(src.reg.into(), self.masm.address_at_reg(addr, 0)?, size, MemOpKind::Normal)?;
 
             self.context.free_reg(addr);
         }
diff --git a/winch/codegen/src/isa/aarch64/masm.rs b/winch/codegen/src/isa/aarch64/masm.rs
index 3f2ad618ced5..17c73d2c5afc 100644
--- a/winch/codegen/src/isa/aarch64/masm.rs
+++ b/winch/codegen/src/isa/aarch64/masm.rs
@@ -175,9 +175,14 @@ impl Masm for MacroAssembler {
         Ok(())
     }
 
-    fn wasm_store(&mut self, src: Reg, dst: Self::Address, size: OperandSize) -> Result<()> {
-        self.asm.str(src, dst, size);
-        Ok(())
+    fn wasm_store(&mut self, src: Reg, dst: Self::Address, size: OperandSize, op_kind: MemOpKind) -> Result<()> {
+        match op_kind {
+            MemOpKind::Atomic => Err(anyhow!(CodeGenError::unimplemented_masm_instruction())),
+            MemOpKind::Normal => {
+                self.asm.str(src, dst, size);
+                Ok(())
+            },
+        }
     }
 
     fn call(
diff --git a/winch/codegen/src/isa/x64/masm.rs b/winch/codegen/src/isa/x64/masm.rs
index f5045aaf7c27..bb358c65ed33 100644
--- a/winch/codegen/src/isa/x64/masm.rs
+++ b/winch/codegen/src/isa/x64/masm.rs
@@ -217,8 +217,13 @@ impl Masm for MacroAssembler {
         self.store_impl(src, dst, size, TRUSTED_FLAGS)
     }
 
-    fn wasm_store(&mut self, src: Reg, dst: Self::Address, size: OperandSize) -> Result<()> {
-        self.store_impl(src.into(), dst, size, UNTRUSTED_FLAGS)
+    fn wasm_store(&mut self, src: Reg, dst: Self::Address, size: OperandSize, op_kind: MemOpKind) -> Result<()> {
+        match op_kind {
+            MemOpKind::Atomic => Err(anyhow!(CodeGenError::unimplemented_masm_instruction())),
+            MemOpKind::Normal => {
+                self.store_impl(src.into(), dst, size, UNTRUSTED_FLAGS)
+            },
+        }
     }
 
     fn pop(&mut self, dst: WritableReg, size: OperandSize) -> Result<()> {
diff --git a/winch/codegen/src/masm.rs b/winch/codegen/src/masm.rs
index d64322c4edd6..c6305c9ae9b4 100644
--- a/winch/codegen/src/masm.rs
+++ b/winch/codegen/src/masm.rs
@@ -739,7 +739,7 @@ pub(crate) trait MacroAssembler {
     /// regards to the endianness depending on the target ISA. For this reason,
     /// [Self::wasm_store], should be explicitly used when emitting WebAssembly
     /// stores.
-    fn wasm_store(&mut self, src: Reg, dst: Self::Address, size: OperandSize) -> Result<()>;
+    fn wasm_store(&mut self, src: Reg, dst: Self::Address, size: OperandSize, op_kind: MemOpKind) -> Result<()>;
 
     /// Perform a zero-extended stack load.
     fn load(&mut self, src: Self::Address, dst: WritableReg, size: OperandSize) -> Result<()>;

From c9b4a0c9c6c947c66ffaa1523e848867472b5a1e Mon Sep 17 00:00:00 2001
From: adhoc <postma.marin@protonmail.com>
Date: Sun, 12 Jan 2025 12:38:57 +0100
Subject: [PATCH 02/17] add fence to x64 asm

---
 winch/codegen/src/isa/x64/asm.rs | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/winch/codegen/src/isa/x64/asm.rs b/winch/codegen/src/isa/x64/asm.rs
index 3e5ba4e30baa..ecf2e2b901a5 100644
--- a/winch/codegen/src/isa/x64/asm.rs
+++ b/winch/codegen/src/isa/x64/asm.rs
@@ -16,10 +16,7 @@ use cranelift_codegen::{
         unwind::UnwindInst,
         x64::{
             args::{
-                self, AluRmiROpcode, Amode, AvxOpcode, CmpOpcode, DivSignedness, ExtMode,
-                FromWritableReg, Gpr, GprMem, GprMemImm, Imm8Gpr, Imm8Reg, RegMem, RegMemImm,
-                ShiftKind as CraneliftShiftKind, SseOpcode, SyntheticAmode, WritableGpr,
-                WritableXmm, Xmm, XmmMem, XmmMemAligned, CC,
+                self, AluRmiROpcode, Amode, AvxOpcode, CmpOpcode, DivSignedness, ExtMode, FenceKind, FromWritableReg, Gpr, GprMem, GprMemImm, Imm8Gpr, Imm8Reg, RegMem, RegMemImm, ShiftKind as CraneliftShiftKind, SseOpcode, SyntheticAmode, WritableGpr, WritableXmm, Xmm, XmmMem, XmmMemAligned, CC
             },
             encoding::rex::{encode_modrm, RexFlags},
             settings as x64_settings, EmitInfo, EmitState, Inst,
@@ -1483,6 +1480,10 @@ impl Assembler {
             dst_hi: dst_hi.to_reg().into(),
         });
     }
+
+    pub fn fence(&mut self, kind: FenceKind) {
+        self.emit(Inst::Fence { kind });
+    }
 }
 
 /// Captures the region in a MachBuffer where an add-with-immediate instruction would be emitted,

From 8454623eb2194ebc997923c4a9a2598aed0d292f Mon Sep 17 00:00:00 2001
From: adhoc <postma.marin@protonmail.com>
Date: Sun, 12 Jan 2025 12:39:09 +0100
Subject: [PATCH 03/17] implement atomic store for x64 masm

---
 winch/codegen/src/isa/x64/masm.rs | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/winch/codegen/src/isa/x64/masm.rs b/winch/codegen/src/isa/x64/masm.rs
index bb358c65ed33..2153485c0934 100644
--- a/winch/codegen/src/isa/x64/masm.rs
+++ b/winch/codegen/src/isa/x64/masm.rs
@@ -33,7 +33,7 @@ use cranelift_codegen::{
     isa::{
         unwind::UnwindInst,
         x64::{
-            args::{ExtMode, CC},
+            args::{ExtMode, FenceKind, CC},
             settings as x64_settings,
         },
     },
@@ -219,7 +219,13 @@ impl Masm for MacroAssembler {
 
     fn wasm_store(&mut self, src: Reg, dst: Self::Address, size: OperandSize, op_kind: MemOpKind) -> Result<()> {
         match op_kind {
-            MemOpKind::Atomic => Err(anyhow!(CodeGenError::unimplemented_masm_instruction())),
+            MemOpKind::Atomic => {
+                // To stay consistent with cranelift, we emit a normal load followed by a mfence,
+                // although, we could probably just emit a xchg.
+                self.store_impl(src.into(), dst, size, UNTRUSTED_FLAGS)?;
+                self.asm.fence(FenceKind::MFence);
+                Ok(())
+            },
             MemOpKind::Normal => {
                 self.store_impl(src.into(), dst, size, UNTRUSTED_FLAGS)
             },

From b36054cc2d923d335ffea48118d4e81fbbfe4b86 Mon Sep 17 00:00:00 2001
From: adhoc <postma.marin@protonmail.com>
Date: Sun, 12 Jan 2025 12:44:34 +0100
Subject: [PATCH 04/17] emit error for 128bits atomic store

---
 winch/codegen/src/isa/x64/masm.rs | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/winch/codegen/src/isa/x64/masm.rs b/winch/codegen/src/isa/x64/masm.rs
index 2153485c0934..dac37293ad03 100644
--- a/winch/codegen/src/isa/x64/masm.rs
+++ b/winch/codegen/src/isa/x64/masm.rs
@@ -220,6 +220,10 @@ impl Masm for MacroAssembler {
     fn wasm_store(&mut self, src: Reg, dst: Self::Address, size: OperandSize, op_kind: MemOpKind) -> Result<()> {
         match op_kind {
             MemOpKind::Atomic => {
+                if size == OperandSize::S128 {
+                    // TODO: we don't support 128-bit atomic store yet.
+                    bail!(CodeGenError::unexpected_operand_size());
+                }
                 // To stay consistent with cranelift, we emit a normal load followed by a mfence,
                 // although, we could probably just emit a xchg.
                 self.store_impl(src.into(), dst, size, UNTRUSTED_FLAGS)?;

From cc27c42e961cc0e3a34ebe76e8ce3cc4a1f07219 Mon Sep 17 00:00:00 2001
From: adhoc <postma.marin@protonmail.com>
Date: Sun, 12 Jan 2025 12:45:13 +0100
Subject: [PATCH 05/17] propagate MemOpKind from Visitor

---
 winch/codegen/src/codegen/mod.rs |  4 ++--
 winch/codegen/src/visitor.rs     | 20 ++++++++++----------
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/winch/codegen/src/codegen/mod.rs b/winch/codegen/src/codegen/mod.rs
index 2373a44e11bf..adad927b9407 100644
--- a/winch/codegen/src/codegen/mod.rs
+++ b/winch/codegen/src/codegen/mod.rs
@@ -868,12 +868,12 @@ where
     }
 
     /// Emit a WebAssembly store.
-    pub fn emit_wasm_store(&mut self, arg: &MemArg, size: OperandSize) -> Result<()> {
+    pub fn emit_wasm_store(&mut self, arg: &MemArg, size: OperandSize, op_kind: MemOpKind) -> Result<()> {
         let src = self.context.pop_to_reg(self.masm, None)?;
         let addr = self.emit_compute_heap_address(&arg, size)?;
         if let Some(addr) = addr {
             self.masm
-                .wasm_store(src.reg.into(), self.masm.address_at_reg(addr, 0)?, size, MemOpKind::Normal)?;
+                .wasm_store(src.reg.into(), self.masm.address_at_reg(addr, 0)?, size, op_kind)?;
 
             self.context.free_reg(addr);
         }
diff --git a/winch/codegen/src/visitor.rs b/winch/codegen/src/visitor.rs
index 0cef2fff4fc8..88926d222681 100644
--- a/winch/codegen/src/visitor.rs
+++ b/winch/codegen/src/visitor.rs
@@ -1983,15 +1983,15 @@ where
     }
 
     fn visit_i32_store(&mut self, memarg: MemArg) -> Self::Output {
-        self.emit_wasm_store(&memarg, OperandSize::S32)
+        self.emit_wasm_store(&memarg, OperandSize::S32, MemOpKind::Normal)
     }
 
     fn visit_i32_store8(&mut self, memarg: MemArg) -> Self::Output {
-        self.emit_wasm_store(&memarg, OperandSize::S8)
+        self.emit_wasm_store(&memarg, OperandSize::S8, MemOpKind::Normal)
     }
 
     fn visit_i32_store16(&mut self, memarg: MemArg) -> Self::Output {
-        self.emit_wasm_store(&memarg, OperandSize::S16)
+        self.emit_wasm_store(&memarg, OperandSize::S16, MemOpKind::Normal)
     }
 
     fn visit_i64_load8_s(&mut self, memarg: MemArg) -> Self::Output {
@@ -2058,19 +2058,19 @@ where
     }
 
     fn visit_i64_store(&mut self, memarg: MemArg) -> Self::Output {
-        self.emit_wasm_store(&memarg, OperandSize::S64)
+        self.emit_wasm_store(&memarg, OperandSize::S64, MemOpKind::Normal)
     }
 
     fn visit_i64_store8(&mut self, memarg: MemArg) -> Self::Output {
-        self.emit_wasm_store(&memarg, OperandSize::S8)
+        self.emit_wasm_store(&memarg, OperandSize::S8, MemOpKind::Normal)
     }
 
     fn visit_i64_store16(&mut self, memarg: MemArg) -> Self::Output {
-        self.emit_wasm_store(&memarg, OperandSize::S16)
+        self.emit_wasm_store(&memarg, OperandSize::S16, MemOpKind::Normal)
     }
 
     fn visit_i64_store32(&mut self, memarg: MemArg) -> Self::Output {
-        self.emit_wasm_store(&memarg, OperandSize::S32)
+        self.emit_wasm_store(&memarg, OperandSize::S32, MemOpKind::Normal)
     }
 
     fn visit_f32_load(&mut self, memarg: MemArg) -> Self::Output {
@@ -2083,7 +2083,7 @@ where
     }
 
     fn visit_f32_store(&mut self, memarg: MemArg) -> Self::Output {
-        self.emit_wasm_store(&memarg, OperandSize::S32)
+        self.emit_wasm_store(&memarg, OperandSize::S32, MemOpKind::Normal)
     }
 
     fn visit_f64_load(&mut self, memarg: MemArg) -> Self::Output {
@@ -2096,7 +2096,7 @@ where
     }
 
     fn visit_f64_store(&mut self, memarg: MemArg) -> Self::Output {
-        self.emit_wasm_store(&memarg, OperandSize::S64)
+        self.emit_wasm_store(&memarg, OperandSize::S64, MemOpKind::Normal)
     }
 
     fn visit_i32_trunc_sat_f32_s(&mut self) -> Self::Output {
@@ -2287,7 +2287,7 @@ where
     }
 
     fn visit_v128_store(&mut self, memarg: MemArg) -> Self::Output {
-        self.emit_wasm_store(&memarg, OperandSize::S128)
+        self.emit_wasm_store(&memarg, OperandSize::S128, MemOpKind::Normal)
     }
 
     fn visit_v128_load8x8_s(&mut self, memarg: MemArg) -> Self::Output {

From 2473ccd9cde849ad87b53cd70abdf8dd1218e603 Mon Sep 17 00:00:00 2001
From: adhoc <postma.marin@protonmail.com>
Date: Sun, 12 Jan 2025 13:07:06 +0100
Subject: [PATCH 06/17] implement atomic store visitor methods

---
 winch/codegen/src/visitor.rs | 50 +++++++++++++++++++++++++++++++-----
 1 file changed, 43 insertions(+), 7 deletions(-)

diff --git a/winch/codegen/src/visitor.rs b/winch/codegen/src/visitor.rs
index 88926d222681..c302354184ab 100644
--- a/winch/codegen/src/visitor.rs
+++ b/winch/codegen/src/visitor.rs
@@ -271,6 +271,14 @@ macro_rules! def_unsupported {
     (emit V128Load16Splat $($rest:tt)*) => {};
     (emit V128Load32Splat $($rest:tt)*) => {};
     (emit V128Load64Splat $($rest:tt)*) => {};
+    (emit I32AtomicStore8 $($rest:tt)*) => {};
+    (emit I32AtomicStore16 $($rest:tt)*) => {};
+    (emit I32AtomicStore $($rest:tt)*) => {};
+    (emit I64AtomicStore8 $($rest:tt)*) => {};
+    (emit I64AtomicStore16 $($rest:tt)*) => {};
+    (emit I64AtomicStore32 $($rest:tt)*) => {};
+    (emit I64AtomicStore $($rest:tt)*) => {};
+
     (emit $unsupported:tt $($rest:tt)*) => {$($rest)*};
 }
 
@@ -2201,7 +2209,7 @@ where
         self.masm.mul_wide(&mut self.context, MulWideKind::Unsigned)
     }
 
-    fn visit_i32_atomic_load8_u(&mut self, memarg: wasmparser::MemArg) -> Self::Output {
+    fn visit_i32_atomic_load8_u(&mut self, memarg: MemArg) -> Self::Output {
         self.emit_wasm_load(
             &memarg,
             WasmValType::I32,
@@ -2210,7 +2218,7 @@ where
         )
     }
 
-    fn visit_i32_atomic_load16_u(&mut self, memarg: wasmparser::MemArg) -> Self::Output {
+    fn visit_i32_atomic_load16_u(&mut self, memarg: MemArg) -> Self::Output {
         self.emit_wasm_load(
             &memarg,
             WasmValType::I32,
@@ -2219,7 +2227,7 @@ where
         )
     }
 
-    fn visit_i32_atomic_load(&mut self, memarg: wasmparser::MemArg) -> Self::Output {
+    fn visit_i32_atomic_load(&mut self, memarg: MemArg) -> Self::Output {
         self.emit_wasm_load(
             &memarg,
             WasmValType::I32,
@@ -2228,7 +2236,7 @@ where
         )
     }
 
-    fn visit_i64_atomic_load8_u(&mut self, memarg: wasmparser::MemArg) -> Self::Output {
+    fn visit_i64_atomic_load8_u(&mut self, memarg: MemArg) -> Self::Output {
         self.emit_wasm_load(
             &memarg,
             WasmValType::I64,
@@ -2237,7 +2245,7 @@ where
         )
     }
 
-    fn visit_i64_atomic_load16_u(&mut self, memarg: wasmparser::MemArg) -> Self::Output {
+    fn visit_i64_atomic_load16_u(&mut self, memarg: MemArg) -> Self::Output {
         self.emit_wasm_load(
             &memarg,
             WasmValType::I64,
@@ -2246,7 +2254,7 @@ where
         )
     }
 
-    fn visit_i64_atomic_load32_u(&mut self, memarg: wasmparser::MemArg) -> Self::Output {
+    fn visit_i64_atomic_load32_u(&mut self, memarg: MemArg) -> Self::Output {
         self.emit_wasm_load(
             &memarg,
             WasmValType::I64,
@@ -2255,7 +2263,7 @@ where
         )
     }
 
-    fn visit_i64_atomic_load(&mut self, memarg: wasmparser::MemArg) -> Self::Output {
+    fn visit_i64_atomic_load(&mut self, memarg: MemArg) -> Self::Output {
         self.emit_wasm_load(
             &memarg,
             WasmValType::I64,
@@ -2264,6 +2272,34 @@ where
         )
     }
 
+    fn visit_i32_atomic_store(&mut self, memarg: MemArg) -> Self::Output {
+        self.emit_wasm_store(&memarg, OperandSize::S32, MemOpKind::Atomic)
+    }
+
+    fn visit_i64_atomic_store(&mut self, memarg: MemArg) -> Self::Output {
+        self.emit_wasm_store(&memarg, OperandSize::S64, MemOpKind::Atomic)
+    }
+
+    fn visit_i32_atomic_store8(&mut self, memarg: MemArg) -> Self::Output {
+        self.emit_wasm_store(&memarg, OperandSize::S8, MemOpKind::Atomic)
+    }
+
+    fn visit_i32_atomic_store16(&mut self, memarg: MemArg) -> Self::Output {
+        self.emit_wasm_store(&memarg, OperandSize::S16, MemOpKind::Atomic)
+    }
+
+    fn visit_i64_atomic_store8(&mut self, memarg: MemArg) -> Self::Output {
+        self.emit_wasm_store(&memarg, OperandSize::S8, MemOpKind::Atomic)
+    }
+
+    fn visit_i64_atomic_store16(&mut self, memarg: MemArg) -> Self::Output {
+        self.emit_wasm_store(&memarg, OperandSize::S16, MemOpKind::Atomic)
+    }
+
+    fn visit_i64_atomic_store32(&mut self, memarg: MemArg) -> Self::Output {
+        self.emit_wasm_store(&memarg, OperandSize::S32, MemOpKind::Atomic)
+    }
+
     wasmparser::for_each_visit_operator!(def_unsupported);
 }
 

From be4ea92b94897fdeed4bd8f2176e80194eb6c989 Mon Sep 17 00:00:00 2001
From: adhoc <postma.marin@protonmail.com>
Date: Sun, 12 Jan 2025 13:07:20 +0100
Subject: [PATCH 07/17] add atomic store tests

---
 .../x64/atomic/store/i32_atomic_store.wat     | 29 +++++++++++++++++++
 .../x64/atomic/store/i32_atomic_store16.wat   | 29 +++++++++++++++++++
 .../x64/atomic/store/i32_atomic_store8.wat    | 29 +++++++++++++++++++
 .../x64/atomic/store/i64_atomic_store.wat     | 29 +++++++++++++++++++
 .../x64/atomic/store/i64_atomic_store16.wat   | 29 +++++++++++++++++++
 .../x64/atomic/store/i64_atomic_store32.wat   | 29 +++++++++++++++++++
 .../x64/atomic/store/i64_atomic_store8.wat    | 29 +++++++++++++++++++
 7 files changed, 203 insertions(+)
 create mode 100644 tests/disas/winch/x64/atomic/store/i32_atomic_store.wat
 create mode 100644 tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat
 create mode 100644 tests/disas/winch/x64/atomic/store/i32_atomic_store8.wat
 create mode 100644 tests/disas/winch/x64/atomic/store/i64_atomic_store.wat
 create mode 100644 tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat
 create mode 100644 tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat
 create mode 100644 tests/disas/winch/x64/atomic/store/i64_atomic_store8.wat

diff --git a/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat b/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat
new file mode 100644
index 000000000000..2ed82872b293
--- /dev/null
+++ b/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat
@@ -0,0 +1,29 @@
+;;! target = "x86_64"
+;;! test = "winch"
+
+(module 
+  (import "env" "memory" (memory 1 1 shared))
+  (func (i32.atomic.store (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      0x4b
+;;   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
+;;       movl    %eax, (%rdx)
+;;       mfence
+;;       addq    $0x10, %rsp
+;;       popq    %rbp
+;;       retq
+;;   4b: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat b/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat
new file mode 100644
index 000000000000..d67b305dfbff
--- /dev/null
+++ b/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat
@@ -0,0 +1,29 @@
+;;! target = "x86_64"
+;;! test = "winch"
+
+(module 
+  (import "env" "memory" (memory 1 1 shared))
+  (func (i32.atomic.store16 (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      0x4c
+;;   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
+;;       movw    %ax, (%rdx)
+;;       mfence
+;;       addq    $0x10, %rsp
+;;       popq    %rbp
+;;       retq
+;;   4c: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i32_atomic_store8.wat b/tests/disas/winch/x64/atomic/store/i32_atomic_store8.wat
new file mode 100644
index 000000000000..1626c3482856
--- /dev/null
+++ b/tests/disas/winch/x64/atomic/store/i32_atomic_store8.wat
@@ -0,0 +1,29 @@
+;;! target = "x86_64"
+;;! test = "winch"
+
+(module 
+  (import "env" "memory" (memory 1 1 shared))
+  (func (i32.atomic.store8 (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      0x4b
+;;   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
+;;       movb    %al, (%rdx)
+;;       mfence
+;;       addq    $0x10, %rsp
+;;       popq    %rbp
+;;       retq
+;;   4b: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat
new file mode 100644
index 000000000000..172f90556ad9
--- /dev/null
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat
@@ -0,0 +1,29 @@
+;;! target = "x86_64"
+;;! test = "winch"
+
+(module 
+  (import "env" "memory" (memory 1 1 shared))
+  (func (i64.atomic.store (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      0x4e
+;;   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
+;;       movq    %rax, (%rdx)
+;;       mfence
+;;       addq    $0x10, %rsp
+;;       popq    %rbp
+;;       retq
+;;   4e: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat
new file mode 100644
index 000000000000..320794a7cd5c
--- /dev/null
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat
@@ -0,0 +1,29 @@
+;;! target = "x86_64"
+;;! test = "winch"
+
+(module 
+  (import "env" "memory" (memory 1 1 shared))
+  (func (i64.atomic.store16 (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      0x4e
+;;   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
+;;       movw    %ax, (%rdx)
+;;       mfence
+;;       addq    $0x10, %rsp
+;;       popq    %rbp
+;;       retq
+;;   4e: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat
new file mode 100644
index 000000000000..08a2307ffb00
--- /dev/null
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat
@@ -0,0 +1,29 @@
+;;! target = "x86_64"
+;;! test = "winch"
+
+(module 
+  (import "env" "memory" (memory 1 1 shared))
+  (func (i64.atomic.store32 (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      0x4d
+;;   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
+;;       movl    %eax, (%rdx)
+;;       mfence
+;;       addq    $0x10, %rsp
+;;       popq    %rbp
+;;       retq
+;;   4d: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store8.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store8.wat
new file mode 100644
index 000000000000..5aadfce63fac
--- /dev/null
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store8.wat
@@ -0,0 +1,29 @@
+;;! target = "x86_64"
+;;! test = "winch"
+
+(module 
+  (import "env" "memory" (memory 1 1 shared))
+  (func (i64.atomic.store8 (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      0x4d
+;;   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
+;;       movb    %al, (%rdx)
+;;       mfence
+;;       addq    $0x10, %rsp
+;;       popq    %rbp
+;;       retq
+;;   4d: ud2

From 8bcd530e83fad8befc16c65a50019d11c4f2fd6e Mon Sep 17 00:00:00 2001
From: adhoc <postma.marin@protonmail.com>
Date: Sun, 12 Jan 2025 13:08:04 +0100
Subject: [PATCH 08/17] fmt

---
 winch/codegen/src/codegen/mod.rs      | 15 ++++++++++++---
 winch/codegen/src/isa/aarch64/masm.rs | 10 ++++++++--
 winch/codegen/src/isa/x64/masm.rs     | 14 +++++++++-----
 winch/codegen/src/masm.rs             |  8 +++++++-
 4 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/winch/codegen/src/codegen/mod.rs b/winch/codegen/src/codegen/mod.rs
index adad927b9407..cd3c4154e022 100644
--- a/winch/codegen/src/codegen/mod.rs
+++ b/winch/codegen/src/codegen/mod.rs
@@ -868,12 +868,21 @@ where
     }
 
     /// Emit a WebAssembly store.
-    pub fn emit_wasm_store(&mut self, arg: &MemArg, size: OperandSize, op_kind: MemOpKind) -> Result<()> {
+    pub fn emit_wasm_store(
+        &mut self,
+        arg: &MemArg,
+        size: OperandSize,
+        op_kind: MemOpKind,
+    ) -> Result<()> {
         let src = self.context.pop_to_reg(self.masm, None)?;
         let addr = self.emit_compute_heap_address(&arg, size)?;
         if let Some(addr) = addr {
-            self.masm
-                .wasm_store(src.reg.into(), self.masm.address_at_reg(addr, 0)?, size, op_kind)?;
+            self.masm.wasm_store(
+                src.reg.into(),
+                self.masm.address_at_reg(addr, 0)?,
+                size,
+                op_kind,
+            )?;
 
             self.context.free_reg(addr);
         }
diff --git a/winch/codegen/src/isa/aarch64/masm.rs b/winch/codegen/src/isa/aarch64/masm.rs
index 17c73d2c5afc..62627c9c980c 100644
--- a/winch/codegen/src/isa/aarch64/masm.rs
+++ b/winch/codegen/src/isa/aarch64/masm.rs
@@ -175,13 +175,19 @@ impl Masm for MacroAssembler {
         Ok(())
     }
 
-    fn wasm_store(&mut self, src: Reg, dst: Self::Address, size: OperandSize, op_kind: MemOpKind) -> Result<()> {
+    fn wasm_store(
+        &mut self,
+        src: Reg,
+        dst: Self::Address,
+        size: OperandSize,
+        op_kind: MemOpKind,
+    ) -> Result<()> {
         match op_kind {
             MemOpKind::Atomic => Err(anyhow!(CodeGenError::unimplemented_masm_instruction())),
             MemOpKind::Normal => {
                 self.asm.str(src, dst, size);
                 Ok(())
-            },
+            }
         }
     }
 
diff --git a/winch/codegen/src/isa/x64/masm.rs b/winch/codegen/src/isa/x64/masm.rs
index dac37293ad03..c9ee4dfe86b1 100644
--- a/winch/codegen/src/isa/x64/masm.rs
+++ b/winch/codegen/src/isa/x64/masm.rs
@@ -217,7 +217,13 @@ impl Masm for MacroAssembler {
         self.store_impl(src, dst, size, TRUSTED_FLAGS)
     }
 
-    fn wasm_store(&mut self, src: Reg, dst: Self::Address, size: OperandSize, op_kind: MemOpKind) -> Result<()> {
+    fn wasm_store(
+        &mut self,
+        src: Reg,
+        dst: Self::Address,
+        size: OperandSize,
+        op_kind: MemOpKind,
+    ) -> Result<()> {
         match op_kind {
             MemOpKind::Atomic => {
                 if size == OperandSize::S128 {
@@ -229,10 +235,8 @@ impl Masm for MacroAssembler {
                 self.store_impl(src.into(), dst, size, UNTRUSTED_FLAGS)?;
                 self.asm.fence(FenceKind::MFence);
                 Ok(())
-            },
-            MemOpKind::Normal => {
-                self.store_impl(src.into(), dst, size, UNTRUSTED_FLAGS)
-            },
+            }
+            MemOpKind::Normal => self.store_impl(src.into(), dst, size, UNTRUSTED_FLAGS),
         }
     }
 
diff --git a/winch/codegen/src/masm.rs b/winch/codegen/src/masm.rs
index c6305c9ae9b4..18c99490aaf1 100644
--- a/winch/codegen/src/masm.rs
+++ b/winch/codegen/src/masm.rs
@@ -739,7 +739,13 @@ pub(crate) trait MacroAssembler {
     /// regards to the endianness depending on the target ISA. For this reason,
     /// [Self::wasm_store], should be explicitly used when emitting WebAssembly
     /// stores.
-    fn wasm_store(&mut self, src: Reg, dst: Self::Address, size: OperandSize, op_kind: MemOpKind) -> Result<()>;
+    fn wasm_store(
+        &mut self,
+        src: Reg,
+        dst: Self::Address,
+        size: OperandSize,
+        op_kind: MemOpKind,
+    ) -> Result<()>;
 
     /// Perform a zero-extended stack load.
     fn load(&mut self, src: Self::Address, dst: WritableReg, size: OperandSize) -> Result<()>;

From c7c5d6a0b8485666070ae8780b1679ecaf0c8156 Mon Sep 17 00:00:00 2001
From: adhoc <postma.marin@protonmail.com>
Date: Sun, 12 Jan 2025 17:00:28 +0100
Subject: [PATCH 09/17] whitelist spec tests

---
 crates/wast-util/src/lib.rs | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/crates/wast-util/src/lib.rs b/crates/wast-util/src/lib.rs
index f537fd5053b5..ea4dc67265ca 100644
--- a/crates/wast-util/src/lib.rs
+++ b/crates/wast-util/src/lib.rs
@@ -488,12 +488,8 @@ impl WastTest {
                 "spec_testsuite/simd_store8_lane.wast",
                 // thread related failures
                 "proposals/threads/atomic.wast",
-                "misc_testsuite/threads/MP_wait.wast",
                 "misc_testsuite/threads/load-store-alignment.wast",
-                "misc_testsuite/threads/MP_atomic.wast",
-                "misc_testsuite/threads/SB_atomic.wast",
                 "misc_testsuite/threads/wait_notify.wast",
-                "misc_testsuite/threads/LB_atomic.wast",
                 "misc_testsuite/threads/atomics_wait_address.wast",
                 "misc_testsuite/threads/atomics_notify.wast",
                 "misc_testsuite/threads/load-store-alignment.wast",

From a9e35fe8d96ea28356c5382481a79df25c2ab591 Mon Sep 17 00:00:00 2001
From: ad hoc <postma.marin@protonmail.com>
Date: Mon, 13 Jan 2025 19:07:45 +0100
Subject: [PATCH 10/17] atomic heap address compute align check

---
 crates/wast-util/src/lib.rs                   |  2 -
 .../winch/x64/atomic/load/i32_atomic_load.wat |  9 +++-
 .../x64/atomic/load/i32_atomic_load16_u.wat   |  9 +++-
 .../winch/x64/atomic/load/i64_atomic_load.wat |  9 +++-
 .../x64/atomic/load/i64_atomic_load16_u.wat   |  9 +++-
 .../x64/atomic/load/i64_atomic_load32_u.wat   |  9 +++-
 .../x64/atomic/store/i32_atomic_store.wat     |  9 +++-
 .../x64/atomic/store/i32_atomic_store16.wat   |  9 +++-
 .../x64/atomic/store/i64_atomic_store.wat     |  9 +++-
 .../x64/atomic/store/i64_atomic_store16.wat   |  9 +++-
 .../x64/atomic/store/i64_atomic_store32.wat   |  9 +++-
 winch/codegen/src/codegen/mod.rs              | 46 +++++++++++++++++--
 winch/codegen/src/masm.rs                     |  6 +++
 13 files changed, 118 insertions(+), 26 deletions(-)

diff --git a/crates/wast-util/src/lib.rs b/crates/wast-util/src/lib.rs
index ea4dc67265ca..432548eee146 100644
--- a/crates/wast-util/src/lib.rs
+++ b/crates/wast-util/src/lib.rs
@@ -488,11 +488,9 @@ impl WastTest {
                 "spec_testsuite/simd_store8_lane.wast",
                 // thread related failures
                 "proposals/threads/atomic.wast",
-                "misc_testsuite/threads/load-store-alignment.wast",
                 "misc_testsuite/threads/wait_notify.wast",
                 "misc_testsuite/threads/atomics_wait_address.wast",
                 "misc_testsuite/threads/atomics_notify.wast",
-                "misc_testsuite/threads/load-store-alignment.wast",
             ];
 
             if unsupported.iter().any(|part| self.path.ends_with(part)) {
diff --git a/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat b/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat
index b561f66e8e44..a671002d8a51 100644
--- a/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat
+++ b/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat
@@ -13,13 +13,17 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x20, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x47
+;;       ja      0x57
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x20, %rsp
 ;;       movq    %rdi, 0x18(%rsp)
 ;;       movq    %rsi, 0x10(%rsp)
 ;;       movl    %edx, 0xc(%rsp)
 ;;       movl    0xc(%rsp), %eax
+;;       andl    $3, %eax
+;;       cmpl    $0, %eax
+;;       jne     0x59
+;;   41: movl    0xc(%rsp), %eax
 ;;       movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rcx
 ;;       addq    %rax, %rcx
@@ -27,4 +31,5 @@
 ;;       addq    $0x20, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   47: ud2
+;;   57: ud2
+;;   59: ud2
diff --git a/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat b/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat
index ed9dbdf8047c..3c156158518b 100644
--- a/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat
+++ b/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat
@@ -13,13 +13,17 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x20, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x49
+;;       ja      0x5b
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x20, %rsp
 ;;       movq    %rdi, 0x18(%rsp)
 ;;       movq    %rsi, 0x10(%rsp)
 ;;       movl    %edx, 0xc(%rsp)
 ;;       movl    0xc(%rsp), %eax
+;;       andw    $1, %ax
+;;       cmpw    $0, %ax
+;;       jne     0x5d
+;;   43: movl    0xc(%rsp), %eax
 ;;       movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rcx
 ;;       addq    %rax, %rcx
@@ -27,4 +31,5 @@
 ;;       addq    $0x20, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   49: ud2
+;;   5b: ud2
+;;   5d: ud2
diff --git a/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat b/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat
index 3aa41b5bb07f..6d67f5534ba1 100644
--- a/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat
+++ b/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat
@@ -14,16 +14,21 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x41
+;;       ja      0x54
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0, %eax
+;;       andq    $7, %rax
+;;       cmpq    $0, %rax
+;;       jne     0x56
+;;   3f: movl    $0, %eax
 ;;       movq    0x60(%r14), %rcx
 ;;       addq    %rax, %rcx
 ;;       movq    (%rcx), %rax
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   41: ud2
+;;   54: ud2
+;;   56: ud2
diff --git a/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat b/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat
index a5a36d758e27..dbd25c81996a 100644
--- a/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat
+++ b/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat
@@ -14,16 +14,21 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x42
+;;       ja      0x55
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0, %eax
+;;       andw    $1, %ax
+;;       cmpw    $0, %ax
+;;       jne     0x57
+;;   3f: movl    $0, %eax
 ;;       movq    0x60(%r14), %rcx
 ;;       addq    %rax, %rcx
 ;;       movzwq  (%rcx), %rax
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   42: ud2
+;;   55: ud2
+;;   57: ud2
diff --git a/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat b/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat
index 4cc24de13d7f..a7ad511e0d37 100644
--- a/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat
+++ b/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat
@@ -14,16 +14,21 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x40
+;;       ja      0x51
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0, %eax
+;;       andl    $3, %eax
+;;       cmpl    $0, %eax
+;;       jne     0x53
+;;   3d: movl    $0, %eax
 ;;       movq    0x60(%r14), %rcx
 ;;       addq    %rax, %rcx
 ;;       movl    (%rcx), %eax
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   40: ud2
+;;   51: ud2
+;;   53: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat b/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat
index 2ed82872b293..b6f98ecffc96 100644
--- a/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat
+++ b/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat
@@ -11,13 +11,17 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x4b
+;;       ja      0x5c
 ;;   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     0x5e
+;;   42: movl    $0, %ecx
 ;;       movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
@@ -26,4 +30,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   4b: ud2
+;;   5c: ud2
+;;   5e: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat b/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat
index d67b305dfbff..3b48acb6bcf6 100644
--- a/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat
+++ b/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat
@@ -11,13 +11,17 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x4c
+;;       ja      0x5f
 ;;   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     0x61
+;;   44: movl    $0, %ecx
 ;;       movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
@@ -26,4 +30,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   4c: ud2
+;;   5f: ud2
+;;   61: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat
index 172f90556ad9..9852a8983c31 100644
--- a/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat
@@ -11,13 +11,17 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x4e
+;;       ja      0x61
 ;;   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     0x63
+;;   46: movl    $0, %ecx
 ;;       movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
@@ -26,4 +30,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   4e: ud2
+;;   61: ud2
+;;   63: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat
index 320794a7cd5c..c13d460fb61b 100644
--- a/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat
@@ -11,13 +11,17 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x4e
+;;       ja      0x61
 ;;   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     0x63
+;;   46: movl    $0, %ecx
 ;;       movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
@@ -26,4 +30,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   4e: ud2
+;;   61: ud2
+;;   63: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat
index 08a2307ffb00..8a4d056aa004 100644
--- a/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat
@@ -11,13 +11,17 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x4d
+;;       ja      0x5e
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movq    $0x2a, %rax
 ;;       movl    $0, %ecx
+;;       andl    $3, %ecx
+;;       cmpl    $0, %ecx
+;;       jne     0x60
+;;   44: movl    $0, %ecx
 ;;       movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
@@ -26,4 +30,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   4d: ud2
+;;   5e: ud2
+;;   60: ud2
diff --git a/winch/codegen/src/codegen/mod.rs b/winch/codegen/src/codegen/mod.rs
index cd3c4154e022..16a1eddfbe30 100644
--- a/winch/codegen/src/codegen/mod.rs
+++ b/winch/codegen/src/codegen/mod.rs
@@ -3,8 +3,7 @@ use crate::{
     codegen::BlockSig,
     isa::reg::{writable, Reg},
     masm::{
-        IntCmpKind, LoadKind, MacroAssembler, MemOpKind, OperandSize, RegImm, SPOffset, ShiftKind,
-        TrapCode,
+        Imm, IntCmpKind, LoadKind, MacroAssembler, MemOpKind, OperandSize, RegImm, SPOffset, ShiftKind, TrapCode
     },
     stack::TypedReg,
 };
@@ -19,7 +18,7 @@ use wasmparser::{
     BinaryReader, FuncValidator, MemArg, Operator, ValidatorResources, VisitOperator,
     VisitSimdOperator,
 };
-use wasmtime_cranelift::{TRAP_BAD_SIGNATURE, TRAP_TABLE_OUT_OF_BOUNDS};
+use wasmtime_cranelift::{TRAP_BAD_SIGNATURE, TRAP_HEAP_MISALIGNED, TRAP_TABLE_OUT_OF_BOUNDS};
 use wasmtime_environ::{
     GlobalIndex, MemoryIndex, PtrSize, TableIndex, Tunables, TypeIndex, WasmHeapType, WasmValType,
     FUNCREF_MASK,
@@ -648,7 +647,12 @@ where
         &mut self,
         memarg: &MemArg,
         access_size: OperandSize,
+        check_align: bool,
     ) -> Result<Option<Reg>> {
+        if check_align {
+            self.check_align(memarg, access_size)?;
+        }
+
         let ptr_size: OperandSize = self.env.ptr_type().try_into()?;
         let enable_spectre_mitigation = self.env.heap_access_spectre_mitigation();
         let add_offset_and_access_size = |offset: ImmOffset, access_size: OperandSize| {
@@ -840,6 +844,40 @@ where
         Ok(addr)
     }
 
+    fn check_align(&mut self, memarg: &MemArg, size: OperandSize) -> Result<()> {
+        if size.bytes() > 1 {
+            let addr = *self.context.stack.peek().unwrap();
+            let effective_addr_reg = self.context.any_gpr(self.masm)?;
+            self.context
+                .move_val_to_reg(&addr, effective_addr_reg, self.masm)?;
+            if memarg.offset != 0 {
+                // self.masm.add(dst, lhs, rhs, size)
+                // self.context.builder.ins().iadd_imm(addr, memarg.offset.signed())
+                self.masm.add(
+                    writable!(effective_addr_reg),
+                    effective_addr_reg,
+                    RegImm::Imm(Imm::I64(memarg.offset)),
+                    size,
+                )?;
+            };
+            self.masm.and(
+                writable!(effective_addr_reg),
+                effective_addr_reg,
+                RegImm::Imm(Imm::I32(size.bytes() - 1)),
+                size,
+            )?;
+
+            self.masm
+                .cmp(effective_addr_reg, RegImm::Imm(Imm::i64(0)), size)?;
+            self.masm.trapif(IntCmpKind::Ne, TRAP_HEAP_MISALIGNED)?;
+
+            // environ.trapnz(builder, f, crate::TRAP_HEAP_MISALIGNED);
+            self.context.free_reg(effective_addr_reg);
+        }
+
+        Ok(())
+    }
+
     /// Emit a WebAssembly load.
     pub fn emit_wasm_load(
         &mut self,
@@ -875,7 +913,7 @@ where
         op_kind: MemOpKind,
     ) -> Result<()> {
         let src = self.context.pop_to_reg(self.masm, None)?;
-        let addr = self.emit_compute_heap_address(&arg, size)?;
+        let addr = self.emit_compute_heap_address(&arg, size, op_kind.is_atomic())?;
         if let Some(addr) = addr {
             self.masm.wasm_store(
                 src.reg.into(),
diff --git a/winch/codegen/src/masm.rs b/winch/codegen/src/masm.rs
index 18c99490aaf1..81200a1774dd 100644
--- a/winch/codegen/src/masm.rs
+++ b/winch/codegen/src/masm.rs
@@ -46,6 +46,12 @@ pub(crate) enum MemOpKind {
     Normal,
 }
 
+impl MemOpKind {
+    pub(crate) fn is_atomic(&self) -> bool {
+        matches!(self, Self::Atomic)
+    }
+}
+
 #[derive(Eq, PartialEq)]
 pub(crate) enum MulWideKind {
     Signed,

From 30cee24a79892e2f2ea903253ca5ad1efad606c6 Mon Sep 17 00:00:00 2001
From: adhoc <postma.marin@protonmail.com>
Date: Mon, 13 Jan 2025 21:07:34 +0100
Subject: [PATCH 11/17] fmt

---
 winch/codegen/src/codegen/mod.rs | 3 ++-
 winch/codegen/src/isa/x64/asm.rs | 5 ++++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/winch/codegen/src/codegen/mod.rs b/winch/codegen/src/codegen/mod.rs
index 16a1eddfbe30..868c67f9fd21 100644
--- a/winch/codegen/src/codegen/mod.rs
+++ b/winch/codegen/src/codegen/mod.rs
@@ -3,7 +3,8 @@ use crate::{
     codegen::BlockSig,
     isa::reg::{writable, Reg},
     masm::{
-        Imm, IntCmpKind, LoadKind, MacroAssembler, MemOpKind, OperandSize, RegImm, SPOffset, ShiftKind, TrapCode
+        Imm, IntCmpKind, LoadKind, MacroAssembler, MemOpKind, OperandSize, RegImm, SPOffset,
+        ShiftKind, TrapCode,
     },
     stack::TypedReg,
 };
diff --git a/winch/codegen/src/isa/x64/asm.rs b/winch/codegen/src/isa/x64/asm.rs
index ecf2e2b901a5..ff56ed658e5a 100644
--- a/winch/codegen/src/isa/x64/asm.rs
+++ b/winch/codegen/src/isa/x64/asm.rs
@@ -16,7 +16,10 @@ use cranelift_codegen::{
         unwind::UnwindInst,
         x64::{
             args::{
-                self, AluRmiROpcode, Amode, AvxOpcode, CmpOpcode, DivSignedness, ExtMode, FenceKind, FromWritableReg, Gpr, GprMem, GprMemImm, Imm8Gpr, Imm8Reg, RegMem, RegMemImm, ShiftKind as CraneliftShiftKind, SseOpcode, SyntheticAmode, WritableGpr, WritableXmm, Xmm, XmmMem, XmmMemAligned, CC
+                self, AluRmiROpcode, Amode, AvxOpcode, CmpOpcode, DivSignedness, ExtMode,
+                FenceKind, FromWritableReg, Gpr, GprMem, GprMemImm, Imm8Gpr, Imm8Reg, RegMem,
+                RegMemImm, ShiftKind as CraneliftShiftKind, SseOpcode, SyntheticAmode, WritableGpr,
+                WritableXmm, Xmm, XmmMem, XmmMemAligned, CC,
             },
             encoding::rex::{encode_modrm, RexFlags},
             settings as x64_settings, EmitInfo, EmitState, Inst,

From 4293461a61635ae4e329ad899ecd15a8be7240b1 Mon Sep 17 00:00:00 2001
From: adhoc <postma.marin@protonmail.com>
Date: Mon, 13 Jan 2025 21:09:07 +0100
Subject: [PATCH 12/17] edit comment

---
 winch/codegen/src/isa/x64/masm.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/winch/codegen/src/isa/x64/masm.rs b/winch/codegen/src/isa/x64/masm.rs
index c9ee4dfe86b1..2aef0f36a3a7 100644
--- a/winch/codegen/src/isa/x64/masm.rs
+++ b/winch/codegen/src/isa/x64/masm.rs
@@ -230,7 +230,7 @@ impl Masm for MacroAssembler {
                     // TODO: we don't support 128-bit atomic store yet.
                     bail!(CodeGenError::unexpected_operand_size());
                 }
-                // To stay consistent with cranelift, we emit a normal load followed by a mfence,
+                // To stay consistent with cranelift, we emit a normal store followed by a mfence,
                 // although, we could probably just emit a xchg.
                 self.store_impl(src.into(), dst, size, UNTRUSTED_FLAGS)?;
                 self.asm.fence(FenceKind::MFence);

From 5864b5fdb31de0cb0d3ba870a4ad7f5e97270479 Mon Sep 17 00:00:00 2001
From: adhoc <postma.marin@protonmail.com>
Date: Mon, 13 Jan 2025 21:20:06 +0100
Subject: [PATCH 13/17] cleanup check align function

---
 winch/codegen/src/codegen/mod.rs | 19 ++++++++-----------
 1 file changed, 8 insertions(+), 11 deletions(-)

diff --git a/winch/codegen/src/codegen/mod.rs b/winch/codegen/src/codegen/mod.rs
index 868c67f9fd21..11dfca8b54cc 100644
--- a/winch/codegen/src/codegen/mod.rs
+++ b/winch/codegen/src/codegen/mod.rs
@@ -848,32 +848,29 @@ where
     fn check_align(&mut self, memarg: &MemArg, size: OperandSize) -> Result<()> {
         if size.bytes() > 1 {
             let addr = *self.context.stack.peek().unwrap();
-            let effective_addr_reg = self.context.any_gpr(self.masm)?;
+            let tmp = self.context.any_gpr(self.masm)?;
             self.context
-                .move_val_to_reg(&addr, effective_addr_reg, self.masm)?;
+                .move_val_to_reg(&addr, tmp, self.masm)?;
             if memarg.offset != 0 {
-                // self.masm.add(dst, lhs, rhs, size)
-                // self.context.builder.ins().iadd_imm(addr, memarg.offset.signed())
                 self.masm.add(
-                    writable!(effective_addr_reg),
-                    effective_addr_reg,
+                    writable!(tmp),
+                    tmp,
                     RegImm::Imm(Imm::I64(memarg.offset)),
                     size,
                 )?;
             };
             self.masm.and(
-                writable!(effective_addr_reg),
-                effective_addr_reg,
+                writable!(tmp),
+                tmp,
                 RegImm::Imm(Imm::I32(size.bytes() - 1)),
                 size,
             )?;
 
             self.masm
-                .cmp(effective_addr_reg, RegImm::Imm(Imm::i64(0)), size)?;
+                .cmp(tmp, RegImm::Imm(Imm::i64(0)), size)?;
             self.masm.trapif(IntCmpKind::Ne, TRAP_HEAP_MISALIGNED)?;
 
-            // environ.trapnz(builder, f, crate::TRAP_HEAP_MISALIGNED);
-            self.context.free_reg(effective_addr_reg);
+            self.context.free_reg(tmp);
         }
 
         Ok(())

From f7044efc38365a3ea5d21f9ea92bbc20af9c4650 Mon Sep 17 00:00:00 2001
From: adhoc <postma.marin@protonmail.com>
Date: Mon, 13 Jan 2025 21:26:49 +0100
Subject: [PATCH 14/17] add doc for check_align

---
 winch/codegen/src/codegen/mod.rs | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/winch/codegen/src/codegen/mod.rs b/winch/codegen/src/codegen/mod.rs
index 11dfca8b54cc..8d4061cdd042 100644
--- a/winch/codegen/src/codegen/mod.rs
+++ b/winch/codegen/src/codegen/mod.rs
@@ -845,12 +845,12 @@ where
         Ok(addr)
     }
 
+    /// Emit checks to ensure that the address at `memarg` is correctly aligned for `size`.
     fn check_align(&mut self, memarg: &MemArg, size: OperandSize) -> Result<()> {
         if size.bytes() > 1 {
             let addr = *self.context.stack.peek().unwrap();
             let tmp = self.context.any_gpr(self.masm)?;
-            self.context
-                .move_val_to_reg(&addr, tmp, self.masm)?;
+            self.context.move_val_to_reg(&addr, tmp, self.masm)?;
             if memarg.offset != 0 {
                 self.masm.add(
                     writable!(tmp),
@@ -866,8 +866,7 @@ where
                 size,
             )?;
 
-            self.masm
-                .cmp(tmp, RegImm::Imm(Imm::i64(0)), size)?;
+            self.masm.cmp(tmp, RegImm::Imm(Imm::i64(0)), size)?;
             self.masm.trapif(IntCmpKind::Ne, TRAP_HEAP_MISALIGNED)?;
 
             self.context.free_reg(tmp);

From 54f24d73a7600651aa62d44a7daaa182ce3bddad Mon Sep 17 00:00:00 2001
From: ad hoc <postma.marin@protonmail.com>
Date: Tue, 14 Jan 2025 15:54:12 +0100
Subject: [PATCH 15/17] review edits

---
 .../winch/x64/atomic/load/i32_atomic_load.wat | 13 +++--
 .../x64/atomic/load/i32_atomic_load16_u.wat   | 13 +++--
 .../winch/x64/atomic/load/i64_atomic_load.wat | 15 +++---
 .../x64/atomic/load/i64_atomic_load16_u.wat   | 15 +++---
 .../x64/atomic/load/i64_atomic_load32_u.wat   | 15 +++---
 .../x64/atomic/store/i32_atomic_store.wat     | 15 +++---
 .../x64/atomic/store/i32_atomic_store16.wat   | 15 +++---
 .../x64/atomic/store/i64_atomic_store.wat     | 15 +++---
 .../x64/atomic/store/i64_atomic_store16.wat   | 15 +++---
 .../x64/atomic/store/i64_atomic_store32.wat   | 15 +++---
 winch/codegen/src/codegen/mod.rs              | 51 ++++++++++++-------
 winch/codegen/src/masm.rs                     |  6 ---
 12 files changed, 102 insertions(+), 101 deletions(-)

diff --git a/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat b/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat
index a671002d8a51..fb6e1f72dc06 100644
--- a/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat
+++ b/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat
@@ -13,23 +13,22 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x20, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x57
+;;       ja      0x55
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x20, %rsp
 ;;       movq    %rdi, 0x18(%rsp)
 ;;       movq    %rsi, 0x10(%rsp)
 ;;       movl    %edx, 0xc(%rsp)
 ;;       movl    0xc(%rsp), %eax
-;;       andl    $3, %eax
-;;       cmpl    $0, %eax
-;;       jne     0x59
-;;   41: movl    0xc(%rsp), %eax
-;;       movq    0x58(%r14), %r11
+;;       andl    $3, %r11d
+;;       cmpl    $0, %r11d
+;;       jne     0x57
+;;   43: movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rcx
 ;;       addq    %rax, %rcx
 ;;       movl    (%rcx), %eax
 ;;       addq    $0x20, %rsp
 ;;       popq    %rbp
 ;;       retq
+;;   55: ud2
 ;;   57: ud2
-;;   59: ud2
diff --git a/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat b/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat
index 3c156158518b..2e4f28ef8a62 100644
--- a/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat
+++ b/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat
@@ -13,23 +13,22 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x20, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x5b
+;;       ja      0x59
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x20, %rsp
 ;;       movq    %rdi, 0x18(%rsp)
 ;;       movq    %rsi, 0x10(%rsp)
 ;;       movl    %edx, 0xc(%rsp)
 ;;       movl    0xc(%rsp), %eax
-;;       andw    $1, %ax
-;;       cmpw    $0, %ax
-;;       jne     0x5d
-;;   43: movl    0xc(%rsp), %eax
-;;       movq    0x58(%r14), %r11
+;;       andw    $1, %r11w
+;;       cmpw    $0, %r11w
+;;       jne     0x5b
+;;   45: movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rcx
 ;;       addq    %rax, %rcx
 ;;       movzwq  (%rcx), %rax
 ;;       addq    $0x20, %rsp
 ;;       popq    %rbp
 ;;       retq
+;;   59: ud2
 ;;   5b: ud2
-;;   5d: ud2
diff --git a/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat b/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat
index 6d67f5534ba1..e8cbbfb071cc 100644
--- a/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat
+++ b/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat
@@ -14,21 +14,20 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x54
+;;       ja      0x4f
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0, %eax
-;;       andq    $7, %rax
-;;       cmpq    $0, %rax
-;;       jne     0x56
-;;   3f: movl    $0, %eax
-;;       movq    0x60(%r14), %rcx
+;;       andq    $7, %r11
+;;       cmpq    $0, %r11
+;;       jne     0x51
+;;   3f: movq    0x60(%r14), %rcx
 ;;       addq    %rax, %rcx
 ;;       movq    (%rcx), %rax
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   54: ud2
-;;   56: ud2
+;;   4f: ud2
+;;   51: ud2
diff --git a/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat b/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat
index dbd25c81996a..1c62828769db 100644
--- a/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat
+++ b/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat
@@ -14,21 +14,20 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x55
+;;       ja      0x52
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0, %eax
-;;       andw    $1, %ax
-;;       cmpw    $0, %ax
-;;       jne     0x57
-;;   3f: movl    $0, %eax
-;;       movq    0x60(%r14), %rcx
+;;       andw    $1, %r11w
+;;       cmpw    $0, %r11w
+;;       jne     0x54
+;;   41: movq    0x60(%r14), %rcx
 ;;       addq    %rax, %rcx
 ;;       movzwq  (%rcx), %rax
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   55: ud2
-;;   57: ud2
+;;   52: ud2
+;;   54: ud2
diff --git a/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat b/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat
index a7ad511e0d37..88b0e886b3e8 100644
--- a/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat
+++ b/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat
@@ -14,21 +14,20 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x51
+;;       ja      0x4e
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0, %eax
-;;       andl    $3, %eax
-;;       cmpl    $0, %eax
-;;       jne     0x53
-;;   3d: movl    $0, %eax
-;;       movq    0x60(%r14), %rcx
+;;       andl    $3, %r11d
+;;       cmpl    $0, %r11d
+;;       jne     0x50
+;;   3f: movq    0x60(%r14), %rcx
 ;;       addq    %rax, %rcx
 ;;       movl    (%rcx), %eax
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   51: ud2
-;;   53: ud2
+;;   4e: ud2
+;;   50: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat b/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat
index b6f98ecffc96..f664c721b7c8 100644
--- a/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat
+++ b/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat
@@ -11,18 +11,17 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x5c
+;;       ja      0x59
 ;;   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     0x5e
-;;   42: movl    $0, %ecx
-;;       movq    0x58(%r14), %r11
+;;       andl    $3, %r11d
+;;       cmpl    $0, %r11d
+;;       jne     0x5b
+;;   44: movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
 ;;       movl    %eax, (%rdx)
@@ -30,5 +29,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   5c: ud2
-;;   5e: ud2
+;;   59: ud2
+;;   5b: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat b/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat
index 3b48acb6bcf6..56b7933cae63 100644
--- a/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat
+++ b/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat
@@ -11,18 +11,17 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x5f
+;;       ja      0x5c
 ;;   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     0x61
-;;   44: movl    $0, %ecx
-;;       movq    0x58(%r14), %r11
+;;       andw    $1, %r11w
+;;       cmpw    $0, %r11w
+;;       jne     0x5e
+;;   46: movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
 ;;       movw    %ax, (%rdx)
@@ -30,5 +29,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   5f: ud2
-;;   61: ud2
+;;   5c: ud2
+;;   5e: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat
index 9852a8983c31..1aa168325a8f 100644
--- a/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat
@@ -11,18 +11,17 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x61
+;;       ja      0x5c
 ;;   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     0x63
-;;   46: movl    $0, %ecx
-;;       movq    0x58(%r14), %r11
+;;       andq    $7, %r11
+;;       cmpq    $0, %r11
+;;       jne     0x5e
+;;   46: movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
 ;;       movq    %rax, (%rdx)
@@ -30,5 +29,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   61: ud2
-;;   63: ud2
+;;   5c: ud2
+;;   5e: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat
index c13d460fb61b..0699c95cb3c6 100644
--- a/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat
@@ -11,18 +11,17 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x61
+;;       ja      0x5e
 ;;   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     0x63
-;;   46: movl    $0, %ecx
-;;       movq    0x58(%r14), %r11
+;;       andw    $1, %r11w
+;;       cmpw    $0, %r11w
+;;       jne     0x60
+;;   48: movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
 ;;       movw    %ax, (%rdx)
@@ -30,5 +29,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   61: ud2
-;;   63: ud2
+;;   5e: ud2
+;;   60: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat
index 8a4d056aa004..8278d67d57a5 100644
--- a/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat
@@ -11,18 +11,17 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x5e
+;;       ja      0x5b
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movq    $0x2a, %rax
 ;;       movl    $0, %ecx
-;;       andl    $3, %ecx
-;;       cmpl    $0, %ecx
-;;       jne     0x60
-;;   44: movl    $0, %ecx
-;;       movq    0x58(%r14), %r11
+;;       andl    $3, %r11d
+;;       cmpl    $0, %r11d
+;;       jne     0x5d
+;;   46: movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
 ;;       movl    %eax, (%rdx)
@@ -30,5 +29,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   5e: ud2
-;;   60: ud2
+;;   5b: ud2
+;;   5d: ud2
diff --git a/winch/codegen/src/codegen/mod.rs b/winch/codegen/src/codegen/mod.rs
index 8d4061cdd042..057b901049c5 100644
--- a/winch/codegen/src/codegen/mod.rs
+++ b/winch/codegen/src/codegen/mod.rs
@@ -648,12 +648,7 @@ where
         &mut self,
         memarg: &MemArg,
         access_size: OperandSize,
-        check_align: bool,
     ) -> Result<Option<Reg>> {
-        if check_align {
-            self.check_align(memarg, access_size)?;
-        }
-
         let ptr_size: OperandSize = self.env.ptr_type().try_into()?;
         let enable_spectre_mitigation = self.env.heap_access_spectre_mitigation();
         let add_offset_and_access_size = |offset: ImmOffset, access_size: OperandSize| {
@@ -846,22 +841,25 @@ where
     }
 
     /// Emit checks to ensure that the address at `memarg` is correctly aligned for `size`.
-    fn check_align(&mut self, memarg: &MemArg, size: OperandSize) -> Result<()> {
+    fn emit_check_align(&mut self, memarg: &MemArg, size: OperandSize) -> Result<()> {
         if size.bytes() > 1 {
-            let addr = *self.context.stack.peek().unwrap();
-            let tmp = self.context.any_gpr(self.masm)?;
-            self.context.move_val_to_reg(&addr, tmp, self.masm)?;
-            if memarg.offset != 0 {
+            let addr = self.context.pop_to_reg(self.masm, None)?;
+            let tmp = scratch!(M);
+            let effective_addr = if memarg.offset != 0 {
                 self.masm.add(
                     writable!(tmp),
-                    tmp,
+                    addr.reg,
                     RegImm::Imm(Imm::I64(memarg.offset)),
                     size,
                 )?;
+                tmp
+            } else {
+                addr.reg
             };
+
             self.masm.and(
                 writable!(tmp),
-                tmp,
+                effective_addr,
                 RegImm::Imm(Imm::I32(size.bytes() - 1)),
                 size,
             )?;
@@ -869,12 +867,21 @@ where
             self.masm.cmp(tmp, RegImm::Imm(Imm::i64(0)), size)?;
             self.masm.trapif(IntCmpKind::Ne, TRAP_HEAP_MISALIGNED)?;
 
-            self.context.free_reg(tmp);
+            self.context.stack.push(addr.into());
         }
 
         Ok(())
     }
 
+    pub fn emit_compute_heap_address_align_checked(
+        &mut self,
+        memarg: &MemArg,
+        access_size: OperandSize,
+    ) -> Result<Option<Reg>> {
+        self.emit_check_align(memarg, access_size)?;
+        self.emit_compute_heap_address(memarg, access_size)
+    }
+
     /// Emit a WebAssembly load.
     pub fn emit_wasm_load(
         &mut self,
@@ -883,8 +890,13 @@ where
         kind: LoadKind,
         op_kind: MemOpKind,
     ) -> Result<()> {
-        if let Some(addr) = self.emit_compute_heap_address(&arg, kind.derive_operand_size())? {
-            let dst = match target_type {
+        let maybe_addr = match op_kind {
+            MemOpKind::Atomic => self.emit_compute_heap_address_align_checked(&arg, kind.derive_operand_size())?,
+            MemOpKind::Normal => self.emit_compute_heap_address(&arg, kind.derive_operand_size())?,
+        };
+
+        if let Some(addr) = maybe_addr {
+            let dst = match target_ty {
                 WasmValType::I32 | WasmValType::I64 => self.context.any_gpr(self.masm)?,
                 WasmValType::F32 | WasmValType::F64 => self.context.any_fpr(self.masm)?,
                 WasmValType::V128 => self.context.reg_for_type(target_type, self.masm)?,
@@ -910,8 +922,13 @@ where
         op_kind: MemOpKind,
     ) -> Result<()> {
         let src = self.context.pop_to_reg(self.masm, None)?;
-        let addr = self.emit_compute_heap_address(&arg, size, op_kind.is_atomic())?;
-        if let Some(addr) = addr {
+
+        let maybe_addr = match op_kind {
+            MemOpKind::Atomic => self.emit_compute_heap_address_align_checked(&arg, size)?,
+            MemOpKind::Normal => self.emit_compute_heap_address(&arg, size)?,
+        };
+
+        if let Some(addr) = maybe_addr {
             self.masm.wasm_store(
                 src.reg.into(),
                 self.masm.address_at_reg(addr, 0)?,
diff --git a/winch/codegen/src/masm.rs b/winch/codegen/src/masm.rs
index 81200a1774dd..18c99490aaf1 100644
--- a/winch/codegen/src/masm.rs
+++ b/winch/codegen/src/masm.rs
@@ -46,12 +46,6 @@ pub(crate) enum MemOpKind {
     Normal,
 }
 
-impl MemOpKind {
-    pub(crate) fn is_atomic(&self) -> bool {
-        matches!(self, Self::Atomic)
-    }
-}
-
 #[derive(Eq, PartialEq)]
 pub(crate) enum MulWideKind {
     Signed,

From e16d7d8485080b349e8eb42f6c823ddbc69d51e7 Mon Sep 17 00:00:00 2001
From: adhoc <postma.marin@protonmail.com>
Date: Tue, 14 Jan 2025 18:30:08 +0100
Subject: [PATCH 16/17] review edits - use scratch register for tmp in
 emit_align_check - pop-push value from stack rather than peeking in
 emit_align_check

---
 .../winch/x64/atomic/load/i32_atomic_load.wat  | 11 ++++++-----
 .../x64/atomic/load/i32_atomic_load16_u.wat    | 11 ++++++-----
 .../winch/x64/atomic/load/i64_atomic_load.wat  | 11 ++++++-----
 .../x64/atomic/load/i64_atomic_load16_u.wat    | 11 ++++++-----
 .../x64/atomic/load/i64_atomic_load32_u.wat    | 11 ++++++-----
 .../x64/atomic/store/i32_atomic_store.wat      | 11 ++++++-----
 .../x64/atomic/store/i32_atomic_store16.wat    | 11 ++++++-----
 .../x64/atomic/store/i64_atomic_store.wat      | 11 ++++++-----
 .../x64/atomic/store/i64_atomic_store16.wat    | 11 ++++++-----
 .../x64/atomic/store/i64_atomic_store32.wat    | 11 ++++++-----
 winch/codegen/src/codegen/mod.rs               | 18 +++++++++---------
 11 files changed, 69 insertions(+), 59 deletions(-)

diff --git a/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat b/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat
index fb6e1f72dc06..0c7553d46556 100644
--- a/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat
+++ b/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat
@@ -13,22 +13,23 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x20, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x55
+;;       ja      0x58
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x20, %rsp
 ;;       movq    %rdi, 0x18(%rsp)
 ;;       movq    %rsi, 0x10(%rsp)
 ;;       movl    %edx, 0xc(%rsp)
 ;;       movl    0xc(%rsp), %eax
+;;       movl    %eax, %r11d
 ;;       andl    $3, %r11d
 ;;       cmpl    $0, %r11d
-;;       jne     0x57
-;;   43: movq    0x58(%r14), %r11
+;;       jne     0x5a
+;;   46: movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rcx
 ;;       addq    %rax, %rcx
 ;;       movl    (%rcx), %eax
 ;;       addq    $0x20, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   55: ud2
-;;   57: ud2
+;;   58: ud2
+;;   5a: ud2
diff --git a/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat b/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat
index 2e4f28ef8a62..0d53faa40058 100644
--- a/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat
+++ b/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat
@@ -13,22 +13,23 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x20, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x59
+;;       ja      0x5c
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x20, %rsp
 ;;       movq    %rdi, 0x18(%rsp)
 ;;       movq    %rsi, 0x10(%rsp)
 ;;       movl    %edx, 0xc(%rsp)
 ;;       movl    0xc(%rsp), %eax
+;;       movl    %eax, %r11d
 ;;       andw    $1, %r11w
 ;;       cmpw    $0, %r11w
-;;       jne     0x5b
-;;   45: movq    0x58(%r14), %r11
+;;       jne     0x5e
+;;   48: movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rcx
 ;;       addq    %rax, %rcx
 ;;       movzwq  (%rcx), %rax
 ;;       addq    $0x20, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   59: ud2
-;;   5b: ud2
+;;   5c: ud2
+;;   5e: ud2
diff --git a/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat b/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat
index e8cbbfb071cc..1259be84e979 100644
--- a/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat
+++ b/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat
@@ -14,20 +14,21 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x4f
+;;       ja      0x52
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0, %eax
+;;       movq    %rax, %r11
 ;;       andq    $7, %r11
 ;;       cmpq    $0, %r11
-;;       jne     0x51
-;;   3f: movq    0x60(%r14), %rcx
+;;       jne     0x54
+;;   42: movq    0x60(%r14), %rcx
 ;;       addq    %rax, %rcx
 ;;       movq    (%rcx), %rax
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   4f: ud2
-;;   51: ud2
+;;   52: ud2
+;;   54: ud2
diff --git a/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat b/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat
index 1c62828769db..dbfd26f07c2c 100644
--- a/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat
+++ b/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat
@@ -14,20 +14,21 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x52
+;;       ja      0x55
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0, %eax
+;;       movl    %eax, %r11d
 ;;       andw    $1, %r11w
 ;;       cmpw    $0, %r11w
-;;       jne     0x54
-;;   41: movq    0x60(%r14), %rcx
+;;       jne     0x57
+;;   44: movq    0x60(%r14), %rcx
 ;;       addq    %rax, %rcx
 ;;       movzwq  (%rcx), %rax
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   52: ud2
-;;   54: ud2
+;;   55: ud2
+;;   57: ud2
diff --git a/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat b/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat
index 88b0e886b3e8..4d647a313f6d 100644
--- a/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat
+++ b/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat
@@ -14,20 +14,21 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x4e
+;;       ja      0x51
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0, %eax
+;;       movl    %eax, %r11d
 ;;       andl    $3, %r11d
 ;;       cmpl    $0, %r11d
-;;       jne     0x50
-;;   3f: movq    0x60(%r14), %rcx
+;;       jne     0x53
+;;   42: movq    0x60(%r14), %rcx
 ;;       addq    %rax, %rcx
 ;;       movl    (%rcx), %eax
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   4e: ud2
-;;   50: ud2
+;;   51: ud2
+;;   53: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat b/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat
index f664c721b7c8..a7750ed52028 100644
--- a/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat
+++ b/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat
@@ -11,17 +11,18 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x59
+;;       ja      0x5c
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0x2a, %eax
 ;;       movl    $0, %ecx
+;;       movl    %ecx, %r11d
 ;;       andl    $3, %r11d
 ;;       cmpl    $0, %r11d
-;;       jne     0x5b
-;;   44: movq    0x58(%r14), %r11
+;;       jne     0x5e
+;;   47: movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
 ;;       movl    %eax, (%rdx)
@@ -29,5 +30,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   59: ud2
-;;   5b: ud2
+;;   5c: ud2
+;;   5e: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat b/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat
index 56b7933cae63..5e77e8cb70fd 100644
--- a/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat
+++ b/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat
@@ -11,17 +11,18 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x5c
+;;       ja      0x5f
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0x2a, %eax
 ;;       movl    $0, %ecx
+;;       movl    %ecx, %r11d
 ;;       andw    $1, %r11w
 ;;       cmpw    $0, %r11w
-;;       jne     0x5e
-;;   46: movq    0x58(%r14), %r11
+;;       jne     0x61
+;;   49: movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
 ;;       movw    %ax, (%rdx)
@@ -29,5 +30,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   5c: ud2
-;;   5e: ud2
+;;   5f: ud2
+;;   61: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat
index 1aa168325a8f..2943236bc1a9 100644
--- a/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat
@@ -11,17 +11,18 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x5c
+;;       ja      0x5f
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movq    $0x2a, %rax
 ;;       movl    $0, %ecx
+;;       movq    %rcx, %r11
 ;;       andq    $7, %r11
 ;;       cmpq    $0, %r11
-;;       jne     0x5e
-;;   46: movq    0x58(%r14), %r11
+;;       jne     0x61
+;;   49: movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
 ;;       movq    %rax, (%rdx)
@@ -29,5 +30,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   5c: ud2
-;;   5e: ud2
+;;   5f: ud2
+;;   61: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat
index 0699c95cb3c6..314cdb9d0c0f 100644
--- a/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat
@@ -11,17 +11,18 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x5e
+;;       ja      0x61
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movq    $0x2a, %rax
 ;;       movl    $0, %ecx
+;;       movl    %ecx, %r11d
 ;;       andw    $1, %r11w
 ;;       cmpw    $0, %r11w
-;;       jne     0x60
-;;   48: movq    0x58(%r14), %r11
+;;       jne     0x63
+;;   4b: movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
 ;;       movw    %ax, (%rdx)
@@ -29,5 +30,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   5e: ud2
-;;   60: ud2
+;;   61: ud2
+;;   63: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat
index 8278d67d57a5..e6d802614a3c 100644
--- a/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat
@@ -11,17 +11,18 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x5b
+;;       ja      0x5e
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movq    $0x2a, %rax
 ;;       movl    $0, %ecx
+;;       movl    %ecx, %r11d
 ;;       andl    $3, %r11d
 ;;       cmpl    $0, %r11d
-;;       jne     0x5d
-;;   46: movq    0x58(%r14), %r11
+;;       jne     0x60
+;;   49: movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
 ;;       movl    %eax, (%rdx)
@@ -29,5 +30,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   5b: ud2
-;;   5d: ud2
+;;   5e: ud2
+;;   60: ud2
diff --git a/winch/codegen/src/codegen/mod.rs b/winch/codegen/src/codegen/mod.rs
index 057b901049c5..ac678410419b 100644
--- a/winch/codegen/src/codegen/mod.rs
+++ b/winch/codegen/src/codegen/mod.rs
@@ -843,31 +843,31 @@ where
     /// Emit checks to ensure that the address at `memarg` is correctly aligned for `size`.
     fn emit_check_align(&mut self, memarg: &MemArg, size: OperandSize) -> Result<()> {
         if size.bytes() > 1 {
+            // Peek addr from top of the stack by popping and pushing.
             let addr = self.context.pop_to_reg(self.masm, None)?;
+            self.context.stack.push(addr.into());
+
             let tmp = scratch!(M);
-            let effective_addr = if memarg.offset != 0 {
+            self.masm.mov(writable!(tmp), RegImm::Reg(addr.reg), size)?;
+
+            if memarg.offset != 0 {
                 self.masm.add(
                     writable!(tmp),
-                    addr.reg,
+                    tmp,
                     RegImm::Imm(Imm::I64(memarg.offset)),
                     size,
                 )?;
-                tmp
-            } else {
-                addr.reg
-            };
+            }
 
             self.masm.and(
                 writable!(tmp),
-                effective_addr,
+                tmp,
                 RegImm::Imm(Imm::I32(size.bytes() - 1)),
                 size,
             )?;
 
             self.masm.cmp(tmp, RegImm::Imm(Imm::i64(0)), size)?;
             self.masm.trapif(IntCmpKind::Ne, TRAP_HEAP_MISALIGNED)?;
-
-            self.context.stack.push(addr.into());
         }
 
         Ok(())

From cd24682db92a6b6b3ba5972f5d5dfb374cd6daf3 Mon Sep 17 00:00:00 2001
From: adhoc <postma.marin@protonmail.com>
Date: Wed, 15 Jan 2025 13:44:46 +0100
Subject: [PATCH 17/17] final edits

---
 .../winch/x64/atomic/load/i32_atomic_load.wat | 16 ++++++-------
 .../x64/atomic/load/i32_atomic_load16_u.wat   | 16 ++++++-------
 .../winch/x64/atomic/load/i64_atomic_load.wat | 14 +++++------
 .../x64/atomic/load/i64_atomic_load16_u.wat   |  8 +++----
 .../x64/atomic/load/i64_atomic_load32_u.wat   |  8 +++----
 .../x64/atomic/store/i32_atomic_store.wat     |  8 +++----
 .../x64/atomic/store/i32_atomic_store16.wat   |  8 +++----
 .../x64/atomic/store/i64_atomic_store.wat     | 14 +++++------
 .../x64/atomic/store/i64_atomic_store16.wat   |  8 +++----
 .../x64/atomic/store/i64_atomic_store32.wat   |  8 +++----
 winch/codegen/src/codegen/mod.rs              | 23 ++++++++++++-------
 11 files changed, 69 insertions(+), 62 deletions(-)

diff --git a/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat b/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat
index 0c7553d46556..a671002d8a51 100644
--- a/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat
+++ b/tests/disas/winch/x64/atomic/load/i32_atomic_load.wat
@@ -13,23 +13,23 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x20, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x58
+;;       ja      0x57
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x20, %rsp
 ;;       movq    %rdi, 0x18(%rsp)
 ;;       movq    %rsi, 0x10(%rsp)
 ;;       movl    %edx, 0xc(%rsp)
 ;;       movl    0xc(%rsp), %eax
-;;       movl    %eax, %r11d
-;;       andl    $3, %r11d
-;;       cmpl    $0, %r11d
-;;       jne     0x5a
-;;   46: movq    0x58(%r14), %r11
+;;       andl    $3, %eax
+;;       cmpl    $0, %eax
+;;       jne     0x59
+;;   41: movl    0xc(%rsp), %eax
+;;       movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rcx
 ;;       addq    %rax, %rcx
 ;;       movl    (%rcx), %eax
 ;;       addq    $0x20, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   58: ud2
-;;   5a: ud2
+;;   57: ud2
+;;   59: ud2
diff --git a/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat b/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat
index 0d53faa40058..3c156158518b 100644
--- a/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat
+++ b/tests/disas/winch/x64/atomic/load/i32_atomic_load16_u.wat
@@ -13,23 +13,23 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x20, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x5c
+;;       ja      0x5b
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x20, %rsp
 ;;       movq    %rdi, 0x18(%rsp)
 ;;       movq    %rsi, 0x10(%rsp)
 ;;       movl    %edx, 0xc(%rsp)
 ;;       movl    0xc(%rsp), %eax
-;;       movl    %eax, %r11d
-;;       andw    $1, %r11w
-;;       cmpw    $0, %r11w
-;;       jne     0x5e
-;;   48: movq    0x58(%r14), %r11
+;;       andw    $1, %ax
+;;       cmpw    $0, %ax
+;;       jne     0x5d
+;;   43: movl    0xc(%rsp), %eax
+;;       movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rcx
 ;;       addq    %rax, %rcx
 ;;       movzwq  (%rcx), %rax
 ;;       addq    $0x20, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   5c: ud2
-;;   5e: ud2
+;;   5b: ud2
+;;   5d: ud2
diff --git a/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat b/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat
index 1259be84e979..6d67f5534ba1 100644
--- a/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat
+++ b/tests/disas/winch/x64/atomic/load/i64_atomic_load.wat
@@ -14,21 +14,21 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x52
+;;       ja      0x54
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0, %eax
-;;       movq    %rax, %r11
-;;       andq    $7, %r11
-;;       cmpq    $0, %r11
-;;       jne     0x54
-;;   42: movq    0x60(%r14), %rcx
+;;       andq    $7, %rax
+;;       cmpq    $0, %rax
+;;       jne     0x56
+;;   3f: movl    $0, %eax
+;;       movq    0x60(%r14), %rcx
 ;;       addq    %rax, %rcx
 ;;       movq    (%rcx), %rax
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   52: ud2
 ;;   54: ud2
+;;   56: ud2
diff --git a/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat b/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat
index dbfd26f07c2c..dbd25c81996a 100644
--- a/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat
+++ b/tests/disas/winch/x64/atomic/load/i64_atomic_load16_u.wat
@@ -20,11 +20,11 @@
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0, %eax
-;;       movl    %eax, %r11d
-;;       andw    $1, %r11w
-;;       cmpw    $0, %r11w
+;;       andw    $1, %ax
+;;       cmpw    $0, %ax
 ;;       jne     0x57
-;;   44: movq    0x60(%r14), %rcx
+;;   3f: movl    $0, %eax
+;;       movq    0x60(%r14), %rcx
 ;;       addq    %rax, %rcx
 ;;       movzwq  (%rcx), %rax
 ;;       addq    $0x10, %rsp
diff --git a/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat b/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat
index 4d647a313f6d..a7ad511e0d37 100644
--- a/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat
+++ b/tests/disas/winch/x64/atomic/load/i64_atomic_load32_u.wat
@@ -20,11 +20,11 @@
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0, %eax
-;;       movl    %eax, %r11d
-;;       andl    $3, %r11d
-;;       cmpl    $0, %r11d
+;;       andl    $3, %eax
+;;       cmpl    $0, %eax
 ;;       jne     0x53
-;;   42: movq    0x60(%r14), %rcx
+;;   3d: movl    $0, %eax
+;;       movq    0x60(%r14), %rcx
 ;;       addq    %rax, %rcx
 ;;       movl    (%rcx), %eax
 ;;       addq    $0x10, %rsp
diff --git a/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat b/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat
index a7750ed52028..b6f98ecffc96 100644
--- a/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat
+++ b/tests/disas/winch/x64/atomic/store/i32_atomic_store.wat
@@ -18,11 +18,11 @@
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0x2a, %eax
 ;;       movl    $0, %ecx
-;;       movl    %ecx, %r11d
-;;       andl    $3, %r11d
-;;       cmpl    $0, %r11d
+;;       andl    $3, %ecx
+;;       cmpl    $0, %ecx
 ;;       jne     0x5e
-;;   47: movq    0x58(%r14), %r11
+;;   42: movl    $0, %ecx
+;;       movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
 ;;       movl    %eax, (%rdx)
diff --git a/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat b/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat
index 5e77e8cb70fd..3b48acb6bcf6 100644
--- a/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat
+++ b/tests/disas/winch/x64/atomic/store/i32_atomic_store16.wat
@@ -18,11 +18,11 @@
 ;;       movq    %rsi, (%rsp)
 ;;       movl    $0x2a, %eax
 ;;       movl    $0, %ecx
-;;       movl    %ecx, %r11d
-;;       andw    $1, %r11w
-;;       cmpw    $0, %r11w
+;;       andw    $1, %cx
+;;       cmpw    $0, %cx
 ;;       jne     0x61
-;;   49: movq    0x58(%r14), %r11
+;;   44: movl    $0, %ecx
+;;       movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
 ;;       movw    %ax, (%rdx)
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat
index 2943236bc1a9..9852a8983c31 100644
--- a/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store.wat
@@ -11,18 +11,18 @@
 ;;       movq    0x10(%r11), %r11
 ;;       addq    $0x10, %r11
 ;;       cmpq    %rsp, %r11
-;;       ja      0x5f
+;;       ja      0x61
 ;;   1c: movq    %rdi, %r14
 ;;       subq    $0x10, %rsp
 ;;       movq    %rdi, 8(%rsp)
 ;;       movq    %rsi, (%rsp)
 ;;       movq    $0x2a, %rax
 ;;       movl    $0, %ecx
-;;       movq    %rcx, %r11
-;;       andq    $7, %r11
-;;       cmpq    $0, %r11
-;;       jne     0x61
-;;   49: movq    0x58(%r14), %r11
+;;       andq    $7, %rcx
+;;       cmpq    $0, %rcx
+;;       jne     0x63
+;;   46: movl    $0, %ecx
+;;       movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
 ;;       movq    %rax, (%rdx)
@@ -30,5 +30,5 @@
 ;;       addq    $0x10, %rsp
 ;;       popq    %rbp
 ;;       retq
-;;   5f: ud2
 ;;   61: ud2
+;;   63: ud2
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat
index 314cdb9d0c0f..c13d460fb61b 100644
--- a/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store16.wat
@@ -18,11 +18,11 @@
 ;;       movq    %rsi, (%rsp)
 ;;       movq    $0x2a, %rax
 ;;       movl    $0, %ecx
-;;       movl    %ecx, %r11d
-;;       andw    $1, %r11w
-;;       cmpw    $0, %r11w
+;;       andw    $1, %cx
+;;       cmpw    $0, %cx
 ;;       jne     0x63
-;;   4b: movq    0x58(%r14), %r11
+;;   46: movl    $0, %ecx
+;;       movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
 ;;       movw    %ax, (%rdx)
diff --git a/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat b/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat
index e6d802614a3c..8a4d056aa004 100644
--- a/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat
+++ b/tests/disas/winch/x64/atomic/store/i64_atomic_store32.wat
@@ -18,11 +18,11 @@
 ;;       movq    %rsi, (%rsp)
 ;;       movq    $0x2a, %rax
 ;;       movl    $0, %ecx
-;;       movl    %ecx, %r11d
-;;       andl    $3, %r11d
-;;       cmpl    $0, %r11d
+;;       andl    $3, %ecx
+;;       cmpl    $0, %ecx
 ;;       jne     0x60
-;;   49: movq    0x58(%r14), %r11
+;;   44: movl    $0, %ecx
+;;       movq    0x58(%r14), %r11
 ;;       movq    (%r11), %rdx
 ;;       addq    %rcx, %rdx
 ;;       movl    %eax, (%rdx)
diff --git a/winch/codegen/src/codegen/mod.rs b/winch/codegen/src/codegen/mod.rs
index ac678410419b..4b05d115c4de 100644
--- a/winch/codegen/src/codegen/mod.rs
+++ b/winch/codegen/src/codegen/mod.rs
@@ -844,11 +844,13 @@ where
     fn emit_check_align(&mut self, memarg: &MemArg, size: OperandSize) -> Result<()> {
         if size.bytes() > 1 {
             // Peek addr from top of the stack by popping and pushing.
-            let addr = self.context.pop_to_reg(self.masm, None)?;
-            self.context.stack.push(addr.into());
-
-            let tmp = scratch!(M);
-            self.masm.mov(writable!(tmp), RegImm::Reg(addr.reg), size)?;
+            let addr = *self
+                .context
+                .stack
+                .peek()
+                .ok_or_else(|| CodeGenError::missing_values_in_stack())?;
+            let tmp = self.context.any_gpr(self.masm)?;
+            self.context.move_val_to_reg(&addr, tmp, self.masm)?;
 
             if memarg.offset != 0 {
                 self.masm.add(
@@ -868,6 +870,7 @@ where
 
             self.masm.cmp(tmp, RegImm::Imm(Imm::i64(0)), size)?;
             self.masm.trapif(IntCmpKind::Ne, TRAP_HEAP_MISALIGNED)?;
+            self.context.free_reg(tmp);
         }
 
         Ok(())
@@ -891,12 +894,16 @@ where
         op_kind: MemOpKind,
     ) -> Result<()> {
         let maybe_addr = match op_kind {
-            MemOpKind::Atomic => self.emit_compute_heap_address_align_checked(&arg, kind.derive_operand_size())?,
-            MemOpKind::Normal => self.emit_compute_heap_address(&arg, kind.derive_operand_size())?,
+            MemOpKind::Atomic => {
+                self.emit_compute_heap_address_align_checked(&arg, kind.derive_operand_size())?
+            }
+            MemOpKind::Normal => {
+                self.emit_compute_heap_address(&arg, kind.derive_operand_size())?
+            }
         };
 
         if let Some(addr) = maybe_addr {
-            let dst = match target_ty {
+            let dst = match target_type {
                 WasmValType::I32 | WasmValType::I64 => self.context.any_gpr(self.masm)?,
                 WasmValType::F32 | WasmValType::F64 => self.context.any_fpr(self.masm)?,
                 WasmValType::V128 => self.context.reg_for_type(target_type, self.masm)?,