Repository: tetratelabs/wazero
Pull request #1964 (2024-01-26 02:35:25)

wazevo: adds support for memory wait
Not too sure what I'm doing so sending out lowering of one atomic op to understand the pattern, hopefully it will be enough to do the rest in the later PR.

I'm trying to compare to wasmtime but not sure if I'm doing something wrong. For this wat file

```
(module
  (import "console" "log" (func $log (param i32)))
  (func
    (i32.const 5)
    (i32.const 0)
    (i64.const 10)
    memory.atomic.wait32
    (call $log)
  )
  (memory 1 1 shared)
  (export "" (func 0))
)
```

I get

```
❯ target/debug/clif-util wasm --target aarch64-apple-darwin test.wat -p -t
function u0:1(i64 vmctx) fast {
    gv0 = vmctx
    gv1 = load.i64 notrap aligned readonly gv0
    sig0 = (i32, i64 vmctx) fast
    fn0 = u0:0 sig0

                                block0(v0: i64):
@0038                               v1 = iconst.i32 5
@003a                               v2 = iconst.i32 0
@003c                               v3 = iconst.i64 10
@003e                               v4 = iconst.i32 0xffff_ffff
@0042                               call fn0(v4, v0)  ; v4 = 0xffff_ffff
@0044                               jump block1

                                block1:
@0044                               return
}
```

I initially used the same wat as my test case but thought some optimization may be happening so tried adding the `call $log`, but still somehow it is just `v4 = iconst.i32 0xffff_ffff`. Any idea what I'm doing wrong?

anuraaga: Thanks @mathetake - can you check this is going in a reasonable direction? I think what's left is defining the trampoline for the wait function
mathetake: Basically you should be able to achieve lowering without adding any new instruction 
anuraaga: Thanks for the patience @mathetake I think this is on the right way now
mathetake: want to add a e2e test case in https://github.com/tetratelabs/wazero/blob/main/internal/engine/wazevo/e2e_test.go#L39 ?
anuraaga: Added a simple e2e for now for the "not equals" case, other ones will be trickier / require notify to be implemented.

And fixed a bug
mathetake: looking good! will take a look tomorrow 
mathetake: lastly could you fix the conflict 
anuraaga: Thanks, resolved conflicts

Patch:
diff --git a/internal/engine/wazevo/call_engine.go b/internal/engine/wazevo/call_engine.go
index 6b97ac32..58e47747 100644
--- a/internal/engine/wazevo/call_engine.go
+++ b/internal/engine/wazevo/call_engine.go
@@ -5,6 +5,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"reflect"
+	"sync/atomic"
 	"unsafe"
 
 	"github.com/tetratelabs/wazero/api"
@@ -82,6 +83,10 @@ type (
 		memmoveAddress uintptr
 		// framePointerBeforeGoCall holds the frame pointer before calling a Go function. Note: only used in amd64.
 		framePointerBeforeGoCall uintptr
+		// memoryWait32TrampolineAddress holds the address of memory_wait32 trampoline function.
+		memoryWait32TrampolineAddress *byte
+		// memoryWait32TrampolineAddress holds the address of memory_wait64 trampoline function.
+		memoryWait64TrampolineAddress *byte
 	}
 )
 
