Repository: tetratelabs/wazero
Pull request #1952 (2024-01-24 19:11:27)

wazevo(amd64): imported func calls / global.get


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index 0138c224..4fb10eb1 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -270,6 +270,17 @@ func (i *instruction) Uses(regs *[]regalloc.VReg) []regalloc.VReg {
 		default:
 			panic(fmt.Sprintf("BUG: invalid operand: %s", i))
 		}
+	case useKindCallInd:
+		op := i.op1
+		switch op.kind {
+		case operandKindReg:
+			*regs = append(*regs, op.r)
+		case operandKindMem:
+			op.amode.uses(regs)
+		default:
+			panic(fmt.Sprintf("BUG: invalid operand: %s", i))
+		}
+		*regs = append(*regs, i.abi.ArgRealRegs...)
 	case useKindCall:
 		*regs = append(*regs, i.abi.ArgRealRegs...)
 	default:
@@ -282,6 +293,19 @@ func (i *instruction) Uses(regs *[]regalloc.VReg) []regalloc.VReg {
 func (i *instruction) AssignUse(index int, v regalloc.VReg) {
 	switch uk := useKinds[i.kind]; uk {
 	case useKindNone:
+	case useKindCallInd:
+		if index != 0 {
+			panic("BUG")
+		}
+		op := &i.op1
+		switch op.kind {
+		case operandKindReg:
+			op.r = v
+		case operandKindMem:
+			op.amode.assignUses(index, v)
+		default:
+			panic("BUG")
+		}
 	case useKindOp1Op2Reg, useKindOp1RegOp2:
 		op, opMustBeReg := &i.op1, &i.op2
 		if uk == useKindOp1RegOp2 {
@@ -1542,6 +1566,7 @@ var defKinds = [instrMax]defKind{
 	movzxRmR:     defKindOp2,
 	gprToXmm:     defKindOp2,
 	call:         defKindCall,
+	callIndirect: defKindCall,
 	ud2:          defKindNone,
 	jmp:          defKindNone,
 	jmpIf:        defKindNone,
@@ -1574,6 +1599,7 @@ const (
 	// useKindOp1RegOp2 is Op1 must be a register, Op2 can be any operand.
 	useKindOp1RegOp2
 	useKindCall
+	useKindCallInd
 )
 
 var useKinds = [instrMax]useKind{
@@ -1590,6 +1616,7 @@ var useKinds = [instrMax]useKind{
 	movzxRmR:     useKindOp1,
 	gprToXmm:     useKindOp1,
 	call:         useKindCall,
+	callIndirect: useKindCallInd,
 	ud2:          useKindNone,
 	jmpIf:        useKindOp1,
 	jmp:          useKindOp1,
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index b758ae7d..1799672f 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -198,7 +198,7 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 	case ssa.OpcodeReturn:
 		panic("BUG: return must be handled by backend.Compiler")
 	case ssa.OpcodeIconst, ssa.OpcodeF32const, ssa.OpcodeF64const: // Constant instructions are inlined.
-	case ssa.OpcodeCall:
+	case ssa.OpcodeCall, ssa.OpcodeCallIndirect:
 		m.lowerCall(instr)
 	case ssa.OpcodeStore:
 		m.lowerStore(instr)
@@ -229,11 +229,34 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 	case ssa.OpcodeExitWithCode:
 		execCtx, code := instr.ExitWithCodeData()
 		m.lowerExitWithCode(m.c.VRegOf(execCtx), code)
+	case ssa.OpcodeLoad:
+		dst := instr.Return()
+		ptr, offset, typ := instr.LoadData()
+		m.lowerLoad(ptr, offset, typ, dst)
 	default:
 		panic("TODO: lowering " + op.String())
 	}
 }
 
+func (m *machine) lowerLoad(ptr ssa.Value, offset uint32, typ ssa.Type, ret ssa.Value) {
+	mem := newOperandMem(m.lowerToAddressMode(ptr, offset))
+	dst := m.c.VRegOf(ret)
+	load := m.allocateInstr()
+	switch typ {
+	case ssa.TypeI32:
+		load.asMovzxRmR(extModeLQ, mem, dst)
+	case ssa.TypeI64:
+		load.asMov64MR(mem, dst)
+	case ssa.TypeF32:
+		load.asXmmUnaryRmR(sseOpcodeMovss, mem, dst)
+	case ssa.TypeF64:
+		load.asXmmUnaryRmR(sseOpcodeMovsd, mem, dst)
+	case ssa.TypeV128:
+		load.asXmmUnaryRmR(sseOpcodeMovdqu, mem, dst)
+	}
+	m.insert(load)
+}
+
 func (m *machine) lowerExitWithCode(execCtx regalloc.VReg, code wazevoapi.ExitCode) {
 	// First we set the exit code in the execution context.
 	exitCodeReg := m.c.AllocateVReg(ssa.TypeI32)
@@ -371,13 +394,14 @@ func (m *machine) lowerStore(si *ssa.Instruction) {
 
 func (m *machine) lowerCall(si *ssa.Instruction) {
 	isDirectCall := si.Opcode() == ssa.OpcodeCall
+	var indirectCalleePtr ssa.Value
 	var directCallee ssa.FuncRef
 	var sigID ssa.SignatureID
 	var args []ssa.Value
 	if isDirectCall {
 		directCallee, sigID, args = si.CallData()
 	} else {
-		panic("TODO")
+		indirectCalleePtr, sigID, args = si.CallIndirectData()
 	}
 	calleeABI := m.c.GetFunctionABI(m.c.SSABuilder().ResolveSignature(sigID))
 
@@ -393,11 +417,12 @@ func (m *machine) lowerCall(si *ssa.Instruction) {
 	}
 
 	if isDirectCall {
-		call := m.allocateInstr()
-		call.asCall(directCallee, calleeABI)
+		call := m.allocateInstr().asCall(directCallee, calleeABI)
 		m.insert(call)
 	} else {
-		panic("TODO")
+		ptrOp := m.getOperand_Mem_Reg(m.c.ValueDefinition(indirectCalleePtr))
+		callInd := m.allocateInstr().asCallIndirect(ptrOp, calleeABI)
+		m.insert(callInd)
 	}
 
 	var index int
diff --git a/internal/engine/wazevo/backend/isa/amd64/operands.go b/internal/engine/wazevo/backend/isa/amd64/operands.go
index 493eba39..abccf553 100644
--- a/internal/engine/wazevo/backend/isa/amd64/operands.go
+++ b/internal/engine/wazevo/backend/isa/amd64/operands.go
@@ -233,6 +233,20 @@ func (a *amode) String() string {
 	panic("BUG: invalid amode kind")
 }
 
+func (m *machine) getOperand_Mem_Reg(def *backend.SSAValueDefinition) (op operand) {
+	if def.IsFromBlockParam() {
+		return newOperandReg(def.BlkParamVReg)
+	}
+	if m.c.MatchInstr(def, ssa.OpcodeLoad) {
+		instr := def.Instr
+		ptr, offset, _ := instr.LoadData()
+		op = newOperandMem(m.lowerToAddressMode(ptr, offset))
+		instr.MarkLowered()
+		return op
+	}
+	return m.getOperand_Reg(def)
+}
+
 func (m *machine) getOperand_Mem_Imm32_Reg(def *backend.SSAValueDefinition) (op operand) {
 	if def.IsFromBlockParam() {
 		return newOperandReg(def.BlkParamVReg)


Test patch:
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index b1b39ff5..c8e69a03 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -122,7 +122,6 @@ func TestE2E(t *testing.T) {
 				},
 			},
 		},
-
 		{
 			name: "many_params_many_results",
 			m:    testcases.ManyParamsManyResults.Module,
@@ -170,10 +169,9 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "imported_function_call",
-			imported:  testcases.ImportedFunctionCall.Imported,
-			m:         testcases.ImportedFunctionCall.Module,
-			skipAMD64: true,
+			name:     "imported_function_call",
+			imported: testcases.ImportedFunctionCall.Imported,
+			m:        testcases.ImportedFunctionCall.Module,
 			calls: []callCase{
 				{params: []uint64{0}, expResults: []uint64{0}},
 				{params: []uint64{2}, expResults: []uint64{2 * 2}},
@@ -234,11 +232,10 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "globals_get",
-			m:         testcases.GlobalsGet.Module,
-			skipAMD64: true,
+			name: "globals_get",
+			m:    testcases.GlobalsGet.Module,
 			calls: []callCase{
-				{expResults: []uint64{0x80000000, 0x8000000000000000, 0x7f7fffff, 0x7fefffffffffffff}},
+				{expResults: []uint64{0x80000000, 0x8000000000000000, 0x7f7fffff, 0x7fefffffffffffff, 1234, 5678}},
 			},
 		},
 		{
diff --git a/internal/engine/wazevo/frontend/frontend_test.go b/internal/engine/wazevo/frontend/frontend_test.go
index a6551dc9..a5d37012 100644
--- a/internal/engine/wazevo/frontend/frontend_test.go
+++ b/internal/engine/wazevo/frontend/frontend_test.go
@@ -1185,7 +1185,8 @@ blk0: (exec_ctx:i64, module_ctx:i64)
 	v3:i64 = Load module_ctx, 0x18
 	v4:f32 = Load module_ctx, 0x28
 	v5:f64 = Load module_ctx, 0x38
-	Jump blk_ret, v2, v3, v4, v5
+	v6:v128 = Load module_ctx, 0x48
+	Jump blk_ret, v2, v3, v4, v5, v6
 `,
 		},
 		{
diff --git a/internal/engine/wazevo/testcases/testcases.go b/internal/engine/wazevo/testcases/testcases.go
index 7e54e753..43f203de 100644
--- a/internal/engine/wazevo/testcases/testcases.go
+++ b/internal/engine/wazevo/testcases/testcases.go
@@ -1291,7 +1291,7 @@ var (
 	GlobalsGet = TestCase{
 		Name: "globals_get",
 		Module: &wasm.Module{
-			TypeSection:     []wasm.FunctionType{{Results: []wasm.ValueType{i32, i64, f32, f64}}},
+			TypeSection:     []wasm.FunctionType{{Results: []wasm.ValueType{i32, i64, f32, f64, v128}}},
 			ExportSection:   []wasm.Export{{Name: ExportedFunctionName, Type: wasm.ExternTypeFunc, Index: 0}},
 			FunctionSection: []wasm.Index{0},
 			GlobalSection: []wasm.Global{
@@ -1311,6 +1311,10 @@ var (
 					Type: wasm.GlobalType{ValType: wasm.ValueTypeF64, Mutable: false},
 					Init: constExprF64(math.MaxFloat64),
 				},
+				{
+					Type: wasm.GlobalType{ValType: wasm.ValueTypeV128, Mutable: false},
+					Init: constExprV128(1234, 5678),
+				},
 			},
 			CodeSection: []wasm.Code{
 				{Body: []byte{
@@ -1318,6 +1322,7 @@ var (
 					wasm.OpcodeGlobalGet, 1,
 					wasm.OpcodeGlobalGet, 2,
 					wasm.OpcodeGlobalGet, 3,
+					wasm.OpcodeGlobalGet, 4,
 					wasm.OpcodeEnd,
 				}},
 			},
@@ -1770,3 +1775,15 @@ func constExprF64(i float64) wasm.ConstantExpression {
 		},
 	}
 }
+
+func constExprV128(lo, hi uint64) wasm.ConstantExpression {
+	return wasm.ConstantExpression{
+		Opcode: wasm.OpcodeVecV128Const,
+		Data: []byte{
+			byte(lo), byte(lo >> 8), byte(lo >> 16), byte(lo >> 24),
+			byte(lo >> 32), byte(lo >> 40), byte(lo >> 48), byte(lo >> 56),
+			byte(hi), byte(hi >> 8), byte(hi >> 16), byte(hi >> 24),
+			byte(hi >> 32), byte(hi >> 40), byte(hi >> 48), byte(hi >> 56),
+		},
+	}
+}
diff --git a/internal/testing/binaryencoding/const_expr.go b/internal/testing/binaryencoding/const_expr.go
index f4c14807..c629c4a2 100644
--- a/internal/testing/binaryencoding/const_expr.go
+++ b/internal/testing/binaryencoding/const_expr.go
@@ -5,6 +5,9 @@ import (
 )
 
 func encodeConstantExpression(expr wasm.ConstantExpression) (ret []byte) {
+	if expr.Opcode == wasm.OpcodeVecV128Const {
+		ret = append(ret, wasm.OpcodeVecPrefix)
+	}
 	ret = append(ret, expr.Opcode)
 	ret = append(ret, expr.Data...)
 	ret = append(ret, wasm.OpcodeEnd)