Repository: tetratelabs/wazero
Pull request #213 (2022-02-08 00:53:42)

jit(arm64): support Add, Mul, Sub operations
part of #187 


Patch:
diff --git a/wasm/interpreter/interpreter.go b/wasm/interpreter/interpreter.go
index 184b4085..3f196271 100644
--- a/wasm/interpreter/interpreter.go
+++ b/wasm/interpreter/interpreter.go
@@ -501,7 +501,9 @@ func (it *interpreter) callHostFunc(ctx context.Context, f *interpreterFunction)
 		raw := it.pop()
 		kind := tp.In(i).Kind()
 		switch kind {
-		case reflect.Float64, reflect.Float32:
+		case reflect.Float32:
+			val.SetFloat(float64(math.Float32frombits(uint32(raw))))
+		case reflect.Float64:
 			val.SetFloat(math.Float64frombits(raw))
 		case reflect.Uint32, reflect.Uint64:
 			val.SetUint(raw)
diff --git a/wasm/jit/engine.go b/wasm/jit/engine.go
index c7abbae8..0213cb8f 100644
--- a/wasm/jit/engine.go
+++ b/wasm/jit/engine.go
@@ -471,7 +471,9 @@ func (e *engine) execHostFunction(f *reflect.Value, ctx *wasm.HostFunctionCallCo
 		raw := e.popValue()
 		kind := tp.In(i).Kind()
 		switch kind {
-		case reflect.Float64, reflect.Float32:
+		case reflect.Float32:
+			val.SetFloat(float64(math.Float32frombits(uint32(raw))))
+		case reflect.Float64:
 			val.SetFloat(math.Float64frombits(raw))
 		case reflect.Uint32, reflect.Uint64:
 			val.SetUint(raw)
diff --git a/wasm/jit/jit_arm64.go b/wasm/jit/jit_arm64.go
index 26193f6c..934266bd 100644
--- a/wasm/jit/jit_arm64.go
+++ b/wasm/jit/jit_arm64.go
@@ -88,9 +88,16 @@ func (c *arm64Compiler) addInstruction(inst *obj.Prog) {
 	c.builder.AddInstruction(inst)
 }
 
-func (c *arm64Compiler) applyConstToRegisterInstruction(insturction obj.As, constValue int64, destinationRegister int16) {
+func (c *arm64Compiler) markRegisterUnused(reg int16) {
+	if !isZeroRegister(reg) {
+		c.locationStack.markRegisterUnused(reg)
+	}
+}
+
+// applyConstToRegisterInstruction adds an instruction where source operand is a constant and destination is a register.
+func (c *arm64Compiler) applyConstToRegisterInstruction(instruction obj.As, constValue int64, destinationRegister int16) {
 	applyConst := c.newProg()
-	applyConst.As = insturction
+	applyConst.As = instruction
 	applyConst.From.Type = obj.TYPE_CONST
 	// Note: in raw arm64 assembly, immediates larger than 16-bits
 	// are not supported, but the assembler takes care of this and
@@ -101,32 +108,96 @@ func (c *arm64Compiler) applyConstToRegisterInstruction(insturction obj.As, cons
 	c.addInstruction(applyConst)
 }
 
-func (c *arm64Compiler) addLoadToRegisterInstruction(insturction obj.As, baseRegister int16, offset int64, destinationRegister int16) {
-	load := c.newProg()
-	load.As = insturction
-	load.From.Type = obj.TYPE_MEM
-	load.From.Reg = baseRegister
-	// Note: in raw arm64 assembly, immediates larger than 16-bits
-	// are not supported, but the assembler takes care of this and
-	// emits corresponding (at most) 4-instructions to load such large constants.
-	load.From.Offset = offset
-	load.To.Type = obj.TYPE_REG
-	load.To.Reg = destinationRegister
-	c.addInstruction(load)
+// applyMemoryToRegisterInstruction adds an instruction where source operand is a memory location and destination is a register.
+// baseRegister is the base absolute address in the memory, and offset is the offset from the absolute address in baseRegister.
+func (c *arm64Compiler) applyMemoryToRegisterInstruction(instruction obj.As, baseRegister int16, offset int64, destinationRegister int16) (err error) {
+	if offset > math.MaxInt16 {
+		// This is a bug in JIT copmiler: caller must check the offset at compilation time, and avoid such a large offset
+		// by loading the const to the register beforehand and then using applyRegisterOffsetMemoryToRegisterInstruction instead.
+		err = fmt.Errorf("memory offset must be smaller than or equal %d, but got %d", math.MaxInt16, offset)
+		return
+	}
+	inst := c.newProg()
+	inst.As = instruction
+	inst.From.Type = obj.TYPE_MEM
+	inst.From.Reg = baseRegister
+	inst.From.Offset = offset
+	inst.To.Type = obj.TYPE_REG
+	inst.To.Reg = destinationRegister
+	c.addInstruction(inst)
+	return
 }
 
-func (c *arm64Compiler) addStoreToMemoryInstruction(insturction obj.As, baseRegister int16, offset int64, source int16) {
-	store := c.newProg()
-	store.As = insturction
-	store.To.Type = obj.TYPE_MEM
-	store.To.Reg = baseRegister
-	// Note: in raw arm64 assembly, immediates larger than 16-bits
-	// are not supported, but the assembler takes care of this and
-	// emits corresponding (at most) 4-instructions to load such large constants.
-	store.To.Offset = offset
-	store.From.Type = obj.TYPE_REG
-	store.From.Reg = source
-	c.addInstruction(store)
+// applyRegisterOffsetMemoryToRegisterInstruction adds an instruction where source operand is a memory location and destination is a register.
+// The difference from applyMemoryToRegisterInstruction is that here we specify the offset by a register rather than offset constant.
+func (c *arm64Compiler) applyRegisterOffsetMemoryToRegisterInstruction(instruction obj.As, baseRegister, offsetRegister, destinationRegister int16) (err error) {
+	inst := c.newProg()
+	inst.As = instruction
+	inst.From.Type = obj.TYPE_MEM
+	inst.From.Reg = baseRegister
+	inst.From.Index = offsetRegister
+	inst.From.Scale = 1
+	inst.To.Type = obj.TYPE_REG
+	inst.To.Reg = destinationRegister
+	c.addInstruction(inst)
+	return nil
+}
+
+// applyRegisterToMemoryInstruction adds an instruction where destination operand is a memory location and source is a register.
+// This is the opposite of applyMemoryToRegisterInstruction.
+func (c *arm64Compiler) applyRegisterToMemoryInstruction(instruction obj.As, baseRegister int16, offset int64, source int16) (err error) {
+	if offset > math.MaxInt16 {
+		// This is a bug in JIT copmiler: caller must check the offset at compilation time, and avoid such a large offset
+		// by loading the const to the register beforehand and then using applyRegisterToRegisterOffsetMemoryInstruction instead.
+		err = fmt.Errorf("memory offset must be smaller than or equal %d, but got %d", math.MaxInt16, offset)
+		return
+	}
+	inst := c.newProg()
+	inst.As = instruction
+	inst.To.Type = obj.TYPE_MEM
+	inst.To.Reg = baseRegister
+	inst.To.Offset = offset
+	inst.From.Type = obj.TYPE_REG
+	inst.From.Reg = source
+	c.addInstruction(inst)
+	return
+}
+
+// applyRegisterToRegisterOffsetMemoryInstruction adds an instruction where destination operand is a memory location and source is a register.
+// The difference from applyRegisterToMemoryInstruction is that here we specify the offset by a register rather than offset constant.
+func (c *arm64Compiler) applyRegisterToRegisterOffsetMemoryInstruction(instruction obj.As, baseRegister, offsetRegister, source int16) {
+	inst := c.newProg()
+	inst.As = instruction
+	inst.To.Type = obj.TYPE_MEM
+	inst.To.Reg = baseRegister
+	inst.To.Index = offsetRegister
+	inst.To.Scale = 1
+	inst.From.Type = obj.TYPE_REG
+	inst.From.Reg = source
+	c.addInstruction(inst)
+}
+
+// applyRegisterToRegisterOffsetMemoryInstruction adds an instruction where both destination and source operands are registers.
+func (c *arm64Compiler) applyRegisterToRegisterInstruction(instruction obj.As, from, to int16) {
+	inst := c.newProg()
+	inst.As = instruction
+	inst.To.Type = obj.TYPE_REG
+	inst.To.Reg = to
+	inst.From.Type = obj.TYPE_REG
+	inst.From.Reg = from
+	c.addInstruction(inst)
+}
+
+// applyRegisterToRegisterOffsetMemoryInstruction adds an instruction which takes two source operands on registers and one destination register operand.
+func (c *arm64Compiler) applyTwoRegistersToRegisterInstruction(instruction obj.As, src1, src2, destination int16) {
+	inst := c.newProg()
+	inst.As = instruction
+	inst.To.Type = obj.TYPE_REG
+	inst.To.Reg = destination
+	inst.From.Type = obj.TYPE_REG
+	inst.From.Reg = src1
+	inst.Reg = src2
+	c.addInstruction(inst)
 }
 
 func (c *arm64Compiler) String() (ret string) { return }
@@ -138,51 +209,72 @@ func (c *arm64Compiler) emitPreamble() error {
 	nop.As = obj.ANOP
 	c.addInstruction(nop)
 
+	// TODO: Push function parameter constants.
+
 	// Before excuting function body, we must initialize the stack base pointer register
 	// so that we can manipulate the memory stack properly.
-	c.initializeReservedStackBasePointerRegister()
-	return nil
+	return c.initializeReservedStackBasePointerRegister()
 }
 
-func (c *arm64Compiler) returnFunction() {
+// returnFunction emits instructions to return from the current function frame.
+// If the current frame is the bottom, the code goes back to the Go code with jitCallStatusCodeReturned status.
+// Otherwise, we jump into the caller's return address (TODO).
+func (c *arm64Compiler) returnFunction() error {
 	// TODO: we don't support function calls yet.
 	// For now the following code just simply returns to Go code.
 
 	// Since we return from the function, we need to decrement the callframe stack pointer, and write it back.
 	callFramePointerReg, _ := c.locationStack.takeFreeRegister(generalPurposeRegisterTypeInt)
-	c.addLoadToRegisterInstruction(arm64.AMOVD, reservedRegisterForEngine, engineGlobalContextCallFrameStackPointerOffset, callFramePointerReg)
+	if err := c.applyMemoryToRegisterInstruction(arm64.AMOVD, reservedRegisterForEngine,
+		engineGlobalContextCallFrameStackPointerOffset, callFramePointerReg); err != nil {
+		return err
+	}
 	c.applyConstToRegisterInstruction(arm64.ASUBS, 1, callFramePointerReg)
-	c.addStoreToMemoryInstruction(arm64.AMOVD, reservedRegisterForEngine, engineGlobalContextCallFrameStackPointerOffset, callFramePointerReg)
+	if err := c.applyRegisterToMemoryInstruction(arm64.AMOVD, reservedRegisterForEngine,
+		engineGlobalContextCallFrameStackPointerOffset, callFramePointerReg); err != nil {
+		return err
+	}
 
-	c.exit(jitCallStatusCodeReturned)
+	return c.exit(jitCallStatusCodeReturned)
 }
 
 // exit adds instructions to give the control back to engine.exec with the given status code.
-func (c *arm64Compiler) exit(status jitCallStatusCode) {
-	tmp, _ := c.locationStack.takeFreeRegister(generalPurposeRegisterTypeInt)
-
+func (c *arm64Compiler) exit(status jitCallStatusCode) error {
 	// Write the current stack pointer to the engine.stackPointer.
-	c.applyConstToRegisterInstruction(arm64.AMOVW, int64(c.locationStack.sp), tmp)
-	c.addStoreToMemoryInstruction(arm64.AMOVW, reservedRegisterForEngine, engineValueStackContextStackPointerOffset, tmp)
+	c.applyConstToRegisterInstruction(arm64.AMOVW, int64(c.locationStack.sp), reservedRegisterForTemporary)
+	if err := c.applyRegisterToMemoryInstruction(arm64.AMOVW, reservedRegisterForEngine,
+		engineValueStackContextStackPointerOffset, reservedRegisterForTemporary); err != nil {
+		return err
+	}
 
 	if status != 0 {
-		c.applyConstToRegisterInstruction(arm64.AMOVW, int64(status), tmp)
-		c.addStoreToMemoryInstruction(arm64.AMOVWU, reservedRegisterForEngine, engineExitContextJITCallStatusCodeOffset, tmp)
+		c.applyConstToRegisterInstruction(arm64.AMOVW, int64(status), reservedRegisterForTemporary)
+		if err := c.applyRegisterToMemoryInstruction(arm64.AMOVWU, reservedRegisterForEngine,
+			engineExitContextJITCallStatusCodeOffset, reservedRegisterForTemporary); err != nil {
+			return err
+		}
 	} else {
 		// If the status == 0, we simply use zero register to store zero.
-		c.addStoreToMemoryInstruction(arm64.AMOVWU, reservedRegisterForEngine, engineExitContextJITCallStatusCodeOffset, zeroRegister)
+		if err := c.applyRegisterToMemoryInstruction(arm64.AMOVWU, reservedRegisterForEngine,
+			engineExitContextJITCallStatusCodeOffset, zeroRegister); err != nil {
+			return err
+		}
 	}
 
 	// The return address to the Go code is stored in archContext.jitReturnAddress which
 	// is embedded in engine. We load the value to the tmpRegister, and then
 	// invoke RET with that register.
-	c.addLoadToRegisterInstruction(arm64.AMOVD, reservedRegisterForEngine, engineArchContextJITCallReturnAddressOffset, tmp)
+	if err := c.applyMemoryToRegisterInstruction(arm64.AMOVD, reservedRegisterForEngine,
+		engineArchContextJITCallReturnAddressOffset, reservedRegisterForTemporary); err != nil {
+		return err
+	}
 
 	ret := c.newProg()
 	ret.As = obj.ARET
 	ret.To.Type = obj.TYPE_REG
-	ret.To.Reg = tmp
+	ret.To.Reg = reservedRegisterForTemporary
 	c.addInstruction(ret)
+	return nil
 }
 
 func (c *arm64Compiler) compileHostFunction(address wasm.FunctionAddress) error {
@@ -247,16 +339,106 @@ func (c *arm64Compiler) compilePick(o *wazeroir.OperationPick) error {
 	return fmt.Errorf("TODO: unsupported on arm64")
 }
 
+// compileAdd implements compiler.compileAdd for the arm64 architecture.
 func (c *arm64Compiler) compileAdd(o *wazeroir.OperationAdd) error {
-	return fmt.Errorf("TODO: unsupported on arm64")
+	x1, x2, err := c.popTwoValuesOnRegisters()
+	if err != nil {
+		return err
+	}
+
+	// Additon can be nop if one of operands is zero.
+	if isZeroRegister(x1.register) {
+		c.locationStack.pushValueOnRegister(x2.register)
+		return nil
+	} else if isZeroRegister(x2.register) {
+		c.locationStack.pushValueOnRegister(x1.register)
+		return nil
+	}
+
+	var inst obj.As
+	switch o.Type {
+	case wazeroir.UnsignedTypeI32:
+		inst = arm64.AADDW
+	case wazeroir.UnsignedTypeI64:
+		inst = arm64.AADD
+	case wazeroir.UnsignedTypeF32:
+		inst = arm64.AFADDS
+	case wazeroir.UnsignedTypeF64:
+		inst = arm64.AFADDD
+	}
+
+	c.applyRegisterToRegisterInstruction(inst, x2.register, x1.register)
+	// The result is placed on a register for x1, so record it.
+	c.locationStack.pushValueOnRegister(x1.register)
+	return nil
 }
 
+// compileSub implements compiler.compileSub for the arm64 architecture.
 func (c *arm64Compiler) compileSub(o *wazeroir.OperationSub) error {
-	return fmt.Errorf("TODO: unsupported on arm64")
+	x1, x2, err := c.popTwoValuesOnRegisters()
+	if err != nil {
+		return err
+	}
+
+	// If both of registers are zeros, this can be nop and simply push the zero register.
+	if isZeroRegister(x1.register) && isZeroRegister(x2.register) {
+		c.locationStack.pushValueOnRegister(zeroRegister)
+		return nil
+	}
+
+	// At this point, at least one of x1 or x2 registers is non zero.
+	// Choose the non-zero register as destination.
+	var destinationReg int16 = x1.register
+	if isZeroRegister(x1.register) {
+		destinationReg = x2.register
+	}
+
+	var inst obj.As
+	switch o.Type {
+	case wazeroir.UnsignedTypeI32:
+		inst = arm64.ASUBW
+	case wazeroir.UnsignedTypeI64:
+		inst = arm64.ASUB
+	case wazeroir.UnsignedTypeF32:
+		inst = arm64.AFSUBS
+	case wazeroir.UnsignedTypeF64:
+		inst = arm64.AFSUBD
+	}
+
+	c.applyTwoRegistersToRegisterInstruction(inst, x2.register, x1.register, destinationReg)
+	c.locationStack.pushValueOnRegister(destinationReg)
+	return nil
 }
 
+// compileMul implements compiler.compileMul for the arm64 architecture.
 func (c *arm64Compiler) compileMul(o *wazeroir.OperationMul) error {
-	return fmt.Errorf("TODO: unsupported on arm64")
+	x1, x2, err := c.popTwoValuesOnRegisters()
+	if err != nil {
+		return err
+	}
+
+	// Multiplcation can be done by putting a zero register if one of operands is zero.
+	if isZeroRegister(x1.register) || isZeroRegister(x2.register) {
+		c.locationStack.pushValueOnRegister(zeroRegister)
+		return nil
+	}
+
+	var inst obj.As
+	switch o.Type {
+	case wazeroir.UnsignedTypeI32:
+		inst = arm64.AMULW
+	case wazeroir.UnsignedTypeI64:
+		inst = arm64.AMUL
+	case wazeroir.UnsignedTypeF32:
+		inst = arm64.AFMULS
+	case wazeroir.UnsignedTypeF64:
+		inst = arm64.AFMULD
+	}
+
+	c.applyRegisterToRegisterInstruction(inst, x2.register, x1.register)
+	// The result is placed on a register for x1, so record it.
+	c.locationStack.pushValueOnRegister(x1.register)
+	return nil
 }
 
 func (c *arm64Compiler) compileClz(o *wazeroir.OperationClz) error {
@@ -486,8 +668,7 @@ func (c *arm64Compiler) emitIntConstant(is32bit bool, value uint64) error {
 		}
 		c.applyConstToRegisterInstruction(inst, int64(value), reg)
 
-		loc := c.locationStack.pushValueOnRegister(reg)
-		loc.setRegisterType(generalPurposeRegisterTypeInt)
+		c.locationStack.pushValueOnRegister(reg)
 	}
 	return nil
 }
@@ -514,14 +695,7 @@ func (c *arm64Compiler) emitFloatConstant(is32bit bool, value uint64) error {
 
 	tmpReg := zeroRegister
 	if value != 0 {
-		// If the target value is not zero, we have to load the constant
-		// temporarily into the integer register since we cannot directly
-		// move the const into float register.
-		tmpReg, err = c.allocateRegister(generalPurposeRegisterTypeInt)
-		if err != nil {
-			return err
-		}
-
+		tmpReg = reservedRegisterForTemporary
 		var inst obj.As
 		if is32bit {
 			inst = arm64.AMOVW
@@ -532,20 +706,15 @@ func (c *arm64Compiler) emitFloatConstant(is32bit bool, value uint64) error {
 	}
 
 	// Use FMOV instruction to move the value on integer register into the float one.
-	mov := c.newProg()
+	var inst obj.As
 	if is32bit {
-		mov.As = arm64.AFMOVS
+		inst = arm64.AFMOVS
 	} else {
-		mov.As = arm64.AFMOVD
+		inst = arm64.AFMOVD
 	}
-	mov.From.Type = obj.TYPE_REG
-	mov.From.Reg = tmpReg
-	mov.To.Type = obj.TYPE_REG
-	mov.To.Reg = reg
-	c.addInstruction(mov)
+	c.applyRegisterToRegisterInstruction(inst, tmpReg, reg)
 
-	loc := c.locationStack.pushValueOnRegister(reg)
-	loc.setRegisterType(generalPurposeRegisterTypeFloat)
+	c.locationStack.pushValueOnRegister(reg)
 	return nil
 }
 
@@ -553,6 +722,67 @@ func (c *arm64Compiler) pushZeroValue() {
 	c.locationStack.pushValueOnRegister(zeroRegister)
 }
 
+// popTwoValuesOnRegisters pops two values from the location stacks, ensures
+// these two values are located on registers, and mark them unused.
+func (c *arm64Compiler) popTwoValuesOnRegisters() (x1, x2 *valueLocation, err error) {
+	x2 = c.locationStack.pop()
+	if err = c.ensureOnGeneralPurposeRegister(x2); err != nil {
+		return
+	}
+
+	x1 = c.locationStack.pop()
+	if err = c.ensureOnGeneralPurposeRegister(x1); err != nil {
+		return
+	}
+
+	c.markRegisterUnused(x1.register)
+	c.markRegisterUnused(x2.register)
+	return
+}
+
+// ensureOnGeneralPurposeRegister emits instructions to ensure that a value is located on a register.
+func (c *arm64Compiler) ensureOnGeneralPurposeRegister(loc *valueLocation) (err error) {
+	if loc.onStack() {
+		err = c.loadValueOnStackToRegister(loc)
+	} else if loc.onConditionalRegister() {
+		err = fmt.Errorf("TODO: support moving conditional to general purpose register")
+	}
+	return
+}
+
+// loadValueOnStackToRegister emits instructions to load the value located on the stack to a register.
+func (c *arm64Compiler) loadValueOnStackToRegister(loc *valueLocation) (err error) {
+	var inst obj.As
+	var reg int16
+	switch loc.regType {
+	case generalPurposeRegisterTypeInt:
+		inst = arm64.AMOVD
+		reg, err = c.allocateRegister(generalPurposeRegisterTypeInt)
+	case generalPurposeRegisterTypeFloat:
+		inst = arm64.AFMOVD
+		reg, err = c.allocateRegister(generalPurposeRegisterTypeFloat)
+	}
+
+	if err != nil {
+		return
+	}
+
+	if offset := int64(loc.stackPointer) * 8; offset > math.MaxInt16 {
+		// The assembler can take care of offsets larger than 2^15-1 by emitting additional instructions to load such large offset,
+		// but it uses "its" temporary register which we cannot track. Therefore, we avoid directly emitting memory load with large offsets,
+		// but instead load the constant manually to "our" temporary register, then emit the load with it.
+		c.applyConstToRegisterInstruction(arm64.AMOVD, offset, reservedRegisterForTemporary)
+		c.applyRegisterOffsetMemoryToRegisterInstruction(inst, reservedRegisterForStackBasePointerAddress, reservedRegisterForTemporary, reg)
+	} else {
+		c.applyMemoryToRegisterInstruction(inst, reservedRegisterForStackBasePointerAddress, offset, reg)
+	}
+
+	// Record that the value holds the register and the register is marked used.
+	loc.setRegister(reg)
+	c.locationStack.markRegisterUsed(reg)
+	return
+}
+
 // allocateRegister returns an unused register of the given type. The register will be taken
 // either from the free register pool or by spilling an used register. If we allocate an used register,
 // this adds an instruction to write the value on a register back to memory stack region.
@@ -574,40 +804,50 @@ func (c *arm64Compiler) allocateRegister(t generalPurposeRegisterType) (reg int1
 
 	// Release the steal target register value onto stack location.
 	reg = stealTarget.register
-	c.releaseRegisterToStack(stealTarget)
+	err = c.releaseRegisterToStack(stealTarget)
 	return
 }
 
 // releaseRegisterToStack adds an instruction to write the value on a register back to memory stack region.
-func (c *arm64Compiler) releaseRegisterToStack(loc *valueLocation) {
-	var inst obj.As
-	switch loc.regType {
-	case generalPurposeRegisterTypeInt:
-		inst = arm64.AMOVD
-	case generalPurposeRegisterTypeFloat:
+func (c *arm64Compiler) releaseRegisterToStack(loc *valueLocation) (err error) {
+	var inst obj.As = arm64.AMOVD
+	if loc.regType == generalPurposeRegisterTypeFloat {
 		inst = arm64.AFMOVD
 	}
 
-	c.addStoreToMemoryInstruction(inst, reservedRegisterForStackBasePointerAddress, int64(loc.stackPointer)*8, loc.register)
+	if offset := int64(loc.stackPointer) * 8; offset > math.MaxInt16 {
+		// The assembler can take care of offsets larger than 2^15-1 by emitting additional instructions to load such large offset,
+		// but it uses "its" temporary register which we cannot track. Therefore, we avoid directly emitting memory load with large offsets,
+		// but instead load the constant manually to "our" temporary register, then emit the load with it.
+		c.applyConstToRegisterInstruction(arm64.AMOVD, offset, reservedRegisterForTemporary)
+		c.applyRegisterToRegisterOffsetMemoryInstruction(inst, reservedRegisterForStackBasePointerAddress, reservedRegisterForTemporary, loc.register)
+	} else {
+		if err = c.applyRegisterToMemoryInstruction(inst, reservedRegisterForStackBasePointerAddress, offset, loc.register); err != nil {
+			return
+		}
+	}
 
 	// Mark the register is free.
 	c.locationStack.releaseRegister(loc)
+	return
 }
 
 // initializeReservedStackBasePointerRegister adds intructions to initialize reservedRegisterForStackBasePointerAddress
 // so that it points to the absolute address of the stack base for this function.
-func (c *arm64Compiler) initializeReservedStackBasePointerRegister() {
-	tmpReg, _ := c.locationStack.takeFreeRegister(generalPurposeRegisterTypeInt)
-
+func (c *arm64Compiler) initializeReservedStackBasePointerRegister() error {
 	// First, load the address of the first element in the value stack into reservedRegisterForStackBasePointerAddress temporarily.
-	c.addLoadToRegisterInstruction(arm64.AMOVD,
+	if err := c.applyMemoryToRegisterInstruction(arm64.AMOVD,
 		reservedRegisterForEngine, engineGlobalContextValueStackElement0AddressOffset,
-		reservedRegisterForStackBasePointerAddress)
+		reservedRegisterForStackBasePointerAddress); err != nil {
+		return err
+	}
 
 	// Next we move the base pointer (engine.stackBasePointer) to the tmp register.
-	c.addLoadToRegisterInstruction(arm64.AMOVD,
+	if err := c.applyMemoryToRegisterInstruction(arm64.AMOVD,
 		reservedRegisterForEngine, engineValueStackContextStackBasePointerOffset,
-		tmpReg)
+		reservedRegisterForTemporary); err != nil {
+		return err
+	}
 
 	// Finally, we calculate "reservedRegisterForStackBasePointerAddress + tmpReg * 8"
 	// where we multiply tmpReg by 8 because stack pointer is an index in the []uint64
@@ -617,8 +857,9 @@ func (c *arm64Compiler) initializeReservedStackBasePointerRegister() {
 	calcStackBasePointerAddress.To.Type = obj.TYPE_REG
 	calcStackBasePointerAddress.To.Reg = reservedRegisterForStackBasePointerAddress
 	// We calculate "tmpReg * 8" as "tmpReg << 3".
-	setLeftShiftedRegister(calcStackBasePointerAddress, tmpReg, 3)
+	setLeftShiftedRegister(calcStackBasePointerAddress, reservedRegisterForTemporary, 3)
 	c.addInstruction(calcStackBasePointerAddress)
+	return nil
 }
 
 // setShiftedRegister modifies the given *obj.Prog so that .From (source operand)
diff --git a/wasm/jit/jit_value_location.go b/wasm/jit/jit_value_location.go
index ef897163..4be89301 100644
--- a/wasm/jit/jit_value_location.go
+++ b/wasm/jit/jit_value_location.go
@@ -137,6 +137,12 @@ func (s *valueLocationStack) pushValueOnRegister(reg int16) (loc *valueLocation)
 			panic("bug in compiler: try pushing a register which is already in use")
 		}
 	}
+
+	if isIntRegister(reg) {
+		loc.setRegisterType(generalPurposeRegisterTypeInt)
+	} else if isFloatRegister(reg) {
+		loc.setRegisterType(generalPurposeRegisterTypeFloat)
+	}
 	s.markRegisterUsed(reg)
 	s.push(loc)
 	return
diff --git a/wasm/jit/jit_value_location_amd64.go b/wasm/jit/jit_value_location_amd64.go
index 33b08b75..fd064545 100644
--- a/wasm/jit/jit_value_location_amd64.go
+++ b/wasm/jit/jit_value_location_amd64.go
@@ -14,7 +14,8 @@ const (
 	// reservedRegisterForStackBasePointerAddress: stack base pointer's address (engine.stackBasePointer) in the current function call.
 	reservedRegisterForStackBasePointerAddress = x86.REG_R14
 	// reservedRegisterForMemory: pointer to the memory slice's data (i.e. &memory.Buffer[0] as uintptr).
-	reservedRegisterForMemory = x86.REG_R15
+	reservedRegisterForMemory    = x86.REG_R15
+	reservedRegisterForTemporary = nilRegister
 )
 
 var (
diff --git a/wasm/jit/jit_value_location_arm64.go b/wasm/jit/jit_value_location_arm64.go
index 4f983a91..4f42445f 100644
--- a/wasm/jit/jit_value_location_arm64.go
+++ b/wasm/jit/jit_value_location_arm64.go
@@ -7,12 +7,13 @@ import "github.com/twitchyliquid64/golang-asm/obj/arm64"
 
 // Reserved registers.
 const (
-	// reservedRegisterForEngine: pointer to engine instance (i.e. *engine as uintptr)
+	// reservedRegisterForEngine holds the pointer to engine instance (i.e. *engine as uintptr)
 	reservedRegisterForEngine = arm64.REG_R0
-	// reservedRegisterForStackBasePointerAddress: stack base pointer's address (engine.stackBasePointer) in the current function call.
+	// reservedRegisterForStackBasePointerAddress holds stack base pointer's address (engine.stackBasePointer) in the current function call.
 	reservedRegisterForStackBasePointerAddress = arm64.REG_R1
-	// reservedRegisterForMemory: pointer to the memory slice's data (i.e. &memory.Buffer[0] as uintptr).
-	reservedRegisterForMemory = arm64.REG_R2
+	// reservedRegisterForMemory holds the pointer to the memory slice's data (i.e. &memory.Buffer[0] as uintptr).
+	reservedRegisterForMemory    = arm64.REG_R2
+	reservedRegisterForTemporary = arm64.REG_R3
 )
 
 // zeroRegister is the alias of the arm64-specific zero register for readability.
@@ -29,12 +30,12 @@ var (
 		arm64.REG_F29, arm64.REG_F30, arm64.REG_F31,
 	}
 	unreservedGeneralPurposeIntRegisters = []int16{
-		arm64.REG_R3, arm64.REG_R4, arm64.REG_R5, arm64.REG_R6, arm64.REG_R7, arm64.REG_R8,
+		arm64.REG_R4, arm64.REG_R5, arm64.REG_R6, arm64.REG_R7, arm64.REG_R8,
 		arm64.REG_R9, arm64.REG_R10, arm64.REG_R11, arm64.REG_R12, arm64.REG_R13,
 		arm64.REG_R14, arm64.REG_R15, arm64.REG_R16, arm64.REG_R17, arm64.REG_R18,
 		arm64.REG_R19, arm64.REG_R20, arm64.REG_R21, arm64.REG_R22, arm64.REG_R23,
 		arm64.REG_R24, arm64.REG_R25, arm64.REG_R26, arm64.REG_R27, arm64.REG_R28,
-		arm64.REG_R29, arm64.REG_R30, arm64.REG_R31,
+		arm64.REG_R29, arm64.REG_R30,
 	}
 )
 
@@ -45,3 +46,7 @@ func isIntRegister(r int16) bool {
 func isFloatRegister(r int16) bool {
 	return generalPurposeFloatRegisters[0] <= r && r <= generalPurposeFloatRegisters[len(generalPurposeFloatRegisters)-1]
 }
+
+func isZeroRegister(r int16) bool {
+	return r == zeroRegister
+}
diff --git a/wasm/module.go b/wasm/module.go
index d08b42d6..6009dfe3 100644
--- a/wasm/module.go
+++ b/wasm/module.go
@@ -226,7 +226,10 @@ const (
 	ValueTypeF64 ValueType = 0x7c
 )
 
-func valueTypeName(t ValueType) string {
+// ValuTypeName returns the type name of the given ValueType as a string.
+// These type names match the names used in the WebAssembly text format.
+// Note that ValueTypeName returns "unknown", if an undefined ValueType value is passed.
+func ValueTypeName(t ValueType) string {
 	switch t {
 	case ValueTypeI32:
 		return "i32"
@@ -242,14 +245,14 @@ func valueTypeName(t ValueType) string {
 
 func (t *FunctionType) String() (ret string) {
 	for _, b := range t.Params {
-		ret += valueTypeName(b)
+		ret += ValueTypeName(b)
 	}
 	if len(t.Params) == 0 {
 		ret += "null"
 	}
 	ret += "_"
 	for _, b := range t.Results {
-		ret += valueTypeName(b)
+		ret += ValueTypeName(b)
 	}
 	if len(t.Results) == 0 {
 		ret += "null"


Test patch:
diff --git a/tests/engine/adhoc_test.go b/tests/engine/adhoc_test.go
index ed1253d3..d998ae79 100644
--- a/tests/engine/adhoc_test.go
+++ b/tests/engine/adhoc_test.go
@@ -1,12 +1,15 @@
 package adhoc
 
 import (
+	"bytes"
 	"context"
 	_ "embed"
+	"math"
 	"reflect"
 	"runtime"
 	"sync"
 	"testing"
+	"text/template"
 
 	"github.com/stretchr/testify/require"
 
@@ -60,6 +63,9 @@ func runTests(t *testing.T, newEngine func() wasm.Engine) {
 	t.Run("imported-and-exported func", func(t *testing.T) {
 		importedAndExportedFunc(t, newEngine)
 	})
+	t.Run("host function with float32 type", func(t *testing.T) {
+		hostFuncWithFloatParam(t, newEngine)
+	})
 }
 
 func fibonacci(t *testing.T, newEngine func() wasm.Engine) {
@@ -194,9 +200,9 @@ func recursiveEntry(t *testing.T, newEngine func() wasm.Engine) {
 	require.NoError(t, err)
 }
 
+// importedAndExportedFunc fails if the engine cannot call an "imported-and-then-exported-back" function
 func importedAndExportedFunc(t *testing.T, newEngine func() wasm.Engine) {
 	ctx := context.Background()
-	// Test that the engine can call "imported-and-then-exported-back" function correctly
 	mod, err := text.DecodeModule([]byte(`(module
 		;; arbitrary function with params
 		(import "env" "add_int" (func $add_int (param i32 i32) (result i32)))
@@ -221,3 +227,80 @@ func importedAndExportedFunc(t *testing.T, newEngine func() wasm.Engine) {
 	require.NoError(t, err)
 	require.Equal(t, uint64(42), results[0])
 }
+
+//  hostFuncWithFloatParam fails if a float parameter corrupts a host function value
+func hostFuncWithFloatParam(t *testing.T, newEngine func() wasm.Engine) {
+	ctx := context.Background()
+	watTemplate, err := template.New("").Parse(`(module
+		;; '{{ . }}' will be replaced by the type name, f32 or f64.
+
+		;; 'identity_float' just returns the given float value as is.
+		(import "test" "identity_float" (func $test.identity_float (param {{ . }}) (result {{ . }})))
+
+		;; 'call->test.identity_float' proxies 'test.identity_float' via call in order to test floats aren't corrupted 
+		;; when the float values are passed through OpCodeCall
+		(func $call->test.identity_float (param {{ . }}) (result {{ . }})
+			local.get 0
+			call $test.identity_float
+		)
+		(export "call->test.identity_float" (func $call->test.identity_float))
+		)`)
+	require.NoError(t, err)
+
+	tests := []struct {
+		testName          string
+		floatType         byte          // wasm.ValueTypeF32 or wasm.ValueTypeF64
+		identityFloatFunc reflect.Value // imported as 'identity_float' to the wasm module
+		floatParam        uint64        // passed to identityFloatFunc as a parameter
+		expectedFloatVal  float64
+	}{
+		{
+			testName:  "host function with f32 param",
+			floatType: wasm.ValueTypeF32,
+			identityFloatFunc: reflect.ValueOf(func(ctx *wasm.HostFunctionCallContext, value float32) float32 {
+				return value
+			}),
+			floatParam:       uint64(math.Float32bits(math.MaxFloat32)), // float bits as a uint32 value, but casted to uint64 to be passed to CallFunction
+			expectedFloatVal: float64(math.MaxFloat32),                  // arbitrary f32 value
+		},
+		{
+			testName:  "host function with f64 param",
+			floatType: wasm.ValueTypeF64,
+			identityFloatFunc: reflect.ValueOf(func(ctx *wasm.HostFunctionCallContext, value float64) float64 {
+				return value
+			}),
+			floatParam:       math.Float64bits(math.MaxFloat64),
+			expectedFloatVal: float64(math.MaxFloat64), // arbitrary f64 value that doesn' fit in f32
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
+		t.Run(tc.testName, func(t *testing.T) {
+			// Build and instantiate the wat for `f32` or `f64`
+			var wat bytes.Buffer
+			err = watTemplate.Execute(&wat, wasm.ValueTypeName(tc.floatType))
+			require.NoError(t, err)
+
+			mod, err := text.DecodeModule(wat.Bytes())
+			require.NoError(t, err)
+
+			store := wasm.NewStore(newEngine())
+
+			err = store.AddHostFunction("test", "identity_float", tc.identityFloatFunc)
+			require.NoError(t, err)
+
+			err = store.Instantiate(mod, "mod")
+			require.NoError(t, err)
+
+			// This call should return the `expectedFloatVal`.
+			// That ensures that the float32 values are not corrupted when they are routed through OpCodeCall.
+			results, resultTypes, err := store.CallFunction(ctx, "mod", "call->test.identity_float", tc.floatParam)
+			require.NoError(t, err)
+			require.Len(t, results, len(resultTypes))
+			require.Equal(t, tc.floatType, resultTypes[0])
+			// Note that the both f32 and f64 result value are expressed as the float64 bits value in wazero.
+			require.Equal(t, tc.expectedFloatVal, math.Float64frombits(results[0]))
+		})
+	}
+}
diff --git a/wasm/jit/jit_arm64_test.go b/wasm/jit/jit_arm64_test.go
index c79e3037..1e90b2c5 100644
--- a/wasm/jit/jit_arm64_test.go
+++ b/wasm/jit/jit_arm64_test.go
@@ -11,6 +11,7 @@ import (
 	"unsafe"
 
 	"github.com/stretchr/testify/require"
+	"github.com/twitchyliquid64/golang-asm/obj/arm64"
 
 	"github.com/tetratelabs/wazero/wasm"
 	"github.com/tetratelabs/wazero/wasm/internal/wazeroir"
@@ -24,20 +25,48 @@ func (j *jitEnv) requireNewCompiler(t *testing.T) *arm64Compiler {
 	return ret
 }
 
+// TODO: delete this as this could be a duplication from other tests especially spectests.
+// Use this until we could run spectests on arm64.
 func TestArm64CompilerEndToEnd(t *testing.T) {
 	ctx := context.Background()
 	for _, tc := range []struct {
 		name string
 		body []byte
+		sig  *wasm.FunctionType
 	}{
-		{name: "empty", body: []byte{wasm.OpcodeEnd}},
-		{name: "br .return", body: []byte{wasm.OpcodeBr, 0, wasm.OpcodeEnd}},
+		{name: "empty", body: []byte{wasm.OpcodeEnd}, sig: &wasm.FunctionType{}},
+		{name: "br .return", body: []byte{wasm.OpcodeBr, 0, wasm.OpcodeEnd}, sig: &wasm.FunctionType{}},
+		{
+			name: "consts",
+			body: []byte{
+				wasm.OpcodeI32Const, 1, wasm.OpcodeI64Const, 1,
+				wasm.OpcodeF32Const, 1, 1, 1, 1, wasm.OpcodeF64Const, 1, 2, 3, 4, 5, 6, 7, 8,
+				wasm.OpcodeEnd,
+			},
+			// We push four constants.
+			sig: &wasm.FunctionType{Results: []wasm.ValueType{wasm.ValueTypeI32, wasm.ValueTypeI64, wasm.ValueTypeF32, wasm.ValueTypeF64}},
+		},
+		{
+			name: "add",
+			body: []byte{wasm.OpcodeI32Const, 1, wasm.OpcodeI32Const, 1, wasm.OpcodeI32Add, wasm.OpcodeEnd},
+			sig:  &wasm.FunctionType{Results: []wasm.ValueType{wasm.ValueTypeI32}},
+		},
+		{
+			name: "sub",
+			body: []byte{wasm.OpcodeI64Const, 1, wasm.OpcodeI64Const, 1, wasm.OpcodeI64Sub, wasm.OpcodeEnd},
+			sig:  &wasm.FunctionType{Results: []wasm.ValueType{wasm.ValueTypeI64}},
+		},
+		{
+			name: "mul",
+			body: []byte{wasm.OpcodeI64Const, 1, wasm.OpcodeI64Const, 1, wasm.OpcodeI64Mul, wasm.OpcodeEnd},
+			sig:  &wasm.FunctionType{Results: []wasm.ValueType{wasm.ValueTypeI64}},
+		},
 	} {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
 			engine := newEngine()
 			f := &wasm.FunctionInstance{
-				FunctionType: &wasm.TypeInstance{Type: &wasm.FunctionType{}},
+				FunctionType: &wasm.TypeInstance{Type: tc.sig},
 				Body:         tc.body,
 			}
 			err := engine.Compile(f)
@@ -57,7 +86,7 @@ func TestArchContextOffsetInEngine(t *testing.T) {
 func TestArm64Compiler_returnFunction(t *testing.T) {
 	env := newJITEnvironment()
 
-	// Build codes.
+	// Build code.
 	compiler := env.requireNewCompiler(t)
 	err := compiler.emitPreamble()
 	require.NoError(t, err)
@@ -87,7 +116,7 @@ func TestArm64Compiler_exit(t *testing.T) {
 
 			env := newJITEnvironment()
 
-			// Build codes.
+			// Build code.
 			compiler := env.requireNewCompiler(t)
 			err := compiler.emitPreamble()
 
@@ -100,7 +129,7 @@ func TestArm64Compiler_exit(t *testing.T) {
 			code, _, _, err := compiler.compile()
 			require.NoError(t, err)
 
-			// Run codes
+			// Run code
 			env.exec(code)
 
 			// JIT status on engine must be updated.
@@ -137,7 +166,7 @@ func TestArm64Compiler_compileConsts(t *testing.T) {
 				t.Run(fmt.Sprintf("0x%x", val), func(t *testing.T) {
 					env := newJITEnvironment()
 
-					// Build codes.
+					// Build code.
 					compiler := env.requireNewCompiler(t)
 					err := compiler.emitPreamble()
 					require.NoError(t, err)
@@ -184,3 +213,309 @@ func TestArm64Compiler_compileConsts(t *testing.T) {
 		})
 	}
 }
+
+func TestArm64Compiler_releaseRegisterToStack(t *testing.T) {
+	const val = 10000
+	for _, tc := range []struct {
+		name         string
+		stackPointer uint64
+		isFloat      bool
+	}{
+		{name: "int", stackPointer: 10, isFloat: false},
+		{name: "float", stackPointer: 10, isFloat: true},
+		{name: "int-huge-height", stackPointer: math.MaxInt16 + 1, isFloat: false},
+		{name: "float-huge-height", stackPointer: math.MaxInt16 + 1, isFloat: true},
+	} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			env := newJITEnvironment()
+
+			// Build code.
+			compiler := env.requireNewCompiler(t)
+			err := compiler.emitPreamble()
+			require.NoError(t, err)
+
+			// Setup the location stack so that we push the const on the specified height.
+			compiler.locationStack.sp = tc.stackPointer
+			compiler.locationStack.stack = make([]*valueLocation, tc.stackPointer)
+
+			if tc.isFloat {
+				err = compiler.compileConstF64(&wazeroir.OperationConstF64{Value: math.Float64frombits(val)})
+			} else {
+				err = compiler.compileConstI64(&wazeroir.OperationConstI64{Value: val})
+			}
+			require.NoError(t, err)
+
+			// Release the register allocated value to the memory stack so that we can see the value after exiting.
+			compiler.releaseRegisterToStack(compiler.locationStack.peek())
+			compiler.returnFunction()
+
+			// Generate the code under test.
+			code, _, _, err := compiler.compile()
+			require.NoError(t, err)
+
+			// Run native code after growing the value stack.
+			env.engine().builtinFunctionGrowValueStack(tc.stackPointer)
+			env.exec(code)
+
+			// JIT status on engine must be returned and stack pointer must end up the specified one.
+			require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+			require.Equal(t, tc.stackPointer+1, env.stackPointer())
+
+			if tc.isFloat {
+				require.Equal(t, math.Float64frombits(val), env.stackTopAsFloat64())
+			} else {
+				require.Equal(t, uint64(val), env.stackTopAsUint64())
+			}
+		})
+	}
+}
+
+func TestArm64Compiler_loadValueOnStackToRegister(t *testing.T) {
+	const val = 123
+	for _, tc := range []struct {
+		name         string
+		stackPointer uint64
+		isFloat      bool
+	}{
+		{name: "int", stackPointer: 10, isFloat: false},
+		{name: "float", stackPointer: 10, isFloat: true},
+		{name: "int-huge-height", stackPointer: math.MaxInt16 + 1, isFloat: false},
+		{name: "float-huge-height", stackPointer: math.MaxInt16 + 1, isFloat: true},
+	} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			env := newJITEnvironment()
+
+			// Build code.
+			compiler := env.requireNewCompiler(t)
+			err := compiler.emitPreamble()
+			require.NoError(t, err)
+
+			// Setup the location stack so that we push the const on the specified height.
+			compiler.locationStack.sp = tc.stackPointer
+			compiler.locationStack.stack = make([]*valueLocation, tc.stackPointer)
+
+			// Record that that top value is on top.
+			require.Len(t, compiler.locationStack.usedRegisters, 0)
+			loc := compiler.locationStack.pushValueOnStack()
+			if tc.isFloat {
+				loc.setRegisterType(generalPurposeRegisterTypeFloat)
+			} else {
+				loc.setRegisterType(generalPurposeRegisterTypeInt)
+			}
+			// At this point the value must be recorded as being on stack.
+			require.True(t, loc.onStack())
+
+			// Release the stack-allocated value to register.
+			compiler.loadValueOnStackToRegister(loc)
+			require.Len(t, compiler.locationStack.usedRegisters, 1)
+			require.True(t, loc.onRegister())
+
+			// To verify the behavior, increment the value on the register.
+			if tc.isFloat {
+				// For float, we cannot add consts, so load the constant first.
+				err = compiler.emitFloatConstant(false, math.Float64bits(1))
+				require.NoError(t, err)
+				// Then, do the increment.
+				compiler.applyRegisterToRegisterInstruction(arm64.AFADDD, compiler.locationStack.peek().register, loc.register)
+				// Delete the loaded const.
+				compiler.locationStack.pop()
+			} else {
+				compiler.applyConstToRegisterInstruction(arm64.AADD, 1, loc.register)
+			}
+
+			// Release the value to the memory stack so that we can see the value after exiting.
+			compiler.releaseRegisterToStack(loc)
+			compiler.returnFunction()
+
+			// Generate the code under test.
+			code, _, _, err := compiler.compile()
+			require.NoError(t, err)
+
+			// Run native code after growing the value stack, and place the original value.
+			env.engine().builtinFunctionGrowValueStack(tc.stackPointer)
+			env.stack()[tc.stackPointer] = val
+			env.exec(code)
+
+			// JIT status on engine must be returned and stack pointer must end up the specified one.
+			require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+			require.Equal(t, tc.stackPointer+1, env.stackPointer())
+
+			if tc.isFloat {
+				require.Equal(t, math.Float64frombits(val)+1, env.stackTopAsFloat64())
+			} else {
+				require.Equal(t, uint64(val)+1, env.stackTopAsUint64())
+			}
+		})
+	}
+}
+
+func TestArm64Compiler_compile_Add_Sub_Mul(t *testing.T) {
+	for _, kind := range []wazeroir.OperationKind{
+		wazeroir.OperationKindAdd,
+		wazeroir.OperationKindSub,
+		wazeroir.OperationKindMul,
+	} {
+		kind := kind
+		t.Run(kind.String(), func(t *testing.T) {
+			for _, unsignedType := range []wazeroir.UnsignedType{
+				wazeroir.UnsignedTypeI32,
+				wazeroir.UnsignedTypeI64,
+				wazeroir.UnsignedTypeF32,
+				wazeroir.UnsignedTypeF64,
+			} {
+				unsignedType := unsignedType
+				t.Run(unsignedType.String(), func(t *testing.T) {
+					for _, values := range [][2]uint64{
+						{0, 0}, {1, 1}, {2, 1}, {100, 1}, {1, 0}, {0, 1}, {math.MaxInt16, math.MaxInt32},
+						{1 << 14, 1 << 21}, {1 << 14, 1 << 21},
+						{0xffff_ffff_ffff_ffff, 0}, {0xffff_ffff_ffff_ffff, 1},
+						{0, 0xffff_ffff_ffff_ffff}, {1, 0xffff_ffff_ffff_ffff},
+						{0, math.Float64bits(math.Inf(1))},
+						{0, math.Float64bits(math.Inf(-1))},
+						{math.Float64bits(math.Inf(1)), 1},
+						{math.Float64bits(math.Inf(-1)), 1},
+						{math.Float64bits(1.11231), math.Float64bits(math.Inf(1))},
+						{math.Float64bits(1.11231), math.Float64bits(math.Inf(-1))},
+						{math.Float64bits(math.Inf(1)), math.Float64bits(1.11231)},
+						{math.Float64bits(math.Inf(-1)), math.Float64bits(1.11231)},
+						{math.Float64bits(math.Inf(1)), math.Float64bits(math.NaN())},
+						{math.Float64bits(math.Inf(-1)), math.Float64bits(math.NaN())},
+						{math.Float64bits(math.NaN()), math.Float64bits(math.Inf(1))},
+						{math.Float64bits(math.NaN()), math.Float64bits(math.Inf(-1))},
+					} {
+						x1, x2 := values[0], values[1]
+						t.Run(fmt.Sprintf("x1=0x%x,x2=0x%x", x1, x2), func(t *testing.T) {
+							env := newJITEnvironment()
+							compiler := env.requireNewCompiler(t)
+							err := compiler.emitPreamble()
+							require.NoError(t, err)
+
+							// Emit consts operands.
+							for _, v := range []uint64{x1, x2} {
+								switch unsignedType {
+								case wazeroir.UnsignedTypeI32:
+									err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: uint32(v)})
+								case wazeroir.UnsignedTypeI64:
+									err = compiler.compileConstI64(&wazeroir.OperationConstI64{Value: v})
+								case wazeroir.UnsignedTypeF32:
+									err = compiler.compileConstF32(&wazeroir.OperationConstF32{Value: math.Float32frombits(uint32(v))})
+								case wazeroir.UnsignedTypeF64:
+									err = compiler.compileConstF64(&wazeroir.OperationConstF64{Value: math.Float64frombits(v)})
+								}
+								require.NoError(t, err)
+							}
+
+							// At this point, two values exist.
+							require.Equal(t, uint64(2), compiler.locationStack.sp)
+
+							// Emit the operation.
+							switch kind {
+							case wazeroir.OperationKindAdd:
+								err = compiler.compileAdd(&wazeroir.OperationAdd{Type: unsignedType})
+							case wazeroir.OperationKindSub:
+								err = compiler.compileSub(&wazeroir.OperationSub{Type: unsignedType})
+							case wazeroir.OperationKindMul:
+								err = compiler.compileMul(&wazeroir.OperationMul{Type: unsignedType})
+							}
+							require.NoError(t, err)
+
+							// We consumed two values, but push the result back.
+							require.Equal(t, uint64(1), compiler.locationStack.sp)
+							resultLocation := compiler.locationStack.peek()
+							// Plus the result must be located on a register.
+							require.True(t, resultLocation.onRegister())
+							// Also, the result must have an appropriate register type.
+							if unsignedType == wazeroir.UnsignedTypeF32 || unsignedType == wazeroir.UnsignedTypeF64 {
+								require.Equal(t, generalPurposeRegisterTypeFloat, resultLocation.regType)
+							} else {
+								require.Equal(t, generalPurposeRegisterTypeInt, resultLocation.regType)
+							}
+
+							// Release the value to the memory stack again to verify the operation, and then return.
+							compiler.releaseRegisterToStack(resultLocation)
+							compiler.returnFunction()
+
+							// Generate the code under test.
+							code, _, _, err := compiler.compile()
+							require.NoError(t, err)
+
+							// Run code.
+							env.exec(code)
+
+							// Check the stack.
+							require.Equal(t, uint64(1), env.stackPointer())
+
+							switch kind {
+							case wazeroir.OperationKindAdd:
+								switch unsignedType {
+								case wazeroir.UnsignedTypeI32:
+									require.Equal(t, uint32(x1)+uint32(x2), env.stackTopAsUint32())
+								case wazeroir.UnsignedTypeI64:
+									require.Equal(t, x1+x2, env.stackTopAsUint64())
+								case wazeroir.UnsignedTypeF32:
+									exp := math.Float32frombits(uint32(x1)) + math.Float32frombits(uint32(x2))
+									if math.IsNaN(float64(exp)) {
+										require.True(t, math.IsNaN(float64(env.stackTopAsFloat32())))
+									} else {
+										require.Equal(t, exp, env.stackTopAsFloat32())
+									}
+								case wazeroir.UnsignedTypeF64:
+									exp := math.Float64frombits(x1) + math.Float64frombits(x2)
+									if math.IsNaN(exp) {
+										require.True(t, math.IsNaN(env.stackTopAsFloat64()))
+									} else {
+										require.Equal(t, exp, env.stackTopAsFloat64())
+									}
+								}
+							case wazeroir.OperationKindSub:
+								switch unsignedType {
+								case wazeroir.UnsignedTypeI32:
+									require.Equal(t, uint32(x1)-uint32(x2), env.stackTopAsUint32())
+								case wazeroir.UnsignedTypeI64:
+									require.Equal(t, x1-x2, env.stackTopAsUint64())
+								case wazeroir.UnsignedTypeF32:
+									exp := math.Float32frombits(uint32(x1)) - math.Float32frombits(uint32(x2))
+									if math.IsNaN(float64(exp)) {
+										require.True(t, math.IsNaN(float64(env.stackTopAsFloat32())))
+									} else {
+										require.Equal(t, exp, env.stackTopAsFloat32())
+									}
+								case wazeroir.UnsignedTypeF64:
+									exp := math.Float64frombits(x1) - math.Float64frombits(x2)
+									if math.IsNaN(exp) {
+										require.True(t, math.IsNaN(env.stackTopAsFloat64()))
+									} else {
+										require.Equal(t, exp, env.stackTopAsFloat64())
+									}
+								}
+							case wazeroir.OperationKindMul:
+								switch unsignedType {
+								case wazeroir.UnsignedTypeI32:
+									require.Equal(t, uint32(x1)*uint32(x2), env.stackTopAsUint32())
+								case wazeroir.UnsignedTypeI64:
+									require.Equal(t, x1*x2, env.stackTopAsUint64())
+								case wazeroir.UnsignedTypeF32:
+									exp := math.Float32frombits(uint32(x1)) * math.Float32frombits(uint32(x2))
+									if math.IsNaN(float64(exp)) {
+										require.True(t, math.IsNaN(float64(env.stackTopAsFloat32())))
+									} else {
+										require.Equal(t, exp, env.stackTopAsFloat32())
+									}
+								case wazeroir.UnsignedTypeF64:
+									exp := math.Float64frombits(x1) * math.Float64frombits(x2)
+									if math.IsNaN(exp) {
+										require.True(t, math.IsNaN(env.stackTopAsFloat64()))
+									} else {
+										require.Equal(t, exp, env.stackTopAsFloat64())
+									}
+								}
+							}
+						})
+					}
+				})
+			}
+		})
+	}
+}
diff --git a/wasm/jit/jit_value_location_test.go b/wasm/jit/jit_value_location_test.go
index 084cd13d..2374cc09 100644
--- a/wasm/jit/jit_value_location_test.go
+++ b/wasm/jit/jit_value_location_test.go
@@ -15,6 +15,7 @@ func TestReservedRegisters(t *testing.T) {
 	require.NotContains(t, unreservedGeneralPurposeIntRegisters, reservedRegisterForEngine)
 	require.NotContains(t, unreservedGeneralPurposeIntRegisters, reservedRegisterForStackBasePointerAddress)
 	require.NotContains(t, unreservedGeneralPurposeIntRegisters, reservedRegisterForMemory)
+	require.NotContains(t, unreservedGeneralPurposeIntRegisters, reservedRegisterForTemporary)
 }
 
 func Test_isIntRegister(t *testing.T) {
diff --git a/wasm/module_test.go b/wasm/module_test.go
index 9cfd7215..7789a25e 100644
--- a/wasm/module_test.go
+++ b/wasm/module_test.go
@@ -99,7 +99,7 @@ func TestValueTypeName(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			require.Equal(t, tc.expected, valueTypeName(tc.input))
+			require.Equal(t, tc.expected, ValueTypeName(tc.input))
 		})
 	}
 }
diff --git a/wasm/store_test.go b/wasm/store_test.go
index 8b8150c5..2b59ce46 100644
--- a/wasm/store_test.go
+++ b/wasm/store_test.go
@@ -439,7 +439,7 @@ func TestStore_executeConstExpression(t *testing.T) {
 	})
 	t.Run("non global expr", func(t *testing.T) {
 		for _, vt := range []ValueType{ValueTypeI32, ValueTypeI64, ValueTypeF32, ValueTypeF64} {
-			t.Run(valueTypeName(vt), func(t *testing.T) {
+			t.Run(ValueTypeName(vt), func(t *testing.T) {
 				t.Run("valid", func(t *testing.T) {
 					// Allocate bytes with enough size for all types.
 					expr := &ConstantExpression{Data: make([]byte, 8)}
@@ -526,7 +526,7 @@ func TestStore_executeConstExpression(t *testing.T) {
 				{valueType: ValueTypeF32, val: uint64(math.Float32bits(634634432.12311))},
 				{valueType: ValueTypeF64, val: math.Float64bits(1.12312311)},
 			} {
-				t.Run(valueTypeName(tc.valueType), func(t *testing.T) {
+				t.Run(ValueTypeName(tc.valueType), func(t *testing.T) {
 					// The index specified in Data equals zero.
 					expr := &ConstantExpression{Data: []byte{0}, Opcode: OpcodeGlobalGet}
 					globals := []*GlobalInstance{{Val: tc.val, Type: &GlobalType{ValType: tc.valueType}}}