@@ -381,6 +386,46 @@ func (c *callEngine) callWithStack(ctx context.Context, paramResultStack []uint6
 			c.execCtx.exitCode = wazevoapi.ExitCodeOK
 			afterGoFunctionCallEntrypoint(c.execCtx.goCallReturnAddress, c.execCtxPtr,
 				uintptr(unsafe.Pointer(c.execCtx.stackPointerBeforeGoCall)), c.execCtx.framePointerBeforeGoCall)
+		case wazevoapi.ExitCodeMemoryWait32:
+			mod := c.callerModuleInstance()
+			mem := mod.MemoryInstance
+			if !mem.Shared {
+				panic(wasmruntime.ErrRuntimeExpectedSharedMemory)
+			}
+
+			s := goCallStackView(c.execCtx.stackPointerBeforeGoCall)
+			timeout, exp, addr := int64(s[0]), uint32(s[1]), uintptr(s[2])
+			base := uintptr(unsafe.Pointer(&mem.Buffer[0]))
+
+			offset := uint32(addr - base)
+			res := mem.Wait32(offset, exp, timeout, func(mem *wasm.MemoryInstance, offset uint32) uint32 {
+				addr := unsafe.Add(unsafe.Pointer(&mem.Buffer[0]), offset)
+				return atomic.LoadUint32((*uint32)(addr))
+			})
+			s[0] = uint64(res)
+			c.execCtx.exitCode = wazevoapi.ExitCodeOK
+			afterGoFunctionCallEntrypoint(c.execCtx.goCallReturnAddress, c.execCtxPtr,
+				uintptr(unsafe.Pointer(c.execCtx.stackPointerBeforeGoCall)), c.execCtx.framePointerBeforeGoCall)
+		case wazevoapi.ExitCodeMemoryWait64:
+			mod := c.callerModuleInstance()
+			mem := mod.MemoryInstance
+			if !mem.Shared {
+				panic(wasmruntime.ErrRuntimeExpectedSharedMemory)
+			}
+
+			s := goCallStackView(c.execCtx.stackPointerBeforeGoCall)
+			timeout, exp, addr := int64(s[0]), uint64(s[1]), uintptr(s[2])
+			base := uintptr(unsafe.Pointer(&mem.Buffer[0]))
+
+			offset := uint32(addr - base)
+			res := mem.Wait64(offset, exp, timeout, func(mem *wasm.MemoryInstance, offset uint32) uint64 {
+				addr := unsafe.Add(unsafe.Pointer(&mem.Buffer[0]), offset)
+				return atomic.LoadUint64((*uint64)(addr))
+			})
+			s[0] = uint64(res)
+			c.execCtx.exitCode = wazevoapi.ExitCodeOK
+			afterGoFunctionCallEntrypoint(c.execCtx.goCallReturnAddress, c.execCtxPtr,
+				uintptr(unsafe.Pointer(c.execCtx.stackPointerBeforeGoCall)), c.execCtx.framePointerBeforeGoCall)
 		case wazevoapi.ExitCodeUnreachable:
 			panic(wasmruntime.ErrRuntimeUnreachable)
 		case wazevoapi.ExitCodeMemoryOutOfBounds:
@@ -397,6 +442,8 @@ func (c *callEngine) callWithStack(ctx context.Context, paramResultStack []uint6
 			panic(wasmruntime.ErrRuntimeIntegerDivideByZero)
 		case wazevoapi.ExitCodeInvalidConversionToInteger:
 			panic(wasmruntime.ErrRuntimeInvalidConversionToInteger)
+		case wazevoapi.ExitCodeUnalignedAtomic:
+			panic(wasmruntime.ErrRuntimeUnalignedAtomic)
 		default:
 			panic("BUG")
 		}
diff --git a/internal/engine/wazevo/engine.go b/internal/engine/wazevo/engine.go
index f2a9b9ea..9bc6d359 100644
--- a/internal/engine/wazevo/engine.go
+++ b/internal/engine/wazevo/engine.go
@@ -56,7 +56,11 @@ type (
 		// tableGrowExecutable is a compiled trampoline executable for table.grow builtin function.
 		tableGrowExecutable []byte
 		// refFuncExecutable is a compiled trampoline executable for ref.func builtin function.
-		refFuncExecutable         []byte
+		refFuncExecutable []byte
+		// memoryWait32Executable is a compiled trampoline executable for memory.wait32 builtin function
+		memoryWait32Executable []byte
+		// memoryWait64Executable is a compiled trampoline executable for memory.wait64 builtin function
+		memoryWait64Executable    []byte
 		listenerBeforeTrampolines map[*wasm.FunctionType][]byte
 		listenerAfterTrampolines  map[*wasm.FunctionType][]byte
 	}
@@ -639,6 +643,36 @@ func (e *engine) compileSharedFunctions() {
 		}
 	}
 
