Repository: tetratelabs/wazero
Pull request #1965 (2024-01-26 12:00:44)

wazevo(amd64): lower Fabs, Fneg, Fadd, Fsub, Fdiv, Fmul, Fsqrt

Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com>


mathetake: nvm, will do it myself.

Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/abi_go_call.go b/internal/engine/wazevo/backend/isa/amd64/abi_go_call.go
index 34c51653..a364044b 100644
--- a/internal/engine/wazevo/backend/isa/amd64/abi_go_call.go
+++ b/internal/engine/wazevo/backend/isa/amd64/abi_go_call.go
@@ -98,7 +98,29 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
 		if arg.Kind == backend.ABIArgKindReg {
 			v = arg.Reg
 		} else {
-			panic("TODO: stack arguments")
+			// We have saved callee saved registers, so we can use them.
+			if arg.Type.IsInt() {
+				v = r15VReg
+			} else {
+				v = xmm15VReg
+			}
+			mem := newOperandMem(newAmodeImmReg(uint32(arg.Offset+16 /* to skip caller_rbp and ret_addr */), rbpVReg))
+			load := m.allocateInstr()
+			switch arg.Type {
+			case ssa.TypeI32:
+				load.asMovzxRmR(extModeLQ, mem, v)
+			case ssa.TypeI64:
+				load.asMov64MR(mem, v)
+			case ssa.TypeF32:
+				load.asXmmUnaryRmR(sseOpcodeMovss, mem, v)
+			case ssa.TypeF64:
+				load.asXmmUnaryRmR(sseOpcodeMovsd, mem, v)
+			case ssa.TypeV128:
+				load.asXmmUnaryRmR(sseOpcodeMovdqu, mem, v)
+			default:
+				panic("BUG")
+			}
+			cur = linkInstr(cur, load)
 		}
 
 		store := m.allocateInstr()
@@ -163,45 +185,89 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
 	// Ready to exit the execution.
 	cur = m.storeReturnAddressAndExit(cur, execCtrPtr)
 
-	// After the call, we need to restore the callee saved registers.
-	cur = m.restoreRegistersInExecutionContext(cur, execCtrPtr, calleeSavedVRegs)
-
 	// We don't need the slice size anymore, so pop it.
 	cur = m.addRSP(8, cur)
 
 	// Ready to set up the results.
 	offsetInGoSlice = 0
