Repository: tetratelabs/wazero
Pull request #456 (2022-04-13 02:24:18)

Fix broken base64 bench
Hasn't been iterated via `b.N` 😄 

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>


Patch:
diff --git a/internal/asm/amd64/impl.go b/internal/asm/amd64/impl.go
index 26a3d860..c2f51498 100644
--- a/internal/asm/amd64/impl.go
+++ b/internal/asm/amd64/impl.go
@@ -724,6 +724,9 @@ func (a *AssemblerImpl) encodeNoneToNone(n *NodeImpl) (err error) {
 	case RET:
 		// https://www.felixcloutier.com/x86/ret
 		err = a.Buf.WriteByte(0xc3)
+	case UD2:
+		// https://mudongliang.github.io/x86/html/file_module_x86_id_318.html
+		_, err = a.Buf.Write([]byte{0x0f, 0x0b})
 	default:
 		err = errorEncodingUnsupported(n)
 	}
diff --git a/internal/wasm/jit/arch.go b/internal/wasm/jit/arch.go
index 0de88b98..714aeb0f 100644
--- a/internal/wasm/jit/arch.go
+++ b/internal/wasm/jit/arch.go
@@ -11,4 +11,4 @@ var (
 // ce is "*callEngine" as uintptr.
 //
 // Note: this is implemented in per-arch Go assembler file. For example, arch_amd64.s implements this for amd64.
-func jitcall(codeSegment, ce uintptr)
+func jitcall(codeSegment, ce uintptr, moduleInstanceAddress uintptr)
diff --git a/internal/wasm/jit/arch_amd64.s b/internal/wasm/jit/arch_amd64.s
index 8cda8954..5208981e 100644
--- a/internal/wasm/jit/arch_amd64.s
+++ b/internal/wasm/jit/arch_amd64.s
@@ -1,8 +1,9 @@
 #include "funcdata.h"
 #include "textflag.h"
 
-// jitcall(codeSegment, ce)
-TEXT ·jitcall(SB),NOSPLIT|NOFRAME,$0-16
-        MOVQ codeSegment+0(FP),AX  // Load the address of native code.
-        MOVQ ce+8(FP),R13          // Load the address of *callEngine.
-        JMP AX                     // Jump to native code.
+// jitcall(codeSegment, ce, moduleInstanceAddress)
+TEXT ·jitcall(SB),NOSPLIT|NOFRAME,$0-24
+        MOVQ ce+8(FP),R13                     // Load the address of *callEngine. into amd64ReservedRegisterForCallEngine.
+        MOVQ moduleInstanceAddress+16(FP),R12 // Load the address of *wasm.ModuleInstance into amd64CallingConventionModuleInstanceAddressRegister.
+        MOVQ codeSegment+0(FP),AX             // Load the address of native code.
+        JMP AX                                // Jump to native code.
diff --git a/internal/wasm/jit/arch_arm64.s b/internal/wasm/jit/arch_arm64.s
index 74fb73f1..f10ccdc8 100644
--- a/internal/wasm/jit/arch_arm64.s
+++ b/internal/wasm/jit/arch_arm64.s
@@ -1,15 +1,17 @@
 #include "funcdata.h"
 #include "textflag.h"
 
-// jitcall(codeSegment, ce)
-TEXT ·jitcall(SB),NOSPLIT|NOFRAME,$0-16
-        // Load the address of native code.
-        MOVD codeSegment+0(FP),R1
-        // Load the address of *callEngine.
+// jitcall(codeSegment, ce, moduleInstanceAddress)
+TEXT ·jitcall(SB),NOSPLIT|NOFRAME,$0-24
+        // Load the address of *callEngine into arm64ReservedRegisterForCallEngine.
         MOVD ce+8(FP),R0
         // In arm64, return address is stored in R30 after jumping into the code.
         // We save the return address value into archContext.jitReturnAddress in Engine.
         // Note that the const 120 drifts after editting Engine or archContext struct. See TestArchContextOffsetInEngine.
         MOVD R30,120(R0)
+        // Load the address of *wasm.ModuleInstance into arm64CallingConventionModuleInstanceAddressRegister.
+        MOVD moduleInstanceAddress+16(FP),R29
+        // Load the address of native code.
+        MOVD codeSegment+0(FP),R1
         // Jump to native code.
         JMP (R1)
diff --git a/internal/wasm/jit/engine.go b/internal/wasm/jit/engine.go
index 8fcfc1cf..5075bbc9 100644
--- a/internal/wasm/jit/engine.go
+++ b/internal/wasm/jit/engine.go
@@ -166,18 +166,20 @@ type (
 	compiledFunction struct {
 		// The following fields are accessed by JITed code.
 
-		// Pre-calculated pointer pointing to the initial byte of .codeSegment slice.
+		// codeInitialAddress is the pre-calculated pointer pointing to the initial byte of .codeSegment slice.
 		// That mean codeInitialAddress always equals uintptr(unsafe.Pointer(&.codeSegment[0]))
 		// and we cache the value (uintptr(unsafe.Pointer(&.codeSegment[0]))) to this field,
 		// so we don't need to repeat the calculation on each function call.
 		codeInitialAddress uintptr
-		// The max of the stack pointer this function can reach. Lazily applied via maybeGrowValueStack.
+		// stackPointerCeil is the max of the stack pointer this function can reach. Lazily applied via maybeGrowValueStack.
 		stackPointerCeil uint64
+		// source is the source function instance from which this is compiled.
+		source *wasm.FunctionInstance
+		// moduleInstanceAddress holds the address of source.ModuleInstance.
+		moduleInstanceAddress uintptr
 
 		// Followings are not accessed by JITed code.
 
-		// The source function instance from which this is compiled.
-		source *wasm.FunctionInstance
 		// codeSegment is holding the compiled native code as a byte slice.
 		codeSegment []byte
 		// See the doc for compiledFunctionStaticData type.
@@ -229,9 +231,10 @@ const (
 	callFrameCompiledFunctionOffset        = 16
 
 	// Offsets for compiledFunction.
-	compiledFunctionCodeInitialAddressOffset = 0
-	compiledFunctionStackPointerCeilOffset   = 8
-	compiledFunctionSourceOffset             = 16
+	compiledFunctionCodeInitialAddressOffset    = 0
+	compiledFunctionStackPointerCeilOffset      = 8
+	compiledFunctionSourceOffset                = 16
+	compiledFunctionModuleInstanceAddressOffset = 24
 
 	// Offsets for wasm.ModuleInstance.
 	moduleInstanceGlobalsOffset = 48
@@ -672,7 +675,9 @@ func (ce *callEngine) execHostFunction(fk wasm.FunctionKind, f *reflect.Value, c
 }
 
 func (ce *callEngine) execWasmFunction(ctx *wasm.ModuleContext, f *compiledFunction) {
-	ce.pushCallFrame(f)
+	// Push the initial callframe.
+	ce.callFrameStack[0] = callFrame{returnAddress: f.codeInitialAddress, compiledFunction: f}
+	ce.globalContext.callFrameStackPointer++
 
 jitentry:
 	{
@@ -683,7 +688,7 @@ jitentry:
 		}
 
 		// Call into the JIT code.
-		jitcall(frame.returnAddress, uintptr(unsafe.Pointer(ce)))
+		jitcall(frame.returnAddress, uintptr(unsafe.Pointer(ce)), f.moduleInstanceAddress)
 
 		// Check the status code from JIT code.
 		switch status := ce.exitContext.statusCode; status {
@@ -722,29 +727,6 @@ jitentry:
 	}
 }
 
-// pushInitialFrame is implemented in assembly as well, but this Go version is used BEFORE jit entry.
-func (ce *callEngine) pushCallFrame(f *compiledFunction) {
-	// Push the new frame to the top of stack.
-	ce.callFrameStack[ce.globalContext.callFrameStackPointer] = callFrame{returnAddress: f.codeInitialAddress, compiledFunction: f}
-	ce.globalContext.callFrameStackPointer++
-
-	// For example, if we have the following state (where "_" means no value pushed),
-	//       base            sp
-	//        |              |
-	// [...., A, B, C, D, E, _, _ ]
-	//
-	// and the target function requires 2 params, we need to pass D and E as arguments.
-	//
-	// Therefore, the target function start executing under the following state:
-	//                base   sp
-	//                 |     |
-	// [...., A, B, C, D, E, _, _ ]
-	//
-	// That means the next stack base pointer is calculated as follows (note stack pointer is relative to base):
-	ce.valueStackContext.stackBasePointer =
-		ce.valueStackContext.stackBasePointer + ce.valueStackContext.stackPointer - uint64(len(f.source.Type.Params))
-}
-
 func (ce *callEngine) builtinFunctionGrowValueStack(stackPointerCeil uint64) {
 	// Extends the valueStack's length to currentLen*2+stackPointerCeil.
 	newLen := ce.globalContext.valueStackLen*2 + (stackPointerCeil)
@@ -809,10 +791,11 @@ func compileHostFunction(f *wasm.FunctionInstance) (*compiledFunction, error) {
 	}
 
 	return &compiledFunction{
-		source:             f,
-		codeSegment:        code,
-		codeInitialAddress: uintptr(unsafe.Pointer(&code[0])),
-		stackPointerCeil:   stackPointerCeil,
+		source:                f,
+		codeSegment:           code,
+		codeInitialAddress:    uintptr(unsafe.Pointer(&code[0])),
+		moduleInstanceAddress: uintptr(unsafe.Pointer(f.Module)),
+		stackPointerCeil:      stackPointerCeil,
 	}, nil
 }
 
@@ -1006,10 +989,11 @@ func compileWasmFunction(enabledFeatures wasm.Features, f *wasm.FunctionInstance
 	}
 
 	return &compiledFunction{
-		source:             f,
-		codeSegment:        code,
-		codeInitialAddress: uintptr(unsafe.Pointer(&code[0])),
-		stackPointerCeil:   stackPointerCeil,
-		staticData:         staticData,
+		source:                f,
+		codeSegment:           code,
+		codeInitialAddress:    uintptr(unsafe.Pointer(&code[0])),
+		moduleInstanceAddress: uintptr(unsafe.Pointer(f.Module)),
+		stackPointerCeil:      stackPointerCeil,
+		staticData:            staticData,
 	}, nil
 }
diff --git a/internal/wasm/jit/jit_impl_amd64.go b/internal/wasm/jit/jit_impl_amd64.go
index 259b5e11..d872a1bb 100644
--- a/internal/wasm/jit/jit_impl_amd64.go
+++ b/internal/wasm/jit/jit_impl_amd64.go
@@ -102,6 +102,13 @@ var (
 	}
 )
 
+var (
+	// amd64CallingConventionModuleInstanceAddressRegister holds *wasm.ModuleInstance of the
+	// next executing function instance. The value is set and used when making function calls
+	// or function returns in the ModuleContextInitialization. See compileModuleContextInitialization.
+	amd64CallingConventionModuleInstanceAddressRegister = amd64.REG_R12
+)
+
 func (c *amd64Compiler) String() string {
 	return c.locationStack.String()
 }
@@ -3405,7 +3412,7 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, compiledFuncti
 	c.locationStack.markRegisterUsed(freeRegs...)
 
 	// Alias these free tmp registers for readability.
-	callFrameStackPointerRegister, tmpRegister, targetAddressRegister,
+	callFrameStackPointerRegister, tmpRegister, targetCompiledFunctionAddressRegister,
 		callFrameStackTopAddressRegister := freeRegs[0], freeRegs[1], freeRegs[2], freeRegs[3]
 
 	// First, we read the current call frame stack pointer.
@@ -3532,14 +3539,14 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, compiledFuncti
 				// Note: FunctionIndex is limited up to 2^27 so this offset never exceeds 32-bit integer.
 				// *8 because the size of *compiledFunction equals 8 bytes.
 				tmpRegister, int64(index)*8,
-				targetAddressRegister,
+				targetCompiledFunctionAddressRegister,
 			)
 		} else {
-			targetAddressRegister = compiledFunctionAddressRegister
+			targetCompiledFunctionAddressRegister = compiledFunctionAddressRegister
 		}
 		// Finally, we are ready to place the address of the target function's *compiledFunction into the new call-frame.
 		// In the example, this is equivalent to set "rc.next".