+	e.be.Init()
+	{
+		src := e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeMemoryWait32, &ssa.Signature{
+			// exec context, timeout, expected, addr
+			Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeI32, ssa.TypeI64},
+			// Returns the status.
+			Results: []ssa.Type{ssa.TypeI32},
+		}, false)
+		e.sharedFunctions.memoryWait32Executable = mmapExecutable(src)
+		if wazevoapi.PerfMapEnabled {
+			exe := e.sharedFunctions.memoryWait32Executable
+			wazevoapi.PerfMap.AddEntry(uintptr(unsafe.Pointer(&exe[0])), uint64(len(exe)), "memory_wait32_trampoline")
+		}
+	}
+
+	e.be.Init()
+	{
+		src := e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeMemoryWait64, &ssa.Signature{
+			// exec context, timeout, expected, addr
+			Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeI64, ssa.TypeI64},
+			// Returns the status.
+			Results: []ssa.Type{ssa.TypeI32},
+		}, false)
+		e.sharedFunctions.memoryWait64Executable = mmapExecutable(src)
+		if wazevoapi.PerfMapEnabled {
+			exe := e.sharedFunctions.memoryWait64Executable
+			wazevoapi.PerfMap.AddEntry(uintptr(unsafe.Pointer(&exe[0])), uint64(len(exe)), "memory_wait64_trampoline")
+		}
+	}
+
 	e.setFinalizer(e.sharedFunctions, sharedFunctionsFinalizer)
 }
 
