Repository: tetratelabs/wazero
Pull request #2052 (2024-02-13 09:39:52)

wazevo(arm64): support for atomic load/store


Patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr.go b/internal/engine/wazevo/backend/isa/arm64/instr.go
index c44df0b7..dd5a0b77 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr.go
@@ -161,6 +161,8 @@ var defKinds = [numInstructionKinds]defKind{
 	emitSourceOffsetInfo: defKindNone,
 	atomicRmw:            defKindRD,
 	atomicCas:            defKindNone,
+	atomicLoad:           defKindRD,
+	atomicStore:          defKindNone,
 }
 
 // Defs returns the list of regalloc.VReg that are defined by the instruction.
@@ -295,6 +297,8 @@ var useKinds = [numInstructionKinds]useKind{
 	emitSourceOffsetInfo: useKindNone,
 	atomicRmw:            useKindRNRM,
 	atomicCas:            useKindRDRewrite,
+	atomicLoad:           useKindRN,
+	atomicStore:          useKindRNRM,
 }
 
 // Uses returns the list of regalloc.VReg that are used by the instruction.
@@ -1476,6 +1480,30 @@ func (i *instruction) String() (str string) {
 			m = m + "b"
 		}
 		str = fmt.Sprintf("%s %s, %s, %s", m, formatVRegSized(i.rd.nr(), size), formatVRegSized(i.rm.nr(), size), formatVRegSized(i.rn.nr(), 64))
+	case atomicLoad:
+		m := "ldar"
+		size := byte(32)
+		switch i.u2 {
+		case 8:
+			size = 64
+		case 2:
+			m = m + "h"
+		case 1:
+			m = m + "b"
+		}
+		str = fmt.Sprintf("%s %s, %s", m, formatVRegSized(i.rd.nr(), size), formatVRegSized(i.rn.nr(), 64))
+	case atomicStore:
+		m := "stlr"
+		size := byte(32)
+		switch i.u2 {
+		case 8:
+			size = 64
+		case 2:
+			m = m + "h"
+		case 1:
+			m = m + "b"
+		}
+		str = fmt.Sprintf("%s %s, %s", m, formatVRegSized(i.rm.nr(), size), formatVRegSized(i.rn.nr(), 64))
 	case udf:
 		str = "udf"
 	case emitSourceOffsetInfo:
@@ -1507,6 +1535,18 @@ func (i *instruction) asAtomicCas(rn, rs, rt operand, size uint64) {
 	i.u2 = size
 }
 
+func (i *instruction) asAtomicLoad(rn, rt operand, size uint64) {
+	i.kind = atomicLoad
+	i.rn, i.rd = rn, rt
+	i.u2 = size
+}
+
+func (i *instruction) asAtomicStore(rn, rt operand, size uint64) {
+	i.kind = atomicStore
+	i.rn, i.rm = rn, rt
+	i.u2 = size
+}
+
 // TODO: delete unnecessary things.
 const (
 	// nop0 represents a no-op of zero size.
@@ -1678,6 +1718,10 @@ const (
 	// atomicCas represents an atomic compare-and-swap operation with three register sources. The value is loaded to
 	// the source register containing the comparison value.
 	atomicCas
+	// atomicLoad represents an atomic load with one source register and a register destination.
+	atomicLoad
+	// atomicStore represents an atomic store with two source registers and no destination.
+	atomicStore
 	// UDF is the undefined instruction. For debugging only.
 	udf
 
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go b/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
index a9d18345..f2efb5e9 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
@@ -397,6 +397,20 @@ func (i *instruction) encode(m *machine) {
 			regNumberInEncoding[i.rn.realReg()],
 			uint32(i.u2),
 		))
+	case atomicLoad:
+		c.Emit4Bytes(encodeAtomicLoadStore(
+			regNumberInEncoding[i.rn.realReg()],
+			regNumberInEncoding[i.rd.realReg()],
+			uint32(i.u2),
+			1,
+		))
+	case atomicStore:
+		c.Emit4Bytes(encodeAtomicLoadStore(
+			regNumberInEncoding[i.rn.realReg()],
+			regNumberInEncoding[i.rm.realReg()],
+			uint32(i.u2),
+			0,
+		))
 	default:
 		panic(i.String())
 	}
@@ -2284,3 +2298,24 @@ func encodeAtomicCas(rs, rt, rn uint32, size uint32) uint32 {
 
 	return _31to21<<21 | rs<<16 | _15to10<<10 | rn<<5 | rt
 }
+
+func encodeAtomicLoadStore(rn, rt, size, l uint32) uint32 {
+	var _31to21, _20to16, _15to10, sz uint32
+
+	switch size {
+	case 8:
+		sz = 0b11
+	case 4:
+		sz = 0b10
+	case 2:
+		sz = 0b01
+	case 1:
+		sz = 0b00
+	}
+
+	_31to21 = 0b00001000_100 | sz<<9 | l<<1
+	_20to16 = 0b11111
+	_15to10 = 0b111111
+
+	return _31to21<<21 | _20to16<<16 | _15to10<<10 | rn<<5 | rt
+}
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
index 543cd6fc..9f902632 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
@@ -711,6 +711,12 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 	case ssa.OpcodeAtomicCas:
 		m.lowerAtomicCas(instr)
 
