Repository: tetratelabs/wazero
Pull request #1968 (2024-01-26 20:22:30)

wazevo(amd64): stack based param/results


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/abi_go_call.go b/internal/engine/wazevo/backend/isa/amd64/abi_go_call.go
index 34c51653..a364044b 100644
--- a/internal/engine/wazevo/backend/isa/amd64/abi_go_call.go
+++ b/internal/engine/wazevo/backend/isa/amd64/abi_go_call.go
@@ -98,7 +98,29 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
 		if arg.Kind == backend.ABIArgKindReg {
 			v = arg.Reg
 		} else {
-			panic("TODO: stack arguments")
+			// We have saved callee saved registers, so we can use them.
+			if arg.Type.IsInt() {
+				v = r15VReg
+			} else {
+				v = xmm15VReg
+			}
+			mem := newOperandMem(newAmodeImmReg(uint32(arg.Offset+16 /* to skip caller_rbp and ret_addr */), rbpVReg))
+			load := m.allocateInstr()
+			switch arg.Type {
+			case ssa.TypeI32:
+				load.asMovzxRmR(extModeLQ, mem, v)
+			case ssa.TypeI64:
+				load.asMov64MR(mem, v)
+			case ssa.TypeF32:
+				load.asXmmUnaryRmR(sseOpcodeMovss, mem, v)
+			case ssa.TypeF64:
+				load.asXmmUnaryRmR(sseOpcodeMovsd, mem, v)
+			case ssa.TypeV128:
+				load.asXmmUnaryRmR(sseOpcodeMovdqu, mem, v)
+			default:
+				panic("BUG")
+			}
+			cur = linkInstr(cur, load)
 		}
 
 		store := m.allocateInstr()
@@ -163,45 +185,89 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
 	// Ready to exit the execution.
 	cur = m.storeReturnAddressAndExit(cur, execCtrPtr)
 
-	// After the call, we need to restore the callee saved registers.
-	cur = m.restoreRegistersInExecutionContext(cur, execCtrPtr, calleeSavedVRegs)
-
 	// We don't need the slice size anymore, so pop it.
 	cur = m.addRSP(8, cur)
 
 	// Ready to set up the results.
 	offsetInGoSlice = 0