@@ -658,6 +692,12 @@ func sharedFunctionsFinalizer(sf *sharedFunctions) {
 	if err := platform.MunmapCodeSegment(sf.refFuncExecutable); err != nil {
 		panic(err)
 	}
+	if err := platform.MunmapCodeSegment(sf.memoryWait32Executable); err != nil {
+		panic(err)
+	}
+	if err := platform.MunmapCodeSegment(sf.memoryWait64Executable); err != nil {
+		panic(err)
+	}
 	for _, f := range sf.listenerBeforeTrampolines {
 		if err := platform.MunmapCodeSegment(f); err != nil {
 			panic(err)
@@ -674,6 +714,8 @@ func sharedFunctionsFinalizer(sf *sharedFunctions) {
 	sf.stackGrowExecutable = nil
 	sf.tableGrowExecutable = nil
 	sf.refFuncExecutable = nil
+	sf.memoryWait32Executable = nil
+	sf.memoryWait64Executable = nil
 	sf.listenerBeforeTrampolines = nil
 	sf.listenerAfterTrampolines = nil
 }
diff --git a/internal/engine/wazevo/frontend/frontend.go b/internal/engine/wazevo/frontend/frontend.go
index 6e7a287f..1c468499 100644
--- a/internal/engine/wazevo/frontend/frontend.go
+++ b/internal/engine/wazevo/frontend/frontend.go
@@ -21,6 +21,8 @@ type Compiler struct {
 	signatures             map[*wasm.FunctionType]*ssa.Signature
 	listenerSignatures     map[*wasm.FunctionType][2]*ssa.Signature
 	memoryGrowSig          ssa.Signature
+	memoryWait32Sig        ssa.Signature
+	memoryWait64Sig        ssa.Signature
 	checkModuleExitCodeSig ssa.Signature
 	tableGrowSig           ssa.Signature
 	refFuncSig             ssa.Signature
@@ -143,6 +145,24 @@ func (c *Compiler) declareSignatures(listenerOn bool) {
 	}
 
 	c.ssaBuilder.DeclareSignature(&c.memmoveSig)
+
+	c.memoryWait32Sig = ssa.Signature{
+		ID: c.memmoveSig.ID + 1,
+		// exec context, timeout, expected, addr
+		Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeI32, ssa.TypeI64},
+		// Returns the status.
+		Results: []ssa.Type{ssa.TypeI32},
+	}
+	c.ssaBuilder.DeclareSignature(&c.memoryWait32Sig)
+
+	c.memoryWait64Sig = ssa.Signature{
+		ID: c.memoryWait32Sig.ID + 1,
+		// exec context, timeout, expected, addr
+		Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeI64, ssa.TypeI64},
+		// Returns the status.
+		Results: []ssa.Type{ssa.TypeI32},
+	}
+	c.ssaBuilder.DeclareSignature(&c.memoryWait64Sig)
 }
 
 // SignatureForWasmFunctionType returns the ssa.Signature for the given wasm.FunctionType.
diff --git a/internal/engine/wazevo/frontend/lower.go b/internal/engine/wazevo/frontend/lower.go
index 64792c9b..fca32c69 100644
--- a/internal/engine/wazevo/frontend/lower.go
+++ b/internal/engine/wazevo/frontend/lower.go
@@ -3151,6 +3151,52 @@ func (c *Compiler) lowerCurrentOpcode() {
 		default:
 			panic("TODO: unsupported vector instruction: " + wasm.VectorInstructionName(vecOp))
 		}
+	case wasm.OpcodeAtomicPrefix:
+		state.pc++
+		atomicOp := c.wasmFunctionBody[state.pc]
+		switch atomicOp {
+		case wasm.OpcodeAtomicMemoryWait32, wasm.OpcodeAtomicMemoryWait64:
+			_, offset := c.readMemArg()
+			if state.unreachable {
+				break
+			}
+
+			c.storeCallerModuleContext()
+
+			var opSize uint64
+			var trampoline wazevoapi.Offset
+			var sig *ssa.Signature
+			switch atomicOp {
+			case wasm.OpcodeAtomicMemoryWait32:
+				opSize = 4
+				trampoline = wazevoapi.ExecutionContextOffsetMemoryWait32TrampolineAddress
+				sig = &c.memoryWait32Sig
+			case wasm.OpcodeAtomicMemoryWait64:
+				opSize = 8
+				trampoline = wazevoapi.ExecutionContextOffsetMemoryWait64TrampolineAddress
+				sig = &c.memoryWait64Sig
+			}
+
+			timeout := state.pop()
+			exp := state.pop()
+			baseAddr := state.pop()
+			addr := c.memOpSetup(baseAddr, uint64(offset), opSize)
+			c.memAlignmentCheck(addr, opSize)
+
+			memoryWaitPtr := builder.AllocateInstruction().
+				AsLoad(c.execCtxPtrValue,
+					trampoline.U32(),
+					ssa.TypeI64,
+				).Insert(builder).Return()
+
+			args := []ssa.Value{c.execCtxPtrValue, timeout, exp, addr}
+			memoryWaitRet := builder.AllocateInstruction().
+				AsCallIndirect(memoryWaitPtr, sig, args).
+				Insert(builder).Return()
+			state.push(memoryWaitRet)
+		default:
+			panic("TODO: unsupported atomic instruction: " + wasm.AtomicInstructionName(atomicOp))
+		}
 	case wasm.OpcodeRefFunc:
 		funcIndex := c.readI32u()
 		if state.unreachable {
@@ -3423,6 +3469,29 @@ func (c *Compiler) memOpSetup(baseAddr ssa.Value, constOffset, operationSizeInBy
 	return
 }
 
+func (c *Compiler) memAlignmentCheck(addr ssa.Value, operationSizeInBytes uint64) {
+	if operationSizeInBytes == 1 {
+		return // No alignment restrictions when accessing a byte
+	}
+	var checkBits uint64
+	switch operationSizeInBytes {
+	case 2:
+		checkBits = 0b1
+	case 4:
+		checkBits = 0b11
+	case 8:
+		checkBits = 0b111
+	}
+
+	builder := c.ssaBuilder
+
+	mask := builder.AllocateInstruction().AsIconst64(checkBits).Insert(builder).Return()
+	masked := builder.AllocateInstruction().AsBand(addr, mask).Insert(builder).Return()
+	zero := builder.AllocateInstruction().AsIconst64(0).Insert(builder).Return()
+	cmp := builder.AllocateInstruction().AsIcmp(masked, zero, ssa.IntegerCmpCondNotEqual).Insert(builder).Return()
+	builder.AllocateInstruction().AsExitIfTrueWithCode(c.execCtxPtrValue, cmp, wazevoapi.ExitCodeUnalignedAtomic).Insert(builder)
+}
+
 func (c *Compiler) callMemmove(dst, src, size ssa.Value) {
 	args := []ssa.Value{dst, src, size} // TODO: reuse the slice.
 	if size.Type() != ssa.TypeI64 {
diff --git a/internal/engine/wazevo/module_engine.go b/internal/engine/wazevo/module_engine.go
index 3a61e331..b753fbe6 100644
--- a/internal/engine/wazevo/module_engine.go
+++ b/internal/engine/wazevo/module_engine.go
@@ -183,6 +183,8 @@ func (m *moduleEngine) NewFunction(index wasm.Index) api.Function {
 	ce.execCtx.checkModuleExitCodeTrampolineAddress = &m.parent.sharedFunctions.checkModuleExitCode[0]
 	ce.execCtx.tableGrowTrampolineAddress = &m.parent.sharedFunctions.tableGrowExecutable[0]
 	ce.execCtx.refFuncTrampolineAddress = &m.parent.sharedFunctions.refFuncExecutable[0]
+	ce.execCtx.memoryWait32TrampolineAddress = &m.parent.sharedFunctions.memoryWait32Executable[0]
+	ce.execCtx.memoryWait64TrampolineAddress = &m.parent.sharedFunctions.memoryWait64Executable[0]
 	ce.execCtx.memmoveAddress = memmovPtr
 	ce.init()
 	return ce
diff --git a/internal/engine/wazevo/ssa/instructions.go b/internal/engine/wazevo/ssa/instructions.go
index 725a7dd7..53edbd22 100644
--- a/internal/engine/wazevo/ssa/instructions.go
+++ b/internal/engine/wazevo/ssa/instructions.go
@@ -1530,11 +1530,12 @@ func (i *Instruction) AsURem(x, y, ctx Value) *Instruction {
 }
 
 // AsBand initializes this instruction as an integer bitwise and instruction with OpcodeBand.
-func (i *Instruction) AsBand(x, amount Value) {
+func (i *Instruction) AsBand(x, amount Value) *Instruction {
 	i.opcode = OpcodeBand
 	i.v = x
 	i.v2 = amount
 	i.typ = x.Type()
+	return i
 }
 
 // AsBor initializes this instruction as an integer bitwise or instruction with OpcodeBor.
@@ -2477,7 +2478,6 @@ func (i *Instruction) Format(b Builder) string {
 		}
 		// Prints Shuffle.[0 1 2 3 4 5 6 7 ...] v2, v3
 		instSuffix = fmt.Sprintf(".%v %s, %s", lanes, i.v.Format(b), i.v2.Format(b))
-
 	default:
 		panic(fmt.Sprintf("TODO: format for %s", i.opcode))
 	}
diff --git a/internal/engine/wazevo/wazevoapi/exitcode.go b/internal/engine/wazevo/wazevoapi/exitcode.go
index fd7c3989..4e9374c1 100644
--- a/internal/engine/wazevo/wazevoapi/exitcode.go
+++ b/internal/engine/wazevo/wazevoapi/exitcode.go
@@ -26,6 +26,9 @@ const (
 	ExitCodeCallGoFunctionWithListener
 	ExitCodeTableGrow
 	ExitCodeRefFunc
+	ExitCodeMemoryWait32
+	ExitCodeMemoryWait64
+	ExitCodeUnalignedAtomic
 	exitCodeMax
 )
 
@@ -46,6 +49,8 @@ func (e ExitCode) String() string {
 		return "unreachable"
 	case ExitCodeMemoryOutOfBounds:
 		return "memory_out_of_bounds"
+	case ExitCodeUnalignedAtomic:
+		return "unaligned_atomic"
 	case ExitCodeTableOutOfBounds:
 		return "table_out_of_bounds"
 	case ExitCodeIndirectCallNullPointer:
@@ -74,6 +79,10 @@ func (e ExitCode) String() string {
 		return "table_grow"
 	case ExitCodeRefFunc:
 		return "ref_func"
+	case ExitCodeMemoryWait32:
+		return "memory_wait32"
+	case ExitCodeMemoryWait64:
+		return "memory_wait64"
 	}
 	panic("TODO")
 }
diff --git a/internal/engine/wazevo/wazevoapi/offsetdata.go b/internal/engine/wazevo/wazevoapi/offsetdata.go
index 29d31e32..0fd3123e 100644
--- a/internal/engine/wazevo/wazevoapi/offsetdata.go
+++ b/internal/engine/wazevo/wazevoapi/offsetdata.go
@@ -47,9 +47,11 @@ const (
 	// ExecutionContextOffsetTableGrowTrampolineAddress is an offset of `tableGrowTrampolineAddress` field in wazevo.executionContext
 	ExecutionContextOffsetTableGrowTrampolineAddress Offset = 1128
 	// ExecutionContextOffsetRefFuncTrampolineAddress is an offset of `refFuncTrampolineAddress` field in wazevo.executionContext
-	ExecutionContextOffsetRefFuncTrampolineAddress Offset = 1136
-	ExecutionContextOffsetMemmoveAddress           Offset = 1144
-	ExecutionContextOffsetFramePointerBeforeGoCall Offset = 1152
+	ExecutionContextOffsetRefFuncTrampolineAddress      Offset = 1136
+	ExecutionContextOffsetMemmoveAddress                Offset = 1144
+	ExecutionContextOffsetFramePointerBeforeGoCall      Offset = 1152
+	ExecutionContextOffsetMemoryWait32TrampolineAddress Offset = 1160
+	ExecutionContextOffsetMemoryWait64TrampolineAddress Offset = 1168
 )
 
 // ModuleContextOffsetData allows the compilers to get the information about offsets to the fields of wazevo.moduleContextOpaque,


Test patch:
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index aade7ac9..29aadac7 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -51,6 +51,7 @@ func TestE2E(t *testing.T) {
 		name        string
 		imported, m *wasm.Module
 		calls       []callCase
+		features    api.CoreFeatures
 		skipAMD64   bool
 	}{
 		{
@@ -606,6 +607,22 @@ func TestE2E(t *testing.T) {
 				{params: []uint64{0x00000000, 0x11111111, 0x11111111, 0xffffffff}, expResults: []uint64{0x11001100, 0xff11ff11}},
 			},
 		},
+		{
+			name:     "memory_wait32",
+			m:        testcases.MemoryWait32.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
+			calls: []callCase{
+				{params: []uint64{0x0, 0xbeef, 0xffffffff}, expResults: []uint64{1}},
+			},
+		},
+		{
+			name:     "memory_wait64",
+			m:        testcases.MemoryWait64.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
+			calls: []callCase{
+				{params: []uint64{0x0, 0xbeef, 0xffffffff}, expResults: []uint64{1}},
+			},
+		},
 		{
 			name: "float_le",
 			m:    testcases.FloatLe.Module,
@@ -633,6 +650,9 @@ func TestE2E(t *testing.T) {
 					cache, err := wazero.NewCompilationCacheWithDir(tmp)
 					require.NoError(t, err)
 					config := opt.NewRuntimeConfigOptimizingCompiler().WithCompilationCache(cache)
+					if tc.features != 0 {
+						config = config.WithCoreFeatures(tc.features)
+					}
 
 					ctx := context.Background()
 					r := wazero.NewRuntimeWithConfig(ctx, config)
diff --git a/internal/engine/wazevo/engine_test.go b/internal/engine/wazevo/engine_test.go
index 241e68e7..6f7679c0 100644
--- a/internal/engine/wazevo/engine_test.go
+++ b/internal/engine/wazevo/engine_test.go
@@ -25,12 +25,18 @@ func Test_sharedFunctionsFinalizer(t *testing.T) {
 	require.NoError(t, err)
 	b7, err := platform.MmapCodeSegment(100)
 	require.NoError(t, err)
+	b8, err := platform.MmapCodeSegment(100)
+	require.NoError(t, err)
+	b9, err := platform.MmapCodeSegment(100)
+	require.NoError(t, err)
 
 	sf.memoryGrowExecutable = b1
 	sf.stackGrowExecutable = b2
 	sf.checkModuleExitCode = b3
 	sf.tableGrowExecutable = b6
 	sf.refFuncExecutable = b7
+	sf.memoryWait32Executable = b8
+	sf.memoryWait64Executable = b9
 
 	sharedFunctionsFinalizer(sf)
 	require.Nil(t, sf.memoryGrowExecutable)
@@ -38,6 +44,8 @@ func Test_sharedFunctionsFinalizer(t *testing.T) {
 	require.Nil(t, sf.checkModuleExitCode)
 	require.Nil(t, sf.tableGrowExecutable)
 	require.Nil(t, sf.refFuncExecutable)
+	require.Nil(t, sf.memoryWait32Executable)
+	require.Nil(t, sf.memoryWait64Executable)
 }
 
 func Test_executablesFinalizer(t *testing.T) {
diff --git a/internal/engine/wazevo/frontend/frontend_test.go b/internal/engine/wazevo/frontend/frontend_test.go
index 826b7cf6..4f3ba849 100644
--- a/internal/engine/wazevo/frontend/frontend_test.go
+++ b/internal/engine/wazevo/frontend/frontend_test.go
@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/experimental"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/testcases"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/wazevoapi"
@@ -27,6 +28,7 @@ func TestCompiler_LowerToSSA(t *testing.T) {
 		exp string
 		// expAfterOpt is not empty when we want to check the result after optimization passes.
 		expAfterOpt string
+		features    api.CoreFeatures
 	}{
 		{
 			name: "empty", m: testcases.Empty.Module,
@@ -1613,6 +1615,62 @@ blk4: () <-- (blk2,blk3)
 blk0: (exec_ctx:i64, module_ctx:i64, v2:v128, v3:v128)
 	v4:v128 = Shuffle.[0 1 2 3 4 5 6 7 24 25 26 27 28 29 30 31] v2, v3
 	Jump blk_ret, v4
+`,
+		},
+		{
+			name:     "MemoryWait32",
+			m:        testcases.MemoryWait32.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
+			exp: `
+signatures:
+	sig6: i64i64i32i64_i32
+
+blk0: (exec_ctx:i64, module_ctx:i64, v2:i32, v3:i32, v4:i64)
+	Store module_ctx, exec_ctx, 0x8
+	v5:i64 = Iconst_64 0x4
+	v6:i64 = UExtend v2, 32->64
+	v7:i64 = Uload32 module_ctx, 0x10
+	v8:i64 = Iadd v6, v5
+	v9:i32 = Icmp lt_u, v7, v8
+	ExitIfTrue v9, exec_ctx, memory_out_of_bounds
+	v10:i64 = Load module_ctx, 0x8
+	v11:i64 = Iadd v10, v6
+	v12:i64 = Iconst_64 0x3
+	v13:i64 = Band v11, v12
+	v14:i64 = Iconst_64 0x0
+	v15:i32 = Icmp neq, v13, v14
+	ExitIfTrue v15, exec_ctx, unaligned_atomic
+	v16:i64 = Load exec_ctx, 0x488
+	v17:i32 = CallIndirect v16:sig6, exec_ctx, v4, v3, v11
+	Jump blk_ret, v17
+`,
+		},
+		{
+			name:     "MemoryWait64",
+			m:        testcases.MemoryWait64.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
+			exp: `
+signatures:
+	sig7: i64i64i64i64_i32
+
+blk0: (exec_ctx:i64, module_ctx:i64, v2:i32, v3:i64, v4:i64)
+	Store module_ctx, exec_ctx, 0x8
+	v5:i64 = Iconst_64 0x8
+	v6:i64 = UExtend v2, 32->64
+	v7:i64 = Uload32 module_ctx, 0x10
+	v8:i64 = Iadd v6, v5
+	v9:i32 = Icmp lt_u, v7, v8
+	ExitIfTrue v9, exec_ctx, memory_out_of_bounds
+	v10:i64 = Load module_ctx, 0x8
+	v11:i64 = Iadd v10, v6
+	v12:i64 = Iconst_64 0x7
+	v13:i64 = Band v11, v12
+	v14:i64 = Iconst_64 0x0
+	v15:i32 = Icmp neq, v13, v14
+	ExitIfTrue v15, exec_ctx, unaligned_atomic
+	v16:i64 = Load exec_ctx, 0x490
+	v17:i32 = CallIndirect v16:sig7, exec_ctx, v4, v3, v11
+	Jump blk_ret, v17
 `,
 		},
 	} {
@@ -1620,7 +1678,11 @@ blk0: (exec_ctx:i64, module_ctx:i64, v2:v128, v3:v128)
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
 			// Just in case let's check the test module is valid.
-			err := tc.m.Validate(api.CoreFeaturesV2)
+			features := tc.features
+			if features == 0 {
+				features = api.CoreFeaturesV2
+			}
+			err := tc.m.Validate(features)
 			require.NoError(t, err, "invalid test case module!")
 
 			b := ssa.NewBuilder()
@@ -1717,6 +1779,8 @@ func TestCompiler_declareSignatures(t *testing.T) {
 			{ID: 6, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32, ssa.TypeI32, ssa.TypeI64}, Results: []ssa.Type{ssa.TypeI32}},
 			{ID: 7, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32}, Results: []ssa.Type{ssa.TypeI64}},
 			{ID: 8, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeI64}},
+			{ID: 9, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeI32, ssa.TypeI64}, Results: []ssa.Type{ssa.TypeI32}},
+			{ID: 10, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeI64, ssa.TypeI64}, Results: []ssa.Type{ssa.TypeI32}},
 		}
 
 		require.Equal(t, len(expected), len(declaredSigs))
@@ -1753,6 +1817,8 @@ func TestCompiler_declareSignatures(t *testing.T) {
 			{ID: 14, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32, ssa.TypeI32, ssa.TypeI64}, Results: []ssa.Type{ssa.TypeI32}},
 			{ID: 15, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32}, Results: []ssa.Type{ssa.TypeI64}},
 			{ID: 16, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeI64}},
+			{ID: 17, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeI32, ssa.TypeI64}, Results: []ssa.Type{ssa.TypeI32}},
+			{ID: 18, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeI64, ssa.TypeI64}, Results: []ssa.Type{ssa.TypeI32}},
 		}
 		require.Equal(t, len(expected), len(declaredSigs))
 		for i := 0; i < len(declaredSigs); i++ {
diff --git a/internal/engine/wazevo/testcases/testcases.go b/internal/engine/wazevo/testcases/testcases.go
index f7b98929..cb26bfb6 100644
--- a/internal/engine/wazevo/testcases/testcases.go
+++ b/internal/engine/wazevo/testcases/testcases.go
@@ -1935,6 +1935,40 @@ var (
 		Name:   "shuffle",
 		Module: VecShuffleWithLane(0, 1, 2, 3, 4, 5, 6, 7, 24, 25, 26, 27, 28, 29, 30, 31),
 	}
+
+	MemoryWait32 = TestCase{
+		Name: "memory_wait32",
+		Module: &wasm.Module{
+			TypeSection:     []wasm.FunctionType{{Params: []wasm.ValueType{i32, i32, i64}, Results: []wasm.ValueType{i32}}},
+			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.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeLocalGet, 2,
+				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicMemoryWait32, 0x1, 0,
+				wasm.OpcodeEnd,
+			}}},
+		},
+	}
+
+	MemoryWait64 = TestCase{
+		Name: "memory_wait64",
+		Module: &wasm.Module{
+			TypeSection:     []wasm.FunctionType{{Params: []wasm.ValueType{i32, i64, i64}, Results: []wasm.ValueType{i32}}},
+			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.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeLocalGet, 2,
+				wasm.OpcodeAtomicPrefix, wasm.OpcodeAtomicMemoryWait64, 0x2, 0,
+				wasm.OpcodeEnd,
+			}}},
+		},
+	}
 )
 
 // VecShuffleWithLane returns a VecShuffle test with a custom 16-bytes immediate (lane indexes).
diff --git a/internal/engine/wazevo/wazevo_test.go b/internal/engine/wazevo/wazevo_test.go
index 2ceb125b..b9d168c7 100644
--- a/internal/engine/wazevo/wazevo_test.go
+++ b/internal/engine/wazevo/wazevo_test.go
@@ -79,4 +79,6 @@ func Test_ExecutionContextOffsets(t *testing.T) {
 	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.refFuncTrampolineAddress)), wazevoapi.ExecutionContextOffsetRefFuncTrampolineAddress)
 	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.memmoveAddress)), wazevoapi.ExecutionContextOffsetMemmoveAddress)
 	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.framePointerBeforeGoCall)), wazevoapi.ExecutionContextOffsetFramePointerBeforeGoCall)
+	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.memoryWait32TrampolineAddress)), wazevoapi.ExecutionContextOffsetMemoryWait32TrampolineAddress)
+	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.memoryWait64TrampolineAddress)), wazevoapi.ExecutionContextOffsetMemoryWait64TrampolineAddress)
 }