Repository: tetratelabs/wazero
Pull request #1987 (2024-01-31 04:58:09)

wazevo: adds support for memory notify


Patch:
diff --git a/internal/engine/wazevo/call_engine.go b/internal/engine/wazevo/call_engine.go
index 58e47747..8aabd4a4 100644
--- a/internal/engine/wazevo/call_engine.go
+++ b/internal/engine/wazevo/call_engine.go
@@ -87,6 +87,8 @@ type (
 		memoryWait32TrampolineAddress *byte
 		// memoryWait32TrampolineAddress holds the address of memory_wait64 trampoline function.
 		memoryWait64TrampolineAddress *byte
+		// memoryNotifyTrampolineAddress holds the address of the memory_notify trampoline function.
+		memoryNotifyTrampolineAddress *byte
 	}
 )
 
@@ -426,6 +428,18 @@ 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.ExitCodeMemoryNotify:
+			mod := c.callerModuleInstance()
+			mem := mod.MemoryInstance
+
+			s := goCallStackView(c.execCtx.stackPointerBeforeGoCall)
+			count, addr := uint32(s[0]), s[1]
+			offset := uint32(uintptr(addr) - uintptr(unsafe.Pointer(&mem.Buffer[0])))
+			res := mem.Notify(offset, count)
+			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:
diff --git a/internal/engine/wazevo/engine.go b/internal/engine/wazevo/engine.go
index 9bc6d359..0fd3b70f 100644
--- a/internal/engine/wazevo/engine.go
+++ b/internal/engine/wazevo/engine.go
@@ -60,7 +60,9 @@ type (
 		// 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
+		memoryWait64Executable []byte
+		// memoryNotifyExecutable is a compiled trampoline executable for memory.notify builtin function
+		memoryNotifyExecutable    []byte
 		listenerBeforeTrampolines map[*wasm.FunctionType][]byte
 		listenerAfterTrampolines  map[*wasm.FunctionType][]byte
 	}
@@ -673,6 +675,21 @@ func (e *engine) compileSharedFunctions() {
 		}
 	}
 
+	e.be.Init()
+	{
+		src := e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeMemoryNotify, &ssa.Signature{
+			// exec context, count, addr
+			Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32, ssa.TypeI64},
+			// Returns the number notified.
+			Results: []ssa.Type{ssa.TypeI32},
+		}, false)
+		e.sharedFunctions.memoryNotifyExecutable = mmapExecutable(src)
+		if wazevoapi.PerfMapEnabled {
+			exe := e.sharedFunctions.memoryNotifyExecutable
+			wazevoapi.PerfMap.AddEntry(uintptr(unsafe.Pointer(&exe[0])), uint64(len(exe)), "memory_notify_trampoline")
+		}
+	}
+
 	e.setFinalizer(e.sharedFunctions, sharedFunctionsFinalizer)
 }
 
