Repository: tetratelabs/wazero
Pull request #1951 (2024-01-24 17:42:06)

wazevo: introduce stack adjustment in stack growth


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index 807b1dc0..0138c224 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -998,11 +998,11 @@ func (u unaryRmROpcode) String() string {
 type shiftROp byte
 
 const (
-	shiftROpRotateLeft           = 0
-	shiftROpRotateRight          = 1
-	shiftROpShiftLeft            = 4
-	shiftROpShiftRightLogical    = 5
-	shiftROpShiftRightArithmetic = 7
+	shiftROpRotateLeft           shiftROp = 0
+	shiftROpRotateRight          shiftROp = 1
+	shiftROpShiftLeft            shiftROp = 4
+	shiftROpShiftRightLogical    shiftROp = 5
+	shiftROpShiftRightArithmetic shiftROp = 7
 )
 
 func (s shiftROp) String() string {
@@ -1535,6 +1535,7 @@ var defKinds = [instrMax]defKind{
 	movRM:        defKindNone,
 	xmmMovRM:     defKindNone,
 	aluRmiR:      defKindNone,
+	shiftR:       defKindNone,
 	imm:          defKindOp2,
 	xmmUnaryRmR:  defKindOp2,
 	mov64MR:      defKindOp2,
@@ -1582,6 +1583,7 @@ var useKinds = [instrMax]useKind{
 	movRM:        useKindOp1RegOp2,
 	xmmMovRM:     useKindOp1RegOp2,
 	aluRmiR:      useKindOp1Op2Reg,
+	shiftR:       useKindOp1Op2Reg,
 	imm:          useKindNone,
 	xmmUnaryRmR:  useKindOp1,
 	mov64MR:      useKindOp1,
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index c6fbb65b..b758ae7d 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -206,6 +206,26 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 		m.lowerAluRmiROp(instr, aluRmiROpcodeAdd)
 	case ssa.OpcodeIsub:
 		m.lowerAluRmiROp(instr, aluRmiROpcodeSub)
+	case ssa.OpcodeImul:
+		m.lowerAluRmiROp(instr, aluRmiROpcodeMul)
+	case ssa.OpcodeBand:
+		m.lowerAluRmiROp(instr, aluRmiROpcodeAnd)
+	case ssa.OpcodeBor:
+		m.lowerAluRmiROp(instr, aluRmiROpcodeOr)
+	case ssa.OpcodeBxor:
+		m.lowerAluRmiROp(instr, aluRmiROpcodeXor)
+	case ssa.OpcodeIshl:
+		m.lowerShiftR(instr, shiftROpShiftLeft)
+	case ssa.OpcodeSshr:
+		m.lowerShiftR(instr, shiftROpShiftRightArithmetic)
+	case ssa.OpcodeRotl:
+		m.lowerShiftR(instr, shiftROpRotateLeft)
+	case ssa.OpcodeRotr:
+		m.lowerShiftR(instr, shiftROpRotateRight)
+	case ssa.OpcodeUshr:
+		m.lowerShiftR(instr, shiftROpShiftRightLogical)
+	case ssa.OpcodeUndefined:
+		m.insert(m.allocateInstr().asUD2())
 	case ssa.OpcodeExitWithCode:
 		execCtx, code := instr.ExitWithCodeData()
 		m.lowerExitWithCode(m.c.VRegOf(execCtx), code)
@@ -277,7 +297,7 @@ func (m *machine) lowerAluRmiROp(si *ssa.Instruction, op aluRmiROpcode) {
 	rd := m.c.VRegOf(si.Return())
 	tmp := m.c.AllocateVReg(si.Return().Type())
 
-	// rm is being overwritten, so we first copy its value to a temp register,
+	// rn is being overwritten, so we first copy its value to a temp register,
 	// in case it is referenced again later.
 	mov := m.allocateInstr()
 	mov.asMovRR(rn.r, tmp, _64)
@@ -293,6 +313,51 @@ func (m *machine) lowerAluRmiROp(si *ssa.Instruction, op aluRmiROpcode) {
 	m.insert(mov2)
 }
 
+func (m *machine) lowerShiftR(si *ssa.Instruction, op shiftROp) {
+	x, amt := si.Arg2()
+	if !x.Type().IsInt() {
+		panic("BUG?")
+	}
+	_64 := x.Type().Bits() == 64
+
+	xDef, amtDef := m.c.ValueDefinition(x), m.c.ValueDefinition(amt)
+
+	opAmt := m.getOperand_Imm32_Reg(amtDef)
+	rx := m.getOperand_Reg(xDef)
+	rd := m.c.VRegOf(si.Return())
+	tmpDst := m.c.AllocateVReg(si.Return().Type())
+
+	// rx is being overwritten, so we first copy its value to a temp register,
+	// in case it is referenced again later.
+	mov := m.allocateInstr()
+	mov.asMovRR(rx.r, tmpDst, _64)
+	m.insert(mov)
+
+	if opAmt.r != regalloc.VRegInvalid {
+		// If opAmt is a register we must copy its value to rcx,
+		// because shiftR encoding mandates that the shift amount is in rcx.
+		mov := m.allocateInstr()
+		mov.asMovRR(opAmt.r, rcxVReg, _64)
+		m.insert(mov)
+
+		alu := m.allocateInstr()
+		alu.asShiftR(op, newOperandReg(rcxVReg), tmpDst, _64)
+		m.insert(alu)
+
+	} else {
+		println("::::::SONO QUI")
+
+		alu := m.allocateInstr()
+		alu.asShiftR(op, opAmt, tmpDst, _64)
+		m.insert(alu)
+	}
+
+	// tmp now contains the result, we copy it to the dest register.
+	mov2 := m.allocateInstr()
+	mov2.asMovRR(tmpDst, rd, _64)
+	m.insert(mov2)
+}
+
 func (m *machine) lowerStore(si *ssa.Instruction) {
 	value, ptr, offset, storeSizeInBits := si.StoreData()
 	rm := m.c.VRegOf(value)
diff --git a/internal/engine/wazevo/backend/isa/amd64/stack.go b/internal/engine/wazevo/backend/isa/amd64/stack.go
new file mode 100644
index 00000000..a695f05b
--- /dev/null
+++ b/internal/engine/wazevo/backend/isa/amd64/stack.go
@@ -0,0 +1,100 @@
+package amd64
+
+import (
+	"encoding/binary"
+	"reflect"
+	"unsafe"
+)
+
+func stackView(rbp, top uintptr) []byte {
+	l := int(top - rbp)
+	var stackBuf []byte
+	{
+		// TODO: use unsafe.Slice after floor version is set to Go 1.20.
+		hdr := (*reflect.SliceHeader)(unsafe.Pointer(&stackBuf))
+		hdr.Data = rbp
+		hdr.Len = l
+		hdr.Cap = l
+	}
+	return stackBuf
+}
+
+// UnwindStack implements wazevo.unwindStack.
+func UnwindStack(_, rbp, top uintptr, returnAddresses []uintptr) []uintptr {
+	stackBuf := stackView(rbp, top)
+
+	for i := uint64(0); i < uint64(len(stackBuf)); {
+		//       (high address)
+		//    +-----------------+
+		//    |     .......     |
+		//    |      ret Y      |
+		//    |     .......     |
+		//    |      ret 0      |
+		//    |      arg X      |
+		//    |     .......     |
+		//    |      arg 1      |
+		//    |      arg 0      |
+		//    |  ReturnAddress  |
+		//    |   Caller_RBP    |
+		//    +-----------------+ <---- Caller_RBP
+		//    |   ...........   |
+		//    |   clobbered  M  |
+		//    |   ............  |
+		//    |   clobbered  0  |
+		//    |   spill slot N  |
+		//    |   ............  |
+		//    |   spill slot 0  |
+		//    |  ReturnAddress  |
+		//    |   Caller_RBP    |
+		//    +-----------------+ <---- RBP
+		//       (low address)
+
+		callerRBP := binary.LittleEndian.Uint64(stackBuf[i:])
+		retAddr := binary.LittleEndian.Uint64(stackBuf[i+8:])
+		returnAddresses = append(returnAddresses, uintptr(retAddr))
+		i = callerRBP - uint64(rbp)
+	}
+	return returnAddresses
+}
+
+// GoCallStackView implements wazevo.goCallStackView.
+func GoCallStackView(stackPointerBeforeGoCall *uint64) []uint64 {
+	panic("implement me")
+}
+
+func AdjustStackAfterGrown(oldRsp, rsp, rbp, top uintptr) {
+	diff := uint64(rsp - oldRsp)
+
+	stackBuf := stackView(rbp, top)
+	for i := uint64(0); i < uint64(len(stackBuf)); {
+		//       (high address)
+		//    +-----------------+
+		//    |     .......     |
+		//    |      ret Y      |
+		//    |     .......     |
+		//    |      ret 0      |
+		//    |      arg X      |
+		//    |     .......     |
+		//    |      arg 1      |
+		//    |      arg 0      |
+		//    |  ReturnAddress  |
+		//    |   Caller_RBP    |
+		//    +-----------------+ <---- Caller_RBP
+		//    |   ...........   |
+		//    |   clobbered  M  |
+		//    |   ............  |
+		//    |   clobbered  0  |
+		//    |   spill slot N  |
+		//    |   ............  |
+		//    |   spill slot 0  |
+		//    |  ReturnAddress  |
+		//    |   Caller_RBP    |
+		//    +-----------------+ <---- RBP
+		//       (low address)
+
+		callerRBP := binary.LittleEndian.Uint64(stackBuf[i:])
+		adjustedCallerRBP := callerRBP + diff
+		binary.LittleEndian.PutUint64(stackBuf[i:], adjustedCallerRBP)
+		i = adjustedCallerRBP - uint64(rbp)
+	}
+}
diff --git a/internal/engine/wazevo/backend/isa/amd64/unwind_stack.go b/internal/engine/wazevo/backend/isa/amd64/unwind_stack.go
deleted file mode 100644
index 727cfeb5..00000000
--- a/internal/engine/wazevo/backend/isa/amd64/unwind_stack.go
+++ /dev/null
@@ -1,59 +0,0 @@
-package amd64
-
-import (
-	"encoding/binary"
-	"reflect"
-	"unsafe"
-)
-
-func UnwindStack(_, fp, top uintptr, returnAddresses []uintptr) []uintptr {
-	l := int(top - fp)
-
-	var stackBuf []byte
-	{
-		// TODO: use unsafe.Slice after floor version is set to Go 1.20.
-		hdr := (*reflect.SliceHeader)(unsafe.Pointer(&stackBuf))
-		hdr.Data = fp
-		hdr.Len = l
-		hdr.Cap = l
-	}
-
-	for i := uint64(0); i < uint64(l); {
-		//       (high address)
-		//    +-----------------+
-		//    |     .......     |
-		//    |      ret Y      |
-		//    |     .......     |
-		//    |      ret 0      |
-		//    |      arg X      |
-		//    |     .......     |
-		//    |      arg 1      |
-		//    |      arg 0      |
-		//    |  ReturnAddress  |
-		//    |   Caller_RBP    |
-		//    +-----------------+ <---- Caller_RBP
-		//    |   ...........   |
-		//    |   clobbered  M  |
-		//    |   ............  |
-		//    |   clobbered  0  |
-		//    |   spill slot N  |
-		//    |   ............  |
-		//    |   spill slot 0  |
-		//    |  ReturnAddress  |
-		//    |   Caller_RBP    |
-		//    +-----------------+ <---- RBP
-		//       (low address)
-
-		callerRBP := binary.LittleEndian.Uint64(stackBuf[i:])
-		retAddr := binary.LittleEndian.Uint64(stackBuf[i+8:])
-		returnAddresses = append(returnAddresses, uintptr(retAddr))
-		i = callerRBP - uint64(fp)
-	}
-	return returnAddresses
-}
-
-// GoCallStackView is a function to get a view of the stack before a Go call, which
-// is the view of the stack allocated in CompileGoFunctionTrampoline.
-func GoCallStackView(stackPointerBeforeGoCall *uint64) []uint64 {
-	panic("implement me")
-}
diff --git a/internal/engine/wazevo/backend/isa/arm64/unwind_stack.go b/internal/engine/wazevo/backend/isa/arm64/unwind_stack.go
index f3e25f2a..adb44b26 100644
--- a/internal/engine/wazevo/backend/isa/arm64/unwind_stack.go
+++ b/internal/engine/wazevo/backend/isa/arm64/unwind_stack.go
@@ -6,8 +6,7 @@ import (
 	"unsafe"
 )
 
-// UnwindStack is a function to unwind the stack, and appends return addresses to `returnAddresses` slice.
-// The implementation must be aligned with the ABI/Calling convention as in machine_pro_epi_logue.go/abi.go.
+// UnwindStack implements wazevo.unwindStack.
 func UnwindStack(sp, _, top uintptr, returnAddresses []uintptr) []uintptr {
 	l := int(top - sp)
 
@@ -60,8 +59,7 @@ func UnwindStack(sp, _, top uintptr, returnAddresses []uintptr) []uintptr {
 	return returnAddresses
 }
 
-// GoCallStackView is a function to get a view of the stack before a Go call, which
-// is the view of the stack allocated in CompileGoFunctionTrampoline.
+// GoCallStackView implements wazevo.goCallStackView.
 func GoCallStackView(stackPointerBeforeGoCall *uint64) []uint64 {
 	//                  (high address)
 	//              +-----------------+ <----+
diff --git a/internal/engine/wazevo/call_engine.go b/internal/engine/wazevo/call_engine.go
index 3dbee422..475dbf5b 100644
--- a/internal/engine/wazevo/call_engine.go
+++ b/internal/engine/wazevo/call_engine.go
@@ -244,15 +244,17 @@ func (c *callEngine) callWithStack(ctx context.Context, paramResultStack []uint6
 		case wazevoapi.ExitCodeOK:
 			return nil
 		case wazevoapi.ExitCodeGrowStack:
-			var newsp uintptr
+			oldsp := uintptr(unsafe.Pointer(c.execCtx.stackPointerBeforeGoCall))
+			var newsp, newfp uintptr
 			if wazevoapi.StackGuardCheckEnabled {
-				newsp, err = c.growStackWithGuarded()
+				newsp, newfp, err = c.growStackWithGuarded()
 			} else {
-				newsp, err = c.growStack()
+				newsp, newfp, err = c.growStack()
 			}
 			if err != nil {
 				return err
 			}
+			adjustStackAfterGrown(oldsp, oldsp, newsp, newfp)
 			c.execCtx.exitCode = wazevoapi.ExitCodeOK
 			afterGoFunctionCallEntrypoint(c.execCtx.goCallReturnAddress, c.execCtxPtr, newsp)
 		case wazevoapi.ExitCodeGrowMemory:
@@ -407,11 +409,11 @@ func opaqueViewFromPtr(ptr uintptr) []byte {
 
 const callStackCeiling = uintptr(50000000) // in uint64 (8 bytes) == 400000000 bytes in total == 400mb.
 
-func (c *callEngine) growStackWithGuarded() (newSP uintptr, err error) {
+func (c *callEngine) growStackWithGuarded() (newSP uintptr, newFP uintptr, err error) {
 	if wazevoapi.StackGuardCheckEnabled {
 		wazevoapi.CheckStackGuardPage(c.stack)
 	}
-	newSP, err = c.growStack()
+	newSP, newFP, err = c.growStack()
 	if err != nil {
 		return
 	}
@@ -422,7 +424,7 @@ func (c *callEngine) growStackWithGuarded() (newSP uintptr, err error) {
 }
 
 // growStack grows the stack, and returns the new stack pointer.
-func (c *callEngine) growStack() (newSP uintptr, err error) {
+func (c *callEngine) growStack() (newSP, newFP uintptr, err error) {
 	currentLen := uintptr(len(c.stack))
 	if callStackCeiling < currentLen {
 		err = wasmruntime.ErrRuntimeStackOverflow
@@ -433,6 +435,7 @@ func (c *callEngine) growStack() (newSP uintptr, err error) {
 	newStack := make([]byte, newLen)
 
 	relSp := c.stackTop - uintptr(unsafe.Pointer(c.execCtx.stackPointerBeforeGoCall))
+	relFp := c.stackTop - c.execCtx.framePointerBeforeGoCall
 
 	// Copy the existing contents in the previous Go-allocated stack into the new one.
 	var prevStackAligned, newStackAligned []byte
@@ -445,6 +448,7 @@ func (c *callEngine) growStack() (newSP uintptr, err error) {
 	newTop := alignedStackTop(newStack)
 	{
 		newSP = newTop - relSp
+		newFP = newTop - relFp
 		sh := (*reflect.SliceHeader)(unsafe.Pointer(&newStackAligned))
 		sh.Data = newSP
 		sh.Len = int(relSp)
diff --git a/internal/engine/wazevo/isa.go b/internal/engine/wazevo/isa.go
new file mode 100644
index 00000000..5f427f2c
--- /dev/null
+++ b/internal/engine/wazevo/isa.go
@@ -0,0 +1,61 @@
+package wazevo
+
+import (
+	"runtime"
+
+	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
+	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/isa/amd64"
+	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/isa/arm64"
+)
+
+func newMachine() backend.Machine {
+	switch runtime.GOARCH {
+	case "arm64":
+		return arm64.NewBackend()
+	case "amd64":
+		return amd64.NewBackend()
+	default:
+		panic("unsupported architecture")
+	}
+}
+
+// unwindStack is a function to unwind the stack, and appends return addresses to `returnAddresses` slice.
+// The implementation must be aligned with the ABI/Calling convention.
+func unwindStack(sp, fp, top uintptr, returnAddresses []uintptr) []uintptr {
+	switch runtime.GOARCH {
+	case "arm64":
+		return arm64.UnwindStack(sp, fp, top, returnAddresses)
+	case "amd64":
+		return amd64.UnwindStack(sp, fp, top, returnAddresses)
+	default:
+		panic("unsupported architecture")
+	}
+}
+
+// goCallStackView is a function to get a view of the stack before a Go call, which
+// is the view of the stack allocated in CompileGoFunctionTrampoline.
+func goCallStackView(stackPointerBeforeGoCall *uint64) []uint64 {
+	switch runtime.GOARCH {
+	case "arm64":
+		return arm64.GoCallStackView(stackPointerBeforeGoCall)
+	case "amd64":
+		return amd64.GoCallStackView(stackPointerBeforeGoCall)
+	default:
+		panic("unsupported architecture")
+	}
+}
+
+// adjustStackAfterGrown is a function to adjust the stack after it is grown.
+// More precisely, absolute addresses (frame pointers) in the stack must be adjusted.
+func adjustStackAfterGrown(oldsp, sp, fp, top uintptr) {
+	switch runtime.GOARCH {
+	case "arm64":
+		// TODO: currently, the frame pointers are not used, and saved old sps are relative to the current stack pointer,
+		//  so no need to adjustment on arm64. However, when we make it absolute, which in my opinion is better perf-wise
+		//  at the expense of slightly costly stack growth, we need to adjust the pushed frame pointers.
+	case "amd64":
+		amd64.AdjustStackAfterGrown(oldsp, sp, fp, top)
+	default:
+		panic("unsupported architecture")
+	}
+}
diff --git a/internal/engine/wazevo/machine.go b/internal/engine/wazevo/machine.go
deleted file mode 100644
index 5319e601..00000000
--- a/internal/engine/wazevo/machine.go
+++ /dev/null
@@ -1,42 +0,0 @@
-package wazevo
-
-import (
-	"runtime"
-
-	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
-	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/isa/amd64"
-	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/isa/arm64"
-)
-
-func newMachine() backend.Machine {
-	switch runtime.GOARCH {
-	case "arm64":
-		return arm64.NewBackend()
-	case "amd64":
-		return amd64.NewBackend()
-	default:
-		panic("unsupported architecture")
-	}
-}
-
-func unwindStack(sp, fp, top uintptr, returnAddresses []uintptr) []uintptr {
-	switch runtime.GOARCH {
-	case "arm64":
-		return arm64.UnwindStack(sp, fp, top, returnAddresses)
-	case "amd64":
-		return amd64.UnwindStack(sp, fp, top, returnAddresses)
-	default:
-		panic("unsupported architecture")
-	}
-}
-
-func goCallStackView(stackPointerBeforeGoCall *uint64) []uint64 {
-	switch runtime.GOARCH {
-	case "arm64":
-		return arm64.GoCallStackView(stackPointerBeforeGoCall)
-	case "amd64":
-		return amd64.GoCallStackView(stackPointerBeforeGoCall)
-	default:
-		panic("unsupported architecture")
-	}
-}


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/unwind_stack_test.go b/internal/engine/wazevo/backend/isa/amd64/stack_test.go
similarity index 53%
rename from internal/engine/wazevo/backend/isa/amd64/unwind_stack_test.go
rename to internal/engine/wazevo/backend/isa/amd64/stack_test.go
index b7442a36..25035029 100644
--- a/internal/engine/wazevo/backend/isa/amd64/unwind_stack_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/stack_test.go
@@ -38,3 +38,32 @@ func TestUnwindStack(t *testing.T) {
 		})
 	}
 }
+
+func addressOf(v *byte) uint64 {
+	return uint64(uintptr(unsafe.Pointer(v)))
+}
+
+func TestAdjustStackAfterGrown(t *testing.T) {
+	// In order to allocate slices on Go heap, we need to allocSlice function.
+	allocSlice := func(size int) []byte {
+		return make([]byte, size)
+	}
+
+	oldStack := allocSlice(512)
+	oldRsp := uintptr(unsafe.Pointer(&oldStack[0]))
+	rbpIndex := uintptr(32)
+	binary.LittleEndian.PutUint64(oldStack[rbpIndex:], addressOf(&oldStack[16+rbpIndex]))
+	binary.LittleEndian.PutUint64(oldStack[rbpIndex+16:], addressOf(&oldStack[32+rbpIndex]))
+	binary.LittleEndian.PutUint64(oldStack[rbpIndex+32:], addressOf(&oldStack[160+rbpIndex]))
+
+	newStack := allocSlice(1024)
+	rsp := uintptr(unsafe.Pointer(&newStack[0]))
+	rbp := rsp + rbpIndex
+	// Coy old stack to new stack which contains the old pointers to the old stack elements.
+	copy(newStack, oldStack)
+
+	AdjustStackAfterGrown(oldRsp, rsp, rbp, uintptr(addressOf(&newStack[len(newStack)-1])))
+	require.Equal(t, addressOf(&newStack[rbpIndex+16]), binary.LittleEndian.Uint64(newStack[rbpIndex:]))
+	require.Equal(t, addressOf(&newStack[rbpIndex+32]), binary.LittleEndian.Uint64(newStack[rbpIndex+16:]))
+	require.Equal(t, addressOf(&newStack[rbpIndex+160]), binary.LittleEndian.Uint64(newStack[rbpIndex+32:]))
+}
diff --git a/internal/engine/wazevo/call_engine_test.go b/internal/engine/wazevo/call_engine_test.go
index b1fd19b7..0a3e2438 100644
--- a/internal/engine/wazevo/call_engine_test.go
+++ b/internal/engine/wazevo/call_engine_test.go
@@ -18,7 +18,7 @@ func TestCallEngine_init(t *testing.T) {
 func TestCallEngine_growStack(t *testing.T) {
 	t.Run("stack overflow", func(t *testing.T) {
 		c := &callEngine{stack: make([]byte, callStackCeiling+1)}
-		_, err := c.growStack()
+		_, _, err := c.growStack()
 		require.Error(t, err)
 	})
 
@@ -33,9 +33,10 @@ func TestCallEngine_growStack(t *testing.T) {
 			execCtx: executionContext{
 				stackGrowRequiredSize:    160,
 				stackPointerBeforeGoCall: (*uint64)(unsafe.Pointer(&s[10])),
+				framePointerBeforeGoCall: uintptr(unsafe.Pointer(&s[14])),
 			},
 		}
-		newSP, err := c.growStack()
+		newSP, newFp, err := c.growStack()
 		require.NoError(t, err)
 		require.Equal(t, 160+32*2+16, len(c.stack))
 
@@ -52,6 +53,7 @@ func TestCallEngine_growStack(t *testing.T) {
 		require.Equal(t, []byte{10, 11, 12, 13, 14}, view)
 		require.True(t, newSP >= uintptr(unsafe.Pointer(c.execCtx.stackBottomPtr)))
 		require.True(t, newSP <= c.stackTop)
+		require.Equal(t, newFp-newSP, uintptr(4))
 	})
 }
 
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index 965d6a68..b1b39ff5 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -102,6 +102,27 @@ func TestE2E(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "arithm_return", m: testcases.ArithmReturn.Module,
+			calls: []callCase{
+				{
+					params: []uint64{
+						21, 10, 0xf0,
+						21, 10, 0xf0,
+					},
+					expResults: []uint64{
+						21 * 10, 21 & 10, 21 | 10, 21 ^ 10,
+						21 << 10, 21 >> 10, 0xf0 >> 10,
+						88080384, 20971520,
+
+						21 * 10, 21 & 10, 21 | 10, 21 ^ 10,
+						21 << 10, 21 >> 10, 0xf0 >> 10,
+						378302368699121664, 20971520,
+					},
+				},
+			},
+		},
+
 		{
 			name: "many_params_many_results",
 			m:    testcases.ManyParamsManyResults.Module,
diff --git a/internal/engine/wazevo/testcases/testcases.go b/internal/engine/wazevo/testcases/testcases.go
index fcc09f0d..7e54e753 100644
--- a/internal/engine/wazevo/testcases/testcases.go
+++ b/internal/engine/wazevo/testcases/testcases.go
@@ -59,7 +59,7 @@ var (
 			wasm.OpcodeI64Const, 6,
 			wasm.OpcodeI64Sub,
 
-			// Large i64 constants are load from register on arm64, amd64.
+			// Large i64 constants are loaded from register on arm64, amd64.
 			wasm.OpcodeI64Const, 3,
 			wasm.OpcodeI64Const, 0xff, 0xff, 0xff, 0xff, 0xff, 0,
 			wasm.OpcodeI64Add,
@@ -80,6 +80,89 @@ var (
 			wasm.OpcodeEnd,
 		}, nil),
 	}
+	ArithmReturn = TestCase{
+		Name: "add_sub_params_return",
+		Module: SingleFunctionModule(
+			wasm.FunctionType{
+				Params: []wasm.ValueType{i32, i32, i32, i64, i64, i64},
+				Results: []wasm.ValueType{
+					i32, i32, i32, i32,
+					i32, i32, i32,
+					i32, i32,
+
+					i64, i64, i64, i64,
+					i64, i64, i64,
+					i64, i64,
+				},
+			},
+			[]byte{
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeI32Mul,
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeI32And,
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeI32Or,
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeI32Xor,
+
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeI32Shl,
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 2,
+				wasm.OpcodeI32ShrS,
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeI32ShrU,
+
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeI32Rotr,
+				// Exercise the case where the amount is an immediate.
+				wasm.OpcodeI32Const, 10,
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeI32Rotl,
+
+				// i64
+
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64Mul,
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64And,
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64Or,
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64Xor,
+
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64Shl,
+				wasm.OpcodeLocalGet, 5,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64ShrS,
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64ShrU,
+
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64Rotr,
+				// Exercise the case where the amount is an immediate.
+				wasm.OpcodeI64Const, 10,
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeI64Rotl,
+
+				wasm.OpcodeEnd,
+			}, nil),
+	}
 	Locals       = TestCase{Name: "locals", Module: SingleFunctionModule(vv, []byte{wasm.OpcodeEnd}, []wasm.ValueType{i32, i64, f32, f64})}
 	LocalsParams = TestCase{
 		Name: "locals_params",