Repository: tetratelabs/wazero
Pull request #1917 (2024-01-12 20:53:06)

wazevo(amd64): pass consts returns


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/abi.go b/internal/engine/wazevo/backend/isa/amd64/abi.go
index 89bf93fb..7b721364 100644
--- a/internal/engine/wazevo/backend/isa/amd64/abi.go
+++ b/internal/engine/wazevo/backend/isa/amd64/abi.go
@@ -71,7 +71,25 @@ func (m *machine) LowerParams(args []ssa.Value) {
 }
 
 // LowerReturns implements backend.Machine.
-func (m *machine) LowerReturns(returns []ssa.Value) {
-	// TODO implement me
-	panic("implement me")
+func (m *machine) LowerReturns(rets []ssa.Value) {
+	a := m.currentABI
+
+	l := len(rets) - 1
+	for i := range rets {
+		// Reverse order in order to avoid overwriting the stack returns existing in the return registers.
+		ret := rets[l-i]
+		r := &a.Rets[l-i]
+		reg := m.c.VRegOf(ret)
+		if def := m.c.ValueDefinition(ret); def.IsFromInstr() {
+			// Constant instructions are inlined.
+			if inst := def.Instr; inst.Constant() {
+				m.InsertLoadConstant(inst, reg)
+			}
+		}
+		if r.Kind == backend.ABIArgKindReg {
+			m.InsertMove(r.Reg, reg, ret.Type())
+		} else {
+			panic("TODO")
+		}
+	}
 }
diff --git a/internal/engine/wazevo/backend/isa/amd64/abi_entry_amd64.s b/internal/engine/wazevo/backend/isa/amd64/abi_entry_amd64.s
index fb56a46e..27691fb2 100644
--- a/internal/engine/wazevo/backend/isa/amd64/abi_entry_amd64.s
+++ b/internal/engine/wazevo/backend/isa/amd64/abi_entry_amd64.s
@@ -7,7 +7,7 @@ TEXT ·entrypoint(SB), NOSPLIT|NOFRAME, $0-48
 	MOVQ functionExectuable+8(FP), DI
 	MOVQ executionContextPtr+16(FP), AX      // First argument is passed in AX.
 	MOVQ moduleContextPtr+24(FP), CX         // Second argument is passed in CX.
-	MOVQ paramResultSlicePtr+32(FP), SI
+	MOVQ paramResultSlicePtr+32(FP), R12
 	MOVQ goAllocatedStackSlicePtr+40(FP), BX
 	JMP  R11
 
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 2cfea86f..7bfb265a 100644
--- a/internal/engine/wazevo/backend/isa/amd64/abi_entry_preamble.go
+++ b/internal/engine/wazevo/backend/isa/amd64/abi_entry_preamble.go
@@ -64,9 +64,16 @@ func (m *machine) compileEntryPreamble(sig *ssa.Signature) *instruction {
 
 	//// ----------------------------------- epilogue ----------------------------------- ////
 
-	// TODO: read the results from regs and the stack, and set them correctly into the paramResultSlicePtr.
-	if len(abi.Rets) > 0 {
-		panic("TODO")
+	// Read the results from regs and the stack, and set them correctly into the paramResultSlicePtr.
+	var offset uint32
+	for i := range abi.Rets {
+		r := &abi.Rets[i]
+		cur = m.goEntryPreamblePassResult(cur, paramResultSlicePtr, offset, r)
+		if r.Type == ssa.TypeV128 {
+			offset += 16
+		} else {
+			offset += 8
+		}
 	}
 
 	// Finally, restore the original RBP and RSP.
@@ -100,3 +107,27 @@ func (m *machine) loadOrStore64AtExecutionCtx(execCtx regalloc.VReg, offset waze
 func (m *machine) linkUD2(cur *instruction) *instruction { //nolint
 	return linkInstr(cur, m.allocateInstr().asUD2())
 }
+
+func (m *machine) goEntryPreamblePassResult(cur *instruction, resultSlicePtr regalloc.VReg, offset uint32, result *backend.ABIArg) *instruction {
+	if result.Kind == backend.ABIArgKindStack {
+		panic("TODO")
+	}
+
+	r := result.Reg
+
+	store := m.allocateInstr()
+	switch result.Type {
+	case ssa.TypeI32:
+		store.asMovRM(r, newOperandMem(newAmodeImmReg(offset, resultSlicePtr)), 4)
+	case ssa.TypeI64:
+		store.asMovRM(r, newOperandMem(newAmodeImmReg(offset, resultSlicePtr)), 8)
+	case ssa.TypeF32:
+		store.asXmmMovRM(sseOpcodeMovss, r, newOperandMem(newAmodeImmReg(offset, resultSlicePtr)))
+	case ssa.TypeF64:
+		store.asXmmMovRM(sseOpcodeMovsd, r, newOperandMem(newAmodeImmReg(offset, resultSlicePtr)))
+	case ssa.TypeV128:
+		panic("TODO")
+	}
+
+	return linkInstr(cur, store)
+}
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index c12614e1..fcf62254 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -49,9 +49,9 @@ func (i *instruction) String() string {
 		return "ret"
 	case imm:
 		if i.b1 {
-			return fmt.Sprintf("movabsq $%d, %s", int64(i.u1), i.op1.format(true))
+			return fmt.Sprintf("movabsq $%d, %s", int64(i.u1), i.op2.format(true))
 		} else {
-			return fmt.Sprintf("movl $%d, %s", int32(i.u1), i.op1.format(false))
+			return fmt.Sprintf("movl $%d, %s", int32(i.u1), i.op2.format(false))
 		}
 	case aluRmiR:
 		return fmt.Sprintf("%s %s, %s", aluRmiROpcode(i.u1), i.op1.format(i.b1), i.op2.format(i.b1))
@@ -117,7 +117,7 @@ func (i *instruction) String() string {
 	case pop64:
 		return fmt.Sprintf("popq %s", i.op1.format(true))
 	case xmmMovRM:
-		panic("TODO")
+		return fmt.Sprintf("%s %s, %s", sseOpcode(i.u1), i.op1.format(true), i.op2.format(true))
 	case xmmLoadConst:
 		panic("TODO")
 	case xmmToGpr:
@@ -155,7 +155,7 @@ func (i *instruction) String() string {
 	case callIndirect:
 		return fmt.Sprintf("callq *%s", i.op1.format(true))
 	default:
-		panic(fmt.Sprintf("BUG: %v", i.kind))
+		panic(fmt.Sprintf("BUG: %d", int(i.kind)))
 	}
 }
 
@@ -165,9 +165,12 @@ func (i *instruction) Defs(regs *[]regalloc.VReg) []regalloc.VReg {
 	switch dk := defKinds[i.kind]; dk {
 	case defKindNone:
 	case defKindOp2:
+		if !i.op2.r.Valid() {
+			panic("BUG" + i.String())
+		}
 		*regs = append(*regs, i.op2.r)
 	default:
-		panic(fmt.Sprintf("BUG: invalid defKind %s for %s", dk, i))
+		panic(fmt.Sprintf("BUG: invalid defKind \"%s\" for %s", dk, i))
 	}
 	return *regs
 }
@@ -177,6 +180,17 @@ func (i *instruction) Uses(regs *[]regalloc.VReg) []regalloc.VReg {
 	*regs = (*regs)[:0]
 	switch uk := useKinds[i.kind]; uk {
 	case useKindNone:
+	case useKindOp1:
+		op := i.op1
+		switch op.kind {
+		case operandKindReg:
+			*regs = append(*regs, op.r)
+		case operandKindMem:
+			op.amode.uses(regs)
+		case operandKindImm32, operandKindLabel:
+		default:
+			panic(fmt.Sprintf("BUG: invalid operand: %s", i))
+		}
 	default:
 		panic(fmt.Sprintf("BUG: invalid useKind %s for %s", uk, i))
 	}
@@ -185,14 +199,41 @@ func (i *instruction) Uses(regs *[]regalloc.VReg) []regalloc.VReg {
 
 // AssignUse implements regalloc.Instr.
 func (i *instruction) AssignUse(index int, v regalloc.VReg) {
-	// TODO implement me
-	panic("implement me")
+	switch uk := useKinds[i.kind]; uk {
+	case useKindNone:
+	case useKindOp1:
+		op := &i.op1
+		switch op.kind {
+		case operandKindReg:
+			if index != 0 {
+				panic("BUG")
+			}
+			if op.r.IsRealReg() {
+				panic("BUG already assigned: " + i.String())
+			}
+			op.r = v
+		case operandKindMem:
+			op.amode.assignUses(index, v)
+		default:
+			panic(fmt.Sprintf("BUG: invalid operand: %s", i))
+		}
+	default:
+		panic(fmt.Sprintf("BUG: invalid useKind %s for %s", uk, i))
+	}
 }
 
 // AssignDef implements regalloc.Instr.
 func (i *instruction) AssignDef(reg regalloc.VReg) {
-	// TODO implement me
-	panic("implement me")
+	switch dk := defKinds[i.kind]; dk {
+	case defKindNone:
+	case defKindOp2:
+		if !i.op2.r.Valid() {
+			panic("BUG already assigned" + i.String())
+		}
+		i.op2.r = reg
+	default:
+		panic(fmt.Sprintf("BUG: invalid defKind \"%s\" for %s", dk, i))
+	}
 }
 
 // IsCopy implements regalloc.Instr.
@@ -323,7 +364,7 @@ const (
 	// value. This is characteristic of mov instructions.
 	xmmUnaryRmR
 
-	// XMM (scalar or vector) unary op (from xmm to reg/mem): stores, movd, movq
+	// XMM (scalar or vector) unary op (from xmm to mem): stores, movd, movq
 	xmmMovRM
 
 	// XMM (vector) unary op (to move a constant value into an xmm register): movups
@@ -554,7 +595,7 @@ func (i *instruction) asRet(abi *backend.FunctionABI) *instruction {
 
 func (i *instruction) asImm(dst regalloc.VReg, value uint64, _64 bool) *instruction {
 	i.kind = imm
-	i.op1 = newOperandReg(dst)
+	i.op2 = newOperandReg(dst)
 	i.u1 = value
 	i.b1 = _64
 	return i
@@ -664,6 +705,17 @@ func (i *instruction) asXmmUnaryRmR(op sseOpcode, rm operand, rd regalloc.VReg,
 	return i
 }
 
+func (i *instruction) asXmmMovRM(op sseOpcode, rm regalloc.VReg, rd operand) *instruction {
+	if rd.kind != operandKindMem {
+		panic("BUG")
+	}
+	i.kind = xmmMovRM
+	i.op1 = newOperandReg(rm)
+	i.op2 = rd
+	i.u1 = uint64(op)
+	return i
+}
+
 func (i *instruction) asPop64(rm regalloc.VReg) *instruction {
 	i.kind = pop64
 	i.op1 = newOperandReg(rm)
@@ -1185,8 +1237,12 @@ const (
 )
 
 var defKinds = [instrMax]defKind{
-	nop0: defKindNone,
-	ret:  defKindNone,
+	nop0:        defKindNone,
+	ret:         defKindNone,
+	movRR:       defKindOp2,
+	imm:         defKindOp2,
+	xmmUnaryRmR: defKindOp2,
+	gprToXmm:    defKindOp2,
 }
 
 // String implements fmt.Stringer.
@@ -1209,8 +1265,12 @@ const (
 )
 
 var useKinds = [instrMax]useKind{
-	nop0: useKindNone,
-	ret:  useKindNone,
+	nop0:        useKindNone,
+	ret:         useKindNone,
+	movRR:       useKindOp1,
+	imm:         useKindNone,
+	xmmUnaryRmR: useKindOp1,
+	gprToXmm:    useKindOp1,
 }
 
 func (u useKind) String() string {
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
index ac7f257e..8f0b76b2 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
@@ -13,7 +13,7 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 	case ret:
 		c.EmitByte(0xc3)
 	case imm:
-		dst := regEncodings[i.op1.r.RealReg()]
+		dst := regEncodings[i.op2.r.RealReg()]
 		con := i.u1
 		if i.b1 { // 64 bit.
 			if lower32willSignExtendTo64(con) {
@@ -597,7 +597,32 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 		c.EmitByte(0x58 | dst.encoding())
 
 	case xmmMovRM:
-		panic("TODO")
+		var legPrefix legacyPrefixes
+		var opcode uint32
+		const opcodeNum = 2
+		switch sseOpcode(i.u1) {
+		case sseOpcodeMovaps:
+			legPrefix, opcode = legacyPrefixesNone, 0x0f29
+		case sseOpcodeMovapd:
+			legPrefix, opcode = legacyPrefixes0x66, 0x0f29
+		case sseOpcodeMovdqa:
+			legPrefix, opcode = legacyPrefixes0x66, 0x0f7f
+		case sseOpcodeMovdqu:
+			legPrefix, opcode = legacyPrefixes0xF3, 0x0f7f
+		case sseOpcodeMovss:
+			legPrefix, opcode = legacyPrefixes0xF3, 0x0f11
+		case sseOpcodeMovsd:
+			legPrefix, opcode = legacyPrefixes0xF2, 0x0f11
+		case sseOpcodeMovups:
+			legPrefix, opcode = legacyPrefixesNone, 0x0f11
+		case sseOpcodeMovupd:
+			legPrefix, opcode = legacyPrefixes0x66, 0x0f11
+		default:
+			panic(fmt.Sprintf("Unsupported sseOpcode: %s", sseOpcode(i.u1)))
+		}
+
+		dst := regEncodings[i.op1.r.RealReg()]
+		encodeRegMem(c, legPrefix, opcode, opcodeNum, dst, i.op2.amode, rexInfo(0).clearW())
 	case xmmLoadConst:
 		panic("TODO")
 	case xmmToGpr:
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index 44f47d7d..bd293d51 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -125,9 +125,16 @@ func (m *machine) LowerConditionalBranch(b *ssa.Instruction) {
 }
 
 // LowerInstr implements backend.Machine.
-func (m *machine) LowerInstr(instruction *ssa.Instruction) {
-	// TODO implement me
-	panic("implement me")
+func (m *machine) LowerInstr(instr *ssa.Instruction) {
+	switch op := instr.Opcode(); op {
+	case ssa.OpcodeBrz, ssa.OpcodeBrnz, ssa.OpcodeJump, ssa.OpcodeBrTable:
+		panic("BUG: branching instructions are handled by LowerBranches")
+	case ssa.OpcodeReturn:
+		panic("BUG: return must be handled by backend.Compiler")
+	case ssa.OpcodeIconst, ssa.OpcodeF32const, ssa.OpcodeF64const: // Constant instructions are inlined.
+	default:
+		panic("TODO: lowering " + op.String())
+	}
 }
 
 // InsertMove implements backend.Machine.
diff --git a/internal/engine/wazevo/backend/isa/amd64/operands.go b/internal/engine/wazevo/backend/isa/amd64/operands.go
index 6c7a0ef1..8a839305 100644
--- a/internal/engine/wazevo/backend/isa/amd64/operands.go
+++ b/internal/engine/wazevo/backend/isa/amd64/operands.go
@@ -31,6 +31,22 @@ const (
 	operandKindLabel
 )
 
+// String implements fmt.Stringer.
+func (o operandKind) String() string {
+	switch o {
+	case operandKindReg:
+		return "reg"
+	case operandKindMem:
+		return "mem"
+	case operandKindImm32:
+		return "imm32"
+	case operandKindLabel:
+		return "label"
+	default:
+		panic("BUG: invalid operand kind")
+	}
+}
+
 func (o *operand) format(_64 bool) string {
 	switch o.kind {
 	case operandKindReg:
@@ -42,7 +58,7 @@ func (o *operand) format(_64 bool) string {
 	case operandKindLabel:
 		return backend.Label(o.imm32).String()
 	default:
-		panic("BUG: invalid operand kind")
+		panic(fmt.Sprintf("BUG: invalid operand: %s", o.kind))
 	}
 }
 
@@ -93,6 +109,41 @@ const (
 	// TODO: there are other addressing modes such as the one without base register.
 )
 
+func (a *amode) uses(rs *[]regalloc.VReg) {
+	switch a.kind {
+	case amodeImmReg:
+		*rs = append(*rs, a.base)
+	case amodeRegRegShift:
+		*rs = append(*rs, a.base, a.index)
+	case amodeRipRelative:
+		// nothing
+	default:
+		panic("BUG: invalid amode kind")
+	}
+}
+
+func (a *amode) assignUses(i int, reg regalloc.VReg) {
+	switch a.kind {
+	case amodeImmReg:
+		if i == 0 {
+			a.base = reg
+		} else {
+			panic("BUG: invalid amode assignment")
+		}
+	case amodeRegRegShift:
+		if i == 0 {
+			a.base = reg
+		} else if i == 1 {
+			a.index = reg
+		} else {
+			panic("BUG: invalid amode assignment")
+		}
+	case amodeRipRelative:
+	default:
+		panic("BUG: invalid amode assignment")
+	}
+}
+
 func newAmodeImmReg(imm32 uint32, base regalloc.VReg) amode {
 	return amode{kind: amodeImmReg, imm32: imm32, base: base}
 }


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 d894796c..f9e8843b 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,28 @@ func TestMachineCompileEntryPreamble(t *testing.T) {
 	movq 16(%rdx), %rbp
 	movq 24(%rdx), %rsp
 	ret
+`,
+		},
+		{
+			name: "only regs rets",
+			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},
+			},
+			exp: `
+	movq %rax, %rdx
+	mov.q %rbp, 16(%rax)
+	mov.q %rsp, 24(%rax)
+	movq %rbx, %rsp
+	callq *%rdi
+	mov.l %rax, (%r12)
+	mov.q %rcx, 8(%r12)
+	movss %xmm0, 16(%r12)
+	movsd %xmm1, 24(%r12)
+	movq 16(%rdx), %rbp
+	movq 24(%rdx), %rsp
+	ret
 `,
 		},
 		// TODO: add Wasm param, result cases.
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 826fa874..18a29e15 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
@@ -1703,6 +1703,78 @@ func TestInstruction_format_encode(t *testing.T) {
 			want:       "48897c087b",
 			wantFormat: "mov.q %rdi, 123(%rax,%rcx,1)",
 		},
+		{
+			setup: func(i *instruction) {
+				a := newOperandMem(newAmodeRegRegShit(123, raxVReg, rcxVReg, 0))
+				i.asXmmMovRM(sseOpcodeMovaps, xmm1VReg, a)
+			},
+			want:       "0f294c087b",
+			wantFormat: "movaps %xmm1, 123(%rax,%rcx,1)",
+		},
+		{
+			setup: func(i *instruction) {
+				a := newOperandMem(newAmodeImmReg(123, rcxVReg))
+				i.asXmmMovRM(sseOpcodeMovaps, xmm15VReg, a)
+			},
+			want:       "440f29797b",
+			wantFormat: "movaps %xmm15, 123(%rcx)",
+		},
+		{
+			setup: func(i *instruction) {
+				a := newOperandMem(newAmodeImmReg(123, rcxVReg))
+				i.asXmmMovRM(sseOpcodeMovapd, xmm15VReg, a)
+			},
+			want:       "66440f29797b",
+			wantFormat: "movapd %xmm15, 123(%rcx)",
+		},
+		{
+			setup: func(i *instruction) {
+				a := newOperandMem(newAmodeImmReg(123, rcxVReg))
+				i.asXmmMovRM(sseOpcodeMovdqa, xmm15VReg, a)
+			},
+			want:       "66440f7f797b",
+			wantFormat: "movdqa %xmm15, 123(%rcx)",
+		},
+		{
+			setup: func(i *instruction) {
+				a := newOperandMem(newAmodeImmReg(123, rcxVReg))
+				i.asXmmMovRM(sseOpcodeMovdqu, xmm15VReg, a)
+			},
+			want:       "f3440f7f797b",
+			wantFormat: "movdqu %xmm15, 123(%rcx)",
+		},
+		{
+			setup: func(i *instruction) {
+				a := newOperandMem(newAmodeImmReg(123, rcxVReg))
+				i.asXmmMovRM(sseOpcodeMovss, xmm15VReg, a)
+			},
+			want:       "f3440f11797b",
+			wantFormat: "movss %xmm15, 123(%rcx)",
+		},
+		{
+			setup: func(i *instruction) {
+				a := newOperandMem(newAmodeImmReg(123, rcxVReg))
+				i.asXmmMovRM(sseOpcodeMovsd, xmm15VReg, a)
+			},
+			want:       "f2440f11797b",
+			wantFormat: "movsd %xmm15, 123(%rcx)",
+		},
+		{
+			setup: func(i *instruction) {
+				a := newOperandMem(newAmodeImmReg(123, rcxVReg))
+				i.asXmmMovRM(sseOpcodeMovups, xmm15VReg, a)
+			},
+			want:       "440f11797b",
+			wantFormat: "movups %xmm15, 123(%rcx)",
+		},
+		{
+			setup: func(i *instruction) {
+				a := newOperandMem(newAmodeImmReg(123, rcxVReg))
+				i.asXmmMovRM(sseOpcodeMovupd, xmm15VReg, a)
+			},
+			want:       "66440f11797b",
+			wantFormat: "movupd %xmm15, 123(%rcx)",
+		},
 	} {
 		tc := tc
 		t.Run(tc.wantFormat, func(t *testing.T) {
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index 5fd5fb78..194c0cfd 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -86,7 +86,6 @@ func TestE2E(t *testing.T) {
 		},
 		{
 			name: "consts", m: testcases.Constants.Module,
-			skipAMD64: true,
 			calls: []callCase{
 				{expResults: []uint64{1, 2, uint64(math.Float32bits(32.0)), math.Float64bits(64.0)}},
 			},