-		c.assembler.CompileRegisterToMemory(amd64.MOVQ, targetAddressRegister, callFrameStackTopAddressRegister, callFrameCompiledFunctionOffset)
+		c.assembler.CompileRegisterToMemory(amd64.MOVQ, targetCompiledFunctionAddressRegister, callFrameStackTopAddressRegister, callFrameCompiledFunctionOffset)
 	}
 
 	// 4) Set ra.1 so that we can return back to this function properly.
@@ -3560,8 +3567,12 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, compiledFuncti
 	// So we increment the call frame stack pointer.
 	c.assembler.CompileNoneToMemory(amd64.INCQ, amd64ReservedRegisterForCallEngine, callEngineGlobalContextCallFrameStackPointerOffset)
 
+	// Also, we have to put the target function's *wasm.ModuleInstance into amd64CallingConventionModuleInstanceAddressRegister.
+	c.assembler.CompileMemoryToRegister(amd64.MOVQ, targetCompiledFunctionAddressRegister, compiledFunctionModuleInstanceAddressOffset,
+		amd64CallingConventionModuleInstanceAddressRegister)
+
 	// And jump into the initial address of the target function.
-	c.assembler.CompileJumpToMemory(amd64.JMP, targetAddressRegister, compiledFunctionCodeInitialAddressOffset)
+	c.assembler.CompileJumpToMemory(amd64.JMP, targetCompiledFunctionAddressRegister, compiledFunctionCodeInitialAddressOffset)
 
 	// All the registers used are temporary so we mark them unused.
 	c.locationStack.markRegisterUnused(freeRegs...)
