Repository: tetratelabs/wazero
Pull request #1934 (2024-01-22 18:14:43)

wazevo(amd64): exported function args


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/abi_entry_preamble.go b/internal/engine/wazevo/backend/isa/amd64/abi_entry_preamble.go
index 4abc41d1..e54e6f5d 100644
--- a/internal/engine/wazevo/backend/isa/amd64/abi_entry_preamble.go
+++ b/internal/engine/wazevo/backend/isa/amd64/abi_entry_preamble.go
@@ -12,19 +12,13 @@ var (
 	// savedExecutionContextPtr also must be a callee-saved reg so that they can be used in the prologue and epilogue.
 	savedExecutionContextPtr = rdxVReg
 	// paramResultSlicePtr must be callee-saved reg so that they can be used in the prologue and epilogue.
-	paramResultSlicePtr = r12VReg //nolint
+	paramResultSlicePtr = r12VReg
 	// goAllocatedStackPtr is not used in the epilogue.
 	goAllocatedStackPtr = rbxVReg
-	// paramResultSliceCopied is not used in the epilogue.
-	paramResultSliceCopied = rsiVReg //nolint
 	// functionExecutable is not used in the epilogue.
 	functionExecutable = rdiVReg
 )
 
-func (m *machine) goEntryPreamblePassArg(cur *instruction, paramSlicePtr regalloc.VReg, arg *backend.ABIArg, argStartOffsetFromSP int64) *instruction {
-	panic("TODO")
-}
-
 // CompileEntryPreamble implements backend.Machine.
 func (m *machine) CompileEntryPreamble(sig *ssa.Signature) []byte {
 	root := m.compileEntryPreamble(sig)
@@ -55,24 +49,26 @@ func (m *machine) compileEntryPreamble(sig *ssa.Signature) *instruction {
 	// 		mov %goAllocatedStackPtr, %rsp
 	cur = m.move64(goAllocatedStackPtr, rspVReg, cur)
 
-	// TODO: read the arguments from paramResultSliceCopied and set them correctly into the regs and stacks.
-	// 	The first two args are module context ptr and execution context ptr, and they are passed in %rax and %rcx.
-	prReg := paramResultSlicePtr
-	if len(abi.Args) > 2 && len(abi.Rets) > 0 {
-		// paramResultSlicePtr is modified during the execution of goEntryPreamblePassArg,
-		// so copy it to another reg.
-		cur = m.move64(paramResultSliceCopied, paramResultSlicePtr, cur)
-		prReg = paramResultSliceCopied
+	if stackSlotSize := abi.AlignedArgResultStackSlotSize(); stackSlotSize > 0 {
+		// Allocate stack slots for the arguments and return values.
+		// 		sub $stackSlotSize, %rsp
+		spDec := m.allocateInstr().asAluRmiR(aluRmiROpcodeSub, newOperandImm32(uint32(stackSlotSize)), rspVReg, true)
+		cur = linkInstr(cur, spDec)
 	}
 
-	stackSlotSize := abi.AlignedArgResultStackSlotSize()
+	var offset uint32
 	for i := range abi.Args {
 		if i < 2 {
 			// module context ptr and execution context ptr are passed in x0 and x1 by the Go assembly function.
 			continue
 		}
 		arg := &abi.Args[i]
-		cur = m.goEntryPreamblePassArg(cur, prReg, arg, -stackSlotSize)
+		cur = m.goEntryPreamblePassArg(cur, paramResultSlicePtr, offset, arg)
+		if arg.Type == ssa.TypeV128 {
+			offset += 16
+		} else {
+			offset += 8
+		}
 	}
 
 	// Now ready to call the real function. Note that at this point stack pointer is already set to the Go-allocated,
@@ -83,7 +79,7 @@ func (m *machine) compileEntryPreamble(sig *ssa.Signature) *instruction {
 	//// ----------------------------------- epilogue ----------------------------------- ////
 
 	// Read the results from regs and the stack, and set them correctly into the paramResultSlicePtr.
-	var offset uint32
+	offset = 0
 	for i := range abi.Rets {
 		r := &abi.Rets[i]
 		cur = m.goEntryPreamblePassResult(cur, paramResultSlicePtr, offset, r)
@@ -126,6 +122,36 @@ func (m *machine) linkUD2(cur *instruction) *instruction { //nolint
 	return linkInstr(cur, m.allocateInstr().asUD2())
 }
 
+func (m *machine) goEntryPreamblePassArg(cur *instruction, paramSlicePtr regalloc.VReg, offsetInParamSlice uint32, arg *backend.ABIArg) *instruction {
+	dst := arg.Reg
+	if arg.Kind == backend.ABIArgKindStack {
+		// Temporary reg.
+		panic("TODO")
+	}
+
+	load := m.allocateInstr()
+	a := newOperandMem(newAmodeImmReg(offsetInParamSlice, paramSlicePtr))
+	switch arg.Type {
+	case ssa.TypeI32:
+		load.asMovzxRmR(extModeLQ, a, dst)
+	case ssa.TypeI64:
+		load.asMov64MR(a, dst)
+	case ssa.TypeF32:
+		load.asXmmUnaryRmR(sseOpcodeMovss, a, dst)
+	case ssa.TypeF64:
+		load.asXmmUnaryRmR(sseOpcodeMovsd, a, dst)
+	case ssa.TypeV128:
+		load.asXmmUnaryRmR(sseOpcodeMovdqu, a, dst)
+	}
+
+	if arg.Kind == backend.ABIArgKindStack {
+		// store back.
+		panic("TODO")
+	}
+
+	return linkInstr(cur, load)
+}
+
 func (m *machine) goEntryPreamblePassResult(cur *instruction, resultSlicePtr regalloc.VReg, offset uint32, result *backend.ABIArg) *instruction {
 	if result.Kind == backend.ABIArgKindStack {
 		panic("TODO")
@@ -134,17 +160,18 @@ func (m *machine) goEntryPreamblePassResult(cur *instruction, resultSlicePtr reg
 	r := result.Reg
 
 	store := m.allocateInstr()
+	a := newOperandMem(newAmodeImmReg(offset, resultSlicePtr))
 	switch result.Type {
 	case ssa.TypeI32:
-		store.asMovRM(r, newOperandMem(newAmodeImmReg(offset, resultSlicePtr)), 4)
+		store.asMovRM(r, a, 4)
 	case ssa.TypeI64:
-		store.asMovRM(r, newOperandMem(newAmodeImmReg(offset, resultSlicePtr)), 8)
+		store.asMovRM(r, a, 8)
 	case ssa.TypeF32:
-		store.asXmmMovRM(sseOpcodeMovss, r, newOperandMem(newAmodeImmReg(offset, resultSlicePtr)))
+		store.asXmmMovRM(sseOpcodeMovss, r, a)
 	case ssa.TypeF64:
-		store.asXmmMovRM(sseOpcodeMovsd, r, newOperandMem(newAmodeImmReg(offset, resultSlicePtr)))
+		store.asXmmMovRM(sseOpcodeMovsd, r, a)
 	case ssa.TypeV128:
-		panic("TODO")
+		store.asXmmMovRM(sseOpcodeMovdqu, r, a)
 	}
 
 	return linkInstr(cur, store)


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/abi_entry_preamble_test.go b/internal/engine/wazevo/backend/isa/amd64/abi_entry_preamble_test.go
index f9e8843b..d429f56b 100644
--- a/internal/engine/wazevo/backend/isa/amd64/abi_entry_preamble_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/abi_entry_preamble_test.go
@@ -28,6 +28,29 @@ func TestMachineCompileEntryPreamble(t *testing.T) {
 	movq 16(%rdx), %rbp
 	movq 24(%rdx), %rsp
 	ret
+`,
+		},
+		{
+			name: "only regs args",
+			sig: &ssa.Signature{
+				// execContext and moduleContext are passed in %rax and %rcx.
+				Params: []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32, ssa.TypeF64, ssa.TypeV128, ssa.TypeI64},
+			},
+			exp: `
+	movq %rax, %rdx
+	mov.q %rbp, 16(%rax)
+	mov.q %rsp, 24(%rax)
+	movq %rbx, %rsp
+	movzx.lq (%r12), %rbx
+	movq 8(%r12), %rsi
+	movss 16(%r12), %xmm0
+	movsd 24(%r12), %xmm1
+	movdqu 32(%r12), %xmm2
+	movq 48(%r12), %rdi
+	callq *%rdi
+	movq 16(%rdx), %rbp
+	movq 24(%rdx), %rsp
+	ret
 `,
 		},
 		{
@@ -35,24 +58,54 @@ func TestMachineCompileEntryPreamble(t *testing.T) {
 			sig: &ssa.Signature{
 				// execContext and moduleContext are passed in %rax and %rcx.
 				Params:  []ssa.Type{ssa.TypeI64, ssa.TypeI64},
-				Results: []ssa.Type{ssa.TypeI32, ssa.TypeI64, ssa.TypeF32, ssa.TypeF64},
+				Results: []ssa.Type{ssa.TypeI32, ssa.TypeV128, ssa.TypeI64, ssa.TypeF32, ssa.TypeF64},
+			},
+			exp: `
+	movq %rax, %rdx
+	mov.q %rbp, 16(%rax)
+	mov.q %rsp, 24(%rax)
+	movq %rbx, %rsp
+	callq *%rdi
+	mov.l %rax, (%r12)
+	movdqu %xmm0, 8(%r12)
+	mov.q %rcx, 24(%r12)
+	movss %xmm1, 32(%r12)
+	movsd %xmm2, 40(%r12)
+	movq 16(%rdx), %rbp
+	movq 24(%rdx), %rsp
+	ret
+`,
+		},
+		{
+			name: "only regs args/rets",
+			sig: &ssa.Signature{
+				// execContext and moduleContext are passed in %rax and %rcx.
+				Params:  []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32, ssa.TypeF64, ssa.TypeV128, ssa.TypeI64},
+				Results: []ssa.Type{ssa.TypeI32, ssa.TypeV128, ssa.TypeI64, ssa.TypeF32, ssa.TypeF64},
 			},
 			exp: `
 	movq %rax, %rdx
 	mov.q %rbp, 16(%rax)
 	mov.q %rsp, 24(%rax)
 	movq %rbx, %rsp
+	movzx.lq (%r12), %rbx
+	movq 8(%r12), %rsi
+	movss 16(%r12), %xmm0
+	movsd 24(%r12), %xmm1
+	movdqu 32(%r12), %xmm2
+	movq 48(%r12), %rdi
 	callq *%rdi
 	mov.l %rax, (%r12)
-	mov.q %rcx, 8(%r12)
-	movss %xmm0, 16(%r12)
-	movsd %xmm1, 24(%r12)
+	movdqu %xmm0, 8(%r12)
+	mov.q %rcx, 24(%r12)
+	movss %xmm1, 32(%r12)
+	movsd %xmm2, 40(%r12)
 	movq 16(%rdx), %rbp
 	movq 24(%rdx), %rsp
 	ret
 `,
 		},
-		// TODO: add Wasm param, result cases.
+		// TODO: add stack based param/results.
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			_, _, m := newSetupWithMockContext()
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go b/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
index 2e5c78a8..41000c94 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
@@ -930,12 +930,26 @@ func TestInstruction_format_encode(t *testing.T) {
 			want:       "f20f10c1",
 			wantFormat: "movsd %xmm1, %xmm0",
 		},
+		{
+			setup: func(i *instruction) {
+				i.asXmmUnaryRmR(sseOpcodeMovsd, newOperandMem(newAmodeImmReg(16, r12VReg)), xmm0VReg)
+			},
+			want:       "f2410f10442410",
+			wantFormat: "movsd 16(%r12), %xmm0",
+		},
 		// movss
 		{
 			setup:      func(i *instruction) { i.asXmmUnaryRmR(sseOpcodeMovss, newOperandReg(xmm1VReg), xmm0VReg) },
 			want:       "f30f10c1",
 			wantFormat: "movss %xmm1, %xmm0",
 		},
+		{
+			setup: func(i *instruction) {
+				i.asXmmUnaryRmR(sseOpcodeMovss, newOperandMem(newAmodeImmReg(160, r12VReg)), xmm15VReg)
+			},
+			want:       "f3450f10bc24a0000000",
+			wantFormat: "movss 160(%r12), %xmm15",
+		},
 		// movups
 		{
 			setup:      func(i *instruction) { i.asXmmUnaryRmR(sseOpcodeMovups, newOperandReg(xmm1VReg), xmm0VReg) },
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index a8605555..0563d667 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -97,7 +97,6 @@ func TestE2E(t *testing.T) {
 		},
 		{
 			name: "add_sub_return", m: testcases.AddSubReturn.Module,
-			skipAMD64: false,
 			calls: []callCase{
 				{
 					params:     []uint64{},
@@ -403,7 +402,6 @@ func TestE2E(t *testing.T) {
 							} else {
 								require.NoError(t, err)
 								require.Equal(t, len(cc.expResults), len(result))
-								require.Equal(t, cc.expResults, result)
 								for i := range cc.expResults {
 									if cc.expResults[i] != result[i] {
 										t.Errorf("result[%d]: exp %d, got %d", i, cc.expResults[i], result[i])