@@ -698,6 +715,9 @@ func sharedFunctionsFinalizer(sf *sharedFunctions) {
 	if err := platform.MunmapCodeSegment(sf.memoryWait64Executable); err != nil {
 		panic(err)
 	}
+	if err := platform.MunmapCodeSegment(sf.memoryNotifyExecutable); err != nil {
+		panic(err)
+	}
 	for _, f := range sf.listenerBeforeTrampolines {
 		if err := platform.MunmapCodeSegment(f); err != nil {
 			panic(err)
@@ -716,6 +736,7 @@ func sharedFunctionsFinalizer(sf *sharedFunctions) {
 	sf.refFuncExecutable = nil
 	sf.memoryWait32Executable = nil
 	sf.memoryWait64Executable = nil
+	sf.memoryNotifyExecutable = nil
 	sf.listenerBeforeTrampolines = nil
 	sf.listenerAfterTrampolines = nil
 }
diff --git a/internal/engine/wazevo/frontend/frontend.go b/internal/engine/wazevo/frontend/frontend.go
index 1c468499..265b642a 100644
--- a/internal/engine/wazevo/frontend/frontend.go
+++ b/internal/engine/wazevo/frontend/frontend.go
@@ -23,6 +23,7 @@ type Compiler struct {
 	memoryGrowSig          ssa.Signature
 	memoryWait32Sig        ssa.Signature
 	memoryWait64Sig        ssa.Signature
+	memoryNotifySig        ssa.Signature
 	checkModuleExitCodeSig ssa.Signature
 	tableGrowSig           ssa.Signature
 	refFuncSig             ssa.Signature
@@ -163,6 +164,15 @@ func (c *Compiler) declareSignatures(listenerOn bool) {
 		Results: []ssa.Type{ssa.TypeI32},
 	}
 	c.ssaBuilder.DeclareSignature(&c.memoryWait64Sig)
+
+	c.memoryNotifySig = ssa.Signature{
+		ID: c.memoryWait64Sig.ID + 1,
+		// exec context, count, addr
+		Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32, ssa.TypeI64},
+		// Returns the number notified.
+		Results: []ssa.Type{ssa.TypeI32},
+	}
+	c.ssaBuilder.DeclareSignature(&c.memoryNotifySig)
 }
 
 // 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 fca32c69..f9d5a19b 100644
--- a/internal/engine/wazevo/frontend/lower.go
+++ b/internal/engine/wazevo/frontend/lower.go
@@ -3194,6 +3194,28 @@ func (c *Compiler) lowerCurrentOpcode() {
 				AsCallIndirect(memoryWaitPtr, sig, args).
 				Insert(builder).Return()
 			state.push(memoryWaitRet)
+		case wasm.OpcodeAtomicMemoryNotify:
+			_, offset := c.readMemArg()
+			if state.unreachable {
+				break
+			}
+
+			c.storeCallerModuleContext()
+			count := state.pop()
+			baseAddr := state.pop()
+			addr := c.memOpSetup(baseAddr, uint64(offset), 4)
+			c.memAlignmentCheck(addr, 4)
+
+			memoryNotifyPtr := builder.AllocateInstruction().
+				AsLoad(c.execCtxPtrValue,
+					wazevoapi.ExecutionContextOffsetMemoryNotifyTrampolineAddress.U32(),
+					ssa.TypeI64,
+				).Insert(builder).Return()
+			args := []ssa.Value{c.execCtxPtrValue, count, addr}
+			memoryNotifyRet := builder.AllocateInstruction().
+				AsCallIndirect(memoryNotifyPtr, &c.memoryNotifySig, args).
+				Insert(builder).Return()
+			state.push(memoryNotifyRet)
 		default:
 			panic("TODO: unsupported atomic instruction: " + wasm.AtomicInstructionName(atomicOp))
 		}
diff --git a/internal/engine/wazevo/module_engine.go b/internal/engine/wazevo/module_engine.go
index b753fbe6..9aee16f2 100644
--- a/internal/engine/wazevo/module_engine.go
+++ b/internal/engine/wazevo/module_engine.go
@@ -185,6 +185,7 @@ func (m *moduleEngine) NewFunction(index wasm.Index) api.Function {
 	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.memoryNotifyTrampolineAddress = &m.parent.sharedFunctions.memoryNotifyExecutable[0]
 	ce.execCtx.memmoveAddress = memmovPtr
 	ce.init()
 	return ce
diff --git a/internal/engine/wazevo/wazevoapi/exitcode.go b/internal/engine/wazevo/wazevoapi/exitcode.go
index 4e9374c1..5ad59498 100644
--- a/internal/engine/wazevo/wazevoapi/exitcode.go
+++ b/internal/engine/wazevo/wazevoapi/exitcode.go
@@ -28,6 +28,7 @@ const (
 	ExitCodeRefFunc
 	ExitCodeMemoryWait32
 	ExitCodeMemoryWait64
+	ExitCodeMemoryNotify
 	ExitCodeUnalignedAtomic
 	exitCodeMax
 )
@@ -83,6 +84,8 @@ func (e ExitCode) String() string {
 		return "memory_wait32"
 	case ExitCodeMemoryWait64:
 		return "memory_wait64"
+	case ExitCodeMemoryNotify:
+		return "memory_notify"
 	}
 	panic("TODO")
 }
diff --git a/internal/engine/wazevo/wazevoapi/offsetdata.go b/internal/engine/wazevo/wazevoapi/offsetdata.go
index 0fd3123e..56a89517 100644
--- a/internal/engine/wazevo/wazevoapi/offsetdata.go
+++ b/internal/engine/wazevo/wazevoapi/offsetdata.go
@@ -52,6 +52,7 @@ const (
 	ExecutionContextOffsetFramePointerBeforeGoCall      Offset = 1152
 	ExecutionContextOffsetMemoryWait32TrampolineAddress Offset = 1160
 	ExecutionContextOffsetMemoryWait64TrampolineAddress Offset = 1168
+	ExecutionContextOffsetMemoryNotifyTrampolineAddress Offset = 1176
 )
 
 // 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 ce5bc3af..3b99cf10 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -612,7 +612,11 @@ func TestE2E(t *testing.T) {
 			m:        testcases.MemoryWait32.Module,
 			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			calls: []callCase{
-				{params: []uint64{0x0, 0xbeef, 0xffffffff}, expResults: []uint64{1}},
+				{params: []uint64{0x0, 0xbeef, 0xffffffff}, expResults: []uint64{1}}, // exp not equal, returns 1
+				{params: []uint64{0x1, 0xbeef, 0xffffffff}, expErr: "unaligned atomic"},
+				{params: []uint64{0x2, 0xbeef, 0xffffffff}, expErr: "unaligned atomic"},
+				{params: []uint64{0x3, 0xbeef, 0xffffffff}, expErr: "unaligned atomic"},
+				{params: []uint64{0x4, 0xbeef, 0xffffffff}, expResults: []uint64{1}}, // exp not equal, returns 1
 			},
 		},
 		{
@@ -620,7 +624,27 @@ func TestE2E(t *testing.T) {
 			m:        testcases.MemoryWait64.Module,
 			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
 			calls: []callCase{
-				{params: []uint64{0x0, 0xbeef, 0xffffffff}, expResults: []uint64{1}},
+				{params: []uint64{0x0, 0xbeef, 0xffffffff}, expResults: []uint64{1}}, // exp not equal, returns 1
+				{params: []uint64{0x1, 0xbeef, 0xffffffff}, expErr: "unaligned atomic"},
+				{params: []uint64{0x2, 0xbeef, 0xffffffff}, expErr: "unaligned atomic"},
+				{params: []uint64{0x3, 0xbeef, 0xffffffff}, expErr: "unaligned atomic"},
+				{params: []uint64{0x4, 0xbeef, 0xffffffff}, expErr: "unaligned atomic"},
+				{params: []uint64{0x5, 0xbeef, 0xffffffff}, expErr: "unaligned atomic"},
+				{params: []uint64{0x6, 0xbeef, 0xffffffff}, expErr: "unaligned atomic"},
+				{params: []uint64{0x7, 0xbeef, 0xffffffff}, expErr: "unaligned atomic"},
+				{params: []uint64{0x8, 0xbeef, 0xffffffff}, expResults: []uint64{1}}, // exp not equal, returns 1
+			},
+		},
+		{
+			name:     "memory_notify",
+			m:        testcases.MemoryNotify.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
+			calls: []callCase{
+				{params: []uint64{0x0, 0x1}, expResults: []uint64{0}}, // no waiters, returns 0
+				{params: []uint64{0x1, 0x1}, expErr: "unaligned atomic"},
+				{params: []uint64{0x2, 0x1}, expErr: "unaligned atomic"},
+				{params: []uint64{0x3, 0x1}, expErr: "unaligned atomic"},
+				{params: []uint64{0x4, 0x1}, expResults: []uint64{0}}, // no waiters, returns 0
 			},
 		},
 		{
diff --git a/internal/engine/wazevo/engine_test.go b/internal/engine/wazevo/engine_test.go
index 6f7679c0..555ef7a5 100644
--- a/internal/engine/wazevo/engine_test.go
+++ b/internal/engine/wazevo/engine_test.go
@@ -29,6 +29,8 @@ func Test_sharedFunctionsFinalizer(t *testing.T) {
 	require.NoError(t, err)
 	b9, err := platform.MmapCodeSegment(100)
 	require.NoError(t, err)
+	b10, err := platform.MmapCodeSegment(100)
+	require.NoError(t, err)
 
 	sf.memoryGrowExecutable = b1
 	sf.stackGrowExecutable = b2
@@ -37,6 +39,7 @@ func Test_sharedFunctionsFinalizer(t *testing.T) {
 	sf.refFuncExecutable = b7
 	sf.memoryWait32Executable = b8
 	sf.memoryWait64Executable = b9
+	sf.memoryNotifyExecutable = b10
 
 	sharedFunctionsFinalizer(sf)
 	require.Nil(t, sf.memoryGrowExecutable)
@@ -46,6 +49,7 @@ func Test_sharedFunctionsFinalizer(t *testing.T) {
 	require.Nil(t, sf.refFuncExecutable)
 	require.Nil(t, sf.memoryWait32Executable)
 	require.Nil(t, sf.memoryWait64Executable)
+	require.Nil(t, sf.memoryNotifyExecutable)
 }
 
 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 4f3ba849..cb331441 100644
--- a/internal/engine/wazevo/frontend/frontend_test.go
+++ b/internal/engine/wazevo/frontend/frontend_test.go
@@ -1671,6 +1671,34 @@ blk0: (exec_ctx:i64, module_ctx:i64, v2:i32, v3:i64, v4:i64)
 	v16:i64 = Load exec_ctx, 0x490
 	v17:i32 = CallIndirect v16:sig7, exec_ctx, v4, v3, v11
 	Jump blk_ret, v17
+`,
+		},
+		{
+			name:     "MemoryNotify",
+			m:        testcases.MemoryNotify.Module,
+			features: api.CoreFeaturesV2 | experimental.CoreFeaturesThreads,
+			exp: `
+signatures:
+	sig8: i64i32i64_i32
+
+blk0: (exec_ctx:i64, module_ctx:i64, v2:i32, v3:i32)
+	Store module_ctx, exec_ctx, 0x8
+	v4:i64 = Iconst_64 0x4
+	v5:i64 = UExtend v2, 32->64
+	v6:i64 = Uload32 module_ctx, 0x10
+	v7:i64 = Iadd v5, v4
+	v8:i32 = Icmp lt_u, v6, v7
+	ExitIfTrue v8, exec_ctx, memory_out_of_bounds
+	v9:i64 = Load module_ctx, 0x8
+	v10:i64 = Iadd v9, v5
+	v11:i64 = Iconst_64 0x3
+	v12:i64 = Band v10, v11
+	v13:i64 = Iconst_64 0x0
+	v14:i32 = Icmp neq, v12, v13
+	ExitIfTrue v14, exec_ctx, unaligned_atomic
+	v15:i64 = Load exec_ctx, 0x498
+	v16:i32 = CallIndirect v15:sig8, exec_ctx, v3, v10
+	Jump blk_ret, v16
 `,
 		},
 	} {
@@ -1781,6 +1809,7 @@ func TestCompiler_declareSignatures(t *testing.T) {
 			{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}},
+			{ID: 11, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32, ssa.TypeI64}, Results: []ssa.Type{ssa.TypeI32}},
 		}
 
 		require.Equal(t, len(expected), len(declaredSigs))
@@ -1819,6 +1848,7 @@ func TestCompiler_declareSignatures(t *testing.T) {
 			{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}},
+			{ID: 19, Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32, 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 cb26bfb6..6b04c9e9 100644
--- a/internal/engine/wazevo/testcases/testcases.go
+++ b/internal/engine/wazevo/testcases/testcases.go
@@ -1969,6 +1969,22 @@ var (
 			}}},
 		},
 	}
+
+	MemoryNotify = TestCase{
+		Name: "memory_notify",
+		Module: &wasm.Module{
+			TypeSection:     []wasm.FunctionType{{Params: []wasm.ValueType{i32, i32}, 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.OpcodeAtomicPrefix, wasm.OpcodeAtomicMemoryNotify, 0x1, 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 b9d168c7..ffe74846 100644
--- a/internal/engine/wazevo/wazevo_test.go
+++ b/internal/engine/wazevo/wazevo_test.go
@@ -81,4 +81,5 @@ func Test_ExecutionContextOffsets(t *testing.T) {
 	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)
+	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.memoryNotifyTrampolineAddress)), wazevoapi.ExecutionContextOffsetMemoryNotifyTrampolineAddress)
 }