mleku/wazero
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

wazevo: adds support for host functions (#1630)

Browse Source 
Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>
This commit is contained in:
Takeshi Yoneda
2023-08-15 12:00:36 +09:00
committed by GitHub
parent 84708fe4e6
commit dbdd3f5541
21 changed files with 584 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
internal/engine/wazevo/backend/isa/arm64/abi.go
View File

@@ -76,19 +76,23 @@ func (a *abiImpl) init(sig *ssa.Signature) {
if len(a.rets) < len(sig.Results) {
a.rets = make([]backend.ABIArg, len(sig.Results))
}
a.rets = a.rets[:len(sig.Results)]
a.retStackSize = a.setABIArgs(a.rets, sig.Results)
if argsNum := len(sig.Params); len(a.args) < argsNum {
a.args = make([]backend.ABIArg, argsNum)
}
a.args = a.args[:len(sig.Params)]
a.argStackSize = a.setABIArgs(a.args, sig.Params)
// Gather the real registers usages in arg/return.
a.retRealRegs = a.retRealRegs[:0]
for i := range a.rets {
r := &a.rets[i]
if r.Kind == backend.ABIArgKindReg {
a.retRealRegs = append(a.retRealRegs, r.Reg)
}
}
a.argRealRegs = a.argRealRegs[:0]
for i := range a.args {
arg := &a.args[i]
if arg.Kind == backend.ABIArgKindReg {

51
internal/engine/wazevo/backend/isa/arm64/abi_go_call.go
View File

@@ -13,7 +13,7 @@ var calleeSavedRegistersPlusLinkRegSorted = []regalloc.VReg{
}
// CompileGoFunctionTrampoline implements backend.Machine.
func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *ssa.Signature) {
func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *ssa.Signature, needModuleContextPtr bool) {
cur := m.allocateInstr()
cur.asNop0()
m.rootInstr = cur
@@ -26,6 +26,25 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
// Next, we need to store all the arguments to the execution context.
argBegin := 1 // Skips exec context by default.
if needModuleContextPtr {
offset := wazevoapi.ExecutionContextOffsets.GoFunctionCallCalleeModuleContextOpaque.I64()
if !offsetFitsInAddressModeKindRegUnsignedImm12(64, offset) {
panic("BUG: too large or un-aligned offset for goFunctionCallCalleeModuleContextOpaque in execution context")
}
// Module context is always the second argument.
moduleCtrPtr := x1VReg
store := m.allocateInstr()
amode := addressMode{kind: addressModeKindRegUnsignedImm12, rn: execCtrPtr, imm: offset}
store.asStore(operandNR(moduleCtrPtr), amode, 64)
store.prev = cur
cur.next = store
cur = store
argBegin++
}
stackPtrReg := x15VReg // Caller save, so we can use it for whatever we want.
imm12Op, ok := asImm12Operand(wazevoapi.ExecutionContextOffsets.GoFunctionCallStackBegin.U64())
if !ok {
@@ -39,7 +58,7 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
cur = goCallStackPtrCalc
abi := m.getOrCreateABIImpl(sig)
for _, arg := range abi.args[1:] { // Skips exec context.
for _, arg := range abi.args[argBegin:] {
if arg.Kind == backend.ABIArgKindReg {
store := m.allocateInstr()
v := arg.Reg
@@ -90,6 +109,7 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
movTmpToReg := m.allocateInstr()
mode := addressMode{kind: addressModeKindPostIndex, rn: stackPtrReg}
tmpRegVRegVec := v17VReg // Caller save, so we can use it for whatever we want.
switch r.Type {
case ssa.TypeI32:
mode.imm = 8 // We use uint64 for all basic types, except SIMD v128.
@@ -101,12 +121,12 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
movTmpToReg.asMove64(r.Reg, tmpRegVReg)
case ssa.TypeF32:
mode.imm = 8 // We use uint64 for all basic types, except SIMD v128.
loadIntoTmp.asFpuLoad(operandNR(tmpRegVReg), mode, 32)
movTmpToReg.asFpuMov64(r.Reg, tmpRegVReg)
loadIntoTmp.asFpuLoad(operandNR(tmpRegVRegVec), mode, 32)
movTmpToReg.asFpuMov64(r.Reg, tmpRegVRegVec)
case ssa.TypeF64:
mode.imm = 8 // We use uint64 for all basic types, except SIMD v128.
loadIntoTmp.asFpuLoad(operandNR(tmpRegVReg), mode, 64)
movTmpToReg.asFpuMov64(r.Reg, tmpRegVReg)
loadIntoTmp.asFpuLoad(operandNR(tmpRegVRegVec), mode, 64)
movTmpToReg.asFpuMov64(r.Reg, tmpRegVRegVec)
}
loadIntoTmp.prev = cur
cur.next = loadIntoTmp
@@ -181,16 +201,19 @@ func (m *machine) restoreRegistersInExecutionContext(cur *instruction, regs []re
}
func (m *machine) setExitCode(cur *instruction, execCtr regalloc.VReg, exitCode wazevoapi.ExitCode) *instruction {
constReg := x17VReg // caller-saved, so we can use it.
m.pendingInstructions = m.pendingInstructions[:0]
m.lowerConstantI32(x17VReg, int32(exitCode))
for _, instr := range m.pendingInstructions {
instr.prev = cur
cur.next = instr
cur = instr
}
// Set the exit status on the execution context.
// movz tmp, #wazevoapi.ExitCodeGrowStack
// str tmp, [exec_context]
loadStatusConst := m.allocateInstrAfterLowering()
loadStatusConst.asMOVZ(tmpRegVReg, uint64(exitCode), 0, true)
loadStatusConst.prev = cur
cur.next = loadStatusConst
cur = loadStatusConst
setExistStatus := m.allocateInstrAfterLowering()
setExistStatus.asStore(operandNR(tmpRegVReg),
setExistStatus.asStore(operandNR(constReg),
addressMode{
kind: addressModeKindRegUnsignedImm12,
rn: execCtr, imm: wazevoapi.ExecutionContextOffsets.ExitCodeOffset.I64(),

176
internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go
View File

@@ -27,11 +27,169 @@ func Test_calleeSavedRegistersPlusLinkRegSorted(t *testing.T) {
func TestMachine_CompileGoFunctionTrampoline(t *testing.T) {
for _, tc := range []struct {
name string
exitCode wazevoapi.ExitCode
sig *ssa.Signature
exp string
name string
exitCode wazevoapi.ExitCode
sig *ssa.Signature
needModuleContextPtr bool
exp string
}{
{
name: "go call",
exitCode: wazevoapi.ExitCodeCallGoFunctionWithIndex(100),
sig: &ssa.Signature{
Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeF64},
Results: []ssa.Type{ssa.TypeI32, ssa.TypeI64, ssa.TypeF32, ssa.TypeF64},
},
needModuleContextPtr: true,
exp: `
str x18, [x0, #0x50]
str x19, [x0, #0x60]
str x20, [x0, #0x70]
str x21, [x0, #0x80]
str x22, [x0, #0x90]
str x23, [x0, #0xa0]
str x24, [x0, #0xb0]
str x25, [x0, #0xc0]
str x26, [x0, #0xd0]
str x28, [x0, #0xe0]
str x30, [x0, #0xf0]
str q18, [x0, #0x100]
str q19, [x0, #0x110]
str q20, [x0, #0x120]
str q21, [x0, #0x130]
str q22, [x0, #0x140]
str q23, [x0, #0x150]
str q24, [x0, #0x160]
str q25, [x0, #0x170]
str q26, [x0, #0x180]
str q27, [x0, #0x190]
str q28, [x0, #0x1a0]
str q29, [x0, #0x1b0]
str q30, [x0, #0x1c0]
str q31, [x0, #0x1d0]
str x1, [x0, #0x450]
add x15, x0, #0x458
str d0, [x15], #0x8
movz w17, #0x6406, LSL 0
str w17, [x0]
mov x27, sp
str x27, [x0, #0x38]
adr x27, #0x1c
str x27, [x0, #0x30]
exit_sequence w0
ldr x18, [x0, #0x50]
ldr x19, [x0, #0x60]
ldr x20, [x0, #0x70]
ldr x21, [x0, #0x80]
ldr x22, [x0, #0x90]
ldr x23, [x0, #0xa0]
ldr x24, [x0, #0xb0]
ldr x25, [x0, #0xc0]
ldr x26, [x0, #0xd0]
ldr x28, [x0, #0xe0]
ldr x30, [x0, #0xf0]
ldr q18, [x0, #0x100]
ldr q19, [x0, #0x110]
ldr q20, [x0, #0x120]
ldr q21, [x0, #0x130]
ldr q22, [x0, #0x140]
ldr q23, [x0, #0x150]
ldr q24, [x0, #0x160]
ldr q25, [x0, #0x170]
ldr q26, [x0, #0x180]
ldr q27, [x0, #0x190]
ldr q28, [x0, #0x1a0]
ldr q29, [x0, #0x1b0]
ldr q30, [x0, #0x1c0]
ldr q31, [x0, #0x1d0]
add x15, x0, #0x458
ldr w27, [x15], #0x8
mov w0, w27
ldr x27, [x15], #0x8
mov x1, x27
ldr s17, [x15], #0x8
mov q0.8b, q17.8b
ldr d17, [x15], #0x8
mov q1.8b, q17.8b
ret
`,
},
{
name: "go call",
exitCode: wazevoapi.ExitCodeCallGoFunctionWithIndex(100),
sig: &ssa.Signature{
Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeF64, ssa.TypeF64, ssa.TypeI32, ssa.TypeI32},
Results: []ssa.Type{},
},
needModuleContextPtr: true,
exp: `
str x18, [x0, #0x50]
str x19, [x0, #0x60]
str x20, [x0, #0x70]
str x21, [x0, #0x80]
str x22, [x0, #0x90]
str x23, [x0, #0xa0]
str x24, [x0, #0xb0]
str x25, [x0, #0xc0]
str x26, [x0, #0xd0]
str x28, [x0, #0xe0]
str x30, [x0, #0xf0]
str q18, [x0, #0x100]
str q19, [x0, #0x110]
str q20, [x0, #0x120]
str q21, [x0, #0x130]
str q22, [x0, #0x140]
str q23, [x0, #0x150]
str q24, [x0, #0x160]
str q25, [x0, #0x170]
str q26, [x0, #0x180]
str q27, [x0, #0x190]
str q28, [x0, #0x1a0]
str q29, [x0, #0x1b0]
str q30, [x0, #0x1c0]
str q31, [x0, #0x1d0]
str x1, [x0, #0x450]
add x15, x0, #0x458
str d0, [x15], #0x8
str d1, [x15], #0x8
str x2, [x15], #0x8
str x3, [x15], #0x8
movz w17, #0x6406, LSL 0
str w17, [x0]
mov x27, sp
str x27, [x0, #0x38]
adr x27, #0x1c
str x27, [x0, #0x30]
exit_sequence w0
ldr x18, [x0, #0x50]
ldr x19, [x0, #0x60]
ldr x20, [x0, #0x70]
ldr x21, [x0, #0x80]
ldr x22, [x0, #0x90]
ldr x23, [x0, #0xa0]
ldr x24, [x0, #0xb0]
ldr x25, [x0, #0xc0]
ldr x26, [x0, #0xd0]
ldr x28, [x0, #0xe0]
ldr x30, [x0, #0xf0]
ldr q18, [x0, #0x100]
ldr q19, [x0, #0x110]
ldr q20, [x0, #0x120]
ldr q21, [x0, #0x130]
ldr q22, [x0, #0x140]
ldr q23, [x0, #0x150]
ldr q24, [x0, #0x160]
ldr q25, [x0, #0x170]
ldr q26, [x0, #0x180]
ldr q27, [x0, #0x190]
ldr q28, [x0, #0x1a0]
ldr q29, [x0, #0x1b0]
ldr q30, [x0, #0x1c0]
ldr q31, [x0, #0x1d0]
add x15, x0, #0x458
ret
`,
},
{
name: "grow memory",
exitCode: wazevoapi.ExitCodeGrowMemory,
@@ -65,10 +223,10 @@ func TestMachine_CompileGoFunctionTrampoline(t *testing.T) {
str q29, [x0, #0x1b0]
str q30, [x0, #0x1c0]
str q31, [x0, #0x1d0]
add x15, x0, #0x450
add x15, x0, #0x458
str x1, [x15], #0x8
movz x27, #0x2, LSL 0
str w27, [x0]
orr w17, wzr, #0x2
str w17, [x0]
mov x27, sp
str x27, [x0, #0x38]
adr x27, #0x1c
@@ -99,7 +257,7 @@ func TestMachine_CompileGoFunctionTrampoline(t *testing.T) {
ldr q29, [x0, #0x1b0]
ldr q30, [x0, #0x1c0]
ldr q31, [x0, #0x1d0]
add x15, x0, #0x450
add x15, x0, #0x458
ldr w27, [x15], #0x8
mov w0, w27
ret
@@ -108,7 +266,7 @@ func TestMachine_CompileGoFunctionTrampoline(t *testing.T) {
} {
t.Run(tc.name, func(t *testing.T) {
_, _, m := newSetupWithMockContext()
m.CompileGoFunctionTrampoline(tc.exitCode, tc.sig)
m.CompileGoFunctionTrampoline(tc.exitCode, tc.sig, tc.needModuleContextPtr)
fmt.Println(m.Format())
require.Equal(t, tc.exp, m.Format())

1
internal/engine/wazevo/backend/isa/arm64/machine.go
View File

@@ -134,6 +134,7 @@ func (m *machine) Reset() {
m.regAllocFn.reset()
m.spillSlotSize = 0
m.unresolvedAddressModes = m.unresolvedAddressModes[:0]
m.rootInstr = nil
}
// InitializeABI implements backend.Machine InitializeABI.

1
internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
View File

@@ -331,6 +331,7 @@ func (m *machine) insertStackBoundsCheck(requiredStackSize int64, cur *instructi
// Set the required stack size and set it to the exec context.
{
// First load the requiredStackSize into the temporary register,
m.pendingInstructions = m.pendingInstructions[:0]
m.lowerConstantI64(tmpRegVReg, requiredStackSize)
// lowerConstantI64 adds instructions into m.pendingInstructions,
// so we manually link them together.

8
internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go
View File

@@ -217,8 +217,8 @@ func TestMachine_insertStackBoundsCheck(t *testing.T) {
str q31, [x0, #0x2c0]
mov x27, sp
str x27, [x0, #0x38]
movz x27, #0x1, LSL 0
str w27, [x0]
orr w17, wzr, #0x1
str w17, [x0]
movz x27, #0xfff0, LSL 0
str x27, [x0, #0x40]
adr x27, #0x1c
@@ -315,8 +315,8 @@ func TestMachine_insertStackBoundsCheck(t *testing.T) {
str q31, [x0, #0x2c0]
mov x27, sp
str x27, [x0, #0x38]
movz x27, #0x1, LSL 0
str w27, [x0]
orr w17, wzr, #0x1
str w17, [x0]
orr x27, xzr, #0x10
str x27, [x0, #0x40]
adr x27, #0x1c

2
internal/engine/wazevo/backend/machine.go
View File

@@ -97,6 +97,6 @@ type (
Encode()
// CompileGoFunctionTrampoline compiles the trampoline function to call a Go function of the given exit code and signature.
CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *ssa.Signature)
CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *ssa.Signature, needModuleContextPtr bool)
}
)

9
internal/engine/wazevo/backend/machine_test.go
View File

@@ -24,10 +24,8 @@ type mockMachine struct {
rinfo *regalloc.RegisterInfo
}
func (m mockMachine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *ssa.Signature) {
// TODO implement me
panic("implement me")
}
// CompileGoFunctionTrampoline implements Machine.CompileGoFunctionTrampoline.
func (m mockMachine) CompileGoFunctionTrampoline(wazevoapi.ExitCode, *ssa.Signature, bool) {}
// Encode implements Machine.Encode.
func (m mockMachine) Encode() {
@@ -36,8 +34,7 @@ func (m mockMachine) Encode() {
}
// ResolveRelocations implements Machine.ResolveRelocations.
func (m mockMachine) ResolveRelocations(refToBinaryOffset map[ssa.FuncRef]int, binary []byte, relocations []RelocationInfo) {
}
func (m mockMachine) ResolveRelocations(map[ssa.FuncRef]int, []byte, []RelocationInfo) {}
// SetupPrologue implements Machine.SetupPrologue.
func (m mockMachine) SetupPrologue() {}

3
internal/engine/wazevo/backend/regalloc/assign.go
View File

@@ -79,7 +79,8 @@ func (a *Allocator) assignRegistersPerInstr(f Function, pc programCounter, instr
return
case 1:
default:
panic("multiple defs (== call instruction) should be special cased")
// multiple defs (== call instruction) can be special cased, and no need to assign (already real regs following the calling convention.
return
}
d := defs[0]

27
internal/engine/wazevo/call_engine.go
View File

@@ -61,11 +61,15 @@ type (
// savedRegisters is the opaque spaces for save/restore registers.
// We want to align 16 bytes for each register, so we use [64][2]uint64.
savedRegisters [64][2]uint64
// goFunctionCallCalleeModuleContextOpaque is the pointer to the target Go function's moduleContextOpaque.
goFunctionCallCalleeModuleContextOpaque uintptr
// goFunctionCallStack is used to pass/receive parameters/results for Go function calls.
goFunctionCallStack [128]uint64
goFunctionCallStack [goFunctionCallStackSize]uint64
}
)
const goFunctionCallStackSize = 128
var initialStackSize uint64 = 512
func (c *callEngine) init() {
@@ -111,7 +115,7 @@ func (c *callEngine) CallWithStack(ctx context.Context, paramResultStack []uint6
entrypoint(c.executable, c.execCtxPtr, c.parent.opaquePtr, paramResultPtr, c.stackTop)
for {
switch c.execCtx.exitCode {
switch ec := c.execCtx.exitCode; ec & wazevoapi.ExitCodeMask {
case wazevoapi.ExitCodeOK:
return nil
case wazevoapi.ExitCodeGrowStack:
@@ -126,7 +130,7 @@ func (c *callEngine) CallWithStack(ctx context.Context, paramResultStack []uint6
case wazevoapi.ExitCodeMemoryOutOfBounds:
return wasmruntime.ErrRuntimeOutOfBoundsMemoryAccess
case wazevoapi.ExitCodeGrowMemory:
mod := moduleInstanceFromPtr(c.execCtx.callerModuleContextPtr)
mod := c.callerModuleInstance()
mem := mod.MemoryInstance
argRes := &c.execCtx.goFunctionCallStack[0]
if res, ok := mem.Grow(uint32(*argRes)); !ok {
@@ -145,14 +149,27 @@ func (c *callEngine) CallWithStack(ctx context.Context, paramResultStack []uint6
}
c.execCtx.exitCode = wazevoapi.ExitCodeOK
afterGoFunctionCallEntrypoint(c.execCtx.goCallReturnAddress, c.execCtxPtr, c.execCtx.stackPointerBeforeGoCall)
case wazevoapi.ExitCodeCallGoFunction:
index := wazevoapi.GoFunctionIndexFromExitCode(ec)
f := hostModuleGoFuncFromOpaque[api.GoFunction](index, c.execCtx.goFunctionCallCalleeModuleContextOpaque)
f.Call(ctx, c.execCtx.goFunctionCallStack[:])
c.execCtx.exitCode = wazevoapi.ExitCodeOK
afterGoFunctionCallEntrypoint(c.execCtx.goCallReturnAddress, c.execCtxPtr, c.execCtx.stackPointerBeforeGoCall)
case wazevoapi.ExitCodeCallGoModuleFunction:
index := wazevoapi.GoFunctionIndexFromExitCode(ec)
f := hostModuleGoFuncFromOpaque[api.GoModuleFunction](index, c.execCtx.goFunctionCallCalleeModuleContextOpaque)
mod := c.callerModuleInstance()
f.Call(ctx, mod, c.execCtx.goFunctionCallStack[:])
c.execCtx.exitCode = wazevoapi.ExitCodeOK
afterGoFunctionCallEntrypoint(c.execCtx.goCallReturnAddress, c.execCtxPtr, c.execCtx.stackPointerBeforeGoCall)
default:
panic("BUG")
}
}
}
func moduleInstanceFromPtr(ptr *byte) *wasm.ModuleInstance {
return *(**wasm.ModuleInstance)(unsafe.Pointer(ptr))
func (c *callEngine) callerModuleInstance() *wasm.ModuleInstance {
return *(**wasm.ModuleInstance)(unsafe.Pointer(c.execCtx.callerModuleContextPtr))
}
func opaqueViewFromPtr(ptr uintptr) []byte {

90
internal/engine/wazevo/e2e_test.go
View File

@@ -17,15 +17,23 @@ import (
"github.com/tetratelabs/wazero/internal/wasm"
)
const (
i32 = wasm.ValueTypeI32
i64 = wasm.ValueTypeI64
f32 = wasm.ValueTypeF32
f64 = wasm.ValueTypeF64
)
func TestE2E(t *testing.T) {
type callCase struct {
params, expResults []uint64
expErr string
}
for _, tc := range []struct {
name string
imported, m *wasm.Module
calls []callCase
name string
imported, m *wasm.Module
needHostModule bool
calls []callCase
}{
{
name: "swap", m: testcases.SwapParamAndReturn.Module,
@@ -224,3 +232,79 @@ func configureWazevo(config wazero.RuntimeConfig) {
// Insert the wazevo implementation.
cm.newEngine = wazevo.NewEngine
}
func TestE2E_host_functions(t *testing.T) {
config := wazero.NewRuntimeConfigCompiler()
// Configure the new optimizing backend!
configureWazevo(config)
ctx := context.Background()
r := wazero.NewRuntimeWithConfig(ctx, config)
defer func() {
require.NoError(t, r.Close(ctx))
}()
var expectedMod api.Module
b := r.NewHostModuleBuilder("env")
b.NewFunctionBuilder().WithFunc(func(ctx2 context.Context, d float64) float64 {
require.Equal(t, ctx, ctx2)
fmt.Printf("%#x\n", math.Float64bits(d))
require.Equal(t, 35.0, d)
return math.Sqrt(d)
}).Export("root")
b.NewFunctionBuilder().WithFunc(func(ctx2 context.Context, mod api.Module, a uint32, b uint64, c float32, d float64) (uint32, uint64, float32, float64) {
require.Equal(t, expectedMod, mod)
require.Equal(t, ctx, ctx2)
require.Equal(t, uint32(2), a)
require.Equal(t, uint64(100), b)
require.Equal(t, float32(15.0), c)
require.Equal(t, 35.0, d)
return a * a, b * b, c * c, d * d
}).Export("square")
_, err := b.Instantiate(ctx)
require.NoError(t, err)
m := &wasm.Module{
ImportFunctionCount: 2,
ImportSection: []wasm.Import{
{Module: "env", Name: "root", Type: wasm.ExternTypeFunc, DescFunc: 0},
{Module: "env", Name: "square", Type: wasm.ExternTypeFunc, DescFunc: 1},
},
TypeSection: []wasm.FunctionType{
{Results: []wasm.ValueType{f64}, Params: []wasm.ValueType{f64}},
{Results: []wasm.ValueType{i32, i64, f32, f64}, Params: []wasm.ValueType{i32, i64, f32, f64}},
{Results: []wasm.ValueType{i32, i64, f32, f64, f64}, Params: []wasm.ValueType{i32, i64, f32, f64}},
},
FunctionSection: []wasm.Index{2},
CodeSection: []wasm.Code{{
Body: []byte{
wasm.OpcodeLocalGet, 0, wasm.OpcodeLocalGet, 1, wasm.OpcodeLocalGet, 2, wasm.OpcodeLocalGet, 3,
wasm.OpcodeCall, 1,
wasm.OpcodeLocalGet, 3,
wasm.OpcodeCall, 0,
wasm.OpcodeEnd,
},
}},
ExportSection: []wasm.Export{{Name: "f", Type: wasm.ExternTypeFunc, Index: 2}},
}
compiled, err := r.CompileModule(ctx, binaryencoding.EncodeModule(m))
require.NoError(t, err)
inst, err := r.InstantiateModule(ctx, compiled, wazero.NewModuleConfig())
require.NoError(t, err)
expectedMod = inst
f := inst.ExportedFunction("f")
res, err := f.Call(ctx, []uint64{2, 100, uint64(math.Float32bits(15.0)), math.Float64bits(35.0)}...)
require.NoError(t, err)
require.Equal(t, []uint64{
2 * 2, 100 * 100, uint64(math.Float32bits(15.0 * 15.0)), math.Float64bits(35.0 * 35.0),
math.Float64bits(math.Sqrt(35.0)),
}, res)
}

122
internal/engine/wazevo/engine.go
View File

@@ -36,17 +36,21 @@ type (
// compiledModule is a compiled variant of a wasm.Module and ready to be used for instantiation.
compiledModule struct {
executable []byte
functionOffsets []compiledFunctionOffset
executable []byte
functionOffsets []compiledFunctionOffset
// The followings are only available for non host modules.
offsets wazevoapi.ModuleContextOffsetData
builtinFunctions *builtinFunctions
}
// compiledFunctionOffset tells us that where in the executable a function begins.
compiledFunctionOffset struct {
// offset is the beggining of the function.
// offset is the beginning of the function.
offset int
// goPreambleSize is the size of Go preamble of the function.
// This is only needed for non host modules.
goPreambleSize int
}
)
@@ -62,10 +66,14 @@ func NewEngine(_ context.Context, _ api.CoreFeatures, _ filecache.Cache) wasm.En
}
// CompileModule implements wasm.Engine.
func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []experimental.FunctionListener, ensureTermination bool) error {
func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []experimental.FunctionListener, _ bool) error {
e.rels = e.rels[:0]
cm := &compiledModule{offsets: wazevoapi.NewModuleContextOffsetData(module)}
if module.IsHostModule {
return e.compileHostModule(module)
}
importedFns, localFns := int(module.ImportFunctionCount), len(module.FunctionSection)
if importedFns+localFns == 0 {
e.addCompiledModule(module, cm)
@@ -181,6 +189,97 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []exper
return nil
}
func (e *engine) compileHostModule(module *wasm.Module) error {
machine := newMachine()
be := backend.NewCompiler(machine, ssa.NewBuilder())
num := len(module.CodeSection)
cm := &compiledModule{}
cm.functionOffsets = make([]compiledFunctionOffset, num)
totalSize := 0 // Total binary size of the executable.
bodies := make([][]byte, num)
var sig ssa.Signature
for i := range module.CodeSection {
totalSize = (totalSize + 15) &^ 15
cm.functionOffsets[i].offset = totalSize
typIndex := module.FunctionSection[i]
typ := &module.TypeSection[typIndex]
if typ.ParamNumInUint64 >= goFunctionCallStackSize || typ.ResultNumInUint64 >= goFunctionCallStackSize {
return fmt.Errorf("too many params or results for a host function (maximum %d): %v",
goFunctionCallStackSize, typ)
}
// We can relax until the index fits together in ExitCode as we do in wazevoapi.ExitCodeCallGoModuleFunctionWithIndex.
// However, 1 << 16 should be large enough for a real use case.
const hostFunctionNumMaximum = 1 << 16
if i >= hostFunctionNumMaximum {
return fmt.Errorf("too many host functions (maximum %d)", hostFunctionNumMaximum)
}
sig.ID = ssa.SignatureID(typIndex) // This is important since we reuse the `machine` which caches the ABI based on the SignatureID.
sig.Params = append(sig.Params[:0],
ssa.TypeI64, // First argument must be exec context.
ssa.TypeI64, // The second argument is the moduleContextOpaque of this host module.
)
for _, t := range typ.Params {
sig.Params = append(sig.Params, frontend.WasmTypeToSSAType(t))
}
sig.Results = sig.Results[:0]
for _, t := range typ.Results {
sig.Results = append(sig.Results, frontend.WasmTypeToSSAType(t))
}
c := &module.CodeSection[i]
if c.GoFunc == nil {
panic("BUG: GoFunc must be set for host module")
}
var exitCode wazevoapi.ExitCode
fn := c.GoFunc
switch fn.(type) {
case api.GoModuleFunction:
exitCode = wazevoapi.ExitCodeCallGoModuleFunctionWithIndex(i)
case api.GoFunction:
exitCode = wazevoapi.ExitCodeCallGoFunctionWithIndex(i)
}
be.Init(false)
machine.CompileGoFunctionTrampoline(exitCode, &sig, true)
be.Encode()
body := be.Buf()
// TODO: optimize as zero copy.
copied := make([]byte, len(body))
copy(copied, body)
bodies[i] = copied
totalSize += len(body)
}
// Allocate executable memory and then copy the generated machine code.
executable, err := platform.MmapCodeSegment(totalSize)
if err != nil {
panic(err)
}
cm.executable = executable
for i, b := range bodies {
offset := cm.functionOffsets[i]
copy(executable[offset.offset:], b)
}
if runtime.GOARCH == "arm64" {
// On arm64, we cannot give all of rwx at the same time, so we change it to exec.
if err = platform.MprotectRX(executable); err != nil {
return err
}
}
e.compiledModules[module.ID] = cm
return nil
}
// Close implements wasm.Engine.
func (e *engine) Close() (err error) {
e.mux.Lock()
@@ -228,10 +327,15 @@ func (e *engine) NewModuleEngine(m *wasm.Module, mi *wasm.ModuleInstance) (wasm.
me.parent = compiled
me.module = mi
if size := compiled.offsets.TotalSize; size != 0 {
opaque := make([]byte, size)
me.opaque = opaque
me.opaquePtr = &opaque[0]
if m.IsHostModule {
me.opaque = buildHostModuleOpaque(m)
me.opaquePtr = &me.opaque[0]
} else {
if size := compiled.offsets.TotalSize; size != 0 {
opaque := make([]byte, size)
me.opaque = opaque
me.opaquePtr = &opaque[0]
}
}
return me, nil
}
@@ -244,7 +348,7 @@ func (e *engine) compileBuiltinFunctions() {
machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeGrowMemory, &ssa.Signature{
Params: []ssa.Type{ssa.TypeI32 /* exec context */, ssa.TypeI32},
Results: []ssa.Type{ssa.TypeI32},
})
}, false)
be.Encode()
src := be.Buf()

14
internal/engine/wazevo/frontend/frontend.go
View File

@@ -64,10 +64,10 @@ func NewFrontendCompiler(m *wasm.Module, ssaBuilder ssa.Builder, offset *wazevoa
sig.Params[0] = executionContextPtrTyp
sig.Params[1] = moduleContextPtrTyp
for j, typ := range wasmSig.Params {
sig.Params[j+2] = wasmToSSA(typ)
sig.Params[j+2] = WasmTypeToSSAType(typ)
}
for j, typ := range wasmSig.Results {
sig.Results[j] = wasmToSSA(typ)
sig.Results[j] = WasmTypeToSSAType(typ)
}
c.signatures[wasmSig] = sig
c.ssaBuilder.DeclareSignature(sig)
@@ -134,7 +134,7 @@ func (c *Compiler) LowerToSSA() error {
builder.AnnotateValue(c.moduleCtxPtrValue, "module_ctx")
for i, typ := range c.wasmFunctionTyp.Params {
st := wasmToSSA(typ)
st := WasmTypeToSSAType(typ)
variable := builder.DeclareVariable(st)
value := entryBlock.AddParam(builder, st)
builder.DefineVariable(variable, value, entryBlock)
@@ -156,7 +156,7 @@ func (c *Compiler) localVariable(index wasm.Index) ssa.Variable {
func (c *Compiler) declareWasmLocals(entry ssa.BasicBlock) {
localCount := wasm.Index(len(c.wasmFunctionTyp.Params))
for i, typ := range c.wasmFunctionLocalTypes {
st := wasmToSSA(typ)
st := WasmTypeToSSAType(typ)
variable := c.ssaBuilder.DeclareVariable(st)
c.wasmLocalToVariable[wasm.Index(i)+localCount] = variable
@@ -227,8 +227,8 @@ func (c *Compiler) declareWasmGlobal(typ wasm.ValueType, mutable bool) {
}
}
// wasmToSSA converts wasm.ValueType to ssa.Type.
func wasmToSSA(vt wasm.ValueType) ssa.Type {
// WasmTypeToSSAType converts wasm.ValueType to ssa.Type.
func WasmTypeToSSAType(vt wasm.ValueType) ssa.Type {
switch vt {
case wasm.ValueTypeI32:
return ssa.TypeI32
@@ -246,7 +246,7 @@ func wasmToSSA(vt wasm.ValueType) ssa.Type {
// addBlockParamsFromWasmTypes adds the block parameters to the given block.
func (c *Compiler) addBlockParamsFromWasmTypes(tps []wasm.ValueType, blk ssa.BasicBlock) {
for _, typ := range tps {
st := wasmToSSA(typ)
st := WasmTypeToSSAType(typ)
blk.AddParam(c.ssaBuilder, st)
}
}

46
internal/engine/wazevo/hostmodule.go Normal file
View File

@@ -0,0 +1,46 @@
package wazevo
import (
"encoding/binary"
"reflect"
"unsafe"
"github.com/tetratelabs/wazero/internal/wasm"
)
func buildHostModuleOpaque(m *wasm.Module) moduleContextOpaque {
size := len(m.CodeSection) * 16
ret := make(moduleContextOpaque, size)
var offset int
for i := range m.CodeSection {
goFn := m.CodeSection[i].GoFunc
writeIface(goFn, ret[offset:])
offset += 16
}
return ret
}
func hostModuleGoFuncFromOpaque[T any](index int, opaqueBegin uintptr) T {
offset := uintptr(index * 16)
ptr := opaqueBegin + offset
var opaqueViewOverFunction []byte
sh := (*reflect.SliceHeader)(unsafe.Pointer(&opaqueViewOverFunction))
sh.Data = ptr
sh.Len = 16
sh.Cap = 16
return readIface(opaqueViewOverFunction).(T)
}
func writeIface(goFn interface{}, buf []byte) {
goFnIface := *(*[2]uint64)(unsafe.Pointer(&goFn))
binary.LittleEndian.PutUint64(buf, goFnIface[0])
binary.LittleEndian.PutUint64(buf[8:], goFnIface[1])
}
func readIface(buf []byte) interface{} {
b := binary.LittleEndian.Uint64(buf)
s := binary.LittleEndian.Uint64(buf[8:])
return *(*interface{})(unsafe.Pointer(&[2]uint64{b, s}))
}

22
internal/engine/wazevo/hostmodule_test.go Normal file
View File

@@ -0,0 +1,22 @@
package wazevo
import (
"context"
"testing"
"github.com/tetratelabs/wazero/api"
"github.com/tetratelabs/wazero/internal/testing/require"
)
func Test_writeIface_readIface(t *testing.T) {
buf := make([]byte, 100)
var called bool
var goFn api.GoFunction = api.GoFunc(func(context.Context, []uint64) {
called = true
})
writeIface(goFn, buf)
got := readIface(buf).(api.GoFunction)
got.Call(context.Background(), nil)
require.True(t, called)
}

13
internal/engine/wazevo/module_engine.go
View File

@@ -26,9 +26,11 @@ type (
// Internally, the buffer is structured as follows:
//
// type moduleContextOpaque struct {
// moduleInstance *wasm.ModuleInstance
// localMemoryBufferPtr *byte (optional)
// localMemoryLength uint64 (optional)
// importedMemoryInstance *wasm.MemoryInstance (optional)
// importedMemoryOwnerOpaqueCtx *byte (optional)
// importedFunctions [importedFunctions] struct { the total size depends on # of imported functions.
// executable *byte
// opaqueCtx *moduleContextOpaque
@@ -38,6 +40,8 @@ type (
// }
//
// See wazevoapi.NewModuleContextOffsetData for the details of the offsets.
//
// Note that for host modules, the structure is entirely different. See buildHostModuleOpaque.
moduleContextOpaque []byte
)
@@ -121,6 +125,11 @@ func (m *moduleEngine) ResolveImportedMemory(importedModuleEngine wasm.ModuleEng
importedME := importedModuleEngine.(*moduleEngine)
inst := importedME.module
if importedME.parent.offsets.ImportedMemoryBegin >= 0 {
// This case can be resolved by recursively resolving the owner.
panic("TODO: support re-exported memory import")
}
offset := m.parent.offsets.ImportedMemoryBegin
b := uint64(uintptr(unsafe.Pointer(inst.MemoryInstance)))
binary.LittleEndian.PutUint64(m.opaque[offset:], b)
@@ -129,7 +138,9 @@ func (m *moduleEngine) ResolveImportedMemory(importedModuleEngine wasm.ModuleEng
// DoneInstantiation implements wasm.ModuleEngine.
func (m *moduleEngine) DoneInstantiation() {
m.setupOpaque()
if !m.module.Source.IsHostModule {
m.setupOpaque()
}
}
// LookupFunction implements wasm.ModuleEngine.

5
internal/engine/wazevo/module_engine_test.go
View File

@@ -69,7 +69,10 @@ func TestModuleEngine_setupOpaque(t *testing.T) {
require.Equal(t, expLen, actualLen)
}
if tc.offset.ImportedMemoryBegin >= 0 {
imported := &moduleEngine{opaque: []byte{1, 2, 3}, module: &wasm.ModuleInstance{MemoryInstance: tc.m.MemoryInstance}}
imported := &moduleEngine{
opaque: []byte{1, 2, 3}, module: &wasm.ModuleInstance{MemoryInstance: tc.m.MemoryInstance},
parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{ImportedMemoryBegin: -1}},
}
imported.opaquePtr = &imported.opaque[0]
m.ResolveImportedMemory(imported)

1
internal/engine/wazevo/wazevo_test.go
View File

@@ -61,5 +61,6 @@ func Test_ExecutionContextOffsets(t *testing.T) {
require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.savedRegisters))%16, wazevoapi.Offset(0),
"SavedRegistersBegin must be aligned to 16 bytes")
require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.savedRegisters)), offsets.SavedRegistersBegin)
require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.goFunctionCallCalleeModuleContextOpaque)), offsets.GoFunctionCallCalleeModuleContextOpaque)
require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.goFunctionCallStack)), offsets.GoFunctionCallStackBegin)
}

24
internal/engine/wazevo/wazevoapi/exitcode.go
View File

@@ -9,8 +9,16 @@ const (
ExitCodeGrowMemory
ExitCodeUnreachable
ExitCodeMemoryOutOfBounds
// ExitCodeCallGoModuleFunction is an exit code for a call to an api.GoModuleFunction.
ExitCodeCallGoModuleFunction
// ExitCodeCallGoFunction is an exit code for a call to an api.GoFunction.
ExitCodeCallGoFunction
exitCodeMax
)
const ExitCodeMask = 0xff
// String implements fmt.Stringer.
func (e ExitCode) String() string {
switch e {
@@ -18,6 +26,10 @@ func (e ExitCode) String() string {
return "ok"
case ExitCodeGrowStack:
return "grow_stack"
case ExitCodeCallGoModuleFunction:
return "call_go_module_function"
case ExitCodeCallGoFunction:
return "call_go_function"
case ExitCodeUnreachable:
return "unreachable"
case ExitCodeMemoryOutOfBounds:
@@ -25,3 +37,15 @@ func (e ExitCode) String() string {
}
panic("TODO")
}
func ExitCodeCallGoModuleFunctionWithIndex(index int) ExitCode {
return ExitCodeCallGoModuleFunction | ExitCode(index<<8)
}
func ExitCodeCallGoFunctionWithIndex(index int) ExitCode {
return ExitCodeCallGoFunction | ExitCode(index<<8)
}
func GoFunctionIndexFromExitCode(exitCode ExitCode) int {
return int(exitCode >> 8)
}

11
internal/engine/wazevo/wazevoapi/exitcode_test.go Normal file
View File

@@ -0,0 +1,11 @@
package wazevoapi
import (
"testing"
"github.com/tetratelabs/wazero/internal/testing/require"
)
func TestExitCode_withinByte(t *testing.T) {
require.True(t, exitCodeMax < ExitCodeMask) //nolint
}

27
internal/engine/wazevo/wazevoapi/offsetdata.go
View File

@@ -3,18 +3,19 @@ package wazevoapi
import "github.com/tetratelabs/wazero/internal/wasm"
var ExecutionContextOffsets = ExecutionContextOffsetData{
ExitCodeOffset: 0,
CallerModuleContextPtr: 8,
OriginalFramePointer: 16,
OriginalStackPointer: 24,
GoReturnAddress: 32,
StackBottomPtr: 40,
GoCallReturnAddress: 48,
StackPointerBeforeGrow: 56,
StackGrowRequiredSize: 64,
MemoryGrowTrampolineAddress: 72,
SavedRegistersBegin: 80,
GoFunctionCallStackBegin: 1104,
ExitCodeOffset: 0,
CallerModuleContextPtr: 8,
OriginalFramePointer: 16,
OriginalStackPointer: 24,
GoReturnAddress: 32,
StackBottomPtr: 40,
GoCallReturnAddress: 48,
StackPointerBeforeGrow: 56,
StackGrowRequiredSize: 64,
MemoryGrowTrampolineAddress: 72,
SavedRegistersBegin: 80,
GoFunctionCallCalleeModuleContextOpaque: 1104,
GoFunctionCallStackBegin: 1112,
}
// ExecutionContextOffsetData allows the compilers to get the information about offsets to the fields of wazevo.executionContext,
@@ -42,6 +43,8 @@ type ExecutionContextOffsetData struct {
MemoryGrowTrampolineAddress Offset
// GoCallReturnAddress is an offset of the first element of `savedRegisters` field in wazevo.executionContext
SavedRegistersBegin Offset
// GoFunctionCallCalleeModuleContextOpaque is an offset of `goFunctionCallCalleeModuleContextOpaque` field in wazevo.executionContext
GoFunctionCallCalleeModuleContextOpaque Offset
// GoFunctionCallStackBegin is an offset of the first element of `goFunctionCallStack` field in wazevo.executionContext
GoFunctionCallStackBegin Offset
}