Repository: tetratelabs/wazero
Pull request #1953 (2024-01-24 20:20:27)

wazevo(amd64): globals.set


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index 4fb10eb1..dcf56ad5 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -211,6 +211,8 @@ func (i *instruction) String() string {
 		}
 	case callIndirect:
 		return fmt.Sprintf("callq *%s", i.op1.format(true))
+	case v128ConstIsland:
+		return fmt.Sprintf("v128ConstIsland (%#x, %#x)", i.u1, i.u2)
 	default:
 		panic(fmt.Sprintf("BUG: %d", int(i.kind)))
 	}
@@ -575,6 +577,9 @@ const (
 	// An instruction that will always trigger the illegal instruction exception.
 	ud2
 
+	// v128ConstIsland is 16 bytes (128-bit) constant that will be loaded into an XMM.
+	v128ConstIsland
+
 	instrMax
 )
 
@@ -662,6 +667,8 @@ func (k instructionKind) String() string {
 		return "jmpTableSeq"
 	case exitSequence:
 		return "exit_sequence"
+	case v128ConstIsland:
+		return "v128ConstIsland"
 	case ud2:
 		return "ud2"
 	default:
@@ -966,6 +973,13 @@ func (i *instruction) asXmmUnaryRmR(op sseOpcode, rm operand, rd regalloc.VReg)
 	return i
 }
 
+func (i *instruction) asV128ConstIsland(lo, hi uint64) *instruction {
+	i.kind = v128ConstIsland
+	i.u1 = lo
+	i.u2 = hi
+	return i
+}
+
 func (i *instruction) asXmmMovRM(op sseOpcode, rm regalloc.VReg, rd operand) *instruction {
 	if rd.kind != operandKindMem {
 		panic("BUG")
@@ -1553,26 +1567,27 @@ const (
 )
 
 var defKinds = [instrMax]defKind{
-	nop0:         defKindNone,
-	ret:          defKindNone,
-	movRR:        defKindOp2,
-	movRM:        defKindNone,
-	xmmMovRM:     defKindNone,
-	aluRmiR:      defKindNone,
-	shiftR:       defKindNone,
-	imm:          defKindOp2,
-	xmmUnaryRmR:  defKindOp2,
-	mov64MR:      defKindOp2,
-	movzxRmR:     defKindOp2,
-	gprToXmm:     defKindOp2,
-	call:         defKindCall,
-	callIndirect: defKindCall,
-	ud2:          defKindNone,
-	jmp:          defKindNone,
-	jmpIf:        defKindNone,
-	cmpRmiR:      defKindNone,
-	exitSequence: defKindNone,
-	lea:          defKindOp2,
+	nop0:            defKindNone,
+	ret:             defKindNone,
+	movRR:           defKindOp2,
+	movRM:           defKindNone,
+	xmmMovRM:        defKindNone,
+	aluRmiR:         defKindNone,
+	shiftR:          defKindNone,
+	imm:             defKindOp2,
+	xmmUnaryRmR:     defKindOp2,
+	mov64MR:         defKindOp2,
+	movzxRmR:        defKindOp2,
+	gprToXmm:        defKindOp2,
+	call:            defKindCall,
+	callIndirect:    defKindCall,
+	ud2:             defKindNone,
+	jmp:             defKindNone,
+	jmpIf:           defKindNone,
+	cmpRmiR:         defKindNone,
+	exitSequence:    defKindNone,
+	lea:             defKindOp2,
+	v128ConstIsland: defKindNone,
 }
 
 // String implements fmt.Stringer.
@@ -1603,26 +1618,27 @@ const (
 )
 
 var useKinds = [instrMax]useKind{
-	nop0:         useKindNone,
-	ret:          useKindNone,
-	movRR:        useKindOp1,
-	movRM:        useKindOp1RegOp2,
-	xmmMovRM:     useKindOp1RegOp2,
-	aluRmiR:      useKindOp1Op2Reg,
-	shiftR:       useKindOp1Op2Reg,
-	imm:          useKindNone,
-	xmmUnaryRmR:  useKindOp1,
-	mov64MR:      useKindOp1,
-	movzxRmR:     useKindOp1,
-	gprToXmm:     useKindOp1,
-	call:         useKindCall,
-	callIndirect: useKindCallInd,
-	ud2:          useKindNone,
-	jmpIf:        useKindOp1,
-	jmp:          useKindOp1,
-	cmpRmiR:      useKindOp1Op2Reg,
-	exitSequence: useKindOp1,
-	lea:          useKindOp1,
+	nop0:            useKindNone,
+	ret:             useKindNone,
+	movRR:           useKindOp1,
+	movRM:           useKindOp1RegOp2,
+	xmmMovRM:        useKindOp1RegOp2,
+	aluRmiR:         useKindOp1Op2Reg,
+	shiftR:          useKindOp1Op2Reg,
+	imm:             useKindNone,
+	xmmUnaryRmR:     useKindOp1,
+	mov64MR:         useKindOp1,
+	movzxRmR:        useKindOp1,
+	gprToXmm:        useKindOp1,
+	call:            useKindCall,
+	callIndirect:    useKindCallInd,
+	ud2:             useKindNone,
+	jmpIf:           useKindOp1,
+	jmp:             useKindOp1,
+	cmpRmiR:         useKindOp1Op2Reg,
+	exitSequence:    useKindOp1,
+	lea:             useKindOp1,
+	v128ConstIsland: useKindNone,
 }
 
 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 b7f2e2c3..06db4a7c 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
@@ -643,7 +643,7 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 		case 8:
 			encodeRegMem(c, legacyPrefixesNone, 0x89, 1, src, m, rex.setW())
 		default:
-			panic("BUG: invalid size")
+			panic(fmt.Sprintf("BUG: invalid size %d: %s", i.u1, i.String()))
 		}
 
 	case shiftR:
@@ -1000,6 +1000,13 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 			panic("BUG: invalid operand kind")
 		}
 
