Repository: tetratelabs/wazero
Pull request #1956 (2024-01-24 22:32:19)

wazevo(amd64): memory loads 


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index dcf56ad5..bd472dd8 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -1577,6 +1577,7 @@ var defKinds = [instrMax]defKind{
 	imm:             defKindOp2,
 	xmmUnaryRmR:     defKindOp2,
 	mov64MR:         defKindOp2,
+	movsxRmR:        defKindOp2,
 	movzxRmR:        defKindOp2,
 	gprToXmm:        defKindOp2,
 	call:            defKindCall,
@@ -1629,6 +1630,7 @@ var useKinds = [instrMax]useKind{
 	xmmUnaryRmR:     useKindOp1,
 	mov64MR:         useKindOp1,
 	movzxRmR:        useKindOp1,
+	movsxRmR:        useKindOp1,
 	gprToXmm:        useKindOp1,
 	call:            useKindCall,
 	callIndirect:    useKindCallInd,
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index c2a9a943..5bcdf00b 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -229,19 +229,65 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 	case ssa.OpcodeExitWithCode:
 		execCtx, code := instr.ExitWithCodeData()
 		m.lowerExitWithCode(m.c.VRegOf(execCtx), code)
+	case ssa.OpcodeExitIfTrueWithCode:
+		execCtx, c, code := instr.ExitIfTrueWithCodeData()
+		m.lowerExitIfTrueWithCode(m.c.VRegOf(execCtx), c, code)
 	case ssa.OpcodeLoad:
-		dst := instr.Return()
 		ptr, offset, typ := instr.LoadData()
+		dst := m.c.VRegOf(instr.Return())
 		m.lowerLoad(ptr, offset, typ, dst)
+	case ssa.OpcodeUload8, ssa.OpcodeUload16, ssa.OpcodeUload32, ssa.OpcodeSload8, ssa.OpcodeSload16, ssa.OpcodeSload32:
+		ptr, offset, _ := instr.LoadData()
+		ret := m.c.VRegOf(instr.Return())
+		m.lowerExtLoad(op, ptr, offset, ret)
 	case ssa.OpcodeVconst:
 		result := instr.Return()
 		lo, hi := instr.VconstData()
 		m.lowerVconst(result, lo, hi)
+	case ssa.OpcodeSExtend, ssa.OpcodeUExtend:
+		from, to, signed := instr.ExtendData()
+		m.lowerExtend(instr.Arg(), instr.Return(), from, to, signed)
 	default:
 		panic("TODO: lowering " + op.String())
 	}
 }
 