+	case ssa.OpcodeAtomicLoad:
+		m.lowerAtomicLoad(instr)
+
+	case ssa.OpcodeAtomicStore:
+		m.lowerAtomicStore(instr)
+
 	default:
 		panic("TODO: lowering " + op.String())
 	}
@@ -2047,7 +2053,7 @@ func (m *machine) lowerAtomicRmwImpl(op atomicRmwOp, rn, rs, rt, tmp operand, si
 
 func (m *machine) lowerAtomicCas(si *ssa.Instruction) {
 	addr, exp, repl := si.Arg3()
-	size := si.AtomicCasData()
+	size := si.AtomicTargetSize()
 
 	addrDef, expDef, replDef := m.compiler.ValueDefinition(addr), m.compiler.ValueDefinition(exp), m.compiler.ValueDefinition(repl)
 	rn := m.getOperand_NR(addrDef, extModeNone)
@@ -2084,6 +2090,41 @@ func (m *machine) lowerAtomicCasImpl(rn, rs, rt operand, size uint64) {
 	m.insert(cas)
 }
 
+func (m *machine) lowerAtomicLoad(si *ssa.Instruction) {
+	addr := si.Arg()
+	size := si.AtomicTargetSize()
+
+	addrDef := m.compiler.ValueDefinition(addr)
+	rn := m.getOperand_NR(addrDef, extModeNone)
+	rt := operandNR(m.compiler.VRegOf(si.Return()))
+
+	m.lowerAtomicLoadImpl(rn, rt, size)
+}
+
+func (m *machine) lowerAtomicLoadImpl(rn, rt operand, size uint64) {
+	ld := m.allocateInstr()
+	ld.asAtomicLoad(rn, rt, size)
+	m.insert(ld)
+}
+
+func (m *machine) lowerAtomicStore(si *ssa.Instruction) {
+	addr, val := si.Arg2()
+	size := si.AtomicTargetSize()
+
+	addrDef := m.compiler.ValueDefinition(addr)
+	valDef := m.compiler.ValueDefinition(val)
+	rn := m.getOperand_NR(addrDef, extModeNone)
+	rt := m.getOperand_NR(valDef, extModeNone)
+
+	m.lowerAtomicStoreImpl(rn, rt, size)
+}
+
+func (m *machine) lowerAtomicStoreImpl(rn, rt operand, size uint64) {
+	ld := m.allocateInstr()
+	ld.asAtomicStore(rn, rt, size)
+	m.insert(ld)
+}
+
 // copyToTmp copies the given regalloc.VReg to a temporary register. This is called before cbr to avoid the regalloc issue
 // e.g. reload happening in the middle of the exit sequence which is not the path the normal path executes
 func (m *machine) copyToTmp(v regalloc.VReg) regalloc.VReg {
diff --git a/internal/engine/wazevo/frontend/lower.go b/internal/engine/wazevo/frontend/lower.go
index 02a75ad7..6f197434 100644
--- a/internal/engine/wazevo/frontend/lower.go
+++ b/internal/engine/wazevo/frontend/lower.go
@@ -3266,8 +3266,7 @@ func (c *Compiler) lowerCurrentOpcode() {
 
 			addr := c.memOpSetup(baseAddr, uint64(offset), size)
 			c.memAlignmentCheck(addr, size)
-			res := builder.AllocateInstruction().AsAtomicStore(addr, val, size).Insert(builder).Return()
-			state.push(res)
+			builder.AllocateInstruction().AsAtomicStore(addr, val, size).Insert(builder)
 		case wasm.OpcodeAtomicI32RmwAdd, wasm.OpcodeAtomicI64RmwAdd, wasm.OpcodeAtomicI32Rmw8AddU, wasm.OpcodeAtomicI32Rmw16AddU, wasm.OpcodeAtomicI64Rmw8AddU, wasm.OpcodeAtomicI64Rmw16AddU, wasm.OpcodeAtomicI64Rmw32AddU,
 			wasm.OpcodeAtomicI32RmwSub, wasm.OpcodeAtomicI64RmwSub, wasm.OpcodeAtomicI32Rmw8SubU, wasm.OpcodeAtomicI32Rmw16SubU, wasm.OpcodeAtomicI64Rmw8SubU, wasm.OpcodeAtomicI64Rmw16SubU, wasm.OpcodeAtomicI64Rmw32SubU,
 			wasm.OpcodeAtomicI32RmwAnd, wasm.OpcodeAtomicI64RmwAnd, wasm.OpcodeAtomicI32Rmw8AndU, wasm.OpcodeAtomicI32Rmw16AndU, wasm.OpcodeAtomicI64Rmw8AndU, wasm.OpcodeAtomicI64Rmw16AndU, wasm.OpcodeAtomicI64Rmw32AndU,
diff --git a/internal/engine/wazevo/ssa/instructions.go b/internal/engine/wazevo/ssa/instructions.go
index 3387fea4..31e14042 100644
--- a/internal/engine/wazevo/ssa/instructions.go
+++ b/internal/engine/wazevo/ssa/instructions.go
@@ -2024,8 +2024,8 @@ func (i *Instruction) AtomicRmwData() (op AtomicRmwOp, size uint64) {
 	return AtomicRmwOp(i.u1), i.u2
 }
 
-// AtomicCasData returns the data for this atomic compare-and-swap instruction.
-func (i *Instruction) AtomicCasData() (size uint64) {
+// AtomicTargetSize returns the target memory size of the atomic instruction.
+func (i *Instruction) AtomicTargetSize() (size uint64) {
 	return i.u1
 }
 


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go b/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go
index cdeac169..df07085c 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go
@@ -1739,6 +1739,48 @@ func TestInstruction_encode(t *testing.T) {
 		{want: "02fce108", setup: func(i *instruction) {
 			i.asAtomicCas(operandNR(x0VReg), operandNR(x1VReg), operandNR(x2VReg), 1)
 		}},
+		{want: "01fcdf88", setup: func(i *instruction) {
+			i.asAtomicLoad(operandNR(x0VReg), operandNR(x1VReg), 4)
+		}},
+		{want: "01fcdf48", setup: func(i *instruction) {
+			i.asAtomicLoad(operandNR(x0VReg), operandNR(x1VReg), 2)
+		}},
+		{want: "01fcdf08", setup: func(i *instruction) {
+			i.asAtomicLoad(operandNR(x0VReg), operandNR(x1VReg), 1)
+		}},
+		{want: "01fcdfc8", setup: func(i *instruction) {
+			i.asAtomicLoad(operandNR(x0VReg), operandNR(x1VReg), 8)
+		}},
+		{want: "01fcdf88", setup: func(i *instruction) {
+			i.asAtomicLoad(operandNR(x0VReg), operandNR(x1VReg), 4)
+		}},
+		{want: "01fcdf48", setup: func(i *instruction) {
+			i.asAtomicLoad(operandNR(x0VReg), operandNR(x1VReg), 2)
+		}},
+		{want: "01fcdf08", setup: func(i *instruction) {
+			i.asAtomicLoad(operandNR(x0VReg), operandNR(x1VReg), 1)
+		}},
+		{want: "01fc9f88", setup: func(i *instruction) {
+			i.asAtomicStore(operandNR(x0VReg), operandNR(x1VReg), 4)
+		}},
+		{want: "01fc9f48", setup: func(i *instruction) {
+			i.asAtomicStore(operandNR(x0VReg), operandNR(x1VReg), 2)
+		}},
+		{want: "01fc9f08", setup: func(i *instruction) {
+			i.asAtomicStore(operandNR(x0VReg), operandNR(x1VReg), 1)
+		}},
+		{want: "01fc9fc8", setup: func(i *instruction) {
+			i.asAtomicStore(operandNR(x0VReg), operandNR(x1VReg), 8)
+		}},
+		{want: "01fc9f88", setup: func(i *instruction) {
+			i.asAtomicStore(operandNR(x0VReg), operandNR(x1VReg), 4)
+		}},
+		{want: "01fc9f48", setup: func(i *instruction) {
+			i.asAtomicStore(operandNR(x0VReg), operandNR(x1VReg), 2)
+		}},
+		{want: "01fc9f08", setup: func(i *instruction) {
+			i.asAtomicStore(operandNR(x0VReg), operandNR(x1VReg), 1)
+		}},
 		{want: "4201231e4201631e4201239e4201639e4201221e4201621e4201229e4201629e", setup: func(i *instruction) {
 			i.asNop0()
 			cur := i
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr_test.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr_test.go
index 350b5760..e87e2760 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr_test.go
@@ -1471,3 +1471,171 @@ casalb w2?, w3?, x1?
 		})
 	}
 }
+
+func TestMachine_lowerAtomicLoad(t *testing.T) {
+	tests := []struct {
+		name   string
+		_64bit bool
+		size   uint64
+		exp    string
+	}{
+		{
+			name: "load 32",
+			size: 4,
+			exp: `
+ldar w2?, x1?
+`,
+		},
+		{
+			name: "load 32_16u",
+			size: 2,
+			exp: `
+ldarh w2?, x1?
+`,
+		},
+		{
+			name: "load 32_8u",
+			size: 1,
+			exp: `
+ldarb w2?, x1?
+`,
+		},
+		{
+			name:   "load 64",
+			size:   8,
+			_64bit: true,
+			exp: `
+ldar x2?, x1?
+`,
+		},
+		{
+			name:   "load 64_32u",
+			size:   4,
+			_64bit: true,
+			exp: `
+ldar w2?, x1?
+`,
+		},
+		{
+			name:   "load 64_16u",
+			size:   2,
+			_64bit: true,
+			exp: `
+ldarh w2?, x1?
+`,
+		},
+		{
+			name:   "load 64_8u",
+			size:   1,
+			_64bit: true,
+			exp: `
+ldarb w2?, x1?
+`,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			_, _, m := newSetupWithMockContext()
+			var typ ssa.Type
+			if tc._64bit {
+				typ = ssa.TypeI64
+			} else {
+				typ = ssa.TypeI32
+			}
+			rn := operandNR(m.compiler.AllocateVReg(ssa.TypeI64))
+			rt := operandNR(m.compiler.AllocateVReg(typ))
+
+			require.Equal(t, 1, int(rn.reg().ID()))
+			require.Equal(t, 2, int(rt.reg().ID()))
+
+			m.lowerAtomicLoadImpl(rn, rt, tc.size)
+			require.Equal(t, tc.exp, "\n"+formatEmittedInstructionsInCurrentBlock(m)+"\n")
+		})
+	}
+}
+
+func TestMachine_lowerAtomicStore(t *testing.T) {
+	tests := []struct {
+		name   string
+		_64bit bool
+		size   uint64
+		exp    string
+	}{
+		{
+			name: "store 32",
+			size: 4,
+			exp: `
+stlr w2?, x1?
+`,
+		},
+		{
+			name: "store 32_16u",
+			size: 2,
+			exp: `
+stlrh w2?, x1?
+`,
+		},
+		{
+			name: "store 32_8u",
+			size: 1,
+			exp: `
+stlrb w2?, x1?
+`,
+		},
+		{
+			name:   "store 64",
+			size:   8,
+			_64bit: true,
+			exp: `
+stlr x2?, x1?
+`,
+		},
+		{
+			name:   "store 64_32u",
+			size:   4,
+			_64bit: true,
+			exp: `
+stlr w2?, x1?
+`,
+		},
+		{
+			name:   "store 64_16u",
+			size:   2,
+			_64bit: true,
+			exp: `
+stlrh w2?, x1?
+`,
+		},
+		{
+			name:   "store 64_8u",
+			size:   1,
+			_64bit: true,
+			exp: `
+stlrb w2?, x1?
+`,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			_, _, m := newSetupWithMockContext()
+			var typ ssa.Type
+			if tc._64bit {
+				typ = ssa.TypeI64
+			} else {
+				typ = ssa.TypeI32
+			}
+			rn := operandNR(m.compiler.AllocateVReg(ssa.TypeI64))
+			rt := operandNR(m.compiler.AllocateVReg(typ))
+
+			require.Equal(t, 1, int(rn.reg().ID()))
+			require.Equal(t, 2, int(rt.reg().ID()))
+
+			m.lowerAtomicStoreImpl(rn, rt, tc.size)
+			require.Equal(t, tc.exp, "\n"+formatEmittedInstructionsInCurrentBlock(m)+"\n")
+		})
+	}
+}
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index a0ab1b91..66c21ecc 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -835,6 +835,16 @@ func TestE2E(t *testing.T) {
 				},
 			},
 		},