+	// To avoid overwriting with the execution context pointer by the result, we need to track the offset,
+	// and defer the restoration of the result to the end of this function.
+	var argOverlapWithExecCtxOffset int32 = -1
 	for i := range abi.Rets {
 		r := &abi.Rets[i]
-		if r.Kind == backend.ABIArgKindReg {
-			v := r.Reg
-			load := m.allocateInstr()
-			mem := newOperandMem(newAmodeImmReg(uint32(offsetInGoSlice), rspVReg))
+		var v regalloc.VReg
+		isRegResult := r.Kind == backend.ABIArgKindReg
+		if isRegResult {
+			v = r.Reg
+			if v.RealReg() == execCtrPtr.RealReg() {
+				argOverlapWithExecCtxOffset = offsetInGoSlice
+				offsetInGoSlice += 8 // always uint64 rep.
+				continue
+			}
+		} else {
+			if r.Type.IsInt() {
+				v = r15VReg
+			} else {
+				v = xmm15VReg
+			}
+		}
+
+		load := m.allocateInstr()
+		mem := newOperandMem(newAmodeImmReg(uint32(offsetInGoSlice), rspVReg))
+		switch r.Type {
+		case ssa.TypeI32:
+			load.asMovzxRmR(extModeLQ, mem, v)
+			offsetInGoSlice += 8 // always uint64 rep.
+		case ssa.TypeI64:
+			load.asMov64MR(mem, v)
+			offsetInGoSlice += 8
+		case ssa.TypeF32:
+			load.asXmmUnaryRmR(sseOpcodeMovss, mem, v)
+			offsetInGoSlice += 8 // always uint64 rep.
+		case ssa.TypeF64:
+			load.asXmmUnaryRmR(sseOpcodeMovsd, mem, v)
+			offsetInGoSlice += 8
+		case ssa.TypeV128:
+			load.asXmmUnaryRmR(sseOpcodeMovdqu, mem, v)
+			offsetInGoSlice += 16
+		default:
+			panic("BUG")
+		}
+		cur = linkInstr(cur, load)
+
+		if !isRegResult {
+			// We need to store it back to the result slot above rbp.
+			store := m.allocateInstr()
+			mem := newOperandMem(newAmodeImmReg(uint32(abi.ArgStackSize+r.Offset+16 /* to skip caller_rbp and ret_addr */), rbpVReg))
 			switch r.Type {
 			case ssa.TypeI32:
-				load.asMovzxRmR(extModeLQ, mem, v)
-				offsetInGoSlice += 8 // always uint64 rep.
+				store.asMovRM(v, mem, 4)
 			case ssa.TypeI64:
-				load.asMov64MR(mem, v)
-				offsetInGoSlice += 8
+				store.asMovRM(v, mem, 8)
 			case ssa.TypeF32:
-				load.asXmmUnaryRmR(sseOpcodeMovss, mem, v)
-				offsetInGoSlice += 8 // always uint64 rep.
+				store.asXmmMovRM(sseOpcodeMovss, v, mem)
 			case ssa.TypeF64:
-				load.asXmmUnaryRmR(sseOpcodeMovsd, mem, v)
-				offsetInGoSlice += 8
+				store.asXmmMovRM(sseOpcodeMovsd, v, mem)
 			case ssa.TypeV128:
-				load.asXmmUnaryRmR(sseOpcodeMovdqu, mem, v)
-				offsetInGoSlice += 16
+				store.asXmmMovRM(sseOpcodeMovdqu, v, mem)
 			default:
 				panic("BUG")
 			}
-			cur = linkInstr(cur, load)
-		} else {
-			panic("TODO: stack results")
+			cur = linkInstr(cur, store)
 		}
 	}
 
+	// Before return, we need to restore the callee saved registers.
+	cur = m.restoreRegistersInExecutionContext(cur, execCtrPtr, calleeSavedVRegs)
+
+	if argOverlapWithExecCtxOffset >= 0 {
+		// At this point execCtt is not used anymore, so we can finally store the
+		// result to the register which overlaps with the execution context pointer.
+		mem := newOperandMem(newAmodeImmReg(uint32(argOverlapWithExecCtxOffset), rspVReg))
+		load := m.allocateInstr().asMov64MR(mem, execCtrPtr)
+		cur = linkInstr(cur, load)
+	}
+
 	// Finally ready to return.
 	cur = m.revertRBPRSP(cur)
 	linkInstr(cur, m.allocateInstr().asRet(nil))
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index f2edb869..51129904 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -2,6 +2,7 @@ package amd64
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
@@ -16,6 +17,7 @@ type instruction struct {
 	b1                  bool
 	addedBeforeRegAlloc bool
 	kind                instructionKind
+	targets             []uint32
 }
 
 // Next implements regalloc.Instr.
@@ -197,8 +199,12 @@ func (i *instruction) String() string {
 		return fmt.Sprintf("jmp %s", i.op1.format(true))
 	case jmpIf:
 		return fmt.Sprintf("j%s %s", cond(i.u1), i.op1.format(true))
-	case jmpTableSequence:
-		panic("TODO")
+	case jmpTableIsland:
+		labels := make([]string, len(i.targets))
+		for index, l := range i.targets {
+			labels[index] = backend.Label(l).String()
+		}
+		return fmt.Sprintf("jump_table_island [%s]", strings.Join(labels, ", "))
 	case exitSequence:
 		return fmt.Sprintf("exit_sequence %s", i.op1.format(true))
 	case ud2:
@@ -565,11 +571,8 @@ const (
 	// Jump conditionally: jcond cond label.
 	jmpIf
 
-	// Jump-table sequence, as one compound instruction (see note in lower.rs for rationale).
-	// The generated code sequence is described in the emit's function match arm for this
-	// instruction.
-	// See comment in lowering about the temporaries signedness.
-	jmpTableSequence
+	// jmpTableIsland is to emit the jump table.
+	jmpTableIsland
 
 	// exitSequence exits the execution and go back to the Go world.
 	exitSequence
@@ -663,8 +666,8 @@ func (k instructionKind) String() string {
 		return "jmpIf"
 	case jmp:
 		return "jmp"
-	case jmpTableSequence:
-		return "jmpTableSeq"
+	case jmpTableIsland:
+		return "jmpTableIsland"
 	case exitSequence:
 		return "exit_sequence"
 	case v128ConstIsland:
@@ -713,6 +716,12 @@ func (i *instruction) asJmpIf(cond cond, target operand) *instruction {
 	return i
 }
 
+func (i *instruction) asJmpTableSequence(targets []uint32) *instruction {
+	i.kind = jmpTableIsland
+	i.targets = targets
+	return i
+}
+
 func (i *instruction) asJmp(target operand) *instruction {
 	i.kind = jmp
 	i.op1 = target
@@ -1576,15 +1585,18 @@ var defKinds = [instrMax]defKind{
 	imm:             defKindOp2,
 	unaryRmR:        defKindOp2,
 	xmmUnaryRmR:     defKindOp2,
+	xmmRmR:          defKindNone,
 	mov64MR:         defKindOp2,
 	movsxRmR:        defKindOp2,
 	movzxRmR:        defKindOp2,
 	gprToXmm:        defKindOp2,
+	cmove:           defKindNone,
 	call:            defKindCall,
 	callIndirect:    defKindCall,
 	ud2:             defKindNone,
 	jmp:             defKindNone,
 	jmpIf:           defKindNone,
+	jmpTableIsland:  defKindNone,
 	cmpRmiR:         defKindNone,
 	exitSequence:    defKindNone,
 	lea:             defKindOp2,
@@ -1624,11 +1636,13 @@ var useKinds = [instrMax]useKind{
 	movRR:           useKindOp1,
 	movRM:           useKindOp1RegOp2,
 	xmmMovRM:        useKindOp1RegOp2,
+	cmove:           useKindOp1Op2Reg,
 	aluRmiR:         useKindOp1Op2Reg,
 	shiftR:          useKindOp1Op2Reg,
 	imm:             useKindNone,
 	unaryRmR:        useKindOp1,
 	xmmUnaryRmR:     useKindOp1,
+	xmmRmR:          useKindOp1Op2Reg,
 	mov64MR:         useKindOp1,
 	movzxRmR:        useKindOp1,
 	movsxRmR:        useKindOp1,
@@ -1642,6 +1656,7 @@ var useKinds = [instrMax]useKind{
 	exitSequence:    useKindOp1,
 	lea:             useKindOp1,
 	v128ConstIsland: useKindNone,
+	jmpTableIsland:  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 fdcc5c56..7fa54545 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
@@ -523,6 +523,7 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 
 	case div:
 		panic("TODO")
+
 	case mulHi:
 		var prefix legacyPrefixes
 		rex := rexInfo(0)
@@ -952,8 +953,11 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 			panic("BUG: invalid operand kind")
 		}
 
-	case jmpTableSequence:
-		panic("TODO")
+	case jmpTableIsland:
+		needsLabelResolution = true
+		for range i.targets {
+			c.Emit8Bytes(0)
+		}
 
 	case exitSequence:
 		execCtx := i.op1.r
@@ -1006,10 +1010,8 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 
 	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))
+		c.Emit8Bytes(lo)
+		c.Emit8Bytes(hi)
 
 	default:
 		panic(fmt.Sprintf("TODO: %v", i.kind))
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index 0f6e5b81..133eb8a1 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -53,7 +53,8 @@ type (
 	}
 
 	labelResolutionPend struct {
-		instr *instruction
+		instr       *instruction
+		instrOffset int64
 		// imm32Offset is the offset of the last 4 bytes of the instruction.
 		imm32Offset int64
 	}
@@ -135,7 +136,8 @@ func (m *machine) LowerSingleBranch(b *ssa.Instruction) {
 		}
 		m.insert(jmp)
 	case ssa.OpcodeBrTable:
-		panic("TODO: implement me")
+		index, target := b.BrTableData()
+		m.lowerBrTable(index, target)
 	default:
 		panic("BUG: unexpected branch opcode" + b.Opcode().String())
 	}
@@ -143,6 +145,50 @@ func (m *machine) LowerSingleBranch(b *ssa.Instruction) {
 
 var condBranchMatches = [...]ssa.Opcode{ssa.OpcodeIcmp, ssa.OpcodeFcmp}
 
+func (m *machine) lowerBrTable(index ssa.Value, targets []ssa.BasicBlock) {
+	_v := m.getOperand_Reg(m.c.ValueDefinition(index))
+	v := m.copyToTmp(_v.r)
+
+	// First, we need to do the bounds check.
+	maxIndex := m.c.AllocateVReg(ssa.TypeI32)
+	m.lowerIconst(maxIndex, uint64(len(targets)-1), false)
+	cmp := m.allocateInstr().asCmpRmiR(true, newOperandReg(maxIndex), v, false)
+	m.insert(cmp)
+
+	// Then do the conditional move maxIndex to v if v > maxIndex.
+	cmov := m.allocateInstr().asCmove(condNB, newOperandReg(maxIndex), v, false)
+	m.insert(cmov)
+
+	// Now that v has the correct index. Load the address of the jump table into the addr.
+	addr := m.c.AllocateVReg(ssa.TypeI64)
+	leaJmpTableAddr := m.allocateInstr()
+	m.insert(leaJmpTableAddr)
+
+	// Then add the target's offset into jmpTableAddr.
+	loadTargetOffsetFromJmpTable := m.allocateInstr().asAluRmiR(aluRmiROpcodeAdd,
+		// Shift by 3 because each entry is 8 bytes.
+		newOperandMem(newAmodeRegRegShift(0, addr, v, 3)), addr, true)
+	m.insert(loadTargetOffsetFromJmpTable)
+
+	// Now ready to jump.
+	jmp := m.allocateInstr().asJmp(newOperandReg(addr))
+	m.insert(jmp)
+
+	jmpTableBegin, jmpTableBeginLabel := m.allocateBrTarget()
+	m.insert(jmpTableBegin)
+	leaJmpTableAddr.asLEA(newAmodeRipRelative(jmpTableBeginLabel), addr)
+
+	jmpTable := m.allocateInstr()
+	// TODO: reuse the slice!
+	labels := make([]uint32, len(targets))
+	for j, target := range targets {
+		labels[j] = uint32(m.ectx.GetOrAllocateSSABlockLabel(target))
+	}
+
+	jmpTable.asJmpTableSequence(labels)
+	m.insert(jmpTable)
+}
+
 // LowerConditionalBranch implements backend.Machine.
 func (m *machine) LowerConditionalBranch(b *ssa.Instruction) {
 	exctx := m.ectx
@@ -234,6 +280,14 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 		m.lowerCtz(instr)
 	case ssa.OpcodePopcnt:
 		m.lowerUnaryRmR(instr, unaryRmROpcodePopcnt)
+	case ssa.OpcodeFadd, ssa.OpcodeFsub, ssa.OpcodeFmul, ssa.OpcodeFdiv:
+		m.lowerXmmRmR(instr)
+	case ssa.OpcodeFabs:
+		m.lowerFabsFneg(instr)
+	case ssa.OpcodeFneg:
+		m.lowerFabsFneg(instr)
+	case ssa.OpcodeSqrt:
+		m.lowerSqrt(instr)
 	case ssa.OpcodeUndefined:
 		m.insert(m.allocateInstr().asUD2())
 	case ssa.OpcodeExitWithCode:
@@ -636,6 +690,116 @@ func (m *machine) lowerShiftR(si *ssa.Instruction, op shiftROp) {
 	m.copyTo(tmpDst, rd)
 }
 
+func (m *machine) lowerXmmRmR(instr *ssa.Instruction) {
+	x, y := instr.Arg2()
+	if !x.Type().IsFloat() {
+		panic("BUG?")
+	}
+	_64 := x.Type().Bits() == 64
+
+	var op sseOpcode
+	if _64 {
+		switch instr.Opcode() {
+		case ssa.OpcodeFadd:
+			op = sseOpcodeAddsd
+		case ssa.OpcodeFsub:
+			op = sseOpcodeSubsd
+		case ssa.OpcodeFmul:
+			op = sseOpcodeMulsd
+		case ssa.OpcodeFdiv:
+			op = sseOpcodeDivsd
+		default:
+			panic("BUG")
+		}
+	} else {
+		switch instr.Opcode() {
+		case ssa.OpcodeFadd:
+			op = sseOpcodeAddss
+		case ssa.OpcodeFsub:
+			op = sseOpcodeSubss
+		case ssa.OpcodeFmul:
+			op = sseOpcodeMulss
+		case ssa.OpcodeFdiv:
+			op = sseOpcodeDivss
+		default:
+			panic("BUG")
+		}
+	}
+
+	xDef, yDef := m.c.ValueDefinition(x), m.c.ValueDefinition(y)
+	rn := m.getOperand_Mem_Reg(yDef)
+	rm := m.getOperand_Reg(xDef)
+	rd := m.c.VRegOf(instr.Return())
+
+	// rm is being overwritten, so we first copy its value to a temp register,
+	// in case it is referenced again later.
+	tmp := m.copyToTmp(rm.r)
+
+	xmm := m.allocateInstr().asXmmRmR(op, rn, tmp)
+	m.insert(xmm)
+
+	m.copyTo(tmp, rd)
+}
+
+func (m *machine) lowerSqrt(instr *ssa.Instruction) {
+	x := instr.Arg()
+	if !x.Type().IsFloat() {
+		panic("BUG")
+	}
+	_64 := x.Type().Bits() == 64
+	var op sseOpcode
+	if _64 {
+		op = sseOpcodeSqrtsd
+	} else {
+		op = sseOpcodeSqrtss
+	}
+
+	xDef := m.c.ValueDefinition(x)
+	rm := m.getOperand_Mem_Reg(xDef)
+	rd := m.c.VRegOf(instr.Return())
+
+	xmm := m.allocateInstr().asXmmUnaryRmR(op, rm, rd)
+	m.insert(xmm)
+}
+
+func (m *machine) lowerFabsFneg(instr *ssa.Instruction) {
+	x := instr.Arg()
+	if !x.Type().IsFloat() {
+		panic("BUG")
+	}
+	_64 := x.Type().Bits() == 64
+	var op sseOpcode
+	var mask uint64
+	if _64 {
+		switch instr.Opcode() {
+		case ssa.OpcodeFabs:
+			mask, op = 0x7fffffffffffffff, sseOpcodeAndpd
+		case ssa.OpcodeFneg:
+			mask, op = 0x8000000000000000, sseOpcodeXorpd
+		}
+	} else {
+		switch instr.Opcode() {
+		case ssa.OpcodeFabs:
+			mask, op = 0x7fffffff, sseOpcodeAndps
+		case ssa.OpcodeFneg:
+			mask, op = 0x80000000, sseOpcodeXorps
+		}
+	}
+
+	tmp := m.c.AllocateVReg(x.Type())
+
+	xDef := m.c.ValueDefinition(x)
+	rm := m.getOperand_Reg(xDef)
+	rd := m.c.VRegOf(instr.Return())
+
+	m.lowerFconst(tmp, mask, _64)
+
+	xmm := m.allocateInstr().asXmmRmR(op, rm, tmp)
+	m.insert(xmm)
+
+	m.copyTo(tmp, rd)
+}
+
 func (m *machine) lowerStore(si *ssa.Instruction) {
 	value, ptr, offset, storeSizeInBits := si.StoreData()
 	rm := m.getOperand_Reg(m.c.ValueDefinition(value))
@@ -677,6 +841,9 @@ func (m *machine) lowerCall(si *ssa.Instruction) {
 		m.maxRequiredStackSizeForCalls = stackSlotSize + 16 // 16 == return address + RBP.
 	}
 
+	// Note: See machine.SetupPrologue for the stack layout.
+	// The stack pointer decrease/increase will be inserted later in the compilation.
+
 	for i, arg := range args {
 		reg := m.c.VRegOf(arg)
 		def := m.c.ValueDefinition(arg)
@@ -707,7 +874,7 @@ func (m *machine) lowerCall(si *ssa.Instruction) {
 
 // callerGenVRegToFunctionArg is the opposite of GenFunctionArgToVReg, which is used to generate the
 // caller side of the function call.
-func (m *machine) callerGenVRegToFunctionArg(a *backend.FunctionABI, argIndex int, reg regalloc.VReg, def *backend.SSAValueDefinition, slotBegin int64) {
+func (m *machine) callerGenVRegToFunctionArg(a *backend.FunctionABI, argIndex int, reg regalloc.VReg, def *backend.SSAValueDefinition, stackSlotSize int64) {
 	arg := &a.Args[argIndex]
 	if def != nil && def.IsFromInstr() {
 		// Constant instructions are inlined.
@@ -718,16 +885,52 @@ func (m *machine) callerGenVRegToFunctionArg(a *backend.FunctionABI, argIndex in
 	if arg.Kind == backend.ABIArgKindReg {
 		m.InsertMove(arg.Reg, reg, arg.Type)
 	} else {
-		panic("TODO")
+		store := m.allocateInstr()
+		mem := newOperandMem(newAmodeImmReg(
+			// -stackSlotSize because the stack pointer is not yet decreased.
+			uint32(arg.Offset-stackSlotSize), rspVReg))
+		switch arg.Type {
+		case ssa.TypeI32:
+			store.asMovRM(reg, mem, 4)
+		case ssa.TypeI64:
+			store.asMovRM(reg, mem, 8)
+		case ssa.TypeF32:
+			store.asXmmMovRM(sseOpcodeMovss, reg, mem)
+		case ssa.TypeF64:
+			store.asXmmMovRM(sseOpcodeMovsd, reg, mem)
+		case ssa.TypeV128:
+			store.asXmmMovRM(sseOpcodeMovdqu, reg, mem)
+		default:
+			panic("BUG")
+		}
+		m.insert(store)
 	}
 }
 
-func (m *machine) callerGenFunctionReturnVReg(a *backend.FunctionABI, retIndex int, reg regalloc.VReg, slotBegin int64) {
+func (m *machine) callerGenFunctionReturnVReg(a *backend.FunctionABI, retIndex int, reg regalloc.VReg, stackSlotSize int64) {
 	r := &a.Rets[retIndex]
 	if r.Kind == backend.ABIArgKindReg {
 		m.InsertMove(reg, r.Reg, r.Type)
 	} else {
-		panic("TODO")
+		load := m.allocateInstr()
+		mem := newOperandMem(newAmodeImmReg(
+			// -stackSlotSize because the stack pointer is not yet decreased.
+			uint32(a.ArgStackSize+r.Offset-stackSlotSize), rspVReg))
+		switch r.Type {
+		case ssa.TypeI32:
+			load.asMovzxRmR(extModeLQ, mem, reg)
+		case ssa.TypeI64:
+			load.asMov64MR(mem, reg)
+		case ssa.TypeF32:
+			load.asXmmUnaryRmR(sseOpcodeMovss, mem, reg)
+		case ssa.TypeF64:
+			load.asXmmUnaryRmR(sseOpcodeMovsd, mem, reg)
+		case ssa.TypeV128:
+			load.asXmmUnaryRmR(sseOpcodeMovdqu, mem, reg)
+		default:
+			panic("BUG")
+		}
+		m.insert(load)
 	}
 }
 
@@ -844,12 +1047,13 @@ func (m *machine) Encode(context.Context) {
 			needLabelResolution := cur.encode(m.c)
 			if needLabelResolution {
 				m.labelResolutionPends = append(m.labelResolutionPends,
-					labelResolutionPend{instr: cur, imm32Offset: int64(len(*bufPtr)) - 4},
+					labelResolutionPend{instr: cur, instrOffset: offset, imm32Offset: int64(len(*bufPtr)) - 4},
 				)
 			}
 		}
 	}
 
+	buf := *bufPtr
 	for i := range m.labelResolutionPends {
 		p := &m.labelResolutionPends[i]
 		switch p.instr.kind {
@@ -858,7 +1062,15 @@ func (m *machine) Encode(context.Context) {
 			targetOffset := ectx.LabelPositions[target].BinaryOffset
 			imm32Offset := p.imm32Offset
 			jmpOffset := int32(targetOffset - (p.imm32Offset + 4)) // +4 because RIP points to the next instruction.
-			binary.LittleEndian.PutUint32((*bufPtr)[imm32Offset:], uint32(jmpOffset))
+			binary.LittleEndian.PutUint32(buf[imm32Offset:], uint32(jmpOffset))
+		case jmpTableIsland:
+			tableBegin := p.instrOffset
+			// Each entry is the offset from the beginning of the jmpTableIsland instruction in 8 bytes.
+			for i, l := range p.instr.targets {
+				targetOffset := ectx.LabelPositions[backend.Label(l)].BinaryOffset
+				jmpOffset := targetOffset - tableBegin
+				binary.LittleEndian.PutUint64(buf[tableBegin+int64(i)*8:], uint64(jmpOffset))
+			}
 		default:
 			panic("BUG")
 		}
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue.go b/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue.go
index 5ba414cb..0e4e22d2 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue.go
@@ -125,6 +125,24 @@ func (m *machine) SetupEpilogue() {
 			continue
 		}
 
+		if abi := cur.abi; abi != nil {
+			// At this point, reg alloc is done, therefore we can safely insert dec/inc RPS instruction
+			// right before/after the call instruction. If this is done before reg alloc, the stack slot
+			// can point to the wrong location and therefore results in a wrong value.
+			// TODO: doing this in `SetupEpilogue` seems weird. Find a better home.
+			call := cur
+			next := call.next
+			if size := uint32(abi.AlignedArgResultStackSlotSize()); size > 0 {
+				dec := m.allocateInstr().asAluRmiR(aluRmiROpcodeSub, newOperandImm32(size), rspVReg, true)
+				linkInstr(call.prev, dec)
+				linkInstr(dec, call)
+				inc := m.allocateInstr().asAluRmiR(aluRmiROpcodeAdd, newOperandImm32(size), rspVReg, true)
+				linkInstr(call, inc)
+				linkInstr(inc, next)
+			}
+			continue
+		}
+
 		// Removes the redundant copy instruction.
 		// TODO: doing this in `SetupEpilogue` seems weird. Find a better home.
 		if cur.IsCopy() && cur.op1.r.RealReg() == cur.op2.r.RealReg() {
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc.go b/internal/engine/wazevo/backend/regalloc/regalloc.go
index 9352cb3c..ace02069 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc.go
@@ -610,13 +610,21 @@ func (a *Allocator) allocBlock(f Function, blk Block) {
 		defs := instr.Defs(&a.vs)
 		switch {
 		case len(defs) > 1:
-			if !call {
-				panic("only call can have multiple defs")
-			}
-			// Call's defining register are all caller-saved registers.
-			// Therefore, we can assume that all of them are allocatable.
+			// Some instructions define multiple values on real registers.
+			// E.g. call instructions (following calling convention) / div instruction on x64 that defines both rax and rdx.
+			//
+			// Note that currently I assume that such instructions define only the pre colored real registers, not the VRegs
+			// that require allocations. If we need to support such case, we need to add the logic to handle it here,
+			// though is there any such instruction?
 			for _, def := range defs {
-				s.useRealReg(def.RealReg(), def)
+				if !def.IsRealReg() {
+					panic("BUG: multiple defs should be on real registers")
+				}
+				r := def.RealReg()
+				if s.regsInUse.has(r) {
+					s.releaseRealReg(r)
+				}
+				s.useRealReg(r, def)
 			}
 		case len(defs) == 1:
 			def := defs[0]
diff --git a/internal/engine/wazevo/ssa/type.go b/internal/engine/wazevo/ssa/type.go
index 4e320db3..e8c8cd9d 100644
--- a/internal/engine/wazevo/ssa/type.go
+++ b/internal/engine/wazevo/ssa/type.go
@@ -48,6 +48,11 @@ func (t Type) IsInt() bool {
 	return t == TypeI32 || t == TypeI64
 }
 
+// IsFloat returns true if the type is a floating point type.
+func (t Type) IsFloat() bool {
+	return t == TypeF32 || t == TypeF64
+}
+
 // Bits returns the number of bits required to represent the type.
 func (t Type) Bits() byte {
 	switch t {


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/abi_go_call_test.go b/internal/engine/wazevo/backend/isa/amd64/abi_go_call_test.go
index 5e401696..3aca2b6b 100644
--- a/internal/engine/wazevo/backend/isa/amd64/abi_go_call_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/abi_go_call_test.go
@@ -76,6 +76,10 @@ L2:
 	mov.q %r12, 48(%rax)
 	exit_sequence %rax
 L3:
+	add $8, %rsp
+	movq 8(%rsp), %rcx
+	movss 16(%rsp), %xmm0
+	movsd 24(%rsp), %xmm1
 	movq 96(%rax), %rdx
 	movq 112(%rax), %r12
 	movq 128(%rax), %r13
@@ -89,11 +93,7 @@ L3:
 	movdqu 256(%rax), %xmm13
 	movdqu 272(%rax), %xmm14
 	movdqu 288(%rax), %xmm15
-	add $8, %rsp
-	movzx.lq (%rsp), %rax
-	movq 8(%rsp), %rcx
-	movss 16(%rsp), %xmm0
-	movsd 24(%rsp), %xmm1
+	movq (%rsp), %rax
 	movq %rbp, %rsp
 	popq %rbp
 	ret
@@ -151,6 +151,7 @@ L2:
 	mov.q %r12, 48(%rax)
 	exit_sequence %rax
 L3:
+	add $8, %rsp
 	movq 96(%rax), %rdx
 	movq 112(%rax), %r12
 	movq 128(%rax), %r13
@@ -164,7 +165,6 @@ L3:
 	movdqu 256(%rax), %xmm13
 	movdqu 272(%rax), %xmm14
 	movdqu 288(%rax), %xmm15
-	add $8, %rsp
 	movq %rbp, %rsp
 	popq %rbp
 	ret
@@ -217,6 +217,7 @@ L2:
 	mov.q %r12, 48(%rax)
 	exit_sequence %rax
 L3:
+	add $8, %rsp
 	movq 96(%rax), %rdx
 	movq 112(%rax), %r12
 	movq 128(%rax), %r13
@@ -230,12 +231,140 @@ L3:
 	movdqu 256(%rax), %xmm13
 	movdqu 272(%rax), %xmm14
 	movdqu 288(%rax), %xmm15
+	movq (%rsp), %rax
+	movq %rbp, %rsp
+	popq %rbp
+	ret
+`,
+		},
+		{
+			name:     "many",
+			exitCode: wazevoapi.ExitCodeCallGoFunctionWithIndex(100, false),
+			sig: &ssa.Signature{
+				Params: []ssa.Type{
+					ssa.TypeI64, ssa.TypeI64, ssa.TypeF64,
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+				},
+				Results: []ssa.Type{
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+					ssa.TypeF64, ssa.TypeV128, ssa.TypeI32, ssa.TypeI64, ssa.TypeF32,
+				},
+			},
+			exp: `
+	pushq %rbp
+	movq %rsp, %rbp
+	sub $248, %rsp
+	cmpq 40(%rax), %rsp
+	jnbe L1
+	add $248, %rsp
+	pushq %r15
+	movabsq $248, %r15
+	mov.q %r15, 64(%rax)
+	popq %r15
+	callq *80(%rax)
+	jmp L2
+L1:
+	add $248, %rsp
+L2:
+	mov.q %rdx, 96(%rax)
+	mov.q %r12, 112(%rax)
+	mov.q %r13, 128(%rax)
+	mov.q %r14, 144(%rax)
+	mov.q %r15, 160(%rax)
+	movdqu %xmm8, 176(%rax)
+	movdqu %xmm9, 192(%rax)
+	movdqu %xmm10, 208(%rax)
+	movdqu %xmm11, 224(%rax)
+	movdqu %xmm12, 240(%rax)
+	movdqu %xmm13, 256(%rax)
+	movdqu %xmm14, 272(%rax)
+	movdqu %xmm15, 288(%rax)
+	mov.q %rcx, 1120(%rax)
+	sub $240, %rsp
+	movsd %xmm0, (%rsp)
+	movsd %xmm1, 8(%rsp)
+	movdqu %xmm2, 16(%rsp)
+	mov.l %rbx, 32(%rsp)
+	mov.q %rsi, 40(%rsp)
+	movss %xmm3, 48(%rsp)
+	movsd %xmm4, 56(%rsp)
+	movdqu %xmm5, 64(%rsp)
+	mov.l %rdi, 80(%rsp)
+	mov.q %r8, 88(%rsp)
+	movss %xmm6, 96(%rsp)
+	movsd %xmm7, 104(%rsp)
+	movdqu 16(%rbp), %xmm15
+	movdqu %xmm15, 112(%rsp)
+	mov.l %r9, 128(%rsp)
+	mov.q %r10, 136(%rsp)
+	movss 32(%rbp), %xmm15
+	movss %xmm15, 144(%rsp)
+	pushq $240
+	movl $25606, %r12d
+	mov.l %r12, (%rax)
+	mov.q %rsp, 56(%rax)
+	mov.q %rbp, 1152(%rax)
+	lea L3(%rip), %r12
+	mov.q %r12, 48(%rax)
+	exit_sequence %rax
+L3:
 	add $8, %rsp
-	movzx.lq (%rsp), %rax
+	movsd (%rsp), %xmm0
+	movdqu 8(%rsp), %xmm1
+	movq 32(%rsp), %rcx
+	movss 40(%rsp), %xmm2
+	movsd 48(%rsp), %xmm3
+	movdqu 56(%rsp), %xmm4
+	movzx.lq 72(%rsp), %rbx
+	movq 80(%rsp), %rsi
+	movss 88(%rsp), %xmm5
+	movsd 96(%rsp), %xmm6
+	movdqu 104(%rsp), %xmm7
+	movzx.lq 120(%rsp), %rdi
+	movq 128(%rsp), %r8
+	movss 136(%rsp), %xmm15
+	movss %xmm15, 40(%rbp)
+	movsd 144(%rsp), %xmm15
+	movsd %xmm15, 48(%rbp)
+	movdqu 152(%rsp), %xmm15
+	movdqu %xmm15, 56(%rbp)
+	movzx.lq 168(%rsp), %r9
+	movq 176(%rsp), %r10
+	movss 184(%rsp), %xmm15
+	movss %xmm15, 72(%rbp)
+	movsd 192(%rsp), %xmm15
+	movsd %xmm15, 80(%rbp)
+	movdqu 200(%rsp), %xmm15
+	movdqu %xmm15, 88(%rbp)
+	movzx.lq 216(%rsp), %r11
+	movq 224(%rsp), %r15
+	mov.q %r15, 104(%rbp)
+	movss 232(%rsp), %xmm15
+	movss %xmm15, 112(%rbp)
+	movq 96(%rax), %rdx
+	movq 112(%rax), %r12
+	movq 128(%rax), %r13
+	movq 144(%rax), %r14
+	movq 160(%rax), %r15
+	movdqu 176(%rax), %xmm8
+	movdqu 192(%rax), %xmm9
+	movdqu 208(%rax), %xmm10
+	movdqu 224(%rax), %xmm11
+	movdqu 240(%rax), %xmm12
+	movdqu 256(%rax), %xmm13
+	movdqu 272(%rax), %xmm14
+	movdqu 288(%rax), %xmm15
+	movq 24(%rsp), %rax
 	movq %rbp, %rsp
 	popq %rbp
 	ret
 `,
+			needModuleContextPtr: true,
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine_test.go b/internal/engine/wazevo/backend/isa/amd64/machine_test.go
index 8bbac962..c9e28f75 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine_test.go
@@ -344,7 +344,7 @@ L2:
 	}
 }
 
-func Test_machine_lowerCtz(t *testing.T) {
+func TestMachine_lowerCtz(t *testing.T) {
 	for _, tc := range []struct {
 		name     string
 		setup    func(*mockCompiler, ssa.Builder, *machine) *backend.SSAValueDefinition
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index 0435029e..fd5de3ad 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -151,6 +151,35 @@ func TestE2E(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "float_arithm", m: testcases.FloatArithm.Module,
+			calls: []callCase{
+				{
+					params: []uint64{
+						math.Float64bits(25), math.Float64bits(5), uint64(math.Float32bits(25)), uint64(math.Float32bits(5)),
+					},
+					expResults: []uint64{
+						math.Float64bits(-25),
+						math.Float64bits(25),
+
+						math.Float64bits(5),
+						math.Float64bits(30),
+						math.Float64bits(20),
+						math.Float64bits(125),
+						math.Float64bits(5),
+
+						uint64(math.Float32bits(-25)),
+						uint64(math.Float32bits(25)),
+
+						uint64(math.Float32bits(5)),
+						uint64(math.Float32bits(30)),
+						uint64(math.Float32bits(20)),
+						uint64(math.Float32bits(125)),
+						uint64(math.Float32bits(5)),
+					},
+				},
+			},
+		},
 		{
 			name: "fibonacci_recursive", m: testcases.FibonacciRecursive.Module,
 			calls: []callCase{
@@ -280,9 +309,8 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "br_table",
-			m:         testcases.BrTable.Module,
-			skipAMD64: true,
+			name: "br_table",
+			m:    testcases.BrTable.Module,
 			calls: []callCase{
 				{params: []uint64{0}, expResults: []uint64{11}},
 				{params: []uint64{1}, expResults: []uint64{12}},
@@ -296,9 +324,8 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "br_table_with_args",
-			m:         testcases.BrTableWithArg.Module,
-			skipAMD64: true,
+			name: "br_table_with_args",
+			m:    testcases.BrTableWithArg.Module,
 			calls: []callCase{
 				{params: []uint64{0, 100}, expResults: []uint64{11 + 100}},
 				{params: []uint64{1, 100}, expResults: []uint64{12 + 100}},
@@ -808,7 +835,6 @@ func TestListener_imported(t *testing.T) {
 }
 
 func TestListener_long(t *testing.T) {
-	skipOnAmd64(t)
 	pickOneParam := binaryencoding.EncodeModule(&wasm.Module{
 		TypeSection: []wasm.FunctionType{{Results: []wasm.ValueType{i32}, Params: []wasm.ValueType{
 			i32, i32, f32, f64, i64, i32, i32, v128, f32,
@@ -858,7 +884,6 @@ func TestListener_long(t *testing.T) {
 }
 
 func TestListener_long_as_is(t *testing.T) {
-	skipOnAmd64(t)
 	params := []wasm.ValueType{
 		i32, i64, i32, i64, i32, i64, i32, i64, i32, i64,
 		i32, i64, i32, i64, i32, i64, i32, i64, i32, i64,
@@ -909,7 +934,6 @@ func TestListener_long_as_is(t *testing.T) {
 }
 
 func TestListener_long_many_consts(t *testing.T) {
-	skipOnAmd64(t)
 	const paramNum = 61
 
 	var exp []uint64
diff --git a/internal/engine/wazevo/testcases/testcases.go b/internal/engine/wazevo/testcases/testcases.go
index 6ec759c2..0a8d3040 100644
--- a/internal/engine/wazevo/testcases/testcases.go
+++ b/internal/engine/wazevo/testcases/testcases.go
@@ -81,7 +81,7 @@ var (
 		}, nil),
 	}
 	ArithmReturn = TestCase{
-		Name: "add_sub_params_return",
+		Name: "arithm return",
 		Module: SingleFunctionModule(
 			wasm.FunctionType{
 				Params: []wasm.ValueType{i32, i32, i32, i64, i64, i64},
@@ -989,6 +989,68 @@ var (
 			wasm.OpcodeEnd,
 		}, []wasm.ValueType{}),
 	}
+	FloatArithm = TestCase{
+		Name: "float_arithm",
+		Module: SingleFunctionModule(wasm.FunctionType{
+			Params:  []wasm.ValueType{f64, f64, f32, f32},
+			Results: []wasm.ValueType{f64, f64, f64, f64, f64, f64, f64, f32, f32, f32, f32, f32, f32, f32},
+		}, []byte{
+			wasm.OpcodeLocalGet, 0,
+			wasm.OpcodeF64Neg,
+
+			wasm.OpcodeLocalGet, 0,
+			wasm.OpcodeF64Neg,
+			wasm.OpcodeF64Abs,
+
+			wasm.OpcodeLocalGet, 0,
+			wasm.OpcodeF64Sqrt,
+
+			wasm.OpcodeLocalGet, 0,
+			wasm.OpcodeLocalGet, 1,
+			wasm.OpcodeF64Add,
+
+			wasm.OpcodeLocalGet, 0,
+			wasm.OpcodeLocalGet, 1,
+			wasm.OpcodeF64Sub,
+
+			wasm.OpcodeLocalGet, 0,
+			wasm.OpcodeLocalGet, 1,
+			wasm.OpcodeF64Mul,
+
+			wasm.OpcodeLocalGet, 0,
+			wasm.OpcodeLocalGet, 1,
+			wasm.OpcodeF64Div,
+
+			// 32-bit floats.
+			wasm.OpcodeLocalGet, 2,
+			wasm.OpcodeF32Neg,
+
+			wasm.OpcodeLocalGet, 2,
+			wasm.OpcodeF32Neg,
+			wasm.OpcodeF32Abs,
+
+			wasm.OpcodeLocalGet, 2,
+			wasm.OpcodeF32Sqrt,
+
+			wasm.OpcodeLocalGet, 2,
+			wasm.OpcodeLocalGet, 3,
+			wasm.OpcodeF32Add,
+
+			wasm.OpcodeLocalGet, 2,
+			wasm.OpcodeLocalGet, 3,
+			wasm.OpcodeF32Sub,
+
+			wasm.OpcodeLocalGet, 2,
+			wasm.OpcodeLocalGet, 3,
+			wasm.OpcodeF32Mul,
+
+			wasm.OpcodeLocalGet, 2,
+			wasm.OpcodeLocalGet, 3,
+			wasm.OpcodeF32Div,
+
+			wasm.OpcodeEnd,
+		}, []wasm.ValueType{}),
+	}
 	FloatConversions = TestCase{
 		Name: "float_conversions",
 		Module: SingleFunctionModule(wasm.FunctionType{