+func (m *machine) lowerExtend(_arg, ret ssa.Value, from, to byte, signed bool) {
+	rd := m.c.VRegOf(ret)
+	arg := m.getOperand_Mem_Reg(m.c.ValueDefinition(_arg))
+
+	ext := m.allocateInstr()
+	switch {
+	case from == 8 && to == 16 && signed:
+		ext.asMovsxRmR(extModeBQ, arg, rd)
+	case from == 8 && to == 16 && !signed:
+		ext.asMovzxRmR(extModeBQ, arg, rd)
+	case from == 8 && to == 32 && signed:
+		ext.asMovsxRmR(extModeBQ, arg, rd)
+	case from == 8 && to == 32 && !signed:
+		ext.asMovzxRmR(extModeBQ, arg, rd)
+	case from == 8 && to == 64 && signed:
+		ext.asMovsxRmR(extModeBQ, arg, rd)
+	case from == 8 && to == 64 && !signed:
+		ext.asMovzxRmR(extModeBQ, arg, rd)
+	case from == 16 && to == 32 && signed:
+		ext.asMovsxRmR(extModeWQ, arg, rd)
+	case from == 16 && to == 32 && !signed:
+		ext.asMovzxRmR(extModeWQ, arg, rd)
+	case from == 16 && to == 64 && signed:
+		ext.asMovsxRmR(extModeWQ, arg, rd)
+	case from == 16 && to == 64 && !signed:
+		ext.asMovzxRmR(extModeWQ, arg, rd)
+	case from == 32 && to == 64 && signed:
+		ext.asMovsxRmR(extModeLQ, arg, rd)
+	case from == 32 && to == 64 && !signed:
+		ext.asMovzxRmR(extModeLQ, arg, rd)
+	default:
+		panic(fmt.Sprintf("BUG: unhandled extend: from=%d, to=%d, signed=%t", from, to, signed))
+	}
+	m.insert(ext)
+}
+
 func (m *machine) lowerVconst(res ssa.Value, lo, hi uint64) {
 	// TODO: use xor when lo == hi == 0.
 
@@ -274,9 +320,8 @@ func (m *machine) lowerVconst(res ssa.Value, lo, hi uint64) {
 	jmp.asJmp(newOperandLabel(afterLoadLabel))
 }
 
-func (m *machine) lowerLoad(ptr ssa.Value, offset uint32, typ ssa.Type, ret ssa.Value) {
+func (m *machine) lowerLoad(ptr ssa.Value, offset uint32, typ ssa.Type, dst regalloc.VReg) {
 	mem := newOperandMem(m.lowerToAddressMode(ptr, offset))
-	dst := m.c.VRegOf(ret)
 	load := m.allocateInstr()
 	switch typ {
 	case ssa.TypeI32:
@@ -289,11 +334,56 @@ func (m *machine) lowerLoad(ptr ssa.Value, offset uint32, typ ssa.Type, ret ssa.
 		load.asXmmUnaryRmR(sseOpcodeMovsd, mem, dst)
 	case ssa.TypeV128:
 		load.asXmmUnaryRmR(sseOpcodeMovdqu, mem, dst)
+	default:
+		panic("BUG")
 	}
 	m.insert(load)
 }
 
-func (m *machine) lowerExitWithCode(execCtx regalloc.VReg, code wazevoapi.ExitCode) {
+func (m *machine) lowerExtLoad(op ssa.Opcode, ptr ssa.Value, offset uint32, dst regalloc.VReg) {
+	mem := newOperandMem(m.lowerToAddressMode(ptr, offset))
+	load := m.allocateInstr()
+	switch op {
+	case ssa.OpcodeUload8:
+		load.asMovzxRmR(extModeBQ, mem, dst)
+	case ssa.OpcodeUload16:
+		load.asMovzxRmR(extModeWQ, mem, dst)
+	case ssa.OpcodeUload32:
+		load.asMovzxRmR(extModeLQ, mem, dst)
+	case ssa.OpcodeSload8:
+		load.asMovsxRmR(extModeBQ, mem, dst)
+	case ssa.OpcodeSload16:
+		load.asMovsxRmR(extModeWQ, mem, dst)
+	case ssa.OpcodeSload32:
+		load.asMovsxRmR(extModeLQ, mem, dst)
+	default:
+		panic("BUG")
+	}
+	m.insert(load)
+}
+
+func (m *machine) lowerExitIfTrueWithCode(execCtx regalloc.VReg, cond ssa.Value, code wazevoapi.ExitCode) {
+	condDef := m.c.ValueDefinition(cond)
+	if !m.c.MatchInstr(condDef, ssa.OpcodeIcmp) {
+		panic("TODO: ExitIfTrue must come after Icmp at the moment: " + condDef.Instr.Opcode().String())
+	}
+	cvalInstr := condDef.Instr
+	cvalInstr.MarkLowered()
+
+	// We need to copy the execution context to a temp register, because if it's spilled,
+	// it might end up being reloaded inside the exiting branch.
+	execCtxTmp := m.copyToTmp(execCtx)
+
+	x, y, c := cvalInstr.IcmpData()
+	m.lowerIcmpToFlag(m.c.ValueDefinition(x), m.c.ValueDefinition(y), x.Type() == ssa.TypeI64)
+
+	jmpIf := m.allocateInstr()
+	m.insert(jmpIf)
+	l := m.lowerExitWithCode(execCtxTmp, code)
+	jmpIf.asJmpIf(condFromSSAIntCmpCond(c).invert(), newOperandLabel(l))
+}
+
+func (m *machine) lowerExitWithCode(execCtx regalloc.VReg, code wazevoapi.ExitCode) (afterLabel backend.Label) {
 	// First we set the exit code in the execution context.
 	exitCodeReg := m.c.AllocateVReg(ssa.TypeI32)
 	m.lowerIconst(exitCodeReg, uint64(code), false)
@@ -338,6 +428,7 @@ func (m *machine) lowerExitWithCode(execCtx regalloc.VReg, code wazevoapi.ExitCo
 	nop, l := m.allocateBrTarget()
 	readRip.asLEA(newAmodeRipRelative(l), ripReg)
 	m.insert(nop)
+	return l
 }
 
 func (m *machine) lowerAluRmiROp(si *ssa.Instruction, op aluRmiROpcode) {
@@ -354,22 +445,17 @@ func (m *machine) lowerAluRmiROp(si *ssa.Instruction, op aluRmiROpcode) {
 	rn := m.getOperand_Reg(xDef)
 	rm := m.getOperand_Mem_Imm32_Reg(yDef)
 	rd := m.c.VRegOf(si.Return())
-	tmp := m.c.AllocateVReg(si.Return().Type())
 
 	// rn is being overwritten, so we first copy its value to a temp register,
 	// in case it is referenced again later.
-	mov := m.allocateInstr()
-	mov.asMovRR(rn.r, tmp, _64)
-	m.insert(mov)
+	tmp := m.copyToTmp(rn.r)
 
 	alu := m.allocateInstr()
 	alu.asAluRmiR(op, rm, tmp, _64)
 	m.insert(alu)
 
 	// tmp now contains the result, we copy it to the dest register.
-	mov2 := m.allocateInstr()
-	mov2.asMovRR(tmp, rd, _64)
-	m.insert(mov2)
+	m.copyTo(tmp, rd)
 }
 
 func (m *machine) lowerShiftR(si *ssa.Instruction, op shiftROp) {
@@ -384,20 +470,15 @@ func (m *machine) lowerShiftR(si *ssa.Instruction, op shiftROp) {
 	opAmt := m.getOperand_Imm32_Reg(amtDef)
 	rx := m.getOperand_Reg(xDef)
 	rd := m.c.VRegOf(si.Return())
-	tmpDst := m.c.AllocateVReg(si.Return().Type())
 
 	// rx is being overwritten, so we first copy its value to a temp register,
 	// in case it is referenced again later.
-	mov := m.allocateInstr()
-	mov.asMovRR(rx.r, tmpDst, _64)
-	m.insert(mov)
+	tmpDst := m.copyToTmp(rx.r)
 
 	if opAmt.r != regalloc.VRegInvalid {
 		// If opAmt is a register we must copy its value to rcx,
 		// because shiftR encoding mandates that the shift amount is in rcx.
-		mov := m.allocateInstr()
-		mov.asMovRR(opAmt.r, rcxVReg, _64)
-		m.insert(mov)
+		m.copyTo(opAmt.r, rcxVReg)
 
 		alu := m.allocateInstr()
 		alu.asShiftR(op, newOperandReg(rcxVReg), tmpDst, _64)
@@ -412,9 +493,7 @@ func (m *machine) lowerShiftR(si *ssa.Instruction, op shiftROp) {
 	}
 
 	// tmp now contains the result, we copy it to the dest register.
-	mov2 := m.allocateInstr()
-	mov2.asMovRR(tmpDst, rd, _64)
-	m.insert(mov2)
+	m.copyTo(tmpDst, rd)
 }
 
 func (m *machine) lowerStore(si *ssa.Instruction) {
@@ -675,3 +754,21 @@ func (m *machine) getVRegSpillSlotOffsetFromSP(id regalloc.VRegID, size byte) in
 	}
 	return offset
 }
+
+func (m *machine) copyTo(src regalloc.VReg, dst regalloc.VReg) {
+	typ := m.c.TypeOf(src)
+	mov := m.allocateInstr()
+	if typ.IsInt() {
+		mov.asMovRR(src, dst, true)
+	} else {
+		mov.asXmmUnaryRmR(sseOpcodeMovdqu, newOperandReg(src), dst)
+	}
+	m.insert(mov)
+}
+
+func (m *machine) copyToTmp(v regalloc.VReg) regalloc.VReg {
+	typ := m.c.TypeOf(v)
+	tmp := m.c.AllocateVReg(typ)
+	m.copyTo(v, tmp)
+	return tmp
+}
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
index 706884e8..dda56171 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
@@ -1884,6 +1884,8 @@ func (m *machine) lowerExitIfTrueWithCode(execCtxVReg regalloc.VReg, cond ssa.Va
 	signed := c.Signed()
 	m.lowerIcmpToFlag(x, y, signed)
 
+	// We need to copy the execution context to a temp register, because if it's spilled,
+	// it might end up being reloaded inside the exiting branch.
 	execCtxTmp := m.copyToTmp(execCtxVReg)
 
 	// We have to skip the entire exit sequence if the condition is false.


Test patch:
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index c5a43df1..d3859d4a 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -194,9 +194,8 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "memory_load_basic",
-			m:         testcases.MemoryLoadBasic.Module,
-			skipAMD64: true,
+			name: "memory_load_basic",
+			m:    testcases.MemoryLoadBasic.Module,
 			calls: []callCase{
 				{params: []uint64{0}, expResults: []uint64{0x03_02_01_00}},
 				{params: []uint64{256}, expResults: []uint64{0x03_02_01_00}},
@@ -206,9 +205,8 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "memory out of bounds",
-			m:         testcases.MemoryLoadBasic.Module,
-			skipAMD64: true,
+			name: "memory out of bounds",
+			m:    testcases.MemoryLoadBasic.Module,
 			calls: []callCase{
 				{params: []uint64{uint64(wasm.MemoryPageSize)}, expErr: "out of bounds memory access"},
 				// We load I32, so we can't load from the last 3 bytes.
@@ -216,11 +214,9 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "memory_loads",
-			m:         testcases.MemoryLoads.Module,
-			skipAMD64: true,
+			name: "memory_loads",
+			m:    testcases.MemoryLoads.Module,
 			calls: []callCase{
-				// These expected results are derived by commenting out `configureWazevo(config)` below to run the old compiler, assuming that it is correct.
 				{params: []uint64{0}, expResults: []uint64{0x3020100, 0x706050403020100, 0x3020100, 0x706050403020100, 0x1211100f, 0x161514131211100f, 0x1211100f, 0x161514131211100f, 0x0, 0xf, 0x0, 0xf, 0x100, 0x100f, 0x100, 0x100f, 0x0, 0xf, 0x0, 0xf, 0x100, 0x100f, 0x100, 0x100f, 0x3020100, 0x1211100f, 0x3020100, 0x1211100f}},
 				{params: []uint64{1}, expResults: []uint64{0x4030201, 0x807060504030201, 0x4030201, 0x807060504030201, 0x13121110, 0x1716151413121110, 0x13121110, 0x1716151413121110, 0x1, 0x10, 0x1, 0x10, 0x201, 0x1110, 0x201, 0x1110, 0x1, 0x10, 0x1, 0x10, 0x201, 0x1110, 0x201, 0x1110, 0x4030201, 0x13121110, 0x4030201, 0x13121110}},
 				{params: []uint64{8}, expResults: []uint64{0xb0a0908, 0xf0e0d0c0b0a0908, 0xb0a0908, 0xf0e0d0c0b0a0908, 0x1a191817, 0x1e1d1c1b1a191817, 0x1a191817, 0x1e1d1c1b1a191817, 0x8, 0x17, 0x8, 0x17, 0x908, 0x1817, 0x908, 0x1817, 0x8, 0x17, 0x8, 0x17, 0x908, 0x1817, 0x908, 0x1817, 0xb0a0908, 0x1a191817, 0xb0a0908, 0x1a191817}},