@@ -3570,13 +3581,12 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, compiledFuncti
 	// This could be reached after returnFunction(), so callEngine.valueStackContext.stackBasePointer
 	// and callEngine.moduleContext.moduleInstanceAddress are changed (See comments in returnFunction()).
 	// Therefore we have to initialize the state according to these changes.
-	//
-	// Due to the change to callEngine.valueStackContext.stackBasePointer.
-	c.compileReservedStackBasePointerInitialization()
-	// Due to the change to callEngine.moduleContext.moduleInstanceAddress.
 	if err := c.compileModuleContextInitialization(); err != nil {
 		return err
 	}
+	// Due to the change to callEngine.valueStackContext.stackBasePointer.
+	c.compileReservedStackBasePointerInitialization()
+
 	// Due to the change to callEngine.moduleContext.moduleInstanceAddress as that might result in
 	// the memory instance manipulation.
 	c.compileReservedMemoryPointerInitialization()
@@ -3594,6 +3604,11 @@ func (c *amd64Compiler) compileReturnFunction() error {
 	// Release all the registers as our calling convention requires the caller-save.
 	c.compileReleaseAllRegistersToStack()
 
+	// amd64CallingConventionModuleInstanceAddressRegister holds the module intstance's address
+	// so mark it used so that it won't be used as a free register.
+	c.locationStack.markRegisterUsed(amd64CallingConventionModuleInstanceAddressRegister)
+	defer c.locationStack.markRegisterUnused(amd64CallingConventionModuleInstanceAddressRegister)
+
 	// Obtain the temporary registers to be used in the followings.
 	regs, found := c.locationStack.takeFreeRegisters(generalPurposeRegisterTypeInt, 3)
 	if !found {
@@ -3657,7 +3672,8 @@ func (c *amd64Compiler) compileReturnFunction() error {
 	//
 	// What we have to do in the following is that
 	//   1) Set callEngine.valueStackContext.stackBasePointer to the value on "rb.caller".
-	//   2) Jump into the address of "ra.caller".
+	//   2) Load rc.caller.moduleInstanceAddress into amd64CallingConventionModuleInstanceAddressRegister.
+	//   3) Jump into the address of "ra.caller".
 
 	// 1) Set callEngine.valueStackContext.stackBasePointer to the value on "rb.caller"
 	c.assembler.CompileMemoryToRegister(amd64.MOVQ,
@@ -3668,7 +3684,18 @@ func (c *amd64Compiler) compileReturnFunction() error {
 	c.assembler.CompileRegisterToMemory(amd64.MOVQ,
 		tmpRegister, amd64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerOffset)
 
-	// 2) Jump into the address of "ra.caller".
+	// 2) Load rc.caller.moduleInstanceAddress into amd64CallingConventionModuleInstanceAddressRegister
+	c.assembler.CompileMemoryToRegister(amd64.MOVQ,
+		// "rc.caller" is BELOW the top address. See the above example for detail.
+		callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameCompiledFunctionOffset),
+		amd64CallingConventionModuleInstanceAddressRegister,
+	)
+	c.assembler.CompileMemoryToRegister(amd64.MOVQ,
+		amd64CallingConventionModuleInstanceAddressRegister, compiledFunctionModuleInstanceAddressOffset,
+		amd64CallingConventionModuleInstanceAddressRegister,
+	)
+
+	// 3) Jump into the address of "ra.caller".
 	c.assembler.CompileMemoryToRegister(amd64.MOVQ,
 		// "ra.caller" is BELOW the top address. See the above example for detail.
 		callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameReturnAddressOffset),
@@ -3783,6 +3810,10 @@ func (c *amd64Compiler) compilePreamble() (err error) {
 	// the caller.
 	c.pushFunctionParams()
 
+	if err := c.compileModuleContextInitialization(); err != nil {
+		return err
+	}
+
 	// Check if it's necessary to grow the value stack by using max stack pointer.
 	if err = c.compileMaybeGrowValueStack(); err != nil {
 		return err
@@ -3790,12 +3821,6 @@ func (c *amd64Compiler) compilePreamble() (err error) {
 
 	c.compileReservedStackBasePointerInitialization()
 
-	// Once the stack base pointer is initialized and the size of stack is ok,
-	// initialize the module context next.
-	if err := c.compileModuleContextInitialization(); err != nil {
-		return err
-	}
-
 	// Finally, we initialize the reserved memory register based on the module context.
 	c.compileReservedMemoryPointerInitialization()
 	return
@@ -3864,9 +3889,13 @@ func (c *amd64Compiler) compileMaybeGrowValueStack() error {
 // callEngine.ModuleContext.ModuleInstanceAddress.
 // This is called in two cases: in function preamble, and on the return from (non-Go) function calls.
 func (c *amd64Compiler) compileModuleContextInitialization() error {
+	// amd64CallingConventionModuleInstanceAddressRegister holds the module intstance's address
+	// so mark it used so that it won't be used as a free register until the module context initialization finishes.
+	c.locationStack.markRegisterUsed(amd64CallingConventionModuleInstanceAddressRegister)
+	defer c.locationStack.markRegisterUnused(amd64CallingConventionModuleInstanceAddressRegister)
 
 	// Obtain the temporary registers to be used in the followings.
-	regs, found := c.locationStack.takeFreeRegisters(generalPurposeRegisterTypeInt, 3)
+	regs, found := c.locationStack.takeFreeRegisters(generalPurposeRegisterTypeInt, 2)
 	if !found {
 		// This in theory never happen as all the registers must be free except indexReg.
 		return fmt.Errorf("could not find enough free registers")
@@ -3874,9 +3903,7 @@ func (c *amd64Compiler) compileModuleContextInitialization() error {
 	c.locationStack.markRegisterUsed(regs...)
 
 	// Alias these free tmp registers for readability.
-	moduleInstanceAddressRegister, tmpRegister, tmpRegister2 := regs[0], regs[1], regs[2]
-
-	c.assembler.CompileConstToRegister(amd64.MOVQ, int64(uintptr(unsafe.Pointer(c.f.Module))), moduleInstanceAddressRegister)
+	tmpRegister, tmpRegister2 := regs[0], regs[1]
 
 	// If the module instance address stays the same, we could skip the entire code below.
 	// The rationale/idea for this is that, in almost all use cases, users instantiate a single
@@ -3884,15 +3911,15 @@ func (c *amd64Compiler) compileModuleContextInitialization() error {
 	// binaries. As a result, this cmp and jmp instruction sequence below must be easy for
 	// x64 CPU to do branch prediction since almost 100% jump happens across function calls.
 	c.assembler.CompileMemoryToRegister(amd64.CMPQ,
-		amd64ReservedRegisterForCallEngine, callEngineModuleContextModuleInstanceAddressOffset, moduleInstanceAddressRegister)
+		amd64ReservedRegisterForCallEngine, callEngineModuleContextModuleInstanceAddressOffset, amd64CallingConventionModuleInstanceAddressRegister)
 	jmpIfModuleNotChange := c.assembler.CompileJump(amd64.JEQ)
 
-	// Otherwise, we need to update fields.
-	// First, save the read module instance address to callEngine.moduleInstanceAddress
-	c.assembler.CompileRegisterToMemory(amd64.MOVQ, moduleInstanceAddressRegister,
+	// If engine.ModuleContext.ModuleInstanceAddress is not euqal the value on amd64CallingConventionModuleInstanceAddressRegister,
+	// we have to put the new value there.
+	c.assembler.CompileRegisterToMemory(amd64.MOVQ, amd64CallingConventionModuleInstanceAddressRegister,
 		amd64ReservedRegisterForCallEngine, callEngineModuleContextModuleInstanceAddressOffset)
 
-	// Otherwise, we have to update the following fields:
+	// Also, we have to update the following fields:
 	// * callEngine.moduleContext.globalElement0Address
 	// * callEngine.moduleContext.tableElement0Address
 	// * callEngine.moduleContext.tableSliceLen
@@ -3909,7 +3936,7 @@ func (c *amd64Compiler) compileModuleContextInitialization() error {
 		// Since ModuleInstance.Globals is []*globalInstance, internally
 		// the address of the first item in the underlying array lies exactly on the globals offset.
 		// See https://go.dev/blog/slices-intro if unfamiliar.
-		c.assembler.CompileMemoryToRegister(amd64.MOVQ, moduleInstanceAddressRegister, moduleInstanceGlobalsOffset, tmpRegister)
+		c.assembler.CompileMemoryToRegister(amd64.MOVQ, amd64CallingConventionModuleInstanceAddressRegister, moduleInstanceGlobalsOffset, tmpRegister)
 
 		c.assembler.CompileRegisterToMemory(amd64.MOVQ, tmpRegister, amd64ReservedRegisterForCallEngine, callEngineModuleContextGlobalElement0AddressOffset)
 	}
@@ -3921,7 +3948,7 @@ func (c *amd64Compiler) compileModuleContextInitialization() error {
 	// why it is ok to skip the initialization if the module's table doesn't exist.
 	if c.f.Module.Table != nil {
 		// First, we need to read the *wasm.Table.
-		c.assembler.CompileMemoryToRegister(amd64.MOVQ, moduleInstanceAddressRegister, moduleInstanceTableOffset, tmpRegister)
+		c.assembler.CompileMemoryToRegister(amd64.MOVQ, amd64CallingConventionModuleInstanceAddressRegister, moduleInstanceTableOffset, tmpRegister)
 
 		// At this point, tmpRegister holds the address of ModuleInstance.Table.
 		// So we are ready to read and put the first item's address stored in Table.Table.
@@ -3946,7 +3973,7 @@ func (c *amd64Compiler) compileModuleContextInitialization() error {
 	// That is ensured by function validation at module instantiation phase, and that's
 	// why it is ok to skip the initialization if the module's memory instance is nil.
 	if c.f.Module.Memory != nil {
-		c.assembler.CompileMemoryToRegister(amd64.MOVQ, moduleInstanceAddressRegister, moduleInstanceMemoryOffset, tmpRegister)
+		c.assembler.CompileMemoryToRegister(amd64.MOVQ, amd64CallingConventionModuleInstanceAddressRegister, moduleInstanceMemoryOffset, tmpRegister)
 
 		// Set length.
 		c.assembler.CompileMemoryToRegister(amd64.MOVQ, tmpRegister, memoryInstanceBufferLenOffset, tmpRegister2)
@@ -3969,7 +3996,7 @@ func (c *amd64Compiler) compileModuleContextInitialization() error {
 		// See the following references for detail:
 		// * https://research.swtch.com/interfaces
 		// * https://github.com/golang/go/blob/release-branch.go1.17/src/runtime/runtime2.go#L207-L210
-		c.assembler.CompileMemoryToRegister(amd64.MOVQ, moduleInstanceAddressRegister, moduleInstanceEngineOffset+interfaceDataOffset, tmpRegister)
+		c.assembler.CompileMemoryToRegister(amd64.MOVQ, amd64CallingConventionModuleInstanceAddressRegister, moduleInstanceEngineOffset+interfaceDataOffset, tmpRegister)
 
 		// "tmpRegister = [tmpRegister + moduleEngineCompiledFunctionsOffset] (== &moduleEngine.compiledFunctions[0])"
 		c.assembler.CompileMemoryToRegister(amd64.MOVQ, tmpRegister, moduleEngineCompiledFunctionsOffset, tmpRegister)
diff --git a/internal/wasm/jit/jit_impl_arm64.go b/internal/wasm/jit/jit_impl_arm64.go
index 21f4ad8d..6f0c93c8 100644
--- a/internal/wasm/jit/jit_impl_arm64.go
+++ b/internal/wasm/jit/jit_impl_arm64.go
@@ -81,6 +81,10 @@ const (
 	arm64ReservedRegisterForTemporary asm.Register = arm64.REG_R27
 )
 
+var (
+	arm64CallingConventionModuleInstanceAddressRegister = arm64.REG_R29
+)
+
 const (
 	// arm64CallEngineArchContextJITCallReturnAddressOffset is the offset of archContext.jitCallReturnAddress in callEngine.
 	arm64CallEngineArchContextJITCallReturnAddressOffset = 120
@@ -201,13 +205,13 @@ func (c *arm64Compiler) compilePreamble() error {
 		return err
 	}
 
-	// We must initialize the stack base pointer register so that we can manipulate the stack properly.
-	c.compileReservedStackBasePointerRegisterInitialization()
-
 	if err := c.compileModuleContextInitialization(); err != nil {
 		return err
 	}
 
+	// We must initialize the stack base pointer register so that we can manipulate the stack properly.
+	c.compileReservedStackBasePointerRegisterInitialization()
+
 	c.compileReservedMemoryRegisterInitialization()
 	return nil
 }
@@ -278,6 +282,11 @@ func (c *arm64Compiler) compileReturnFunction() error {
 		return err
 	}
 
+	// arm64CallingConventionModuleInstanceAddressRegister holds the module intstance's address
+	// so mark it used so that it won't be used as a free register.
+	c.locationStack.markRegisterUsed(arm64CallingConventionModuleInstanceAddressRegister)
+	defer c.locationStack.markRegisterUnused(arm64CallingConventionModuleInstanceAddressRegister)
+
 	tmpRegs, found := c.locationStack.takeFreeRegisters(generalPurposeRegisterTypeInt, 3)
 	if !found {
 		return fmt.Errorf("BUG: all the registers should be free at this point")
@@ -331,7 +340,8 @@ func (c *arm64Compiler) compileReturnFunction() error {
 	//
 	// What we have to do in the following is that
 	//   1) Set ce.valueStackContext.stackBasePointer to the value on "rb.caller".
-	//   2) Jump into the address of "ra.caller".
+	//   2) Load rc.caller.moduleInstanceAddress into arm64CallingConventionModuleInstanceAddressRegister.
+	//   3) Jump into the address of "ra.caller".
 
 	// 1) Set ce.valueStackContext.stackBasePointer to the value on "rb.caller".
 	c.assembler.CompileMemoryToRegister(arm64.MOVD,
@@ -342,7 +352,16 @@ func (c *arm64Compiler) compileReturnFunction() error {
 		tmpReg,
 		arm64ReservedRegisterForCallEngine, callEngineValueStackContextStackBasePointerOffset)
 
-	// 2) Branch into the address of "ra.caller".
+	// 2) Load rc.caller.moduleInstanceAddress into arm64CallingConventionModuleInstanceAddressRegister.
+	c.assembler.CompileMemoryToRegister(arm64.MOVD,
+		// "rb.caller" is below the top address.
+		callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameCompiledFunctionOffset),
+		arm64CallingConventionModuleInstanceAddressRegister)
+	c.assembler.CompileMemoryToRegister(arm64.MOVD,
+		arm64CallingConventionModuleInstanceAddressRegister, compiledFunctionModuleInstanceAddressOffset,
+		arm64CallingConventionModuleInstanceAddressRegister)
+
+	// 3) Branch into the address of "ra.caller".
 	c.assembler.CompileMemoryToRegister(arm64.MOVD,
 		// "rb.caller" is below the top address.
 		callFrameStackTopAddressRegister, -(callFrameDataSize - callFrameReturnAddressOffset),
@@ -990,6 +1009,12 @@ func (c *arm64Compiler) compileCallImpl(index wasm.Index, compiledFunctionAddres
 		tmp,
 		arm64ReservedRegisterForCallEngine, callEngineGlobalContextCallFrameStackPointerOffset)
 
+	// Also, we have to put the compiledFunction's moduleinstance address into arm64CallingConventionModuleInstanceAddressRegister.
+	c.assembler.CompileMemoryToRegister(arm64.MOVD,
+		compiledFunctionRegister, compiledFunctionModuleInstanceAddressOffset,
+		arm64CallingConventionModuleInstanceAddressRegister,
+	)
+
 	// Then, br into the target function's initial address.
 	c.assembler.CompileMemoryToRegister(arm64.MOVD,
 		compiledFunctionRegister, compiledFunctionCodeInitialAddressOffset,
@@ -1016,13 +1041,13 @@ func (c *arm64Compiler) compileCallImpl(index wasm.Index, compiledFunctionAddres
 		}
 	}
 
-	// On the function return, we initialize the state for this function.
-	c.compileReservedStackBasePointerRegisterInitialization()
-
 	if err := c.compileModuleContextInitialization(); err != nil {
 		return err
 	}
 
+	// On the function return, we initialize the state for this function.
+	c.compileReservedStackBasePointerRegisterInitialization()
+
 	c.compileReservedMemoryRegisterInitialization()
 	return nil
 }
@@ -3030,27 +3055,32 @@ func (c *arm64Compiler) compileReservedMemoryRegisterInitialization() {
 // ce.ModuleContext.ModuleInstanceAddress.
 // This is called in two cases: in function preamble, and on the return from (non-Go) function calls.
 func (c *arm64Compiler) compileModuleContextInitialization() error {
-	regs, found := c.locationStack.takeFreeRegisters(generalPurposeRegisterTypeInt, 3)
+	c.markRegisterUsed(arm64CallingConventionModuleInstanceAddressRegister)
+	defer c.markRegisterUnused(arm64CallingConventionModuleInstanceAddressRegister)
+
+	regs, found := c.locationStack.takeFreeRegisters(generalPurposeRegisterTypeInt, 2)
 	if !found {
 		return fmt.Errorf("BUG: all the registers should be free at this point")
 	}
 	c.markRegisterUsed(regs...)
 
 	// Alias these free registers for readability.
-	moduleInstanceAddressRegister, tmpX, tmpY := regs[0], regs[1], regs[2]
-
-	// Load the absolute address of the current function's module instance.
-	// Note: this should be modified to support Clone() functionality per #179.
-	c.assembler.CompileConstToRegister(arm64.MOVD, int64(uintptr(unsafe.Pointer(c.f.Module))), moduleInstanceAddressRegister)
+	tmpX, tmpY := regs[0], regs[1]
 
 	// "tmpX = ce.ModuleInstanceAddress"
 	c.assembler.CompileMemoryToRegister(arm64.MOVD, arm64ReservedRegisterForCallEngine, callEngineModuleContextModuleInstanceAddressOffset, tmpX)
 
 	// If the module instance address stays the same, we could skip the entire code below.
-	c.assembler.CompileTwoRegistersToNone(arm64.CMP, moduleInstanceAddressRegister, tmpX)
+	c.assembler.CompileTwoRegistersToNone(arm64.CMP, arm64CallingConventionModuleInstanceAddressRegister, tmpX)
 	brIfModuleUnchanged := c.assembler.CompileJump(arm64.BEQ)
 
-	// Otherwise, we have to update the following fields:
+	// Otherwise, update the moduleEngine.moduleContext.ModuleInstanceAddress.
+	c.assembler.CompileRegisterToMemory(arm64.MOVD,
+		arm64CallingConventionModuleInstanceAddressRegister,
+		arm64ReservedRegisterForCallEngine, callEngineModuleContextModuleInstanceAddressOffset,
+	)
+
+	// Also, we have to update the following fields:
 	// * callEngine.moduleContext.globalElement0Address
 	// * callEngine.moduleContext.memoryElement0Address
 	// * callEngine.moduleContext.memorySliceLen
@@ -3066,7 +3096,7 @@ func (c *arm64Compiler) compileModuleContextInitialization() error {
 	if len(c.f.Module.Globals) > 0 {
 		// "tmpX = &moduleInstance.Globals[0]"
 		c.assembler.CompileMemoryToRegister(arm64.MOVD,
-			moduleInstanceAddressRegister, moduleInstanceGlobalsOffset,
+			arm64CallingConventionModuleInstanceAddressRegister, moduleInstanceGlobalsOffset,
 			tmpX,
 		)
 
@@ -3086,7 +3116,7 @@ func (c *arm64Compiler) compileModuleContextInitialization() error {
 		// "tmpX = moduleInstance.Memory"
 		c.assembler.CompileMemoryToRegister(
 			arm64.MOVD,
-			moduleInstanceAddressRegister, moduleInstanceMemoryOffset,
+			arm64CallingConventionModuleInstanceAddressRegister, moduleInstanceMemoryOffset,
 			tmpX,
 		)
 
@@ -3130,7 +3160,7 @@ func (c *arm64Compiler) compileModuleContextInitialization() error {
 		// "tmpX = &tables[0] (type of **wasm.Table)"
 		c.assembler.CompileMemoryToRegister(
 			arm64.MOVD,
-			moduleInstanceAddressRegister, moduleInstanceTableOffset,
+			arm64CallingConventionModuleInstanceAddressRegister, moduleInstanceTableOffset,
 			tmpX,
 		)
 
@@ -3175,7 +3205,7 @@ func (c *arm64Compiler) compileModuleContextInitialization() error {
 		// * https://github.com/golang/go/blob/release-branch.go1.17/src/runtime/runtime2.go#L207-L210
 		c.assembler.CompileMemoryToRegister(
 			arm64.MOVD,
-			moduleInstanceAddressRegister, moduleInstanceEngineOffset+interfaceDataOffset,
+			arm64CallingConventionModuleInstanceAddressRegister, moduleInstanceEngineOffset+interfaceDataOffset,
 			tmpX,
 		)
 


Test patch:
diff --git a/internal/wasm/jit/engine_test.go b/internal/wasm/jit/engine_test.go
index 4261b096..1222ee86 100644
--- a/internal/wasm/jit/engine_test.go
+++ b/internal/wasm/jit/engine_test.go
@@ -60,6 +60,7 @@ func TestJIT_VerifyOffsetValue(t *testing.T) {
 	require.Equal(t, int(unsafe.Offsetof(compiledFunc.codeInitialAddress)), compiledFunctionCodeInitialAddressOffset)
 	require.Equal(t, int(unsafe.Offsetof(compiledFunc.stackPointerCeil)), compiledFunctionStackPointerCeilOffset)
 	require.Equal(t, int(unsafe.Offsetof(compiledFunc.source)), compiledFunctionSourceOffset)
+	require.Equal(t, int(unsafe.Offsetof(compiledFunc.moduleInstanceAddress)), compiledFunctionModuleInstanceAddressOffset)
 
 	// Offsets for wasm.ModuleInstance.
 	var moduleInstance wasm.ModuleInstance
diff --git a/internal/wasm/jit/jit_controlflow_test.go b/internal/wasm/jit/jit_controlflow_test.go
index 62e96ad1..f56a2745 100644
--- a/internal/wasm/jit/jit_controlflow_test.go
+++ b/internal/wasm/jit/jit_controlflow_test.go
@@ -31,7 +31,10 @@ func TestCompiler_compileHostFunction(t *testing.T) {
 	require.Equal(t, jitCallStatusCodeCallHostFunction, env.jitStatus())
 
 	// Re-enter the return address.
-	jitcall(env.callFrameStackPeek().returnAddress, uintptr(unsafe.Pointer(env.callEngine())))
+	jitcall(env.callFrameStackPeek().returnAddress,
+		uintptr(unsafe.Pointer(env.callEngine())),
+		uintptr(unsafe.Pointer(env.module())),
+	)
 
 	// After that, the code must exit with returned status.
 	require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
@@ -665,8 +668,9 @@ func TestCompiler_compileCallIndirect(t *testing.T) {
 						require.NoError(t, err)
 
 						cf := &compiledFunction{
-							codeSegment:        code,
-							codeInitialAddress: uintptr(unsafe.Pointer(&code[0])),
+							codeSegment:           code,
+							codeInitialAddress:    uintptr(unsafe.Pointer(&code[0])),
+							moduleInstanceAddress: uintptr(unsafe.Pointer(env.moduleInstance)),
 							source: &wasm.FunctionInstance{
 								TypeID: targetTypeID,
 							},
@@ -717,7 +721,10 @@ func TestCompiler_compileCallIndirect(t *testing.T) {
 							// Grow the callFrame stack, and exec again from the return address.
 							ce := env.callEngine()
 							ce.builtinFunctionGrowCallFrameStack()
-							jitcall(env.callFrameStackPeek().returnAddress, uintptr(unsafe.Pointer(ce)))
+							jitcall(
+								env.callFrameStackPeek().returnAddress, uintptr(unsafe.Pointer(ce)),
+								uintptr(unsafe.Pointer(env.module())),
+							)
 						}
 
 						require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
@@ -774,8 +781,9 @@ func TestCompiler_compileCall(t *testing.T) {
 					require.NoError(t, err)
 					index := wasm.Index(i)
 					me.compiledFunctions = append(me.compiledFunctions, &compiledFunction{
-						codeSegment:        code,
-						codeInitialAddress: uintptr(unsafe.Pointer(&code[0])),
+						codeSegment:           code,
+						codeInitialAddress:    uintptr(unsafe.Pointer(&code[0])),
+						moduleInstanceAddress: uintptr(unsafe.Pointer(env.moduleInstance)),
 					})
 					env.module().Functions = append(env.module().Functions,
 						&wasm.FunctionInstance{Type: targetFunctionType, Index: index})
@@ -816,7 +824,10 @@ func TestCompiler_compileCall(t *testing.T) {
 				// Grow the callFrame stack, and exec again from the return address.
 				ce := env.callEngine()
 				ce.builtinFunctionGrowCallFrameStack()
-				jitcall(env.callFrameStackPeek().returnAddress, uintptr(unsafe.Pointer(ce)))
+				jitcall(
+					env.callFrameStackPeek().returnAddress, uintptr(unsafe.Pointer(ce)),
+					uintptr(unsafe.Pointer(env.module())),
+				)
 			}
 
 			// Check status and returned values.
@@ -879,7 +890,10 @@ func TestCompiler_returnFunction(t *testing.T) {
 				require.NoError(t, err)
 
 				// Compiles and adds to the engine.
-				compiledFunction := &compiledFunction{codeSegment: code, codeInitialAddress: uintptr(unsafe.Pointer(&code[0]))}
+				compiledFunction := &compiledFunction{
+					codeSegment: code, codeInitialAddress: uintptr(unsafe.Pointer(&code[0])),
+					moduleInstanceAddress: uintptr(unsafe.Pointer(env.moduleInstance)),
+				}
 				moduleEngine.compiledFunctions = append(moduleEngine.compiledFunctions, compiledFunction)
 
 				// Pushes the frame whose return address equals the beginning of the function just compiled.
diff --git a/internal/wasm/jit/jit_initialization_test.go b/internal/wasm/jit/jit_initialization_test.go
index cc0fd6eb..dfacc3e5 100644
--- a/internal/wasm/jit/jit_initialization_test.go
+++ b/internal/wasm/jit/jit_initialization_test.go
@@ -177,7 +177,10 @@ func TestCompiler_compileMaybeGrowValueStack(t *testing.T) {
 		// Reenter from the return address.
 		returnAddress := env.callFrameStackPeek().returnAddress
 		require.NotZero(t, returnAddress)
-		jitcall(returnAddress, uintptr(unsafe.Pointer(env.callEngine())))
+		jitcall(
+			returnAddress, uintptr(unsafe.Pointer(env.callEngine())),
+			uintptr(unsafe.Pointer(env.module())),
+		)
 
 		// Check the result. This should be "Returned".
 		require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
diff --git a/internal/wasm/jit/jit_memory_test.go b/internal/wasm/jit/jit_memory_test.go
index 2a84404f..8a3f4fc4 100644
--- a/internal/wasm/jit/jit_memory_test.go
+++ b/internal/wasm/jit/jit_memory_test.go
@@ -40,7 +40,10 @@ func TestCompiler_compileMemoryGrow(t *testing.T) {
 	require.Equal(t, builtinFunctionIndexMemoryGrow, env.builtinFunctionCallAddress())
 
 	// Reenter from the return address.
-	jitcall(env.callFrameStackPeek().returnAddress, uintptr(unsafe.Pointer(env.callEngine())))
+	jitcall(
+		env.callFrameStackPeek().returnAddress, uintptr(unsafe.Pointer(env.callEngine())),
+		uintptr(unsafe.Pointer(env.module())),
+	)
 
 	// Check if the code successfully executed the code after builtin function call.
 	require.Equal(t, expValue, env.stackTopAsUint32())
diff --git a/internal/wasm/jit/jit_test.go b/internal/wasm/jit/jit_test.go
index 6bc7a4f0..6d51822a 100644
--- a/internal/wasm/jit/jit_test.go
+++ b/internal/wasm/jit/jit_test.go
@@ -120,8 +120,9 @@ func (j *jitEnv) callEngine() *callEngine {
 
 func (j *jitEnv) exec(code []byte) {
 	compiledFunction := &compiledFunction{
-		codeSegment:        code,
-		codeInitialAddress: uintptr(unsafe.Pointer(&code[0])),
+		codeSegment:           code,
+		codeInitialAddress:    uintptr(unsafe.Pointer(&code[0])),
+		moduleInstanceAddress: uintptr(unsafe.Pointer(j.moduleInstance)),
 		source: &wasm.FunctionInstance{
 			Kind:   wasm.FunctionKindWasm,
 			Type:   &wasm.FunctionType{},
@@ -129,11 +130,13 @@ func (j *jitEnv) exec(code []byte) {
 		},
 	}
 
-	j.ce.pushCallFrame(compiledFunction)
+	j.ce.callFrameStack[j.ce.globalContext.callFrameStackPointer] = callFrame{compiledFunction: compiledFunction}
+	j.ce.globalContext.callFrameStackPointer++
 
 	jitcall(
 		uintptr(unsafe.Pointer(&code[0])),
 		uintptr(unsafe.Pointer(j.ce)),
+		uintptr(unsafe.Pointer(j.moduleInstance)),
 	)
 }
 
diff --git a/tests/bench/bench_test.go b/tests/bench/bench_test.go
index ad9f6ae4..f67da6a6 100644
--- a/tests/bench/bench_test.go
+++ b/tests/bench/bench_test.go
@@ -22,7 +22,7 @@ func BenchmarkEngines(b *testing.B) {
 		defer m.Close()
 		runAllBenches(b, m)
 	})
-	if runtime.GOARCH == "amd64" {
+	if runtime.GOARCH == "amd64" || runtime.GOARCH == "arm64" {
 		b.Run("jit", func(b *testing.B) {
 			m := instantiateHostFunctionModuleWithEngine(b, wazero.NewRuntimeConfigJIT())
 			defer m.Close()
@@ -46,8 +46,10 @@ func runBase64Benches(b *testing.B, m api.Module) {
 		numPerExec := uint64(numPerExec)
 		b.ResetTimer()
 		b.Run(fmt.Sprintf("base64_%d_per_exec", numPerExec), func(b *testing.B) {
-			if _, err := base64.Call(nil, numPerExec); err != nil {
-				b.Fatal(err)
+			for i := 0; i < b.N; i++ {
+				if _, err := base64.Call(nil, numPerExec); err != nil {
+					b.Fatal(err)
+				}
 			}
 		})
 	}