+	case v128ConstIsland:
+		lo, hi := i.u1, i.u2
+		c.Emit4Bytes(uint32(lo))
+		c.Emit4Bytes(uint32(lo >> 32))
+		c.Emit4Bytes(uint32(hi))
+		c.Emit4Bytes(uint32(hi >> 32))
+
 	default:
 		panic(fmt.Sprintf("TODO: %v", i.kind))
 	}
diff --git a/internal/engine/wazevo/backend/isa/amd64/lower_constant.go b/internal/engine/wazevo/backend/isa/amd64/lower_constant.go
index 9d32392b..c4be8b4c 100644
--- a/internal/engine/wazevo/backend/isa/amd64/lower_constant.go
+++ b/internal/engine/wazevo/backend/isa/amd64/lower_constant.go
@@ -32,7 +32,7 @@ func (m *machine) InsertLoadConstant(instr *ssa.Instruction, vr regalloc.VReg) {
 	case ssa.TypeI32, ssa.TypeI64:
 		m.lowerIconst(vr, v, bits == 64)
 	default:
-		panic("TODO")
+		panic("BUG")
 	}
 }
 
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index 1799672f..0b244863 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -233,11 +233,47 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 		dst := instr.Return()
 		ptr, offset, typ := instr.LoadData()
 		m.lowerLoad(ptr, offset, typ, dst)
+	case ssa.OpcodeVconst:
+		result := instr.Return()
+		lo, hi := instr.VconstData()
+		m.lowerVconst(result, lo, hi)
 	default:
 		panic("TODO: lowering " + op.String())
 	}
 }
 
+func (m *machine) lowerVconst(res ssa.Value, lo, hi uint64) {
+	// TODO: use xor when lo == hi == 0.
+
+	dst := m.c.VRegOf(res)
+
+	islandAddr := m.c.AllocateVReg(ssa.TypeI64)
+	lea := m.allocateInstr()
+	load := m.allocateInstr().asXmmUnaryRmR(sseOpcodeMovdqu, newOperandMem(newAmodeImmReg(0, islandAddr)), dst)
+	jmp := m.allocateInstr()
+
+	constLabelNop, constLabel := m.allocateBrTarget()
+	constIsland := m.allocateInstr().asV128ConstIsland(lo, hi)
+	afterLoadNop, afterLoadLabel := m.allocateBrTarget()
+
+	// 		lea    constLabel(%rip), %islandAddr
+	// 		movdqu (%islandAddr), %dst
+	//		jmp    afterConst
+	// constLabel:
+	//		constIsland $lo, $hi
+	// afterConst:
+
+	m.insert(lea)
+	m.insert(load)
+	m.insert(jmp)
+	m.insert(constLabelNop)
+	m.insert(constIsland)
+	m.insert(afterLoadNop)
+
+	lea.asLEA(newAmodeRipRelative(constLabel), islandAddr)
+	jmp.asJmp(newOperandLabel(afterLoadLabel))
+}
+
 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)