+	// To avoid overwriting with the execution context pointer by the result, we need to track the offset,
+	// and defer the restoration of the result to the end of this function.
+	var argOverlapWithExecCtxOffset int32 = -1
 	for i := range abi.Rets {
 		r := &abi.Rets[i]
-		if r.Kind == backend.ABIArgKindReg {
-			v := r.Reg
-			load := m.allocateInstr()
-			mem := newOperandMem(newAmodeImmReg(uint32(offsetInGoSlice), rspVReg))
+		var v regalloc.VReg
+		isRegResult := r.Kind == backend.ABIArgKindReg
+		if isRegResult {
+			v = r.Reg
+			if v.RealReg() == execCtrPtr.RealReg() {
+				argOverlapWithExecCtxOffset = offsetInGoSlice
+				offsetInGoSlice += 8 // always uint64 rep.
+				continue
+			}
+		} else {
+			if r.Type.IsInt() {
+				v = r15VReg
+			} else {
+				v = xmm15VReg
+			}
+		}
+
+		load := m.allocateInstr()
+		mem := newOperandMem(newAmodeImmReg(uint32(offsetInGoSlice), rspVReg))
+		switch r.Type {
+		case ssa.TypeI32:
+			load.asMovzxRmR(extModeLQ, mem, v)
+			offsetInGoSlice += 8 // always uint64 rep.
+		case ssa.TypeI64:
+			load.asMov64MR(mem, v)
+			offsetInGoSlice += 8
+		case ssa.TypeF32:
+			load.asXmmUnaryRmR(sseOpcodeMovss, mem, v)
+			offsetInGoSlice += 8 // always uint64 rep.
+		case ssa.TypeF64:
+			load.asXmmUnaryRmR(sseOpcodeMovsd, mem, v)
+			offsetInGoSlice += 8
+		case ssa.TypeV128:
+			load.asXmmUnaryRmR(sseOpcodeMovdqu, mem, v)
+			offsetInGoSlice += 16
+		default:
+			panic("BUG")
+		}
+		cur = linkInstr(cur, load)
+
+		if !isRegResult {
+			// We need to store it back to the result slot above rbp.
+			store := m.allocateInstr()
+			mem := newOperandMem(newAmodeImmReg(uint32(abi.ArgStackSize+r.Offset+16 /* to skip caller_rbp and ret_addr */), rbpVReg))
 			switch r.Type {
 			case ssa.TypeI32:
-				load.asMovzxRmR(extModeLQ, mem, v)
-				offsetInGoSlice += 8 // always uint64 rep.
+				store.asMovRM(v, mem, 4)
 			case ssa.TypeI64:
-				load.asMov64MR(mem, v)
-				offsetInGoSlice += 8
+				store.asMovRM(v, mem, 8)
 			case ssa.TypeF32:
-				load.asXmmUnaryRmR(sseOpcodeMovss, mem, v)
-				offsetInGoSlice += 8 // always uint64 rep.
+				store.asXmmMovRM(sseOpcodeMovss, v, mem)
 			case ssa.TypeF64:
-				load.asXmmUnaryRmR(sseOpcodeMovsd, mem, v)
-				offsetInGoSlice += 8
+				store.asXmmMovRM(sseOpcodeMovsd, v, mem)
 			case ssa.TypeV128:
-				load.asXmmUnaryRmR(sseOpcodeMovdqu, mem, v)
-				offsetInGoSlice += 16
+				store.asXmmMovRM(sseOpcodeMovdqu, v, mem)
 			default:
 				panic("BUG")
 			}
-			cur = linkInstr(cur, load)
-		} else {
-			panic("TODO: stack results")
+			cur = linkInstr(cur, store)
 		}
 	}
 
+	// Before return, we need to restore the callee saved registers.
+	cur = m.restoreRegistersInExecutionContext(cur, execCtrPtr, calleeSavedVRegs)
+
+	if argOverlapWithExecCtxOffset >= 0 {
+		// At this point execCtt is not used anymore, so we can finally store the
+		// result to the register which overlaps with the execution context pointer.
+		mem := newOperandMem(newAmodeImmReg(uint32(argOverlapWithExecCtxOffset), rspVReg))
+		load := m.allocateInstr().asMov64MR(mem, execCtrPtr)
+		cur = linkInstr(cur, load)
+	}
+
 	// Finally ready to return.
 	cur = m.revertRBPRSP(cur)
 	linkInstr(cur, m.allocateInstr().asRet(nil))
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index 0f6e5b81..38099b81 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -677,6 +677,9 @@ func (m *machine) lowerCall(si *ssa.Instruction) {
 		m.maxRequiredStackSizeForCalls = stackSlotSize + 16 // 16 == return address + RBP.
 	}
 
+	// Note: See machine.SetupPrologue for the stack layout.
+	// The stack pointer decrease/increase will be inserted later in the compilation.
+
 	for i, arg := range args {
 		reg := m.c.VRegOf(arg)
 		def := m.c.ValueDefinition(arg)
@@ -707,7 +710,7 @@ func (m *machine) lowerCall(si *ssa.Instruction) {
 
 // callerGenVRegToFunctionArg is the opposite of GenFunctionArgToVReg, which is used to generate the
 // caller side of the function call.
-func (m *machine) callerGenVRegToFunctionArg(a *backend.FunctionABI, argIndex int, reg regalloc.VReg, def *backend.SSAValueDefinition, slotBegin int64) {
+func (m *machine) callerGenVRegToFunctionArg(a *backend.FunctionABI, argIndex int, reg regalloc.VReg, def *backend.SSAValueDefinition, stackSlotSize int64) {
 	arg := &a.Args[argIndex]
 	if def != nil && def.IsFromInstr() {
 		// Constant instructions are inlined.
@@ -718,16 +721,52 @@ func (m *machine) callerGenVRegToFunctionArg(a *backend.FunctionABI, argIndex in
 	if arg.Kind == backend.ABIArgKindReg {
 		m.InsertMove(arg.Reg, reg, arg.Type)
 	} else {
-		panic("TODO")
+		store := m.allocateInstr()
+		mem := newOperandMem(newAmodeImmReg(
+			// -stackSlotSize because the stack pointer is not yet decreased.
+			uint32(arg.Offset-stackSlotSize), rspVReg))
+		switch arg.Type {
+		case ssa.TypeI32:
+			store.asMovRM(reg, mem, 4)
+		case ssa.TypeI64:
+			store.asMovRM(reg, mem, 8)
+		case ssa.TypeF32:
+			store.asXmmMovRM(sseOpcodeMovss, reg, mem)
+		case ssa.TypeF64:
+			store.asXmmMovRM(sseOpcodeMovsd, reg, mem)
+		case ssa.TypeV128:
+			store.asXmmMovRM(sseOpcodeMovdqu, reg, mem)
+		default:
+			panic("BUG")
+		}
+		m.insert(store)
 	}
 }
 
-func (m *machine) callerGenFunctionReturnVReg(a *backend.FunctionABI, retIndex int, reg regalloc.VReg, slotBegin int64) {
+func (m *machine) callerGenFunctionReturnVReg(a *backend.FunctionABI, retIndex int, reg regalloc.VReg, stackSlotSize int64) {
 	r := &a.Rets[retIndex]
 	if r.Kind == backend.ABIArgKindReg {
 		m.InsertMove(reg, r.Reg, r.Type)
 	} else {
-		panic("TODO")
+		load := m.allocateInstr()
+		mem := newOperandMem(newAmodeImmReg(
+			// -stackSlotSize because the stack pointer is not yet decreased.
+			uint32(a.ArgStackSize+r.Offset-stackSlotSize), rspVReg))
+		switch r.Type {
+		case ssa.TypeI32:
+			load.asMovzxRmR(extModeLQ, mem, reg)
+		case ssa.TypeI64:
+			load.asMov64MR(mem, reg)
+		case ssa.TypeF32:
+			load.asXmmUnaryRmR(sseOpcodeMovss, mem, reg)
+		case ssa.TypeF64:
+			load.asXmmUnaryRmR(sseOpcodeMovsd, mem, reg)
+		case ssa.TypeV128:
+			load.asXmmUnaryRmR(sseOpcodeMovdqu, mem, reg)
+		default:
+			panic("BUG")
+		}
+		m.insert(load)
 	}
 }
 
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue.go b/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue.go
index 5ba414cb..0e4e22d2 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue.go
@@ -125,6 +125,24 @@ func (m *machine) SetupEpilogue() {
 			continue
 		}
 
+		if abi := cur.abi; abi != nil {
+			// At this point, reg alloc is done, therefore we can safely insert dec/inc RPS instruction
+			// right before/after the call instruction. If this is done before reg alloc, the stack slot
+			// can point to the wrong location and therefore results in a wrong value.
+			// TODO: doing this in `SetupEpilogue` seems weird. Find a better home.
+			call := cur
+			next := call.next
+			if size := uint32(abi.AlignedArgResultStackSlotSize()); size > 0 {
+				dec := m.allocateInstr().asAluRmiR(aluRmiROpcodeSub, newOperandImm32(size), rspVReg, true)
+				linkInstr(call.prev, dec)
+				linkInstr(dec, call)
+				inc := m.allocateInstr().asAluRmiR(aluRmiROpcodeAdd, newOperandImm32(size), rspVReg, true)
+				linkInstr(call, inc)
+				linkInstr(inc, next)
+			}
+			continue
+		}
+
 		// Removes the redundant copy instruction.
 		// TODO: doing this in `SetupEpilogue` seems weird. Find a better home.
 		if cur.IsCopy() && cur.op1.r.RealReg() == cur.op2.r.RealReg() {


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/abi_go_call_test.go b/internal/engine/wazevo/backend/isa/amd64/abi_go_call_test.go
index 5e401696..3aca2b6b 100644
--- a/internal/engine/wazevo/backend/isa/amd64/abi_go_call_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/abi_go_call_test.go
@@ -76,6 +76,10 @@ L2:
 	mov.q %r12, 48(%rax)
 	exit_sequence %rax
 L3:
+	add $8, %rsp
+	movq 8(%rsp), %rcx
+	movss 16(%rsp), %xmm0
+	movsd 24(%rsp), %xmm1
 	movq 96(%rax), %rdx
 	movq 112(%rax), %r12
 	movq 128(%rax), %r13
@@ -89,11 +93,7 @@ L3:
 	movdqu 256(%rax), %xmm13
 	movdqu 272(%rax), %xmm14
 	movdqu 288(%rax), %xmm15
-	add $8, %rsp
-	movzx.lq (%rsp), %rax
-	movq 8(%rsp), %rcx
-	movss 16(%rsp), %xmm0
-	movsd 24(%rsp), %xmm1
+	movq (%rsp), %rax
 	movq %rbp, %rsp
 	popq %rbp
 	ret
@@ -151,6 +151,7 @@ L2:
 	mov.q %r12, 48(%rax)
 	exit_sequence %rax
 L3:
+	add $8, %rsp
 	movq 96(%rax), %rdx
 	movq 112(%rax), %r12
 	movq 128(%rax), %r13
@@ -164,7 +165,6 @@ L3:
 	movdqu 256(%rax), %xmm13
 	movdqu 272(%rax), %xmm14
 	movdqu 288(%rax), %xmm15
-	add $8, %rsp
 	movq %rbp, %rsp
 	popq %rbp
 	ret
@@ -217,6 +217,7 @@ L2:
 	mov.q %r12, 48(%rax)
 	exit_sequence %rax
 L3:
+	add $8, %rsp
 	movq 96(%rax), %rdx
 	movq 112(%rax), %r12
 	movq 128(%rax), %r13
@@ -230,12 +231,140 @@ L3:
 	movdqu 256(%rax), %xmm13
 	movdqu 272(%rax), %xmm14
 	movdqu 288(%rax), %xmm15
+	movq (%rsp), %rax
+	movq %rbp, %rsp
+	popq %rbp
+	ret
+`,
+		},
+		{
+			name:     "many",
+			exitCode: wazevoapi.ExitCodeCallGoFunctionWithIndex(100, false),
+			sig: &ssa.Signature{
+				Params: []ssa.Type{
+					ssa.TypeI64, ssa.TypeI64, ssa.TypeF64,
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+				},
+				Results: []ssa.Type{
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+				},
+			},
+			exp: `
+	pushq %rbp
+	movq %rsp, %rbp
+	sub $248, %rsp
+	cmpq 40(%rax), %rsp
+	jnbe L1
+	add $248, %rsp
+	pushq %r15
+	movabsq $248, %r15
+	mov.q %r15, 64(%rax)
+	popq %r15
+	callq *80(%rax)
+	jmp L2
+L1:
+	add $248, %rsp
+L2:
+	mov.q %rdx, 96(%rax)
+	mov.q %r12, 112(%rax)
+	mov.q %r13, 128(%rax)
+	mov.q %r14, 144(%rax)
+	mov.q %r15, 160(%rax)
+	movdqu %xmm8, 176(%rax)
+	movdqu %xmm9, 192(%rax)
+	movdqu %xmm10, 208(%rax)
+	movdqu %xmm11, 224(%rax)
+	movdqu %xmm12, 240(%rax)
+	movdqu %xmm13, 256(%rax)
+	movdqu %xmm14, 272(%rax)
+	movdqu %xmm15, 288(%rax)
+	mov.q %rcx, 1120(%rax)
+	sub $240, %rsp
+	movsd %xmm0, (%rsp)
+	movsd %xmm1, 8(%rsp)
+	movdqu %xmm2, 16(%rsp)
+	mov.l %rbx, 32(%rsp)
+	mov.q %rsi, 40(%rsp)
+	movss %xmm3, 48(%rsp)
+	movsd %xmm4, 56(%rsp)
+	movdqu %xmm5, 64(%rsp)
+	mov.l %rdi, 80(%rsp)
+	mov.q %r8, 88(%rsp)
+	movss %xmm6, 96(%rsp)
+	movsd %xmm7, 104(%rsp)
+	movdqu 16(%rbp), %xmm15
+	movdqu %xmm15, 112(%rsp)
+	mov.l %r9, 128(%rsp)
+	mov.q %r10, 136(%rsp)
+	movss 32(%rbp), %xmm15
+	movss %xmm15, 144(%rsp)
+	pushq $240
+	movl $25606, %r12d
+	mov.l %r12, (%rax)
+	mov.q %rsp, 56(%rax)
+	mov.q %rbp, 1152(%rax)
+	lea L3(%rip), %r12
+	mov.q %r12, 48(%rax)
+	exit_sequence %rax
+L3:
 	add $8, %rsp
-	movzx.lq (%rsp), %rax
+	movsd (%rsp), %xmm0
+	movdqu 8(%rsp), %xmm1
+	movq 32(%rsp), %rcx
+	movss 40(%rsp), %xmm2
+	movsd 48(%rsp), %xmm3
+	movdqu 56(%rsp), %xmm4
+	movzx.lq 72(%rsp), %rbx
+	movq 80(%rsp), %rsi
+	movss 88(%rsp), %xmm5
+	movsd 96(%rsp), %xmm6
+	movdqu 104(%rsp), %xmm7
+	movzx.lq 120(%rsp), %rdi
+	movq 128(%rsp), %r8
+	movss 136(%rsp), %xmm15
+	movss %xmm15, 40(%rbp)
+	movsd 144(%rsp), %xmm15
+	movsd %xmm15, 48(%rbp)
+	movdqu 152(%rsp), %xmm15
+	movdqu %xmm15, 56(%rbp)
+	movzx.lq 168(%rsp), %r9
+	movq 176(%rsp), %r10
+	movss 184(%rsp), %xmm15
+	movss %xmm15, 72(%rbp)
+	movsd 192(%rsp), %xmm15
+	movsd %xmm15, 80(%rbp)
+	movdqu 200(%rsp), %xmm15
+	movdqu %xmm15, 88(%rbp)
+	movzx.lq 216(%rsp), %r11
+	movq 224(%rsp), %r15
+	mov.q %r15, 104(%rbp)
+	movss 232(%rsp), %xmm15
+	movss %xmm15, 112(%rbp)
+	movq 96(%rax), %rdx
+	movq 112(%rax), %r12
+	movq 128(%rax), %r13
+	movq 144(%rax), %r14
+	movq 160(%rax), %r15
+	movdqu 176(%rax), %xmm8
+	movdqu 192(%rax), %xmm9
+	movdqu 208(%rax), %xmm10
+	movdqu 224(%rax), %xmm11
+	movdqu 240(%rax), %xmm12
+	movdqu 256(%rax), %xmm13
+	movdqu 272(%rax), %xmm14
+	movdqu 288(%rax), %xmm15
+	movq 24(%rsp), %rax
 	movq %rbp, %rsp
 	popq %rbp
 	ret
 `,
+			needModuleContextPtr: true,
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index 0435029e..5c08b76f 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -808,7 +808,6 @@ func TestListener_imported(t *testing.T) {
 }
 
 func TestListener_long(t *testing.T) {
-	skipOnAmd64(t)
 	pickOneParam := binaryencoding.EncodeModule(&wasm.Module{
 		TypeSection: []wasm.FunctionType{{Results: []wasm.ValueType{i32}, Params: []wasm.ValueType{
 			i32, i32, f32, f64, i64, i32, i32, v128, f32,
@@ -858,7 +857,6 @@ func TestListener_long(t *testing.T) {
 }
 
 func TestListener_long_as_is(t *testing.T) {
-	skipOnAmd64(t)
 	params := []wasm.ValueType{
 		i32, i64, i32, i64, i32, i64, i32, i64, i32, i64,
 		i32, i64, i32, i64, i32, i64, i32, i64, i32, i64,
@@ -909,7 +907,6 @@ func TestListener_long_as_is(t *testing.T) {
 }
 
 func TestListener_long_many_consts(t *testing.T) {
-	skipOnAmd64(t)
 	const paramNum = 61
 
 	var exp []uint64