+		{
+			name:      "atomic_store_load",
+			m:         testcases.AtomicStoreLoad.Module,
+			features:  api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
+			skipAMD64: true,
+			calls: []callCase{
+				{params: []uint64{1, 2, 3, 4, 5, 6, 7}, expResults: []uint64{1, 2, 3, 4, 5, 6, 7}},
+				{params: []uint64{10, 20, 30, 40, 50, 60, 70}, expResults: []uint64{10, 20, 30, 40, 50, 60, 70}},
+			},
+		},
 		{
 			name: "float_le",
 			m:    testcases.FloatLe.Module,
diff --git a/internal/engine/wazevo/frontend/frontend_test.go b/internal/engine/wazevo/frontend/frontend_test.go
index fa4dcb35..c251f4eb 100644
--- a/internal/engine/wazevo/frontend/frontend_test.go
+++ b/internal/engine/wazevo/frontend/frontend_test.go
@@ -2254,100 +2254,8 @@ blk0: (exec_ctx:i64, module_ctx:i64, v2:i32, v3:i32, v4:i32, v5:i64, v6:i64, v7:
 `,
 		},
 		{
-			name:     "AtomicLoad",
-			m:        testcases.AtomicLoad.Module,
-			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
-			exp: `
-blk0: (exec_ctx:i64, module_ctx:i64)
-	v2:i32 = Iconst_32 0x0
-	v3:i64 = Iconst_64 0x1
-	v4:i64 = UExtend v2, 32->64
-	v5:i64 = Uload32 module_ctx, 0x10
-	v6:i64 = Iadd v4, v3
-	v7:i32 = Icmp lt_u, v5, v6
-	ExitIfTrue v7, exec_ctx, memory_out_of_bounds
-	v8:i64 = Load module_ctx, 0x8
-	v9:i64 = Iadd v8, v4
-	v10:i32 = AtomicLoad_8, v9
-	v11:i32 = Iconst_32 0x8
-	v12:i64 = Iconst_64 0x2
-	v13:i64 = UExtend v11, 32->64
-	v14:i64 = Iadd v13, v12
-	v15:i32 = Icmp lt_u, v5, v14
-	ExitIfTrue v15, exec_ctx, memory_out_of_bounds
-	v16:i64 = Iadd v8, v13
-	v17:i64 = Iconst_64 0x1
-	v18:i64 = Band v16, v17
-	v19:i64 = Iconst_64 0x0
-	v20:i32 = Icmp neq, v18, v19
-	ExitIfTrue v20, exec_ctx, unaligned_atomic
-	v21:i32 = AtomicLoad_16, v16
-	v22:i32 = Iconst_32 0x10
-	v23:i64 = Iconst_64 0x4
-	v24:i64 = UExtend v22, 32->64
-	v25:i64 = Iadd v24, v23
-	v26:i32 = Icmp lt_u, v5, v25
-	ExitIfTrue v26, exec_ctx, memory_out_of_bounds
-	v27:i64 = Iadd v8, v24
-	v28:i64 = Iconst_64 0x3
-	v29:i64 = Band v27, v28
-	v30:i64 = Iconst_64 0x0
-	v31:i32 = Icmp neq, v29, v30
-	ExitIfTrue v31, exec_ctx, unaligned_atomic
-	v32:i32 = AtomicLoad_32, v27
-	v33:i32 = Iconst_32 0x18
-	v34:i64 = Iconst_64 0x1
-	v35:i64 = UExtend v33, 32->64
-	v36:i64 = Iadd v35, v34
-	v37:i32 = Icmp lt_u, v5, v36
-	ExitIfTrue v37, exec_ctx, memory_out_of_bounds
-	v38:i64 = Iadd v8, v35
-	v39:i64 = AtomicLoad_8, v38
-	v40:i32 = Iconst_32 0x20
-	v41:i64 = Iconst_64 0x2
-	v42:i64 = UExtend v40, 32->64
-	v43:i64 = Iadd v42, v41
-	v44:i32 = Icmp lt_u, v5, v43
-	ExitIfTrue v44, exec_ctx, memory_out_of_bounds
-	v45:i64 = Iadd v8, v42
-	v46:i64 = Iconst_64 0x1
-	v47:i64 = Band v45, v46
-	v48:i64 = Iconst_64 0x0
-	v49:i32 = Icmp neq, v47, v48
-	ExitIfTrue v49, exec_ctx, unaligned_atomic
-	v50:i64 = AtomicLoad_16, v45
-	v51:i32 = Iconst_32 0x28
-	v52:i64 = Iconst_64 0x4
-	v53:i64 = UExtend v51, 32->64
-	v54:i64 = Iadd v53, v52
-	v55:i32 = Icmp lt_u, v5, v54
-	ExitIfTrue v55, exec_ctx, memory_out_of_bounds
-	v56:i64 = Iadd v8, v53
-	v57:i64 = Iconst_64 0x3
-	v58:i64 = Band v56, v57
-	v59:i64 = Iconst_64 0x0
-	v60:i32 = Icmp neq, v58, v59
-	ExitIfTrue v60, exec_ctx, unaligned_atomic
-	v61:i64 = AtomicLoad_32, v56
-	v62:i32 = Iconst_32 0x30
-	v63:i64 = Iconst_64 0x8
-	v64:i64 = UExtend v62, 32->64
-	v65:i64 = Iadd v64, v63
-	v66:i32 = Icmp lt_u, v5, v65
-	ExitIfTrue v66, exec_ctx, memory_out_of_bounds
-	v67:i64 = Iadd v8, v64
-	v68:i64 = Iconst_64 0x7
-	v69:i64 = Band v67, v68
-	v70:i64 = Iconst_64 0x0
-	v71:i32 = Icmp neq, v69, v70
-	ExitIfTrue v71, exec_ctx, unaligned_atomic
-	v72:i64 = AtomicLoad_64, v67
-	Jump blk_ret, v10, v21, v32, v39, v50, v61, v72
-`,
-		},
-		{
-			name:     "AtomicStore",
-			m:        testcases.AtomicStore.Module,
+			name:     "AtomicStoreLoad",
+			m:        testcases.AtomicStoreLoad.Module,
 			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			exp: `
 blk0: (exec_ctx:i64, module_ctx:i64, v2:i32, v3:i32, v4:i32, v5:i64, v6:i64, v7:i64, v8:i64)
@@ -2361,80 +2269,161 @@ blk0: (exec_ctx:i64, module_ctx:i64, v2:i32, v3:i32, v4:i32, v5:i64, v6:i64, v7:
 	v15:i64 = Load module_ctx, 0x8
 	v16:i64 = Iadd v15, v11
 	AtomicStore_8, v16, v2
-	v17:i32 = Iconst_32 0x8
-	v18:i64 = Iconst_64 0x2
+	v17:i32 = Iconst_32 0x0
+	v18:i64 = Iconst_64 0x1
 	v19:i64 = UExtend v17, 32->64
 	v20:i64 = Iadd v19, v18
 	v21:i32 = Icmp lt_u, v12, v20
 	ExitIfTrue v21, exec_ctx, memory_out_of_bounds
 	v22:i64 = Iadd v15, v19
-	v23:i64 = Iconst_64 0x1
-	v24:i64 = Band v22, v23
-	v25:i64 = Iconst_64 0x0
-	v26:i32 = Icmp neq, v24, v25
-	ExitIfTrue v26, exec_ctx, unaligned_atomic
-	AtomicStore_16, v22, v3
-	v27:i32 = Iconst_32 0x10
-	v28:i64 = Iconst_64 0x4
-	v29:i64 = UExtend v27, 32->64
-	v30:i64 = Iadd v29, v28
-	v31:i32 = Icmp lt_u, v12, v30
-	ExitIfTrue v31, exec_ctx, memory_out_of_bounds
-	v32:i64 = Iadd v15, v29
-	v33:i64 = Iconst_64 0x3
-	v34:i64 = Band v32, v33
-	v35:i64 = Iconst_64 0x0
-	v36:i32 = Icmp neq, v34, v35
-	ExitIfTrue v36, exec_ctx, unaligned_atomic
-	AtomicStore_32, v32, v4
-	v37:i32 = Iconst_32 0x18
-	v38:i64 = Iconst_64 0x1
-	v39:i64 = UExtend v37, 32->64
-	v40:i64 = Iadd v39, v38
-	v41:i32 = Icmp lt_u, v12, v40
-	ExitIfTrue v41, exec_ctx, memory_out_of_bounds
-	v42:i64 = Iadd v15, v39
-	AtomicStore_8, v42, v5
-	v43:i32 = Iconst_32 0x20
-	v44:i64 = Iconst_64 0x2
-	v45:i64 = UExtend v43, 32->64
-	v46:i64 = Iadd v45, v44
-	v47:i32 = Icmp lt_u, v12, v46
-	ExitIfTrue v47, exec_ctx, memory_out_of_bounds
-	v48:i64 = Iadd v15, v45
-	v49:i64 = Iconst_64 0x1
-	v50:i64 = Band v48, v49
-	v51:i64 = Iconst_64 0x0
-	v52:i32 = Icmp neq, v50, v51
-	ExitIfTrue v52, exec_ctx, unaligned_atomic
-	AtomicStore_16, v48, v6
-	v53:i32 = Iconst_32 0x28
-	v54:i64 = Iconst_64 0x4
-	v55:i64 = UExtend v53, 32->64
-	v56:i64 = Iadd v55, v54
-	v57:i32 = Icmp lt_u, v12, v56
-	ExitIfTrue v57, exec_ctx, memory_out_of_bounds
-	v58:i64 = Iadd v15, v55
-	v59:i64 = Iconst_64 0x3
-	v60:i64 = Band v58, v59
-	v61:i64 = Iconst_64 0x0
-	v62:i32 = Icmp neq, v60, v61
-	ExitIfTrue v62, exec_ctx, unaligned_atomic
-	AtomicStore_32, v58, v7
-	v63:i32 = Iconst_32 0x30
-	v64:i64 = Iconst_64 0x8
-	v65:i64 = UExtend v63, 32->64
-	v66:i64 = Iadd v65, v64
-	v67:i32 = Icmp lt_u, v12, v66
-	ExitIfTrue v67, exec_ctx, memory_out_of_bounds
-	v68:i64 = Iadd v15, v65
-	v69:i64 = Iconst_64 0x7
-	v70:i64 = Band v68, v69
-	v71:i64 = Iconst_64 0x0
-	v72:i32 = Icmp neq, v70, v71
-	ExitIfTrue v72, exec_ctx, unaligned_atomic
-	AtomicStore_64, v68, v8
-	Jump blk_ret
+	v23:i32 = AtomicLoad_8, v22
+	v24:i32 = Iconst_32 0x8
+	v25:i64 = Iconst_64 0x2
+	v26:i64 = UExtend v24, 32->64
+	v27:i64 = Iadd v26, v25
+	v28:i32 = Icmp lt_u, v12, v27
+	ExitIfTrue v28, exec_ctx, memory_out_of_bounds
+	v29:i64 = Iadd v15, v26
+	v30:i64 = Iconst_64 0x1
+	v31:i64 = Band v29, v30
+	v32:i64 = Iconst_64 0x0
+	v33:i32 = Icmp neq, v31, v32
+	ExitIfTrue v33, exec_ctx, unaligned_atomic
+	AtomicStore_16, v29, v3
+	v34:i32 = Iconst_32 0x8
+	v35:i64 = Iconst_64 0x2
+	v36:i64 = UExtend v34, 32->64
+	v37:i64 = Iadd v36, v35
+	v38:i32 = Icmp lt_u, v12, v37
+	ExitIfTrue v38, exec_ctx, memory_out_of_bounds
+	v39:i64 = Iadd v15, v36
+	v40:i64 = Iconst_64 0x1
+	v41:i64 = Band v39, v40
+	v42:i64 = Iconst_64 0x0
+	v43:i32 = Icmp neq, v41, v42
+	ExitIfTrue v43, exec_ctx, unaligned_atomic
+	v44:i32 = AtomicLoad_16, v39
+	v45:i32 = Iconst_32 0x10
+	v46:i64 = Iconst_64 0x4
+	v47:i64 = UExtend v45, 32->64
+	v48:i64 = Iadd v47, v46
+	v49:i32 = Icmp lt_u, v12, v48
+	ExitIfTrue v49, exec_ctx, memory_out_of_bounds
+	v50:i64 = Iadd v15, v47
+	v51:i64 = Iconst_64 0x3
+	v52:i64 = Band v50, v51
+	v53:i64 = Iconst_64 0x0
+	v54:i32 = Icmp neq, v52, v53
+	ExitIfTrue v54, exec_ctx, unaligned_atomic
+	AtomicStore_32, v50, v4
+	v55:i32 = Iconst_32 0x10
+	v56:i64 = Iconst_64 0x4
+	v57:i64 = UExtend v55, 32->64
+	v58:i64 = Iadd v57, v56
+	v59:i32 = Icmp lt_u, v12, v58
+	ExitIfTrue v59, exec_ctx, memory_out_of_bounds
+	v60:i64 = Iadd v15, v57
+	v61:i64 = Iconst_64 0x3
+	v62:i64 = Band v60, v61
+	v63:i64 = Iconst_64 0x0
+	v64:i32 = Icmp neq, v62, v63
+	ExitIfTrue v64, exec_ctx, unaligned_atomic
+	v65:i32 = AtomicLoad_32, v60
+	v66:i32 = Iconst_32 0x18
+	v67:i64 = Iconst_64 0x1
+	v68:i64 = UExtend v66, 32->64
+	v69:i64 = Iadd v68, v67
+	v70:i32 = Icmp lt_u, v12, v69
+	ExitIfTrue v70, exec_ctx, memory_out_of_bounds
+	v71:i64 = Iadd v15, v68
+	AtomicStore_8, v71, v5
+	v72:i32 = Iconst_32 0x18
+	v73:i64 = Iconst_64 0x1
+	v74:i64 = UExtend v72, 32->64
+	v75:i64 = Iadd v74, v73
+	v76:i32 = Icmp lt_u, v12, v75
+	ExitIfTrue v76, exec_ctx, memory_out_of_bounds
+	v77:i64 = Iadd v15, v74
+	v78:i64 = AtomicLoad_8, v77
+	v79:i32 = Iconst_32 0x20
+	v80:i64 = Iconst_64 0x2
+	v81:i64 = UExtend v79, 32->64
+	v82:i64 = Iadd v81, v80
+	v83:i32 = Icmp lt_u, v12, v82
+	ExitIfTrue v83, exec_ctx, memory_out_of_bounds
+	v84:i64 = Iadd v15, v81
+	v85:i64 = Iconst_64 0x1
+	v86:i64 = Band v84, v85
+	v87:i64 = Iconst_64 0x0
+	v88:i32 = Icmp neq, v86, v87
+	ExitIfTrue v88, exec_ctx, unaligned_atomic
+	AtomicStore_16, v84, v6
+	v89:i32 = Iconst_32 0x20
+	v90:i64 = Iconst_64 0x2
+	v91:i64 = UExtend v89, 32->64
+	v92:i64 = Iadd v91, v90
+	v93:i32 = Icmp lt_u, v12, v92
+	ExitIfTrue v93, exec_ctx, memory_out_of_bounds
+	v94:i64 = Iadd v15, v91
+	v95:i64 = Iconst_64 0x1
+	v96:i64 = Band v94, v95
+	v97:i64 = Iconst_64 0x0
+	v98:i32 = Icmp neq, v96, v97
+	ExitIfTrue v98, exec_ctx, unaligned_atomic
+	v99:i64 = AtomicLoad_16, v94
+	v100:i32 = Iconst_32 0x28
+	v101:i64 = Iconst_64 0x4
+	v102:i64 = UExtend v100, 32->64
+	v103:i64 = Iadd v102, v101
+	v104:i32 = Icmp lt_u, v12, v103
+	ExitIfTrue v104, exec_ctx, memory_out_of_bounds
+	v105:i64 = Iadd v15, v102
+	v106:i64 = Iconst_64 0x3
+	v107:i64 = Band v105, v106
+	v108:i64 = Iconst_64 0x0
+	v109:i32 = Icmp neq, v107, v108
+	ExitIfTrue v109, exec_ctx, unaligned_atomic
+	AtomicStore_32, v105, v7
+	v110:i32 = Iconst_32 0x28
+	v111:i64 = Iconst_64 0x4
+	v112:i64 = UExtend v110, 32->64
+	v113:i64 = Iadd v112, v111
+	v114:i32 = Icmp lt_u, v12, v113
+	ExitIfTrue v114, exec_ctx, memory_out_of_bounds
+	v115:i64 = Iadd v15, v112
+	v116:i64 = Iconst_64 0x3
+	v117:i64 = Band v115, v116
+	v118:i64 = Iconst_64 0x0
+	v119:i32 = Icmp neq, v117, v118
+	ExitIfTrue v119, exec_ctx, unaligned_atomic
+	v120:i64 = AtomicLoad_32, v115
+	v121:i32 = Iconst_32 0x30
+	v122:i64 = Iconst_64 0x8
+	v123:i64 = UExtend v121, 32->64
+	v124:i64 = Iadd v123, v122
+	v125:i32 = Icmp lt_u, v12, v124
+	ExitIfTrue v125, exec_ctx, memory_out_of_bounds
+	v126:i64 = Iadd v15, v123
+	v127:i64 = Iconst_64 0x7
+	v128:i64 = Band v126, v127
+	v129:i64 = Iconst_64 0x0
+	v130:i32 = Icmp neq, v128, v129
+	ExitIfTrue v130, exec_ctx, unaligned_atomic
+	AtomicStore_64, v126, v8
+	v131:i32 = Iconst_32 0x30
+	v132:i64 = Iconst_64 0x8
+	v133:i64 = UExtend v131, 32->64
+	v134:i64 = Iadd v133, v132
+	v135:i32 = Icmp lt_u, v12, v134
+	ExitIfTrue v135, exec_ctx, memory_out_of_bounds
+	v136:i64 = Iadd v15, v133
+	v137:i64 = Iconst_64 0x7
+	v138:i64 = Band v136, v137
+	v139:i64 = Iconst_64 0x0
+	v140:i32 = Icmp neq, v138, v139
+	ExitIfTrue v140, exec_ctx, unaligned_atomic
+	v141:i64 = AtomicLoad_64, v136
+	Jump blk_ret, v23, v44, v65, v78, v99, v120, v141
 `,
 		},
 		{
diff --git a/internal/engine/wazevo/testcases/testcases.go b/internal/engine/wazevo/testcases/testcases.go
index 5193b01b..aec54c46 100644
--- a/internal/engine/wazevo/testcases/testcases.go
+++ b/internal/engine/wazevo/testcases/testcases.go
@@ -2208,42 +2208,12 @@ var (
 		},
 	}
 
-	AtomicLoad = TestCase{
-		Name: "atomic_load",
-		Module: &wasm.Module{
-			TypeSection: []wasm.FunctionType{{
-				Params:  []wasm.ValueType{},
-				Results: []wasm.ValueType{i32, i32, i32, i64, i64, i64, i64},
-			}},
-			ExportSection:   []wasm.Export{{Name: ExportedFunctionName, Type: wasm.ExternTypeFunc, Index: 0}},
-			MemorySection:   &wasm.Memory{Min: 1, Max: 1, IsMaxEncoded: true, IsShared: true},
-			FunctionSection: []wasm.Index{0},
-			CodeSection: []wasm.Code{{Body: []byte{
-				wasm.OpcodeI32Const, 0,
-				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Load8U, 0x0, 0,
-				wasm.OpcodeI32Const, 8,
-				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Load16U, 0x1, 0,
-				wasm.OpcodeI32Const, 16,
-				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Load, 0x2, 0,
-				wasm.OpcodeI32Const, 24,
-				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Load8U, 0x0, 0,
-				wasm.OpcodeI32Const, 32,
-				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Load16U, 0x1, 0,
-				wasm.OpcodeI32Const, 40,
-				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Load32U, 0x2, 0,
-				wasm.OpcodeI32Const, 48,
-				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Load, 0x3, 0,
-				wasm.OpcodeEnd,
-			}}},
-		},
-	}
-
-	AtomicStore = TestCase{
-		Name: "atomic_store",
+	AtomicStoreLoad = TestCase{
+		Name: "atomic_store_load",
 		Module: &wasm.Module{
 			TypeSection: []wasm.FunctionType{{
 				Params:  []wasm.ValueType{i32, i32, i32, i64, i64, i64, i64},
-				Results: []wasm.ValueType{},
+				Results: []wasm.ValueType{i32, i32, i32, i64, i64, i64, i64},
 			}},
 			ExportSection:   []wasm.Export{{Name: ExportedFunctionName, Type: wasm.ExternTypeFunc, Index: 0}},
 			MemorySection:   &wasm.Memory{Min: 1, Max: 1, IsMaxEncoded: true, IsShared: true},
@@ -2252,24 +2222,38 @@ var (
 				wasm.OpcodeI32Const, 0,
 				wasm.OpcodeLocalGet, 0,
 				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Store8, 0x0, 0,
+				wasm.OpcodeI32Const, 0,
+				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Load8U, 0x0, 0,
 				wasm.OpcodeI32Const, 8,
 				wasm.OpcodeLocalGet, 1,
 				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Store16, 0x1, 0,
+				wasm.OpcodeI32Const, 8,
+				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Load16U, 0x1, 0,
 				wasm.OpcodeI32Const, 16,
 				wasm.OpcodeLocalGet, 2,
 				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Store, 0x2, 0,
+				wasm.OpcodeI32Const, 16,
+				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI32Load, 0x2, 0,
 				wasm.OpcodeI32Const, 24,
 				wasm.OpcodeLocalGet, 3,
 				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Store8, 0x0, 0,
+				wasm.OpcodeI32Const, 24,
+				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Load8U, 0x0, 0,
 				wasm.OpcodeI32Const, 32,
 				wasm.OpcodeLocalGet, 4,
 				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Store16, 0x1, 0,
+				wasm.OpcodeI32Const, 32,
+				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Load16U, 0x1, 0,
 				wasm.OpcodeI32Const, 40,
 				wasm.OpcodeLocalGet, 5,
 				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Store32, 0x2, 0,
+				wasm.OpcodeI32Const, 40,
+				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Load32U, 0x2, 0,
 				wasm.OpcodeI32Const, 48,
 				wasm.OpcodeLocalGet, 6,
 				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Store, 0x3, 0,
+				wasm.OpcodeI32Const, 48,
+				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicI64Load, 0x3, 0,
 				wasm.OpcodeEnd,
 			}}},
 		},