wazevo: adds support for host functions (#1630)

Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>

internal/engine/wazevo/backend/isa/arm64/abi.go

@@ -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 {

internal/engine/wazevo/backend/isa/arm64/abi_go_call.go

@@ -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(),

internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go

@@ -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())

internal/engine/wazevo/backend/isa/arm64/machine.go

@@ -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.

internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go

@@ -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.

internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go

@@ -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

internal/engine/wazevo/backend/machine.go

@@ -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)
 	}
 )

internal/engine/wazevo/backend/machine_test.go

@@ -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() {}