@@ -387,8 +423,19 @@ func (m *machine) lowerStore(si *ssa.Instruction) {
 	base := m.c.VRegOf(ptr)
 
 	store := m.allocateInstr()
-	// TODO: optimization to find whether we could fit newAmodeRegRegShift.
-	store.asMovRM(rm, newOperandMem(newAmodeImmReg(offset, base)), storeSizeInBits/8)
+	mem := newOperandMem(newAmodeImmReg(offset, base))
+	switch value.Type() {
+	case ssa.TypeI32:
+		store.asMovRM(rm, mem, storeSizeInBits/8)
+	case ssa.TypeI64:
+		store.asMovRM(rm, mem, storeSizeInBits/8)
+	case ssa.TypeF32:
+		store.asXmmMovRM(sseOpcodeMovss, rm, mem)
+	case ssa.TypeF64:
+		store.asXmmMovRM(sseOpcodeMovsd, rm, mem)
+	case ssa.TypeV128:
+		store.asXmmMovRM(sseOpcodeMovdqu, rm, mem)
+	}
 	m.insert(store)
 }
 


Test patch:
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 60547552..fe716659 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
@@ -2779,6 +2779,11 @@ func TestInstruction_format_encode(t *testing.T) {
 			want:       "498b6f10498b6718c3",
 			wantFormat: "exit_sequence %r15",
 		},
+		{
+			setup:      func(i *instruction) { i.asV128ConstIsland(0xffffffff_eeeeeeee, 0xaaaaaaaa_bbbbbbbb) },
+			want:       "eeeeeeeeffffffffbbbbbbbbaaaaaaaa",
+			wantFormat: "v128ConstIsland (0xffffffffeeeeeeee, 0xaaaaaaaabbbbbbbb)",
+		},
 	} {
 		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 c8e69a03..72db916d 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -239,10 +239,11 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "globals_set",
-			m:         testcases.GlobalsSet.Module,
-			skipAMD64: true,
-			calls:     []callCase{{expResults: []uint64{1, 2, uint64(math.Float32bits(3.0)), math.Float64bits(4.0)}}},
+			name: "globals_set",
+			m:    testcases.GlobalsSet.Module,
+			calls: []callCase{{expResults: []uint64{
+				1, 2, uint64(math.Float32bits(3.0)), math.Float64bits(4.0), 10, 20,
+			}}},
 		},
 		{
 			name:      "globals_mutable",
@@ -714,7 +715,6 @@ wasm stack trace:
 }
 
 func TestStackUnwind_unreachable(t *testing.T) {
-	skipOnAmd64(t)
 	unreachable := &wasm.Module{
 		TypeSection:     []wasm.FunctionType{{}},
 		ExportSection:   []wasm.Export{{Name: "main", Type: wasm.ExternTypeFunc, Index: 0}},
diff --git a/internal/engine/wazevo/frontend/frontend_test.go b/internal/engine/wazevo/frontend/frontend_test.go
index a5d37012..598241d8 100644
--- a/internal/engine/wazevo/frontend/frontend_test.go
+++ b/internal/engine/wazevo/frontend/frontend_test.go
@@ -1202,7 +1202,9 @@ blk0: (exec_ctx:i64, module_ctx:i64)
 	Store v4, module_ctx, 0x28
 	v5:f64 = F64const 4.000000
 	Store v5, module_ctx, 0x38
-	Jump blk_ret, v2, v3, v4, v5
+	v6:v128 = Vconst 000000000000000a 0000000000000014
+	Store v6, 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 43f203de..6ec759c2 100644
--- a/internal/engine/wazevo/testcases/testcases.go
+++ b/internal/engine/wazevo/testcases/testcases.go
@@ -1330,9 +1330,9 @@ var (
 	}
 
 	GlobalsSet = TestCase{
-		Name: "globals_get",
+		Name: "globals_set",
 		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{
@@ -1352,6 +1352,10 @@ var (
 					Type: wasm.GlobalType{ValType: wasm.ValueTypeF64, Mutable: true},
 					Init: constExprF64(0),
 				},
+				{
+					Type: wasm.GlobalType{ValType: wasm.ValueTypeV128, Mutable: true},
+					Init: constExprV128(0, 0),
+				},
 			},
 			CodeSection: []wasm.Code{
 				{Body: []byte{
@@ -1367,6 +1371,10 @@ var (
 					wasm.OpcodeF64Const, 0, 0, 0, 0, 0, 0, 16, 64, // 4.0
 					wasm.OpcodeGlobalSet, 3,
 					wasm.OpcodeGlobalGet, 3,
+					wasm.OpcodeVecPrefix, wasm.OpcodeVecV128Const,
+					10, 0, 0, 0, 0, 0, 0, 0, 20, 0, 0, 0, 0, 0, 0, 0,
+					wasm.OpcodeGlobalSet, 4,
+					wasm.OpcodeGlobalGet, 4,
 					wasm.OpcodeEnd,
 				}},
 			},