Repository: tetratelabs/wazero
Pull request #781 (2022-08-30 02:15:56)

compiler: hold stack base pointers in bytes, not index to the stack slice
To reduce the redundant instructions and register usage on the hot path.

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>


Patch:
diff --git a/README.md b/README.md
index 1958ca6a..3b7db38a 100644
--- a/README.md
+++ b/README.md
@@ -17,16 +17,34 @@ Import wazero and extend your Go application with code written in any language!
 
 The best way to learn wazero is by trying one of our [examples](examples).
 
-For the impatient, here's how invoking a factorial function looks in wazero:
+For the impatient, here's a peek of a general flow with wazero:
 
-```golang
+First, you need to compile your code into the WebAssembly Binary Format (Wasm).
+
+Here's source in [TinyGo](https://wazero.io/languages/tinygo), which exports an
+"add" function:
+```go
+package main
+
+//export add
+func add(x, y uint32) uint32 {
+	return x + y
+}
+```
+
+Here's the minimal command to build a `%.wasm` binary.
+```bash
+tinygo build -o add.wasm -target=wasi add.go
+```
+
+Finally, you can run that inside your Go application.
+```go
 func main() {
 	// Choose the context to use for function calls.
 	ctx := context.Background()
 
-	// Read a WebAssembly binary containing an exported "fac" function.
-	// * Ex. (func (export "fac") (param i64) (result i64) ...
-	wasm, err := os.ReadFile("./path/to/fac.wasm")
+	// Read a WebAssembly binary containing an exported "add" function.
+	wasm, err := os.ReadFile("./path/to/add.wasm")
 	if err != nil {
 		log.Panicln(err)
 	}
@@ -41,14 +59,17 @@ func main() {
 		log.Panicln(err)
 	}
 
-	// Discover 7! is 5040
-	fmt.Println(module.ExportedFunction("fac").Call(ctx, 7))
+	// Discover 1+2=3
+	fmt.Println(module.ExportedFunction("add").Call(ctx, 1, 2))
 }
 ```
 
-Note: `fac.wasm` was compiled from [fac.wat][3], in the [WebAssembly 1.0][1]
-Text Format, it could have been written in another language that compiles to
-(targets) WebAssembly, such as AssemblyScript, C, C++, Rust, TinyGo or Zig.
+Notes:
+
+* The Wasm binary is often called the "guest" in WebAssembly.
+* The embedding application is often called the "host" in WebAssembly.
+* Many languages compile to (target) Wasm including AssemblyScript, C, C++,
+  Rust, TinyGo and Zig!
 
 ## Deeper dive
 
@@ -194,7 +215,6 @@ wazero is a registered trademark of Tetrate.io, Inc. in the United States and/or
 
 [1]: https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/
 [2]: https://www.w3.org/TR/2022/WD-wasm-core-2-20220419/
-[3]: ./internal/integration_test/vs/testdata/fac.wat
 [4]: https://github.com/WebAssembly/meetings/blob/main/process/subgroups.md
 [5]: https://github.com/WebAssembly/WASI
 [6]: https://pkg.go.dev/golang.org/x/sys/unix
diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go
index 44e83f3e..c0fed4bc 100644
--- a/internal/engine/compiler/engine.go
+++ b/internal/engine/compiler/engine.go
@@ -79,9 +79,9 @@ type (
 		// valueStackElement0Address is &engine.valueStack[0] as uintptr.
 		// Note: this is updated when growing the stack in builtinFunctionGrowValueStack.
 		valueStackElement0Address uintptr
-		// valueStackLen is len(engine.valueStack[0]).
+		// valueStackLenInBytes is len(engine.valueStack[0]) * 8 (bytes).
 		// Note: this is updated when growing the stack in builtinFunctionGrowValueStack.
-		valueStackLen uint64
+		valueStackLenInBytes uint64
 
 		// callFrameStackElementZeroAddress is &engine.callFrameStack[0] as uintptr.
 		// Note: this is updated when growing the stack in builtinFunctionGrowCallFrameStack.
@@ -130,14 +130,14 @@ type (
 
 	// valueStackContext stores the data to access engine.valueStack.
 	valueStackContext struct {
-		// stackPointer on .valueStack field which is accessed by [stackBasePointer] + [stackPointer].
+		// stackPointer on .valueStack field which is accessed by valueStack[stackBasePointer+stackBasePointerInBytes*8].
 		//
 		// Note: stackPointer is not used in assembly since the native code knows exact position of
 		// each variable in the value stack from the info from compilation.
 		// Therefore, only updated when native code exit from the Compiler world and go back to the Go function.
 		stackPointer uint64
 
-		// stackBasePointer is updated whenever we make function calls.
+		// stackBasePointerInBytes is updated whenever we make function calls.
 		// Background: Functions might be compiled as if they use the stack from the bottom.
 		// However, in reality, they have to use it from the middle of the stack depending on
 		// when these function calls are made. So instead of accessing stack via stackPointer alone,
@@ -148,7 +148,7 @@ type (
 		// Note: This is saved on callFrameTop().returnStackBasePointer whenever making function call.
 		// Also, this is changed whenever we make function call or return from functions where we execute jump instruction.
 		// In either case, the caller of "jmp" instruction must set this field properly.
-		stackBasePointer uint64
+		stackBasePointerInBytes uint64
 	}
 
 	// exitContext will be manipulated whenever compiled native code returns into the Go function.
@@ -164,14 +164,14 @@ type (
 	// callFrame holds the information to which the caller function can return.
 	// callFrame is created for currently executed function frame as well,
 	// so some fields are not yet set when native code is currently executing it.
-	// That is, callFrameTop().returnAddress or returnStackBasePointer are not set
+	// That is, callFrameTop().returnAddress or returnStackBasePointerInBytes are not set
 	// until it makes a function call.
 	callFrame struct {
 		// Set when making function call from this function frame, or for the initial function frame to call from
 		// callEngine.execWasmFunction.
 		returnAddress uintptr
 		// Set when making function call from this function frame.
-		returnStackBasePointer uint64
+		returnStackBasePointerInBytes uint64
 		// Set when making function call to this function frame.
 		function *function
 		// _ is a necessary padding to make the size of callFrame struct a power of 2.
@@ -230,7 +230,7 @@ const (
 
 	// Offsets for callEngine globalContext.
 	callEngineGlobalContextValueStackElement0AddressOffset     = 0
-	callEngineGlobalContextValueStackLenOffset                 = 8
+	callEngineGlobalContextValueStackLenInBytesOffset          = 8
 	callEngineGlobalContextCallFrameStackElement0AddressOffset = 16
 	callEngineGlobalContextCallFrameStackLenOffset             = 24
 	callEngineGlobalContextCallFrameStackPointerOffset         = 32
@@ -247,19 +247,19 @@ const (
 	callEngineModuleContextElementInstancesElement0AddressOffset = 104
 
 	// Offsets for callEngine valueStackContext.
-	callEngineValueStackContextStackPointerOffset     = 112
-	callEngineValueStackContextStackBasePointerOffset = 120
+	callEngineValueStackContextStackPointerOffset            = 112
+	callEngineValueStackContextStackBasePointerInBytesOffset = 120
 
 	// Offsets for callEngine exitContext.
 	callEngineExitContextNativeCallStatusCodeOffset       = 128
 	callEngineExitContextBuiltinFunctionCallAddressOffset = 132
 
 	// Offsets for callFrame.
-	callFrameDataSize                      = 32
-	callFrameDataSizeMostSignificantSetBit = 5
-	callFrameReturnAddressOffset           = 0
-	callFrameReturnStackBasePointerOffset  = 8
-	callFrameFunctionOffset                = 16
+	callFrameDataSize                            = 32
+	callFrameDataSizeMostSignificantSetBit       = 5
+	callFrameReturnAddressOffset                 = 0
+	callFrameReturnStackBasePointerInBytesOffset = 8
+	callFrameFunctionOffset                      = 16
 
 	// Offsets for function.
 	functionCodeInitialAddressOffset    = 0
@@ -385,7 +385,7 @@ func (s nativeCallStatusCode) String() (ret string) {
 func (c *callFrame) String() string {
 	return fmt.Sprintf(
 		"[%s: return address=0x%x, return stack base pointer=%d]",
-		c.function.source.FunctionDefinition.DebugName(), c.returnAddress, c.returnStackBasePointer,
+		c.function.source.FunctionDefinition.DebugName(), c.returnAddress, c.returnStackBasePointerInBytes>>3,
 	)
 }
 
@@ -652,7 +652,7 @@ func (e *moduleEngine) newCallEngine(source *wasm.FunctionInstance, compiled *fu
 	callFrameStackHeader := (*reflect.SliceHeader)(unsafe.Pointer(&ce.callFrameStack))
 	ce.globalContext = globalContext{
 		valueStackElement0Address:        valueStackHeader.Data,
-		valueStackLen:                    uint64(valueStackHeader.Len),
+		valueStackLenInBytes:             uint64(valueStackHeader.Len) << 3,
 		callFrameStackElementZeroAddress: callFrameStackHeader.Data,
 		callFrameStackLen:                uint64(callFrameStackHeader.Len),
 		callFrameStackPointer:            0,
@@ -681,7 +681,7 @@ func (ce *callEngine) callFrameAt(depth uint64) *callFrame {
 }
 
 func (ce *callEngine) valueStackTopIndex() uint64 {
-	return ce.valueStackContext.stackBasePointer + ce.valueStackContext.stackPointer
+	return ce.valueStackContext.stackPointer + (ce.valueStackContext.stackBasePointerInBytes >> 3)
 }
 
 const (
@@ -703,7 +703,7 @@ entry:
 		frame := ce.callFrameTop()
 		if buildoptions.IsDebugMode {
 			fmt.Printf("callframe=%s, stackBasePointer: %d, stackPointer: %d\n",
-				frame.String(), ce.valueStackContext.stackBasePointer, ce.valueStackContext.stackPointer)
+				frame.String(), ce.valueStackContext.stackBasePointerInBytes>>3, ce.valueStackContext.stackPointer)
 		}
 
 		// Call into the native code.
@@ -758,14 +758,15 @@ entry:
 
 func (ce *callEngine) builtinFunctionGrowValueStack(stackPointerCeil uint64) {
 	// Extends the valueStack's length to currentLen*2+stackPointerCeil.
-	newLen := ce.globalContext.valueStackLen*2 + (stackPointerCeil)
+	newLen := uint64(len(ce.valueStack))<<1 + (stackPointerCeil)
 	newStack := make([]uint64, newLen)
-	top := ce.valueStackContext.stackBasePointer + ce.valueStackContext.stackPointer
+	ce.valueStackTopIndex()
+	top := ce.valueStackTopIndex()
 	copy(newStack[:top], ce.valueStack[:top])
 	ce.valueStack = newStack
 	valueStackHeader := (*reflect.SliceHeader)(unsafe.Pointer(&ce.valueStack))
 	ce.globalContext.valueStackElement0Address = valueStackHeader.Data
-	ce.globalContext.valueStackLen = uint64(valueStackHeader.Len)
+	ce.globalContext.valueStackLenInBytes = newLen << 3
 }
 
 var callStackCeiling = uint64(buildoptions.CallStackCeiling)
diff --git a/internal/engine/compiler/impl_amd64.go b/internal/engine/compiler/impl_amd64.go
index 27db2efa..9f77e1c5 100644
--- a/internal/engine/compiler/impl_amd64.go
+++ b/internal/engine/compiler/impl_amd64.go
@@ -4604,11 +4604,11 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, functionAddres
 	{
 		// We must save the current stack base pointer (which lives on callEngine.valueStackContext.stackPointer)
 		// to the call frame stack. In the example, this is equivalent to writing the value into "rb.1".
-		c.assembler.CompileMemoryToRegister(amd64.MOVQ, amd64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerOffset, tmpRegister)
+		c.assembler.CompileMemoryToRegister(amd64.MOVQ, amd64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerInBytesOffset, tmpRegister)
 
 		c.assembler.CompileRegisterToMemory(amd64.MOVQ, tmpRegister,
 			// "rb.1" is BELOW the top address. See the above example for detail.
-			callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameReturnStackBasePointerOffset),
+			callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameReturnStackBasePointerInBytesOffset),
 		)
 	}
 
@@ -4617,10 +4617,10 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, functionAddres
 		// At this point, tmpRegister holds the old stack base pointer. We could get the new frame's
 		// stack base pointer by "old stack base pointer + old stack pointer - # of function params"
 		// See the comments in callEngine.pushCallFrame which does exactly the same calculation in Go.
-		c.assembler.CompileConstToRegister(amd64.ADDQ, offset, tmpRegister)
+		c.assembler.CompileConstToRegister(amd64.ADDQ, offset<<3, tmpRegister)
 
 		// Write the calculated value to callEngine.valueStackContext.stackBasePointer.
-		c.assembler.CompileRegisterToMemory(amd64.MOVQ, tmpRegister, amd64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerOffset)
+		c.assembler.CompileRegisterToMemory(amd64.MOVQ, tmpRegister, amd64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerInBytesOffset)
 	}
 
 	// 3) Set rc.next to specify which function is executed on the current call frame (needs to make builtin function calls).
@@ -4791,11 +4791,11 @@ func (c *amd64Compiler) compileReturnFunction() error {
 	// 1) Set callEngine.valueStackContext.stackBasePointer to the value on "rb.caller"
 	c.assembler.CompileMemoryToRegister(amd64.MOVQ,
 		// "rb.caller" is BELOW the top address. See the above example for detail.
-		callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameReturnStackBasePointerOffset),
+		callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameReturnStackBasePointerInBytesOffset),
 		tmpRegister,
 	)
 	c.assembler.CompileRegisterToMemory(amd64.MOVQ,
-		tmpRegister, amd64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerOffset)
+		tmpRegister, amd64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerInBytesOffset)
 
 	// 2) Load rc.caller.moduleInstanceAddress into amd64CallingConventionModuleInstanceAddressRegister
 	c.assembler.CompileMemoryToRegister(amd64.MOVQ,
@@ -4935,10 +4935,12 @@ func (c *amd64Compiler) compileReleaseRegisterToStack(loc *runtimeValueLocation)
 }
 
 func (c *amd64Compiler) compileExitFromNativeCode(status nativeCallStatusCode) {
-	c.assembler.CompileConstToMemory(amd64.MOVB, int64(status), amd64ReservedRegisterForCallEngine, callEngineExitContextNativeCallStatusCodeOffset)
+	c.assembler.CompileConstToMemory(amd64.MOVB, int64(status),
+		amd64ReservedRegisterForCallEngine, callEngineExitContextNativeCallStatusCodeOffset)
 
 	// Write back the cached SP to the actual eng.stackPointer.
-	c.assembler.CompileConstToMemory(amd64.MOVQ, int64(c.locationStack.sp), amd64ReservedRegisterForCallEngine, callEngineValueStackContextStackPointerOffset)
+	c.assembler.CompileConstToMemory(amd64.MOVQ, int64(c.locationStack.sp),
+		amd64ReservedRegisterForCallEngine, callEngineValueStackContextStackPointerOffset)
 
 	c.assembler.CompileStandAlone(amd64.RET)
 }
@@ -4970,19 +4972,9 @@ func (c *amd64Compiler) compileReservedStackBasePointerInitialization() {
 		amd64ReservedRegisterForCallEngine, callEngineGlobalContextValueStackElement0AddressOffset,
 		amd64ReservedRegisterForStackBasePointerAddress)
 
-	// Since initializeReservedRegisters is called at the beginning of function
-	// calls (or right after they return), we have free registers at this point.
-	tmpReg, _ := c.locationStack.takeFreeRegister(registerTypeGeneralPurpose)
-
 	// next we move the base pointer (callEngine.stackBasePointer) to the tmp register.
-	c.assembler.CompileMemoryToRegister(amd64.MOVQ,
-		amd64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerOffset,
-		tmpReg,
-	)
-
-	c.assembler.CompileMemoryWithIndexToRegister(
-		amd64.LEAQ,
-		amd64ReservedRegisterForStackBasePointerAddress, 0, tmpReg, 8,
+	c.assembler.CompileMemoryToRegister(amd64.ADDQ,
+		amd64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerInBytesOffset,
 		amd64ReservedRegisterForStackBasePointerAddress,
 	)
 }
@@ -5002,13 +4994,13 @@ func (c *amd64Compiler) compileReservedMemoryPointerInitialization() {
 func (c *amd64Compiler) compileMaybeGrowValueStack() error {
 	tmpRegister, _ := c.allocateRegister(registerTypeGeneralPurpose)
 
-	c.assembler.CompileMemoryToRegister(amd64.MOVQ, amd64ReservedRegisterForCallEngine, callEngineGlobalContextValueStackLenOffset, tmpRegister)
-	c.assembler.CompileMemoryToRegister(amd64.SUBQ, amd64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerOffset, tmpRegister)
+	c.assembler.CompileMemoryToRegister(amd64.MOVQ, amd64ReservedRegisterForCallEngine, callEngineGlobalContextValueStackLenInBytesOffset, tmpRegister)
+	c.assembler.CompileMemoryToRegister(amd64.SUBQ, amd64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerInBytesOffset, tmpRegister)
 
 	// If stack base pointer + max stack pointer > valueStackLen, we need to grow the stack.
 	cmpWithStackPointerCeil := c.assembler.CompileRegisterToConst(amd64.CMPQ, tmpRegister, 0)
 	c.onStackPointerCeilDeterminedCallBack = func(stackPointerCeil uint64) {
-		cmpWithStackPointerCeil.AssignDestinationConstant(int64(stackPointerCeil))
+		cmpWithStackPointerCeil.AssignDestinationConstant(int64(stackPointerCeil) << 3)
 	}
 
 	// Jump if we have no need to grow.
diff --git a/internal/engine/compiler/impl_arm64.go b/internal/engine/compiler/impl_arm64.go
index 60b76deb..7b16b9ff 100644
--- a/internal/engine/compiler/impl_arm64.go
+++ b/internal/engine/compiler/impl_arm64.go
@@ -230,14 +230,14 @@ func (c *arm64Compiler) compileMaybeGrowValueStack() error {
 	// "tmpX = len(ce.valueStack)"
 	c.assembler.CompileMemoryToRegister(
 		arm64.LDRD,
-		arm64ReservedRegisterForCallEngine, callEngineGlobalContextValueStackLenOffset,
+		arm64ReservedRegisterForCallEngine, callEngineGlobalContextValueStackLenInBytesOffset,
 		tmpX,
 	)
 
 	// "tmpY = ce.stackBasePointer"
 	c.assembler.CompileMemoryToRegister(
 		arm64.LDRD,
-		arm64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerOffset,
+		arm64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerInBytesOffset,
 		tmpY,
 	)
 
@@ -256,7 +256,9 @@ func (c *arm64Compiler) compileMaybeGrowValueStack() error {
 	)
 	// At this point of compilation, we don't know the value of stack point ceil,
 	// so we lazily resolve the value later.
-	c.onStackPointerCeilDeterminedCallBack = func(stackPointerCeil uint64) { loadStackPointerCeil.AssignSourceConstant(int64(stackPointerCeil)) }
+	c.onStackPointerCeilDeterminedCallBack = func(stackPointerCeil uint64) {
+		loadStackPointerCeil.AssignSourceConstant(int64(stackPointerCeil) << 3)
+	}
 
 	// Compare tmpX (len(ce.valueStack) - ce.stackBasePointer) and tmpY (ce.stackPointerCeil)
 	c.assembler.CompileTwoRegistersToNone(arm64.CMP, tmpX, tmpY)
@@ -347,11 +349,11 @@ func (c *arm64Compiler) compileReturnFunction() error {
 	// 1) Set ce.valueStackContext.stackBasePointer to the value on "rb.caller".
 	c.assembler.CompileMemoryToRegister(arm64.LDRD,
 		// "rb.caller" is below the top address.
-		callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameReturnStackBasePointerOffset),
+		callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameReturnStackBasePointerInBytesOffset),
 		tmpReg)
 	c.assembler.CompileRegisterToMemory(arm64.STRD,
 		tmpReg,
-		arm64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerOffset)
+		arm64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerInBytesOffset)
 
 	// 2) Load rc.caller.moduleInstanceAddress into arm64CallingConventionModuleInstanceAddressRegister.
 	c.assembler.CompileMemoryToRegister(arm64.LDRD,
@@ -982,12 +984,12 @@ func (c *arm64Compiler) compileCallImpl(index wasm.Index, targetFunctionAddressR
 
 	// 1) Set rb.current so that we can return back to this function properly.
 	c.assembler.CompileMemoryToRegister(arm64.LDRD,
-		arm64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerOffset,
+		arm64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerInBytesOffset,
 		oldStackBasePointer)
 	c.assembler.CompileRegisterToMemory(arm64.STRD,
 		oldStackBasePointer,
 		// "rb.current" is BELOW the top address. See the above example for detail.
-		callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameReturnStackBasePointerOffset))
+		callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameReturnStackBasePointerInBytesOffset))
 
 	// 2) Set ce.valueStackContext.stackBasePointer for the next function.
 	//
@@ -995,10 +997,10 @@ func (c *arm64Compiler) compileCallImpl(index wasm.Index, targetFunctionAddressR
 	// stack base pointer by "old stack base pointer + old stack pointer - # of function params"
 	// See the comments in ce.pushCallFrame which does exactly the same calculation in Go.
 	if offset := int64(c.locationStack.sp) - int64(functype.ParamNumInUint64); offset > 0 {
-		c.assembler.CompileConstToRegister(arm64.ADD, offset, oldStackBasePointer)
+		c.assembler.CompileConstToRegister(arm64.ADD, offset<<3, oldStackBasePointer)
 		c.assembler.CompileRegisterToMemory(arm64.STRD,
 			oldStackBasePointer,
-			arm64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerOffset)
+			arm64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerInBytesOffset)
 	}
 
 	// 3) Set rc.next to specify which function is executed on the current call frame.
@@ -4098,15 +4100,11 @@ func (c *arm64Compiler) compileReservedStackBasePointerRegisterInitialization()
 
 	// next we move the base pointer (ce.stackBasePointer) to arm64ReservedRegisterForTemporary.
 	c.assembler.CompileMemoryToRegister(arm64.LDRD,
-		arm64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerOffset,
+		arm64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerInBytesOffset,
 		arm64ReservedRegisterForTemporary)
 
-	// Finally, we calculate "arm64ReservedRegisterForStackBasePointerAddress + arm64ReservedRegisterForTemporary << 3"
-	// where we shift tmpReg by 3 because stack pointer is an index in the []uint64
-	// so we must multiply the value by the size of uint64 = 8 bytes.
-	c.assembler.CompileLeftShiftedRegisterToRegister(arm64.ADD,
-		arm64ReservedRegisterForTemporary, 3, arm64ReservedRegisterForStackBasePointerAddress,
-		arm64ReservedRegisterForStackBasePointerAddress)
+	// Finally, we calculate "callEngineValueStackContextStackBasePointerInBytesOffset + arm64ReservedRegisterForTemporary"
+	c.assembler.CompileRegisterToRegister(arm64.ADD, arm64ReservedRegisterForTemporary, arm64ReservedRegisterForStackBasePointerAddress)
 }
 
 func (c *arm64Compiler) compileReservedMemoryRegisterInitialization() {


Test patch:
diff --git a/internal/engine/compiler/compiler_controlflow_test.go b/internal/engine/compiler/compiler_controlflow_test.go
index f1642831..5040582c 100644
--- a/internal/engine/compiler/compiler_controlflow_test.go
+++ b/internal/engine/compiler/compiler_controlflow_test.go
@@ -965,12 +965,12 @@ func TestCompiler_returnFunction(t *testing.T) {
 				// Set the return address to the beginning of the function so that we can execute the constI32 above.
 				returnAddress: f.codeInitialAddress,
 				// Note: return stack base pointer is set to funcaddr*5 and this is where the const should be pushed.
-				returnStackBasePointer: uint64(funcIndex) * 5,
-				function:               f,
+				returnStackBasePointerInBytes: (uint64(funcIndex) * 5) << 3,
+				function:                      f,
 			}
 			ce.callFrameStack[ce.globalContext.callFrameStackPointer] = frame
 			ce.globalContext.callFrameStackPointer++
-			stackPointerToExpectedValue[frame.returnStackBasePointer] = expValue
+			stackPointerToExpectedValue[frame.returnStackBasePointerInBytes>>3] = expValue
 		}
 
 		require.Equal(t, uint64(callFrameNums), env.callFrameStackPointer())
diff --git a/internal/engine/compiler/compiler_test.go b/internal/engine/compiler/compiler_test.go
index c2bbc51c..eaef2e70 100644
--- a/internal/engine/compiler/compiler_test.go
+++ b/internal/engine/compiler/compiler_test.go
@@ -37,7 +37,7 @@ func init() {
 	var ce callEngine
 	// Offsets for callEngine.globalContext.
 	requireEqual(int(unsafe.Offsetof(ce.valueStackElement0Address)), callEngineGlobalContextValueStackElement0AddressOffset, "callEngineGlobalContextValueStackElement0AddressOffset")
-	requireEqual(int(unsafe.Offsetof(ce.valueStackLen)), callEngineGlobalContextValueStackLenOffset, "callEngineGlobalContextValueStackLenOffset")
+	requireEqual(int(unsafe.Offsetof(ce.valueStackLenInBytes)), callEngineGlobalContextValueStackLenInBytesOffset, "callEngineGlobalContextValueStackLenInBytesOffset")
 	requireEqual(int(unsafe.Offsetof(ce.callFrameStackElementZeroAddress)), callEngineGlobalContextCallFrameStackElement0AddressOffset, "callEngineGlobalContextCallFrameStackElement0AddressOffset")
 	requireEqual(int(unsafe.Offsetof(ce.callFrameStackLen)), callEngineGlobalContextCallFrameStackLenOffset, "callEngineGlobalContextCallFrameStackLenOffset")
 	requireEqual(int(unsafe.Offsetof(ce.callFrameStackPointer)), callEngineGlobalContextCallFrameStackPointerOffset, "callEngineGlobalContextCallFrameStackPointerOffset")
@@ -55,7 +55,7 @@ func init() {
 
 	// Offsets for callEngine.valueStackContext
 	requireEqual(int(unsafe.Offsetof(ce.stackPointer)), callEngineValueStackContextStackPointerOffset, "callEngineValueStackContextStackPointerOffset")
-	requireEqual(int(unsafe.Offsetof(ce.stackBasePointer)), callEngineValueStackContextStackBasePointerOffset, "callEngineValueStackContextStackBasePointerOffset")
+	requireEqual(int(unsafe.Offsetof(ce.stackBasePointerInBytes)), callEngineValueStackContextStackBasePointerInBytesOffset, "callEngineValueStackContextStackBasePointerInBytesOffset")
 
 	// Offsets for callEngine.exitContext.
 	requireEqual(int(unsafe.Offsetof(ce.statusCode)), callEngineExitContextNativeCallStatusCodeOffset, "callEngineExitContextNativeCallStatusCodeOffset")
@@ -68,7 +68,7 @@ func init() {
 	requireEqual(0, callFrameDataSize&(callFrameDataSize-1), "callFrameDataSize&(callFrameDataSize-1)")
 	requireEqual(math.Ilogb(float64(callFrameDataSize)), callFrameDataSizeMostSignificantSetBit, "callFrameDataSizeMostSignificantSetBit")
 	requireEqual(int(unsafe.Offsetof(frame.returnAddress)), callFrameReturnAddressOffset, "callFrameReturnAddressOffset")
-	requireEqual(int(unsafe.Offsetof(frame.returnStackBasePointer)), callFrameReturnStackBasePointerOffset, "callFrameReturnStackBasePointerOffset")
+	requireEqual(int(unsafe.Offsetof(frame.returnStackBasePointerInBytes)), callFrameReturnStackBasePointerInBytesOffset, "callFrameReturnStackBasePointerInBytesOffset")
 	requireEqual(int(unsafe.Offsetof(frame.function)), callFrameFunctionOffset, "callFrameFunctionOffset")
 
 	// Offsets for code.
@@ -173,7 +173,7 @@ func (j *compilerEnv) stackPointer() uint64 {
 }
 
 func (j *compilerEnv) stackBasePointer() uint64 {
-	return j.ce.valueStackContext.stackBasePointer
+	return j.ce.valueStackContext.stackBasePointerInBytes >> 3
 }
 
 func (j *compilerEnv) setStackPointer(sp uint64) {
@@ -201,7 +201,7 @@ func (j *compilerEnv) callFrameStackPointer() uint64 {
 }
 
 func (j *compilerEnv) setValueStackBasePointer(sp uint64) {
-	j.ce.valueStackContext.stackBasePointer = sp
+	j.ce.valueStackContext.stackBasePointerInBytes = sp << 3
 }
 
 func (j *compilerEnv) setCallFrameStackPointerLen(l uint64) {
diff --git a/internal/engine/compiler/engine_test.go b/internal/engine/compiler/engine_test.go
index ec4b92e0..efc85a74 100644
--- a/internal/engine/compiler/engine_test.go
+++ b/internal/engine/compiler/engine_test.go
@@ -332,7 +332,7 @@ wasm stack trace:
 
 	// After recover, the stack pointers must be reset, but the underlying slices must be intact
 	// for the subsequent calls.
-	require.Equal(t, uint64(0), ce.stackPointer)
+	require.Equal(t, uint64(0), ce.stackBasePointerInBytes)
 	require.Equal(t, uint64(0), ce.callFrameStackPointer)
 	require.Equal(t, beforeRecoverValueStack, ce.valueStack)
 	require.Equal(t, beforeRecoverCallFrameStack, ce.callFrameStack)