Align preambles, trampolines to 16 bytes. (#2432)

Signed-off-by: Nuno Cruces <ncruces@users.noreply.github.com>

Author: ncruces

internal/engine/wazevo/engine.go

@@ -187,7 +187,9 @@ func (exec *executables) compileEntryPreambles(m *wasm.Module, machine backend.M
 		be.Init()
 		buf := machine.CompileEntryPreamble(&sig)
 		preambles = append(preambles, buf...)
-		sizes[i] = len(buf)
+		align := 15 & -len(preambles) // Align 16-bytes boundary.
+		preambles = append(preambles, make([]byte, align)...)
+		sizes[i] = len(buf) + align
 	}
 
 	exec.entryPreambles = mmapExecutable(preambles)
@@ -725,88 +727,70 @@ func (e *engine) compileSharedFunctions() {
 	var sizes [8]int
 	var trampolines []byte
 
+	addTrampoline := func(i int, buf []byte) {
+		trampolines = append(trampolines, buf...)
+		align := 15 & -len(trampolines) // Align 16-bytes boundary.
+		trampolines = append(trampolines, make([]byte, align)...)
+		sizes[i] = len(buf) + align
+	}
+
 	e.be.Init()
-	{
-		src := e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeGrowMemory, &ssa.Signature{
+	addTrampoline(0,
+		e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeGrowMemory, &ssa.Signature{
 			Params:  []ssa.Type{ssa.TypeI64 /* exec context */, ssa.TypeI32},
 			Results: []ssa.Type{ssa.TypeI32},
-		}, false)
-		trampolines = append(trampolines, src...)
-		sizes[0] = len(src)
-	}
+		}, false))
 
 	e.be.Init()
-	{
-		src := e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeTableGrow, &ssa.Signature{
+	addTrampoline(1,
+		e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeTableGrow, &ssa.Signature{
 			Params:  []ssa.Type{ssa.TypeI64 /* exec context */, ssa.TypeI32 /* table index */, ssa.TypeI32 /* num */, ssa.TypeI64 /* ref */},
 			Results: []ssa.Type{ssa.TypeI32},
-		}, false)
-		trampolines = append(trampolines, src...)
-		sizes[1] = len(src)
-	}
+		}, false))
 
 	e.be.Init()
-	{
-		src := e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeCheckModuleExitCode, &ssa.Signature{
+	addTrampoline(2,
+		e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeCheckModuleExitCode, &ssa.Signature{
 			Params:  []ssa.Type{ssa.TypeI32 /* exec context */},
 			Results: []ssa.Type{ssa.TypeI32},
-		}, false)
-		trampolines = append(trampolines, src...)
-		sizes[2] = len(src)
-	}
+		}, false))
 
 	e.be.Init()
-	{
-		src := e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeRefFunc, &ssa.Signature{
+	addTrampoline(3,
+		e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeRefFunc, &ssa.Signature{
 			Params:  []ssa.Type{ssa.TypeI64 /* exec context */, ssa.TypeI32 /* function index */},
 			Results: []ssa.Type{ssa.TypeI64}, // returns the function reference.
-		}, false)
-		trampolines = append(trampolines, src...)
-		sizes[3] = len(src)
-	}
+		}, false))
 
 	e.be.Init()
-	{
-		src := e.machine.CompileStackGrowCallSequence()
-		trampolines = append(trampolines, src...)
-		sizes[4] = len(src)
-	}
+	addTrampoline(4, e.machine.CompileStackGrowCallSequence())
 
 	e.be.Init()
-	{
-		src := e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeMemoryWait32, &ssa.Signature{
+	addTrampoline(5,
+		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)
-		trampolines = append(trampolines, src...)
-		sizes[5] = len(src)
-	}
+		}, false))
 
 	e.be.Init()
-	{
-		src := e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeMemoryWait64, &ssa.Signature{
+	addTrampoline(6,
+		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)
-		trampolines = append(trampolines, src...)
-		sizes[6] = len(src)
-	}
+		}, false))
 
 	e.be.Init()
-	{
-		src := e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeMemoryNotify, &ssa.Signature{
+	addTrampoline(7,
+		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)
-		trampolines = append(trampolines, src...)
-		sizes[7] = len(src)
-	}
+		}, false))
 
 	fns := &sharedFunctions{
 		executable:          mmapExecutable(trampolines),
@@ -916,7 +900,9 @@ func (e *engine) getListenerTrampolineForType(functionType *wasm.FunctionType) (
 		buf := e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeCallListenerBefore, beforeSig, false)
 		executable = append(executable, buf...)
 
-		offset := len(buf)
+		align := 15 & -len(executable) // Align 16-bytes boundary.
+		executable = append(executable, make([]byte, align)...)
+		offset := len(executable)
 
 		e.be.Init()
 		buf = e.machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeCallListenerAfter, afterSig, false)

internal/engine/wazevo/engine_test.go

@@ -80,6 +80,51 @@ func TestEngine_CompileModule(t *testing.T) {
 	}
 }
 
+func TestEngine_CompileModule_alignment(t *testing.T) {
+	ctx := context.Background()
+	e := NewEngine(ctx, 0, nil).(*engine)
+
+	okModule := &wasm.Module{
+		TypeSection:     []wasm.FunctionType{{}},
+		FunctionSection: []wasm.Index{0, 0, 0, 0},
+		CodeSection: []wasm.Code{
+			{Body: []byte{wasm.OpcodeEnd}},
+			{Body: []byte{wasm.OpcodeEnd}},
+			{Body: []byte{wasm.OpcodeEnd}},
+			{Body: []byte{wasm.OpcodeEnd}},
+		},
+		ID: wasm.ModuleID{},
+	}
+
+	err := e.CompileModule(ctx, okModule, nil, false)
+	require.NoError(t, err)
+
+	cm, ok := e.getCompiledModuleFromMemory(okModule)
+	require.True(t, ok)
+
+	for _, offset := range cm.functionOffsets {
+		require.True(t, offset&15 == 0)
+	}
+
+	for _, ptr := range cm.entryPreamblesPtrs {
+		require.True(t, uintptr(unsafe.Pointer(ptr))&15 == 0)
+	}
+
+	shared := cm.sharedFunctions
+	require.True(t, uintptr(unsafe.Pointer(shared.memoryGrowAddress))&15 == 0)
+	require.True(t, uintptr(unsafe.Pointer(shared.checkModuleExitCodeAddress))&15 == 0)
+	require.True(t, uintptr(unsafe.Pointer(shared.stackGrowAddress))&15 == 0)
+	require.True(t, uintptr(unsafe.Pointer(shared.tableGrowAddress))&15 == 0)
+	require.True(t, uintptr(unsafe.Pointer(shared.refFuncAddress))&15 == 0)
+	require.True(t, uintptr(unsafe.Pointer(shared.memoryWait32Address))&15 == 0)
+	require.True(t, uintptr(unsafe.Pointer(shared.memoryWait64Address))&15 == 0)
+	require.True(t, uintptr(unsafe.Pointer(shared.memoryNotifyAddress))&15 == 0)
+	for _, trampoline := range shared.listenerTrampolines {
+		require.True(t, uintptr(unsafe.Pointer(trampoline.before))&15 == 0)
+		require.True(t, uintptr(unsafe.Pointer(trampoline.after))&15 == 0)
+	}
+}
+
 func TestEngine_sortedCompiledModules(t *testing.T) {
 	requireEqualExisting := func(t *testing.T, e *engine, expected []uintptr) {
 		actual := make([]uintptr, 0)