Repository: tetratelabs/wazero
Pull request #111 (2022-01-02 02:35:04)

jit: support for mul
Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>

codefromthecrypt: yay! I found the long form tedious

> Message ID: ***@***.***>
>


Patch:
diff --git a/wasm/jit/compiler.go b/wasm/jit/compiler.go
index a553594a..df03d693 100644
--- a/wasm/jit/compiler.go
+++ b/wasm/jit/compiler.go
@@ -25,6 +25,7 @@ type compiler interface {
 	compilePick(o *wazeroir.OperationPick) error
 	compileAdd(o *wazeroir.OperationAdd) error
 	compileSub(o *wazeroir.OperationSub) error
+	compileMul(o *wazeroir.OperationMul) error
 	compileEq(o *wazeroir.OperationEq) error
 	compileNe(o *wazeroir.OperationNe) error
 	compileEqz(o *wazeroir.OperationEqz) error
diff --git a/wasm/jit/engine.go b/wasm/jit/engine.go
index a4f0bc58..8277135a 100644
--- a/wasm/jit/engine.go
+++ b/wasm/jit/engine.go
@@ -537,7 +537,7 @@ func (e *engine) compileWasmFunction(f *wasm.FunctionInstance) (*compiledWasmFun
 		case *wazeroir.OperationSub:
 			err = compiler.compileSub(o)
 		case *wazeroir.OperationMul:
-			err = fmt.Errorf("unsupported operation")
+			err = compiler.compileMul(o)
 		case *wazeroir.OperationClz:
 			err = fmt.Errorf("unsupported operation")
 		case *wazeroir.OperationCtz:
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index c77b216a..c776fdfd 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -6,6 +6,8 @@ package jit
 // This file implements the compiler for amd64/x86_64 target.
 // Please refer to https://www.felixcloutier.com/x86/index.html
 // if unfamiliar with amd64 instructions used here.
+// Note that x86 pkg used here prefixes all the instructions with "A"
+// e.g. MOVQ will be given as x86.AMOVQ.
 
 import (
 	"encoding/binary"
@@ -133,7 +135,7 @@ func (c *amd64Compiler) compileSwap(o *wazeroir.OperationSwap) error {
 	// If x1 is on the conditional register, we must move it to a gp
 	// register before swap.
 	if x1.onConditionalRegister() {
-		if err := c.moveConditionalToGeneralPurposeRegister(x1); err != nil {
+		if err := c.moveConditionalToFreeGeneralPurposeRegister(x1); err != nil {
 			return err
 		}
 	}
@@ -572,7 +574,7 @@ func (c *amd64Compiler) emitDropRange(r *wazeroir.InclusiveRange) error {
 			if topIsConditional {
 				// If the top is conditional, and it's not target of drop,
 				// we must assign it to the register before we emit any instructions here.
-				if err := c.moveConditionalToGeneralPurposeRegister(top); err != nil {
+				if err := c.moveConditionalToFreeGeneralPurposeRegister(top); err != nil {
 					return err
 				}
 				topIsConditional = false
@@ -590,7 +592,7 @@ func (c *amd64Compiler) emitDropRange(r *wazeroir.InclusiveRange) error {
 	return nil
 }
 
-// handleSelect uses top three values on the stack:
+// compileSelect uses top three values on the stack:
 // Assume we have stack as [..., x1, x2, c], if the value of c
 // equals zero, then the stack results in [..., x1]
 // otherwise, [..., x2].
@@ -766,7 +768,7 @@ func (c *amd64Compiler) compileSub(o *wazeroir.OperationSub) error {
 		return err
 	}
 
-	// x1 += x2.
+	// x1 -= x2.
 	prog := c.newProg()
 	prog.From.Type = obj.TYPE_REG
 	prog.From.Reg = x2.register
@@ -781,6 +783,129 @@ func (c *amd64Compiler) compileSub(o *wazeroir.OperationSub) error {
 	return nil
 }
 
+// compileMul adds instructions to multiply two operands which may be on the stack or registers.
+// After execution, the result of multiplication pushed onto the stack.
+func (c *amd64Compiler) compileMul(o *wazeroir.OperationMul) (err error) {
+	switch o.Type {
+	case wazeroir.UnsignedTypeI32:
+		err = c.compileMulForInts(true, x86.AMULL)
+	case wazeroir.UnsignedTypeI64:
+		err = c.compileMulForInts(false, x86.AMULQ)
+	case wazeroir.UnsignedTypeF32:
+		err = c.compileMulForFloats(x86.AMULSS)
+	case wazeroir.UnsignedTypeF64:
+		err = c.compileMulForFloats(x86.AMULSD)
+	}
+	return
+}
+
+func (c *amd64Compiler) compileMulForInts(is32Bit bool, mulInstruction obj.As) error {
+	// See https://www.felixcloutier.com/x86/mul if unfamiliar with the convention for
+	// integer multiplication. In summary, the destination operand must be on the AX
+	// register and the overflow info is stored in RDX. So, we have to ensure that
+	// 1) Previously located value on RDX must be saved to memory stack.
+	// 2) One of the operands (x1 or x2) must be on RAX (the AX register).
+	const (
+		resultRegister   = x86.REG_AX
+		reservedRegister = x86.REG_DX
+	)
+
+	x2 := c.locationStack.pop()
+	x1 := c.locationStack.pop()
+
+	var valueOnAX *valueLocation
+	if x1.register == resultRegister {
+		valueOnAX = x1
+	} else if x2.register == resultRegister {
+		valueOnAX = x2
+	} else {
+		valueOnAX = x2
+		// This case we  move x2 to AX register.
+		c.onValueReleaseRegisterToStack(resultRegister)
+		if x2.onConditionalRegister() {
+			c.moveConditionalToGeneralPurposeRegister(x2, resultRegister)
+		} else if x2.onStack() {
+			x2.setRegister(resultRegister)
+			c.moveStackToRegister(x2)
+			c.locationStack.markRegisterUsed(resultRegister)
+		} else {
+			moveX2ToAX := c.newProg()
+			if is32Bit {
+				moveX2ToAX.As = x86.AMOVL
+			} else {
+				moveX2ToAX.As = x86.AMOVQ
+			}
+			moveX2ToAX.To.Reg = resultRegister
+			moveX2ToAX.To.Type = obj.TYPE_REG
+			moveX2ToAX.From.Reg = x2.register
+			moveX2ToAX.From.Type = obj.TYPE_REG
+			c.addInstruction(moveX2ToAX)
+			// We no longer uses the prev register of x2.
+			c.locationStack.releaseRegister(x2)
+			x2.setRegister(resultRegister)
+			c.locationStack.markRegisterUsed(resultRegister)
+		}
+	}
+
+	// We have to make sure that at this point the operands must be on registers.
+	if err := c.ensureOnGeneralPurposeRegister(x2); err != nil {
+		return err
+	}
+	if err := c.ensureOnGeneralPurposeRegister(x1); err != nil {
+		return err
+	}
+
+	// We have to save the existing value on RDX.
+	if x1.register != reservedRegister && x2.register != reservedRegister {
+		c.onValueReleaseRegisterToStack(reservedRegister)
+	}
+
+	// Now ready to emit the mul instruction.
+	mul := c.newProg()
+	mul.As = mulInstruction
+	mul.To.Type = obj.TYPE_NONE
+	mul.From.Type = obj.TYPE_REG
+	if x1 == valueOnAX {
+		mul.From.Reg = x2.register
+		c.locationStack.markRegisterUnused(x2.register)
+	} else {
+		mul.From.Reg = x1.register
+		c.locationStack.markRegisterUnused(x1.register)
+	}
+	c.addInstruction(mul)
+
+	// Now we have the result in the AX register,
+	// so we record it.
+	c.locationStack.pushValueOnRegister(resultRegister)
+	return nil
+}
+
+func (c *amd64Compiler) compileMulForFloats(instruction obj.As) error {
+	x2 := c.locationStack.pop()
+	if err := c.ensureOnGeneralPurposeRegister(x2); err != nil {
+		return err
+	}
+
+	x1 := c.locationStack.peek() // Note this is peek!
+	if err := c.ensureOnGeneralPurposeRegister(x1); err != nil {
+		return err
+	}
+
+	// x1 *= x2.
+	prog := c.newProg()
+	prog.From.Type = obj.TYPE_REG
+	prog.From.Reg = x2.register
+	prog.To.Type = obj.TYPE_REG
+	prog.To.Reg = x1.register
+	prog.As = instruction
+	c.addInstruction(prog)
+
+	// We no longer need x2 register after MUL operation here,
+	// so we release it.
+	c.locationStack.releaseRegister(x2)
+	return nil
+}
+
 func (c *amd64Compiler) compileEq(o *wazeroir.OperationEq) error {
 	return c.emitEqOrNe(o.Type, true)
 }
@@ -1532,7 +1657,7 @@ func (c *amd64Compiler) moveStackToRegister(loc *valueLocation) {
 	c.addInstruction(prog)
 }
 
-func (c *amd64Compiler) moveConditionalToGeneralPurposeRegister(loc *valueLocation) error {
+func (c *amd64Compiler) moveConditionalToFreeGeneralPurposeRegister(loc *valueLocation) error {
 	// Get the free register.
 	reg, ok := c.locationStack.takeFreeRegister(generalPurposeRegisterTypeInt)
 	if !ok {
@@ -1542,10 +1667,16 @@ func (c *amd64Compiler) moveConditionalToGeneralPurposeRegister(loc *valueLocati
 		return fmt.Errorf("conditional register mov requires a free register")
 	}
 
+	c.moveConditionalToGeneralPurposeRegister(loc, reg)
+	return nil
+}
+
+func (c *amd64Compiler) moveConditionalToGeneralPurposeRegister(loc *valueLocation, reg int16) {
 	// Set the flag bit to the destination.
 	prog := c.newProg()
 	prog.To.Type = obj.TYPE_REG
 	prog.To.Reg = reg
+
 	// See
 	// - https://c9x.me/x86/html/file_module_x86_id_288.html
 	// - https://github.com/golang/go/blob/master/src/cmd/internal/obj/x86/asm6.go#L1453-L1468
@@ -1590,7 +1721,6 @@ func (c *amd64Compiler) moveConditionalToGeneralPurposeRegister(loc *valueLocati
 	// Mark it uses the register.
 	loc.setRegister(reg)
 	c.locationStack.markRegisterUsed(reg)
-	return nil
 }
 
 // allocateRegister returns an unused register of the given type. The register will be taken
@@ -1683,7 +1813,7 @@ func (c *amd64Compiler) callFunctionFromRegisterIndex(reg int16) {
 
 func (c *amd64Compiler) releaseAllRegistersToStack() {
 	used := len(c.locationStack.usedRegisters)
-	for i := len(c.locationStack.stack) - 1; i >= 0 && used > 0; i-- {
+	for i := uint64(0); i < c.locationStack.sp && used > 0; i++ {
 		if loc := c.locationStack.stack[i]; loc.onRegister() {
 			c.releaseRegisterToStack(loc)
 			used--
@@ -1752,6 +1882,15 @@ func (c *amd64Compiler) setFunctionCallIndexFromConst(index int64) {
 	c.addInstruction(prog)
 }
 
+func (c *amd64Compiler) onValueReleaseRegisterToStack(reg int16) {
+	prevValue := c.locationStack.findValueForRegister(reg)
+	if prevValue == nil {
+		// This case the target register is not used by any value.
+		return
+	}
+	c.releaseRegisterToStack(prevValue)
+}
+
 func (c *amd64Compiler) releaseRegisterToStack(loc *valueLocation) {
 	// Push value.
 	prog := c.newProg()
@@ -1861,7 +2000,7 @@ func (c *amd64Compiler) ensureOnGeneralPurposeRegister(loc *valueLocation) error
 			return err
 		}
 	} else if loc.onConditionalRegister() {
-		if err := c.moveConditionalToGeneralPurposeRegister(loc); err != nil {
+		if err := c.moveConditionalToFreeGeneralPurposeRegister(loc); err != nil {
 			return err
 		}
 	}
diff --git a/wasm/jit/jit_value_location_amd64.go b/wasm/jit/jit_value_location_amd64.go
index 14954412..472f37b4 100644
--- a/wasm/jit/jit_value_location_amd64.go
+++ b/wasm/jit/jit_value_location_amd64.go
@@ -250,3 +250,15 @@ func (s *valueLocationStack) takeStealTargetFromUsedRegister(tp generalPurposeRe
 	}
 	return nil, false
 }
+
+// findValueForRegister returns the valueLocation of the given register or nil if not found.
+// If not found, return nil.
+func (s *valueLocationStack) findValueForRegister(reg int16) *valueLocation {
+	for i := uint64(0); i < s.sp; i++ {
+		loc := s.stack[i]
+		if loc.register == reg {
+			return loc
+		}
+	}
+	return nil
+}


Test patch:
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index 221cdf6c..a328383a 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -1142,7 +1142,7 @@ func TestAmd64Compiler_emitEqOrNe(t *testing.T) {
 						// to the stack.
 						top := compiler.locationStack.peek()
 						require.True(t, top.onConditionalRegister() && !top.onRegister())
-						err = compiler.moveConditionalToGeneralPurposeRegister(top)
+						err = compiler.moveConditionalToFreeGeneralPurposeRegister(top)
 						require.NoError(t, err)
 						require.True(t, !top.onConditionalRegister() && top.onRegister())
 						compiler.releaseRegisterToStack(top)
@@ -1209,7 +1209,7 @@ func TestAmd64Compiler_emitEqOrNe(t *testing.T) {
 						// to the stack.
 						top := compiler.locationStack.peek()
 						require.True(t, top.onConditionalRegister() && !top.onRegister())
-						err = compiler.moveConditionalToGeneralPurposeRegister(top)
+						err = compiler.moveConditionalToFreeGeneralPurposeRegister(top)
 						require.NoError(t, err)
 						require.True(t, !top.onConditionalRegister() && top.onRegister())
 						compiler.releaseRegisterToStack(top)
@@ -1283,7 +1283,7 @@ func TestAmd64Compiler_emitEqOrNe(t *testing.T) {
 						// to the stack.
 						flag := compiler.locationStack.peek()
 						require.True(t, flag.onConditionalRegister() && !flag.onRegister())
-						err = compiler.moveConditionalToGeneralPurposeRegister(flag)
+						err = compiler.moveConditionalToFreeGeneralPurposeRegister(flag)
 						require.NoError(t, err)
 						require.True(t, !flag.onConditionalRegister() && flag.onRegister())
 						compiler.releaseRegisterToStack(flag)
@@ -1356,7 +1356,7 @@ func TestAmd64Compiler_emitEqOrNe(t *testing.T) {
 						// to the stack.
 						flag := compiler.locationStack.peek()
 						require.True(t, flag.onConditionalRegister() && !flag.onRegister())
-						err = compiler.moveConditionalToGeneralPurposeRegister(flag)
+						err = compiler.moveConditionalToFreeGeneralPurposeRegister(flag)
 						require.NoError(t, err)
 						require.True(t, !flag.onConditionalRegister() && flag.onRegister())
 						compiler.releaseRegisterToStack(flag)
@@ -1414,7 +1414,7 @@ func TestAmd64Compiler_compileEqz(t *testing.T) {
 				// to the stack.
 				top := compiler.locationStack.peek()
 				require.True(t, top.onConditionalRegister() && !top.onRegister())
-				err = compiler.moveConditionalToGeneralPurposeRegister(top)
+				err = compiler.moveConditionalToFreeGeneralPurposeRegister(top)
 				require.NoError(t, err)
 				require.True(t, !top.onConditionalRegister() && top.onRegister())
 				compiler.releaseRegisterToStack(top)
@@ -1462,7 +1462,7 @@ func TestAmd64Compiler_compileEqz(t *testing.T) {
 				// to the stack.
 				top := compiler.locationStack.peek()
 				require.True(t, top.onConditionalRegister() && !top.onRegister())
-				err = compiler.moveConditionalToGeneralPurposeRegister(top)
+				err = compiler.moveConditionalToFreeGeneralPurposeRegister(top)
 				require.NoError(t, err)
 				require.True(t, !top.onConditionalRegister() && top.onRegister())
 				compiler.releaseRegisterToStack(top)
@@ -1544,7 +1544,7 @@ func TestAmd64Compiler_compileLe_or_Lt(t *testing.T) {
 						// to the stack.
 						top := compiler.locationStack.peek()
 						require.True(t, top.onConditionalRegister() && !top.onRegister())
-						err = compiler.moveConditionalToGeneralPurposeRegister(top)
+						err = compiler.moveConditionalToFreeGeneralPurposeRegister(top)
 						require.NoError(t, err)
 						require.True(t, !top.onConditionalRegister() && top.onRegister())
 						compiler.releaseRegisterToStack(top)
@@ -1627,7 +1627,7 @@ func TestAmd64Compiler_compileLe_or_Lt(t *testing.T) {
 						// to the stack.
 						top := compiler.locationStack.peek()
 						require.True(t, top.onConditionalRegister() && !top.onRegister())
-						err = compiler.moveConditionalToGeneralPurposeRegister(top)
+						err = compiler.moveConditionalToFreeGeneralPurposeRegister(top)
 						require.NoError(t, err)
 						require.True(t, !top.onConditionalRegister() && top.onRegister())
 						compiler.releaseRegisterToStack(top)
@@ -1702,7 +1702,7 @@ func TestAmd64Compiler_compileLe_or_Lt(t *testing.T) {
 						// to the stack.
 						flag := compiler.locationStack.peek()
 						require.True(t, flag.onConditionalRegister() && !flag.onRegister())
-						err = compiler.moveConditionalToGeneralPurposeRegister(flag)
+						err = compiler.moveConditionalToFreeGeneralPurposeRegister(flag)
 						require.NoError(t, err)
 						require.True(t, !flag.onConditionalRegister() && flag.onRegister())
 						compiler.releaseRegisterToStack(flag)
@@ -1773,7 +1773,7 @@ func TestAmd64Compiler_compileLe_or_Lt(t *testing.T) {
 						// to the stack.
 						flag := compiler.locationStack.peek()
 						require.True(t, flag.onConditionalRegister() && !flag.onRegister())
-						err = compiler.moveConditionalToGeneralPurposeRegister(flag)
+						err = compiler.moveConditionalToFreeGeneralPurposeRegister(flag)
 						require.NoError(t, err)
 						require.True(t, !flag.onConditionalRegister() && flag.onRegister())
 						compiler.releaseRegisterToStack(flag)
@@ -1861,7 +1861,7 @@ func TestAmd64Compiler_compileGe_or_Gt(t *testing.T) {
 						// to the stack.
 						top := compiler.locationStack.peek()
 						require.True(t, top.onConditionalRegister() && !top.onRegister())
-						err = compiler.moveConditionalToGeneralPurposeRegister(top)
+						err = compiler.moveConditionalToFreeGeneralPurposeRegister(top)
 						require.NoError(t, err)
 						require.True(t, !top.onConditionalRegister() && top.onRegister())
 						compiler.releaseRegisterToStack(top)
@@ -1945,7 +1945,7 @@ func TestAmd64Compiler_compileGe_or_Gt(t *testing.T) {
 						// to the stack.
 						top := compiler.locationStack.peek()
 						require.True(t, top.onConditionalRegister() && !top.onRegister())
-						err = compiler.moveConditionalToGeneralPurposeRegister(top)
+						err = compiler.moveConditionalToFreeGeneralPurposeRegister(top)
 						require.NoError(t, err)
 						require.True(t, !top.onConditionalRegister() && top.onRegister())
 						compiler.releaseRegisterToStack(top)
@@ -2020,7 +2020,7 @@ func TestAmd64Compiler_compileGe_or_Gt(t *testing.T) {
 						// to the stack.
 						flag := compiler.locationStack.peek()
 						require.True(t, flag.onConditionalRegister() && !flag.onRegister())
-						err = compiler.moveConditionalToGeneralPurposeRegister(flag)
+						err = compiler.moveConditionalToFreeGeneralPurposeRegister(flag)
 						require.NoError(t, err)
 						require.True(t, !flag.onConditionalRegister() && flag.onRegister())
 						compiler.releaseRegisterToStack(flag)
@@ -2093,7 +2093,7 @@ func TestAmd64Compiler_compileGe_or_Gt(t *testing.T) {
 						// to the stack.
 						flag := compiler.locationStack.peek()
 						require.True(t, flag.onConditionalRegister() && !flag.onRegister())
-						err = compiler.moveConditionalToGeneralPurposeRegister(flag)
+						err = compiler.moveConditionalToFreeGeneralPurposeRegister(flag)
 						require.NoError(t, err)
 						require.True(t, !flag.onConditionalRegister() && flag.onRegister())
 						compiler.releaseRegisterToStack(flag)
@@ -2301,6 +2301,327 @@ func TestAmd64Compiler_compileSub(t *testing.T) {
 	})
 }
 
+func TestAmd64Compiler_compileMul(t *testing.T) {
+	t.Run("int32", func(t *testing.T) {
+		for _, tc := range []struct {
+			name string
+			// Interpret -1 as stack.
+			x1Reg, x2Reg int16
+		}{
+			{
+				name:  "x1:rax,x2:random_reg",
+				x1Reg: x86.REG_AX,
+				x2Reg: x86.REG_R10,
+			},
+			{
+				name:  "x1:rax,x2:stack",
+				x1Reg: x86.REG_AX,
+				x2Reg: -1,
+			},
+			{
+				name:  "x1:random_reg,x2:rax",
+				x1Reg: x86.REG_R10,
+				x2Reg: x86.REG_AX,
+			},
+			{
+				name:  "x1:staack,x2:rax",
+				x1Reg: -1,
+				x2Reg: x86.REG_AX,
+			},
+			{
+				name:  "x1:random_reg,x2:random_reg",
+				x1Reg: x86.REG_R10,
+				x2Reg: x86.REG_R9,
+			},
+			{
+				name:  "x1:stack,x2:random_reg",
+				x1Reg: -1,
+				x2Reg: x86.REG_R9,
+			},
+			{
+				name:  "x1:random_reg,x2:stack",
+				x1Reg: x86.REG_R9,
+				x2Reg: -1,
+			},
+			{
+				name:  "x1:stack,x2:stack",
+				x1Reg: -1,
+				x2Reg: -1,
+			},
+		} {
+			tc := tc
+			t.Run(tc.name, func(t *testing.T) {
+				const x1Value uint32 = 1 << 11
+				const x2Value uint32 = 51
+				const dxValue uint64 = 111111
+
+				eng := newEngine()
+				compiler := requireNewCompiler(t)
+				compiler.initializeReservedRegisters()
+
+				// Pretend there was an existing value on the DX register. We expect compileMul to save this to the stack.
+				compiler.movIntConstToRegister(int64(dxValue), x86.REG_DX)
+				compiler.locationStack.pushValueOnRegister(x86.REG_DX)
+
+				// Setup values.
+				if tc.x1Reg != -1 {
+					compiler.movIntConstToRegister(int64(x1Value), tc.x1Reg)
+					compiler.locationStack.pushValueOnRegister(tc.x1Reg)
+				} else {
+					loc := compiler.locationStack.pushValueOnStack()
+					eng.stack[loc.stackPointer] = uint64(x1Value)
+				}
+				if tc.x2Reg != -1 {
+					compiler.movIntConstToRegister(int64(x2Value), tc.x2Reg)
+					compiler.locationStack.pushValueOnRegister(tc.x2Reg)
+				} else {
+					loc := compiler.locationStack.pushValueOnStack()
+					eng.stack[loc.stackPointer] = uint64(x2Value)
+				}
+
+				err := compiler.compileMul(&wazeroir.OperationMul{Type: wazeroir.UnsignedTypeI32})
+				require.NoError(t, err)
+				require.Equal(t, int16(x86.REG_AX), compiler.locationStack.peek().register)
+				require.Equal(t, uint64(2), compiler.locationStack.sp)
+				require.Len(t, compiler.locationStack.usedRegisters, 1)
+
+				// To verify the behavior, we push the value
+				// to the stack.
+				compiler.releaseAllRegistersToStack()
+				compiler.returnFunction()
+
+				// Generate the code under test.
+				code, _, err := compiler.generate()
+				require.NoError(t, err)
+				// Run code.
+				jitcall(
+					uintptr(unsafe.Pointer(&code[0])),
+					uintptr(unsafe.Pointer(eng)),
+					0,
+				)
+				fmt.Println(eng.stack[:3])
+
+				// Check the stack.
+				require.Equal(t, uint64(2), eng.stackPointer)
+				require.Equal(t, uint64(x1Value*x2Value), eng.stack[eng.stackPointer-1])
+				// Verify the previous DX register value was saved to the stack.
+				require.Equal(t, dxValue, eng.stack[eng.stackPointer-2])
+			})
+		}
+	})
+	t.Run("int64", func(t *testing.T) {
+		for _, tc := range []struct {
+			name string
+			// Interpret -1 as stack.
+			x1Reg, x2Reg int16
+		}{
+			{
+				name:  "x1:rax,x2:random_reg",
+				x1Reg: x86.REG_AX,
+				x2Reg: x86.REG_R10,
+			},
+			{
+				name:  "x1:rax,x2:stack",
+				x1Reg: x86.REG_AX,
+				x2Reg: -1,
+			},
+			{
+				name:  "x1:random_reg,x2:rax",
+				x1Reg: x86.REG_R10,
+				x2Reg: x86.REG_AX,
+			},
+			{
+				name:  "x1:stack,x2:rax",
+				x1Reg: -1,
+				x2Reg: x86.REG_AX,
+			},
+			{
+				name:  "x1:random_reg,x2:random_reg",
+				x1Reg: x86.REG_R10,
+				x2Reg: x86.REG_R9,
+			},
+			{
+				name:  "x1:stack,x2:random_reg",
+				x1Reg: -1,
+				x2Reg: x86.REG_R9,
+			},
+			{
+				name:  "x1:random_reg,x2:stack",
+				x1Reg: x86.REG_R9,
+				x2Reg: -1,
+			},
+			{
+				name:  "x1:stack,x2:stack",
+				x1Reg: -1,
+				x2Reg: -1,
+			},
+		} {
+			tc := tc
+			t.Run(tc.name, func(t *testing.T) {
+				const x1Value uint64 = 1 << 35
+				const x2Value uint64 = 51
+				const dxValue uint64 = 111111
+
+				eng := newEngine()
+				compiler := requireNewCompiler(t)
+				compiler.initializeReservedRegisters()
+
+				// Pretend there was an existing value on the DX register. We expect compileMul to save this to the stack.
+				compiler.movIntConstToRegister(int64(dxValue), x86.REG_DX)
+				compiler.locationStack.pushValueOnRegister(x86.REG_DX)
+
+				// Setup values.
+				if tc.x1Reg != -1 {
+					compiler.movIntConstToRegister(int64(x1Value), tc.x1Reg)
+					compiler.locationStack.pushValueOnRegister(tc.x1Reg)
+				} else {
+					loc := compiler.locationStack.pushValueOnStack()
+					eng.stack[loc.stackPointer] = uint64(x1Value)
+				}
+				if tc.x2Reg != -1 {
+					compiler.movIntConstToRegister(int64(x2Value), tc.x2Reg)
+					compiler.locationStack.pushValueOnRegister(tc.x2Reg)
+				} else {
+					loc := compiler.locationStack.pushValueOnStack()
+					eng.stack[loc.stackPointer] = uint64(x2Value)
+				}
+
+				err := compiler.compileMul(&wazeroir.OperationMul{Type: wazeroir.UnsignedTypeI64})
+				require.NoError(t, err)
+				require.Equal(t, int16(x86.REG_AX), compiler.locationStack.peek().register)
+				require.Equal(t, uint64(2), compiler.locationStack.sp)
+				require.Len(t, compiler.locationStack.usedRegisters, 1)
+
+				// To verify the behavior, we push the value
+				// to the stack.
+				compiler.releaseAllRegistersToStack()
+				compiler.returnFunction()
+
+				// Generate the code under test.
+				code, _, err := compiler.generate()
+				require.NoError(t, err)
+				// Run code.
+				jitcall(
+					uintptr(unsafe.Pointer(&code[0])),
+					uintptr(unsafe.Pointer(eng)),
+					0,
+				)
+				fmt.Println(eng.stack[:3])
+
+				// Check the stack.
+				require.Equal(t, uint64(2), eng.stackPointer)
+				require.Equal(t, uint64(x1Value*x2Value), eng.stack[eng.stackPointer-1])
+				// Verify the previous DX register value was saved to the stack.
+				require.Equal(t, dxValue, eng.stack[eng.stackPointer-2])
+			})
+		}
+	})
+	t.Run("float32", func(t *testing.T) {
+		for i, tc := range []struct {
+			x1, x2 float32
+		}{
+			{x1: 100, x2: -1.1},
+			{x1: -1, x2: 100},
+			{x1: 100, x2: 200},
+			{x1: 100.01234124, x2: 100.01234124},
+			{x1: 100.01234124, x2: -100.01234124},
+			{x1: 200.12315, x2: 100},
+			{x1: float32(math.Inf(1)), x2: 100},
+			{x1: float32(math.Inf(-1)), x2: 100},
+		} {
+			tc := tc
+			t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+				compiler := requireNewCompiler(t)
+				compiler.initializeReservedRegisters()
+				err := compiler.compileConstF32(&wazeroir.OperationConstF32{Value: tc.x1})
+				require.NoError(t, err)
+				x1 := compiler.locationStack.peek()
+				err = compiler.compileConstF32(&wazeroir.OperationConstF32{Value: tc.x2})
+				require.NoError(t, err)
+				x2 := compiler.locationStack.peek()
+
+				err = compiler.compileMul(&wazeroir.OperationMul{Type: wazeroir.UnsignedTypeF32})
+				require.NoError(t, err)
+				require.Contains(t, compiler.locationStack.usedRegisters, x1.register)
+				require.NotContains(t, compiler.locationStack.usedRegisters, x2.register)
+
+				// To verify the behavior, we push the value
+				// to the stack.
+				compiler.releaseRegisterToStack(x1)
+				compiler.returnFunction()
+
+				// Generate the code under test.
+				code, _, err := compiler.generate()
+				require.NoError(t, err)
+				// Run code.
+				eng := newEngine()
+				mem := newMemoryInst()
+				jitcall(
+					uintptr(unsafe.Pointer(&code[0])),
+					uintptr(unsafe.Pointer(eng)),
+					uintptr(unsafe.Pointer(&mem.Buffer[0])),
+				)
+				// Check the stack.
+				require.Equal(t, uint64(1), eng.stackPointer)
+				require.Equal(t, tc.x1*tc.x2, math.Float32frombits(uint32(eng.stack[eng.stackPointer-1])))
+			})
+		}
+	})
+	t.Run("float64", func(t *testing.T) {
+		for i, tc := range []struct {
+			x1, x2 float64
+		}{
+			{x1: 100, x2: -1.1},
+			{x1: -1, x2: 100},
+			{x1: 100, x2: 200},
+			{x1: 100.01234124, x2: 100.01234124},
+			{x1: 100.01234124, x2: -100.01234124},
+			{x1: 200.12315, x2: 100},
+			{x1: 6.8719476736e+10 /* = 1 << 36 */, x2: 100},
+			{x1: 6.8719476736e+10 /* = 1 << 36 */, x2: 1.37438953472e+11 /* = 1 << 37*/},
+			{x1: math.Inf(1), x2: 100},
+			{x1: math.Inf(-1), x2: 100},
+		} {
+			tc := tc
+			t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+				compiler := requireNewCompiler(t)
+				compiler.initializeReservedRegisters()
+				err := compiler.compileConstF64(&wazeroir.OperationConstF64{Value: tc.x1})
+				require.NoError(t, err)
+				x1 := compiler.locationStack.peek()
+				err = compiler.compileConstF64(&wazeroir.OperationConstF64{Value: tc.x2})
+				require.NoError(t, err)
+				x2 := compiler.locationStack.peek()
+
+				err = compiler.compileMul(&wazeroir.OperationMul{Type: wazeroir.UnsignedTypeF64})
+				require.NoError(t, err)
+				require.Contains(t, compiler.locationStack.usedRegisters, x1.register)
+				require.NotContains(t, compiler.locationStack.usedRegisters, x2.register)
+
+				// To verify the behavior, we push the value
+				// to the stack.
+				compiler.releaseRegisterToStack(x1)
+				compiler.returnFunction()
+
+				// Generate the code under test.
+				code, _, err := compiler.generate()
+				require.NoError(t, err)
+				// Run code.
+				eng := newEngine()
+				mem := newMemoryInst()
+				jitcall(
+					uintptr(unsafe.Pointer(&code[0])),
+					uintptr(unsafe.Pointer(eng)),
+					uintptr(unsafe.Pointer(&mem.Buffer[0])),
+				)
+				// Check the stack.
+				require.Equal(t, uint64(1), eng.stackPointer)
+				require.Equal(t, tc.x1*tc.x2, math.Float64frombits(eng.stack[eng.stackPointer-1]))
+			})
+		}
+	})
+}
+
 func TestAmd64Compiler_compileCall(t *testing.T) {
 	t.Run("host function", func(t *testing.T) {
 		const functionIndex = 5