Repository: tetratelabs/wazero
Pull request #376 (2022-03-16 00:35:07)

jit: isolate and document per-build tag specific definitions 
This is the continuation of refactor on jit pkg. Notably, 
this commit isolates per-build tag definition into arch_*
pkg and provide build-tag free documentation interface
on `arch.go`.


mathetake: windows failure seems persistent..
mathetake: Seems like this is due to the notorious `TestModuleEngine_Close_Concurrent`.. I will open an issue 

```
fatal error: unexpected signal during runtime execution
[signal 0xc0000005 code=0x1 addr=0x2cbf0fedd28 pc=0x8fc66a]

goroutine 22 [running]:
runtime.throw({0xbbcef2, 0x2cbf0fedd[18](https://github.com/tetratelabs/wazero/runs/5563644494?check_suite_focus=true#step:5:18)})
	C:/hostedtoolcache/windows/go/1.17.7/x64/src/runtime/panic.go:1[19](https://github.com/tetratelabs/wazero/runs/5563644494?check_suite_focus=true#step:5:19)8 +0x76 fp=0xc0005397c0 sp=0xc000539790 pc=0x908eb6
runtime.sigpanic()
	C:/hostedtoolcache/windows/go/1.17.7/x64/src/runtime/signal_windows.go:248 +0x22b fp=0xc000539808 sp=0xc0005397c0 pc=0x91d60b
runtime.newBucket(0x1, 0x9)
	C:/hostedtoolcache/windows/go/1.17.7/x64/src/runtime/mprof.go:175 +0x6a fp=0xc000539838 sp=0xc000539808 pc=0x8fc66a
runtime.stkbucket(0x1, 0xd0, {0xc0005398e8, 0xc0005396f0, 0x[20](https://github.com/tetratelabs/wazero/runs/5563644494?check_suite_focus=true#step:5:20)}, 0x1)
	C:/hostedtoolcache/windows/go/1.17.7/x64/src/runtime/mprof.go:240 +0x1f4 fp=0xc0005398a0 sp=0xc000539838 pc=0x8fc974
runtime.mProf_Malloc(0xc000022340, 0xd0)
	C:/hostedtoolcache/windows/go/1.17.7/x64/src/runtime/mprof.go:344 +0xb1 fp=0xc000539a10 sp=0xc0005398a0 pc=0x8fce91
runtime.profilealloc(0x2cbcbac4b0d, 0xd0, 0x8e6e2f)
	C:/hostedtoolcache/windows/go/1.17.7/x64/src/runtime/malloc.go:1270 +0x85 fp=0xc000539a48 sp=0xc000539a10 pc=0x8ddb05
runtime.mallocgc(0xd0, 0xb80120, 0x1)
	C:/hostedtoolcache/windows/go/1.17.7/x64/src/runtime/malloc.go:1149 +0x72a fp=0xc000539ac8 sp=0xc000539a48 pc=0x8dd6aa
runtime.newobject(0x8dd634)
	C:/hostedtoolcache/windows/go/1.17.7/x64/src/runtime/malloc.go:1234 +0x27 fp=0xc000539af0 sp=0xc000539ac8 pc=0x8dd987
runtime.CallersFrames(...)
	C:/hostedtoolcache/windows/go/1.17.7/x64/src/runtime/symtab.go:66
runtime.Caller(0x2)
	C:/hostedtoolcache/windows/go/1.17.7/x64/src/runtime/extern.go:203 +0x7e fp=0xc000539c00 sp=0xc000539af0 pc=0x8d9ebe
github.com/stretchr/testify/assert.CallerInfo()
	C:/Users/runneradmin/go/pkg/mod/github.com/stretchr/testify@v1.7.0/assert/assertions.go:115 +0x6c fp=0xc000539cc8 sp=0xc000539c00 pc=0xaf5fac
github.com/stretchr/testify/assert.Fail({0x2cbf0e8e2c8, 0xc00037e340}, {0xbb8ff4, 0x[21](https://github.com/tetratelabs/wazero/runs/5563644494?check_suite_focus=true#step:5:21)}, {0x0, 0x0, 0x0})
	C:/Users/runneradmin/go/pkg/mod/github.com/stretchr/testify@v1.7.0/assert/assertions.go:[24](https://github.com/tetratelabs/wazero/runs/5563644494?check_suite_focus=true#step:5:24)6 +0x89 fp=0xc000539d98 sp=0xc000539cc8 pc=0xaf6929
github.com/stretchr/testify/assert.Error({0x2cbf0e8e2c8, 0xc00037e340}, {0x0, 0x0}, {0x0, 0x0, 0x0})
	C:/Users/runneradmin/go/pkg/mod/github.com/stretchr/testify@v1.7.0/assert/assertions.go:1349 +0x9d fp=0xc000539de0 sp=0xc000539d98 pc=0xaf9d7d
github.com/stretchr/testify/require.Error({0xc19[25](https://github.com/tetratelabs/wazero/runs/5563644494?check_suite_focus=true#step:5:25)8, 0xc00037e340}, {0x0, 0x0}, {0x0, 0x0, 0x0})
	C:/Users/runneradmin/go/pkg/mod/github.com/stretchr/testify@v1.7.0/require/require.go:253 +0x96 fp=0xc000539e[28](https://github.com/tetratelabs/wazero/runs/5563644494?check_suite_focus=true#step:5:28) sp=0xc000539de0 pc=0xafb476
github.com/tetratelabs/wazero/internal/wasm/jit.TestModuleEngine_Close_Concurrent.func2(0xc00037e340)
	D:/a/wazero/wazero/internal/wasm/jit/engine_test.go:361 +0xaa7 fp=0xc000539f70 sp=0xc000539e28 pc=0xb15507
testing.tRunner(0xc00037e340, 0xc00004c780)
	C:/hostedtoolcache/windows/go/1.17.7/x64/src/testing/testing.go:1259 +0x102 fp=0xc000539fc0 sp=0xc000539f70 pc=0x9a9202
testing.(*T).Run·dwrap·21()
	C:/hostedtoolcache/windows/go/1.17.7/x64/src/testing/testing.go:1[30](https://github.com/tetratelabs/wazero/runs/5563644494?check_suite_focus=true#step:5:30)6 +0x2a fp=0xc000539fe0 sp=0xc000539fc0 pc=0x9a9f0a
runtime.goexit()
	C:/hostedtoolcache/windows/go/1.17.7/x64/src/runtime/asm_amd64.s:1581 +0x1 fp=0xc000539fe8 sp=0xc000539fe0 pc=0x9396e1
created by testing.(*T).Run
	C:/hostedtoolcache/windows/go/1.17.7/x64/src/testing/testing.go:1306 +0x[35](https://github.com/tetratelabs/wazero/runs/5563644494?check_suite_focus=true#step:5:35)a
```

Patch:
diff --git a/internal/wasm/jit/arch.go b/internal/wasm/jit/arch.go
new file mode 100644
index 00000000..16fb8da3
--- /dev/null
+++ b/internal/wasm/jit/arch.go
@@ -0,0 +1,29 @@
+package jit
+
+import (
+	wasm "github.com/tetratelabs/wazero/internal/wasm"
+	"github.com/tetratelabs/wazero/internal/wazeroir"
+)
+
+// Following variables must be initialized in init() function in architecture-specific arch_*.go file which
+// is guarded by build tag. This file exists as build-tag free in order to provide documentation interface
+// over these per-architecture definitions of the following symbols.
+var (
+	// newArchContext returns a new archContext which is architecture-specific type to be embedded in callEngine.
+	newArchContext func() archContext
+
+	// jitcall is used by callEngine.execWasmFunction and the entrypoint to enter the JITed native code.
+	// codeSegment is the pointer to the initial instruction of the compiled native code.
+	// ce is "*callEngine" as uintptr.
+	jitcall func(codeSegment, ce uintptr)
+
+	// newCompiler returns a new compiler interface which can be used to compile the given function instance.
+	// Note: ir param can be nil for host functions.
+	newCompiler func(f *wasm.FunctionInstance, ir *wazeroir.CompilationResult) (compiler, error)
+
+	// unreservedGeneralPurposeIntRegisters contains unreserved general purpose registers of integer type.
+	unreservedGeneralPurposeIntRegisters []int16
+
+	// unreservedGeneralPurposeFloatRegisters contains unreserved general purpose registers of scalar float type.
+	unreservedGeneralPurposeFloatRegisters []int16
+)
diff --git a/internal/wasm/jit/arch_amd64.go b/internal/wasm/jit/arch_amd64.go
new file mode 100644
index 00000000..fb4b2290
--- /dev/null
+++ b/internal/wasm/jit/arch_amd64.go
@@ -0,0 +1,63 @@
+package jit
+
+import (
+	"fmt"
+
+	asm "github.com/twitchyliquid64/golang-asm"
+	"github.com/twitchyliquid64/golang-asm/obj/x86"
+
+	wasm "github.com/tetratelabs/wazero/internal/wasm"
+	"github.com/tetratelabs/wazero/internal/wazeroir"
+)
+
+// init initializes variables for the amd64 architecture
+func init() {
+	jitcall = jitcallImpl
+	newCompiler = newCompilerImpl
+	newArchContext = newArchContextImpl
+	unreservedGeneralPurposeFloatRegisters = []int16{
+		x86.REG_X0, x86.REG_X1, x86.REG_X2, x86.REG_X3,
+		x86.REG_X4, x86.REG_X5, x86.REG_X6, x86.REG_X7,
+		x86.REG_X8, x86.REG_X9, x86.REG_X10, x86.REG_X11,
+		x86.REG_X12, x86.REG_X13, x86.REG_X14, x86.REG_X15,
+	}
+	// Note that we never invoke "call" instruction,
+	// so we don't need to care about the calling convention.
+	// TODO: Maybe it is safe just save rbp, rsp somewhere
+	// in Go-allocated variables, and reuse these registers
+	// in JITed functions and write them back before returns.
+	unreservedGeneralPurposeIntRegisters = []int16{
+		x86.REG_AX, x86.REG_CX, x86.REG_DX, x86.REG_BX,
+		x86.REG_SI, x86.REG_DI, x86.REG_R8, x86.REG_R9,
+		x86.REG_R10, x86.REG_R11, x86.REG_R12,
+	}
+}
+
+// jitcallImpl implements jitcall for amd64 architecture.
+// Note: this function's body is defined in arch_amd64.s
+func jitcallImpl(codeSegment, ce uintptr)
+
+// newCompilerImpl implements newCompiler for amd64 architecture.
+func newCompilerImpl(f *wasm.FunctionInstance, ir *wazeroir.CompilationResult) (compiler, error) {
+	b, err := asm.NewBuilder("amd64", 1024)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create a new assembly builder: %w", err)
+	}
+
+	compiler := &amd64Compiler{
+		f:             f,
+		builder:       b,
+		locationStack: newValueLocationStack(),
+		currentLabel:  wazeroir.EntrypointLabel,
+		ir:            ir,
+		labels:        map[string]*labelInfo{},
+	}
+	return compiler, nil
+}
+
+// archContext is embedded in callEngine in order to store architecture-specific data.
+// For amd64, this is empty.
+type archContext struct{}
+
+// newArchContextImpl implements newArchContext for amd64 architecture.
+func newArchContextImpl() (ret archContext) { return }
diff --git a/internal/wasm/jit/jit_amd64.s b/internal/wasm/jit/arch_amd64.s
similarity index 77%
rename from internal/wasm/jit/jit_amd64.s
rename to internal/wasm/jit/arch_amd64.s
index 5b84dc72..a1ceb0d9 100644
--- a/internal/wasm/jit/jit_amd64.s
+++ b/internal/wasm/jit/arch_amd64.s
@@ -3,8 +3,8 @@
 #include "funcdata.h"
 #include "textflag.h"
 
-// jitcall(codeSegment, ce)
-TEXT ·jitcall(SB),NOSPLIT|NOFRAME,$0-16
+// jitcallImpl(codeSegment, ce)
+TEXT ·jitcallImpl(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.
diff --git a/internal/wasm/jit/arch_arm64.go b/internal/wasm/jit/arch_arm64.go
new file mode 100644
index 00000000..e4cbe169
--- /dev/null
+++ b/internal/wasm/jit/arch_arm64.go
@@ -0,0 +1,98 @@
+package jit
+
+import (
+	"fmt"
+	"math"
+
+	asm "github.com/twitchyliquid64/golang-asm"
+	"github.com/twitchyliquid64/golang-asm/obj/arm64"
+
+	wasm "github.com/tetratelabs/wazero/internal/wasm"
+	"github.com/tetratelabs/wazero/internal/wazeroir"
+)
+
+// init initializes variables for the arm64 architecture
+func init() {
+	jitcall = jitcallImpl
+	newCompiler = newCompilerImpl
+	newArchContext = newArchContextImpl
+	unreservedGeneralPurposeFloatRegisters = []int16{
+		arm64.REG_F0, arm64.REG_F1, arm64.REG_F2, arm64.REG_F3,
+		arm64.REG_F4, arm64.REG_F5, arm64.REG_F6, arm64.REG_F7, arm64.REG_F8,
+		arm64.REG_F9, arm64.REG_F10, arm64.REG_F11, arm64.REG_F12, arm64.REG_F13,
+		arm64.REG_F14, arm64.REG_F15, arm64.REG_F16, arm64.REG_F17, arm64.REG_F18,
+		arm64.REG_F19, arm64.REG_F20, arm64.REG_F21, arm64.REG_F22, arm64.REG_F23,
+		arm64.REG_F24, arm64.REG_F25, arm64.REG_F26, arm64.REG_F27, arm64.REG_F28,
+		arm64.REG_F29, arm64.REG_F30, arm64.REG_F31,
+	}
+	// Note (see arm64 section in https://go.dev/doc/asm):
+	// * REG_R18 is reserved as a platform register, and we don't use it in JIT.
+	// * REG_R28 is reserved for Goroutine by Go runtime, and we don't use it in JIT.
+	unreservedGeneralPurposeIntRegisters = []int16{
+		arm64.REG_R4, arm64.REG_R5, arm64.REG_R6, arm64.REG_R7, arm64.REG_R8,
+		arm64.REG_R9, arm64.REG_R10, arm64.REG_R11, arm64.REG_R12, arm64.REG_R13,
+		arm64.REG_R14, arm64.REG_R15, arm64.REG_R16, arm64.REG_R17, arm64.REG_R19,
+		arm64.REG_R20, arm64.REG_R21, arm64.REG_R22, arm64.REG_R23, arm64.REG_R24,
+		arm64.REG_R25, arm64.REG_R26, arm64.REG_R27, arm64.REG_R29, arm64.REG_R30,
+	}
+}
+
+// jitcallImpl implements jitcallfor arm64 architecture.
+// Note: this function's body is defined in arch_arm64.s
+func jitcallImpl(codeSegment, ce uintptr)
+
+// newCompilerImpl implements newCompiler for amd64 architecture.
+func newCompilerImpl(f *wasm.FunctionInstance, ir *wazeroir.CompilationResult) (compiler, error) {
+	b, err := asm.NewBuilder("arm64", 1024)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create a new assembly builder: %w", err)
+	}
+
+	compiler := &arm64Compiler{
+		f:             f,
+		builder:       b,
+		locationStack: newValueLocationStack(),
+		ir:            ir,
+		labels:        map[string]*labelInfo{},
+	}
+	return compiler, nil
+}
+
+// archContext is embedded in callEngine in order to store architecture-specific data.
+type archContext struct {
+	// jitCallReturnAddress holds the absolute return address for jitcall.
+	// The value is set whenever jitcall is executed and done in jit_arm64.s
+	// Native code can return back to the ce.execWasmFunction's main loop back by
+	// executing "ret" instruction with this value. See arm64Compiler.exit.
+	// Note: this is only used by JIT code so mark this as nolint.
+	jitCallReturnAddress uint64 //nolint
+
+	// Loading large constants in arm64 is a bit costly, so we place the following
+	// consts on callEngine struct so that we can quickly access them during various operations.
+
+	// minimum32BitSignedInt is used for overflow check for 32-bit signed division.
+	// Note: this can be obtained by moving $1 and doing left-shift with 31, but it is
+	// slower than directly loading fron this location.
+	minimum32BitSignedInt int32
+	// Note: this can be obtained by moving $1 and doing left-shift with 63, but it is
+	// slower than directly loading fron this location.
+	// minimum64BitSignedInt is used for overflow check for 64-bit signed division.
+	minimum64BitSignedInt int64
+}
+
+const (
+	// callEngineArchContextJITCallReturnAddressOffset is the offset of archContext.jitCallReturnAddress in callEngine.
+	callEngineArchContextJITCallReturnAddressOffset = 120
+	// callEngineArchContextMinimum32BitSignedIntOffset is the offset of archContext.minimum32BitSignedIntAddress in callEngine.
+	callEngineArchContextMinimum32BitSignedIntOffset = 128
+	// callEngineArchContextMinimum64BitSignedIntOffset is the offset of archContext.minimum64BitSignedIntAddress in callEngine.
+	callEngineArchContextMinimum64BitSignedIntOffset = 136
+)
+
+// newArchContextImpl implements newArchContext for amd64 architecture.
+func newArchContextImpl() archContext {
+	return archContext{
+		minimum32BitSignedInt: math.MinInt32,
+		minimum64BitSignedInt: math.MinInt64,
+	}
+}
diff --git a/internal/wasm/jit/jit_arm64.s b/internal/wasm/jit/arch_arm64.s
similarity index 88%
rename from internal/wasm/jit/jit_arm64.s
rename to internal/wasm/jit/arch_arm64.s
index b59294fa..6af929e3 100644
--- a/internal/wasm/jit/jit_arm64.s
+++ b/internal/wasm/jit/arch_arm64.s
@@ -3,8 +3,8 @@
 #include "funcdata.h"
 #include "textflag.h"
 
-// jitcall(codeSegment, ce)
-TEXT ·jitcall(SB),NOSPLIT|NOFRAME,$0-16
+// jitcallImpl(codeSegment, ce)
+TEXT ·jitcallImpl(SB),NOSPLIT|NOFRAME,$0-16
         // Load the address of native code.
         MOVD codeSegment+0(FP),R1
         // Load the address of *callEngine.
diff --git a/internal/wasm/jit/arch_other.go b/internal/wasm/jit/arch_other.go
new file mode 100644
index 00000000..5f21edbc
--- /dev/null
+++ b/internal/wasm/jit/arch_other.go
@@ -0,0 +1,6 @@
+//go:build !amd64 && !arm64
+
+package jit
+
+// archContext is empty on an unsupported architecture.
+type archContext struct{}
diff --git a/internal/wasm/jit/engine.go b/internal/wasm/jit/engine.go
index 12c20d71..771d7478 100644
--- a/internal/wasm/jit/engine.go
+++ b/internal/wasm/jit/engine.go
@@ -779,10 +779,15 @@ func (ce *callEngine) builtinFunctionMemoryGrow(mem *wasm.MemoryInstance) {
 	ce.moduleContext.memoryElement0Address = bufSliceHeader.Data
 }
 
+// golang-asm is not goroutine-safe so we take lock until we complete the compilation.
+// TODO: delete after https://github.com/tetratelabs/wazero/issues/233
+var assemblerMutex = &sync.Mutex{}
+
 func compileHostFunction(f *wasm.FunctionInstance) (*compiledFunction, error) {
-	compiler, done, err := newCompiler(f, nil)
-	defer done()
+	assemblerMutex.Lock()
+	defer assemblerMutex.Unlock()
 
+	compiler, err := newCompiler(f, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -810,13 +815,15 @@ func compileHostFunction(f *wasm.FunctionInstance) (*compiledFunction, error) {
 }
 
 func compileWasmFunction(f *wasm.FunctionInstance) (*compiledFunction, error) {
+	assemblerMutex.Lock()
+	defer assemblerMutex.Unlock()
+
 	ir, err := wazeroir.Compile(f)
 	if err != nil {
 		return nil, fmt.Errorf("failed to lower to wazeroir: %w", err)
 	}
 
-	compiler, done, err := newCompiler(f, ir)
-	defer done()
+	compiler, err := newCompiler(f, ir)
 	if err != nil {
 		return nil, fmt.Errorf("failed to initialize assembly builder: %w", err)
 	}
diff --git a/internal/wasm/jit/jit_amd64.go b/internal/wasm/jit/jit_amd64.go
index 88622510..c3e4ed1f 100644
--- a/internal/wasm/jit/jit_amd64.go
+++ b/internal/wasm/jit/jit_amd64.go
@@ -11,7 +11,6 @@ import (
 	"fmt"
 	"math"
 	"runtime"
-	"sync"
 	"unsafe"
 
 	asm "github.com/twitchyliquid64/golang-asm"
@@ -74,51 +73,29 @@ func init() {
 	float64ForMaximumSigned64bitIntPlusOneAddress = uintptr(unsafe.Pointer(&float64ForMaximumSigned64bitIntPlusOne))
 }
 
-// jitcall is implemented in jit_amd64.s as a Go Assembler function.
-// This is used by callEngine.execWasmFunction and the entrypoint to enter the JITed native code.
-// codeSegment is the pointer to the initial instruction of the compiled native code.
-// ce is the pointer to the *callEngine as uintptr.
-func jitcall(codeSegment, ce uintptr)
-
-// archContext is embedded in callEngine in order to store architecture-specific data.
-// For amd64, this is empty.
-type archContext struct{}
-
-func newArchContext() (ret archContext) { return }
-
-// golang-asm is not goroutine-safe so we take lock until we complete the compilation.
-// TODO: delete after https://github.com/tetratelabs/wazero/issues/233
-var assemblerMutex = &sync.Mutex{}
-
-func unlockAssembler() {
-	assemblerMutex.Unlock()
-}
-
-// newCompiler returns a new compiler interface which can be used to compile the given function instance.
-// The function returned must be invoked when finished compiling, so use `defer` to ensure this.
-// Note: ir param can be nil for host functions.
-func newCompiler(f *wasm.FunctionInstance, ir *wazeroir.CompilationResult) (compiler, func(), error) {
-	// golang-asm is not goroutine-safe so we take lock until we complete the compilation.
-	// TODO: delete after https://github.com/tetratelabs/wazero/issues/233
-	assemblerMutex.Lock()
-
-	// We can choose arbitrary number instead of 1024 which indicates the cache size in the compiler.
-	// TODO: optimize the number.
-	b, err := asm.NewBuilder("amd64", 1024)
-	if err != nil {
-		return nil, unlockAssembler, fmt.Errorf("failed to create a new assembly builder: %w", err)
-	}
+const (
+	conditionalRegisterStateE  = conditionalRegisterStateUnset + 1 + iota // ZF equal to zero
+	conditionalRegisterStateNE                                            //˜ZF not equal to zero
+	conditionalRegisterStateS                                             // SF negative
+	conditionalRegisterStateNS                                            // ˜SF non-negative
+	conditionalRegisterStateG                                             // ˜(SF xor OF) & ˜ ZF greater (signed >)
+	conditionalRegisterStateGE                                            // ˜(SF xor OF) greater or equal (signed >=)
+	conditionalRegisterStateL                                             // SF xor OF less (signed <)
+	conditionalRegisterStateLE                                            // (SF xor OF) | ZF less or equal (signed <=)
+	conditionalRegisterStateA                                             // ˜CF & ˜ZF above (unsigned >)
+	conditionalRegisterStateAE                                            // ˜CF above or equal (unsigned >=)
+	conditionalRegisterStateB                                             // CF below (unsigned <)
+	conditionalRegisterStateBE                                            // CF | ZF below or equal (unsigned <=)
+)
 
-	compiler := &amd64Compiler{
-		f:             f,
-		builder:       b,
-		locationStack: newValueLocationStack(),
-		currentLabel:  wazeroir.EntrypointLabel,
-		ir:            ir,
-		labels:        map[string]*labelInfo{},
-	}
-	return compiler, unlockAssembler, nil
-}
+const (
+	// reservedRegisterForCallEngine: pointer to callEngine (i.e. *callEngine as uintptr)
+	reservedRegisterForCallEngine = x86.REG_R13
+	// reservedRegisterForStackBasePointerAddress: stack base pointer's address (callEngine.stackBasePointer) in the current function call.
+	reservedRegisterForStackBasePointerAddress = x86.REG_R14
+	// reservedRegisterForMemory: pointer to the memory slice's data (i.e. &memory.Buffer[0] as uintptr).
+	reservedRegisterForMemory = x86.REG_R15
+)
 
 func (c *amd64Compiler) String() string {
 	return c.locationStack.String()
@@ -161,6 +138,12 @@ func (c *amd64Compiler) addStaticData(d []byte) {
 	c.staticData = append(c.staticData, d)
 }
 
+func (c *amd64Compiler) pushValueLocationOnRegister(reg int16) (ret *valueLocation) {
+	ret = c.locationStack.pushValueLocationOnRegister(reg)
+	c.locationStack.markRegisterUsed(reg)
+	return
+}
+
 type labelInfo struct {
 	// initialInstruction is the initial instruction for this label so other block can jump into it.
 	initialInstruction *obj.Prog
@@ -527,7 +510,7 @@ func (c *amd64Compiler) compileSwap(o *wazeroir.OperationSwap) error {
 		reg := x1.register
 		c.locationStack.markRegisterUnused(reg)
 		// Save x1's value to the temporary top of the stack.
-		tmpStackLocation := c.locationStack.pushValueLocationOnRegister(reg)
+		tmpStackLocation := c.pushValueLocationOnRegister(reg)
 		c.compileReleaseRegisterToStack(tmpStackLocation)
 		// Then move the x2's value to the x1's register location.
 		x2.register = reg
@@ -547,7 +530,7 @@ func (c *amd64Compiler) compileSwap(o *wazeroir.OperationSwap) error {
 		reg := x2.register
 		c.locationStack.markRegisterUnused(reg)
 		// Save x2's value to the temporary top of the stack.
-		tmpStackLocation := c.locationStack.pushValueLocationOnRegister(reg)
+		tmpStackLocation := c.pushValueLocationOnRegister(reg)
 		c.compileReleaseRegisterToStack(tmpStackLocation)
 		// Then move the x1's value to the x2's register location.
 		x1.register = reg
@@ -572,7 +555,7 @@ func (c *amd64Compiler) compileSwap(o *wazeroir.OperationSwap) error {
 		x2.setRegister(reg)
 		c.compileLoadValueOnStackToRegister(x2)
 		// Save x2's value to the temporary top of the stack.
-		tmpStackLocation := c.locationStack.pushValueLocationOnRegister(reg)
+		tmpStackLocation := c.pushValueLocationOnRegister(reg)
 		c.compileReleaseRegisterToStack(tmpStackLocation)
 		// Then move the x1's value to the x2's register location.
 		x1.register = reg
@@ -625,7 +608,7 @@ func (c *amd64Compiler) compileGlobalGet(o *wazeroir.OperationGlobalGet) error {
 	c.compileMemoryToRegisterInstruction(x86.AMOVQ, intReg, globalInstanceValueOffset, valueReg)
 
 	// Record that the retrieved global value on the top of the stack is now in a register.
-	loc := c.locationStack.pushValueLocationOnRegister(valueReg)
+	loc := c.pushValueLocationOnRegister(valueReg)
 	switch wasmType {
 	case wasm.ValueTypeI32, wasm.ValueTypeI64:
 		loc.setRegisterType(generalPurposeRegisterTypeInt)
@@ -1242,7 +1225,7 @@ func (c *amd64Compiler) compilePick(o *wazeroir.OperationPick) error {
 	}
 	// Now we already placed the picked value on the register,
 	// so push the location onto the stack.
-	loc := c.locationStack.pushValueLocationOnRegister(reg)
+	loc := c.pushValueLocationOnRegister(reg)
 	loc.setRegisterType(pickTarget.registerType())
 	return nil
 }
@@ -1414,7 +1397,7 @@ func (c *amd64Compiler) compileMulForInts(is32Bit bool, mulInstruction obj.As) e
 
 	// Now we have the result in the AX register,
 	// so we record it.
-	result := c.locationStack.pushValueLocationOnRegister(resultRegister)
+	result := c.pushValueLocationOnRegister(resultRegister)
 	result.setRegisterType(generalPurposeRegisterTypeInt)
 	return nil
 }
@@ -1498,7 +1481,7 @@ func (c *amd64Compiler) compileClz(o *wazeroir.OperationClz) error {
 
 	// We reused the same register of target for the result.
 	c.locationStack.markRegisterUnused(target.register)
-	result := c.locationStack.pushValueLocationOnRegister(target.register)
+	result := c.pushValueLocationOnRegister(target.register)
 	result.setRegisterType(generalPurposeRegisterTypeInt)
 	return nil
 }
@@ -1551,7 +1534,7 @@ func (c *amd64Compiler) compileCtz(o *wazeroir.OperationCtz) error {
 
 	// We reused the same register of target for the result.
 	c.locationStack.markRegisterUnused(target.register)
-	result := c.locationStack.pushValueLocationOnRegister(target.register)
+	result := c.pushValueLocationOnRegister(target.register)
 	result.setRegisterType(generalPurposeRegisterTypeInt)
 	return nil
 }
@@ -1571,7 +1554,7 @@ func (c *amd64Compiler) compilePopcnt(o *wazeroir.OperationPopcnt) error {
 
 	// We reused the same register of target for the result.
 	c.locationStack.markRegisterUnused(target.register)
-	result := c.locationStack.pushValueLocationOnRegister(target.register)
+	result := c.pushValueLocationOnRegister(target.register)
 	result.setRegisterType(generalPurposeRegisterTypeInt)
 	return nil
 }
@@ -1605,7 +1588,7 @@ func (c *amd64Compiler) compileDivForInts(is32Bit bool, signed bool) error {
 	}
 	// Now we have the quotient of the division result in the AX register,
 	// so we record it.
-	result := c.locationStack.pushValueLocationOnRegister(x86.REG_AX)
+	result := c.pushValueLocationOnRegister(x86.REG_AX)
 	result.setRegisterType(generalPurposeRegisterTypeInt)
 	return nil
 }
@@ -1628,7 +1611,7 @@ func (c *amd64Compiler) compileRem(o *wazeroir.OperationRem) (err error) {
 
 	// Now we have the remainder of the division result in the DX register,
 	// so we record it.
-	result := c.locationStack.pushValueLocationOnRegister(x86.REG_DX)
+	result := c.pushValueLocationOnRegister(x86.REG_DX)
 	result.setRegisterType(generalPurposeRegisterTypeInt)
 	return
 }
@@ -1881,7 +1864,7 @@ func (c *amd64Compiler) compileSimpleBinaryOp(instruction obj.As) error {
 	// We already stored the result in the register used by x1
 	// so we record it.
 	c.locationStack.markRegisterUnused(x1.register)
-	result := c.locationStack.pushValueLocationOnRegister(x1.register)
+	result := c.pushValueLocationOnRegister(x1.register)
 	result.setRegisterType(x1.registerType())
 	return nil
 }
@@ -2190,7 +2173,7 @@ func (c *amd64Compiler) compileMinOrMax(is32Bit bool, minOrMaxInstruction obj.As
 	// Record that we consumed the x2 and placed the minOrMax result in the x1's register.
 	c.locationStack.markRegisterUnused(x2.register)
 	c.locationStack.markRegisterUnused(x1.register)
-	c.locationStack.pushValueLocationOnRegister(x1.register)
+	c.pushValueLocationOnRegister(x1.register)
 	return nil
 }
 
@@ -2249,7 +2232,7 @@ func (c *amd64Compiler) compileCopysign(o *wazeroir.OperationCopysign) error {
 	// Record that we consumed the x2 and placed the copysign result in the x1's register.
 	c.locationStack.markRegisterUnused(x2.register)
 	c.locationStack.markRegisterUnused(x1.register)
-	c.locationStack.pushValueLocationOnRegister(x1.register)
+	c.pushValueLocationOnRegister(x1.register)
 	return nil
 }
 
@@ -2389,7 +2372,7 @@ func (c *amd64Compiler) emitUnsignedI32TruncFromFloat(isFloat32Bit bool) error {
 	// We consumed the source's register and placed the conversion result
 	// in the result register.
 	c.locationStack.markRegisterUnused(source.register)
-	loc := c.locationStack.pushValueLocationOnRegister(result)
+	loc := c.pushValueLocationOnRegister(result)
 	loc.setRegisterType(generalPurposeRegisterTypeInt)
 	return nil
 }
@@ -2476,7 +2459,7 @@ func (c *amd64Compiler) emitUnsignedI64TruncFromFloat(isFloat32Bit bool) error {
 	// We consumed the source's register and placed the conversion result
 	// in the result register.
 	c.locationStack.markRegisterUnused(source.register)
-	loc := c.locationStack.pushValueLocationOnRegister(result)
+	loc := c.pushValueLocationOnRegister(result)
 	loc.setRegisterType(generalPurposeRegisterTypeInt)
 	return nil
 }
@@ -2560,7 +2543,7 @@ func (c *amd64Compiler) emitSignedI32TruncFromFloat(isFloat32Bit bool) error {
 	// We consumed the source's register and placed the conversion result
 	// in the result register.
 	c.locationStack.markRegisterUnused(source.register)
-	loc := c.locationStack.pushValueLocationOnRegister(result)
+	loc := c.pushValueLocationOnRegister(result)
 	loc.setRegisterType(generalPurposeRegisterTypeInt)
 	return nil
 }
@@ -2637,7 +2620,7 @@ func (c *amd64Compiler) emitSignedI64TruncFromFloat(isFloat32Bit bool) error {
 	// We consumed the source's register and placed the conversion result
 	// in the result register.
 	c.locationStack.markRegisterUnused(source.register)
-	loc := c.locationStack.pushValueLocationOnRegister(result)
+	loc := c.pushValueLocationOnRegister(result)
 	loc.setRegisterType(generalPurposeRegisterTypeInt)
 	return nil
 }
@@ -2773,7 +2756,7 @@ func (c *amd64Compiler) emitUnsignedInt64ToFloatConversion(isFloat32bit bool) er
 	// We consumed the origin's register and placed the conversion result
 	// in the dest register.
 	c.locationStack.markRegisterUnused(origin.register)
-	loc := c.locationStack.pushValueLocationOnRegister(dest)
+	loc := c.pushValueLocationOnRegister(dest)
 	loc.setRegisterType(generalPurposeRegisterTypeFloat)
 	return nil
 }
@@ -2794,7 +2777,7 @@ func (c *amd64Compiler) compileSimpleConversion(convInstruction obj.As, destinat
 	c.compileRegisterToRegister(convInstruction, origin.register, dest)
 
 	c.locationStack.markRegisterUnused(origin.register)
-	loc := c.locationStack.pushValueLocationOnRegister(dest)
+	loc := c.pushValueLocationOnRegister(dest)
 	loc.setRegisterType(destinationRegisterType)
 	return nil
 }
@@ -3011,7 +2994,7 @@ func (c *amd64Compiler) compileEqOrNeForFloats(x1Reg, x2Reg int16, cmpInstructio
 	c.compileRegisterToRegister(x86.AMOVBLZX, cmpResultReg, cmpResultReg)
 
 	// Now we have the result in cmpResultReg register, so we record it.
-	loc := c.locationStack.pushValueLocationOnRegister(cmpResultReg)
+	loc := c.pushValueLocationOnRegister(cmpResultReg)
 	loc.setRegisterType(generalPurposeRegisterTypeInt)
 	// Also, we no longer need nanFragRegister.
 	c.locationStack.markRegisterUnused(nanFragReg)
@@ -3263,7 +3246,7 @@ func (c *amd64Compiler) compileLoad(o *wazeroir.OperationLoad) error {
 			// we access memory as memory.Buffer[ceil-targetSizeInBytes: ceil].
 			reservedRegisterForMemory, -targetSizeInBytes, reg, 1,
 			reg)
-		top := c.locationStack.pushValueLocationOnRegister(reg)
+		top := c.pushValueLocationOnRegister(reg)
 		top.setRegisterType(generalPurposeRegisterTypeInt)
 	} else {
 		// For float types, we read the value to the float register.
@@ -3275,7 +3258,7 @@ func (c *amd64Compiler) compileLoad(o *wazeroir.OperationLoad) error {
 			// we access memory as memory.Buffer[ceil-targetSizeInBytes: ceil].
 			reservedRegisterForMemory, -targetSizeInBytes, reg, 1,
 			floatReg)
-		top := c.locationStack.pushValueLocationOnRegister(floatReg)
+		top := c.pushValueLocationOnRegister(floatReg)
 		top.setRegisterType(generalPurposeRegisterTypeFloat)
 		// We no longer need the int register so mark it unused.
 		c.locationStack.markRegisterUnused(reg)
@@ -3310,7 +3293,7 @@ func (c *amd64Compiler) compileLoad8(o *wazeroir.OperationLoad8) error {
 		reservedRegisterForMemory, -targetSizeInBytes, reg, 1,
 		reg)
 
-	top := c.locationStack.pushValueLocationOnRegister(reg)
+	top := c.pushValueLocationOnRegister(reg)
 
 	// The result of load8 is always int type.
 	top.setRegisterType(generalPurposeRegisterTypeInt)
@@ -3343,7 +3326,7 @@ func (c *amd64Compiler) compileLoad16(o *wazeroir.OperationLoad16) error {
 		reservedRegisterForMemory, -targetSizeInBytes, reg, 1,
 		reg)
 
-	top := c.locationStack.pushValueLocationOnRegister(reg)
+	top := c.pushValueLocationOnRegister(reg)
 	// The result of load16 is always int type.
 	top.setRegisterType(generalPurposeRegisterTypeInt)
 	return nil
@@ -3368,7 +3351,7 @@ func (c *amd64Compiler) compileLoad32(o *wazeroir.OperationLoad32) error {
 		// We access memory as memory.Buffer[ceil-targetSizeInBytes: ceil].
 		reservedRegisterForMemory, -targetSizeInBytes, reg, 1,
 		reg)
-	top := c.locationStack.pushValueLocationOnRegister(reg)
+	top := c.pushValueLocationOnRegister(reg)
 
 	// The result of load32 is always int type.
 	top.setRegisterType(generalPurposeRegisterTypeInt)
@@ -3486,7 +3469,7 @@ func (c *amd64Compiler) compileMemorySize() error {
 	if err != nil {
 		return err
 	}
-	loc := c.locationStack.pushValueLocationOnRegister(reg)
+	loc := c.pushValueLocationOnRegister(reg)
 
 	c.compileMemoryToRegisterInstruction(x86.AMOVQ, reservedRegisterForCallEngine, callEngineModuleContextMemorySliceLenOffset, loc.register)
 
@@ -3505,7 +3488,7 @@ func (c *amd64Compiler) compileConstI32(o *wazeroir.OperationConstI32) error {
 	if err != nil {
 		return err
 	}
-	loc := c.locationStack.pushValueLocationOnRegister(reg)
+	loc := c.pushValueLocationOnRegister(reg)
 	loc.setRegisterType(generalPurposeRegisterTypeInt)
 
 	c.compileConstToRegisterInstruction(x86.AMOVL, int64(o.Value), reg)
@@ -3520,7 +3503,7 @@ func (c *amd64Compiler) compileConstI64(o *wazeroir.OperationConstI64) error {
 	if err != nil {
 		return err
 	}
-	loc := c.locationStack.pushValueLocationOnRegister(reg)
+	loc := c.pushValueLocationOnRegister(reg)
 	loc.setRegisterType(generalPurposeRegisterTypeInt)
 
 	c.compileConstToRegisterInstruction(x86.AMOVQ, int64(o.Value), reg)
@@ -3535,7 +3518,7 @@ func (c *amd64Compiler) compileConstF32(o *wazeroir.OperationConstF32) error {
 	if err != nil {
 		return err
 	}
-	loc := c.locationStack.pushValueLocationOnRegister(reg)
+	loc := c.pushValueLocationOnRegister(reg)
 	loc.setRegisterType(generalPurposeRegisterTypeFloat)
 
 	// We cannot directly load the value from memory to float regs,
@@ -3558,7 +3541,7 @@ func (c *amd64Compiler) compileConstF64(o *wazeroir.OperationConstF64) error {
 	if err != nil {
 		return err
 	}
-	loc := c.locationStack.pushValueLocationOnRegister(reg)
+	loc := c.pushValueLocationOnRegister(reg)
 	loc.setRegisterType(generalPurposeRegisterTypeFloat)
 
 	// We cannot directly load the value from memory to float regs,
@@ -3719,7 +3702,7 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, compiledFuncti
 	if !isNilRegister(compiledFunctionAddressRegister) {
 		// If we need to get the target funcaddr from register (call_indirect case), we must save it before growing the
 		// call-frame stack, as the register is not saved across function calls.
-		savedOffsetLocation := c.locationStack.pushValueLocationOnRegister(compiledFunctionAddressRegister)
+		savedOffsetLocation := c.pushValueLocationOnRegister(compiledFunctionAddressRegister)
 		c.compileReleaseRegisterToStack(savedOffsetLocation)
 	}
 
@@ -4045,12 +4028,13 @@ func (c *amd64Compiler) compileReleaseAllRegistersToStack() {
 }
 
 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
+	for i := uint64(0); i < c.locationStack.sp; i++ {
+		prevValue := c.locationStack.stack[i]
+		if prevValue.register == reg {
+			c.compileReleaseRegisterToStack(prevValue)
+			break
+		}
 	}
-	c.compileReleaseRegisterToStack(prevValue)
 }
 
 func (c *amd64Compiler) compileReleaseRegisterToStack(loc *valueLocation) {
diff --git a/internal/wasm/jit/jit_arm64.go b/internal/wasm/jit/jit_arm64.go
index 4c8b55d7..f0309845 100644
--- a/internal/wasm/jit/jit_arm64.go
+++ b/internal/wasm/jit/jit_arm64.go
@@ -13,7 +13,6 @@ import (
 	"encoding/binary"
 	"fmt"
 	"math"
-	"sync"
 	"unsafe"
 
 	asm "github.com/twitchyliquid64/golang-asm"
@@ -24,83 +23,6 @@ import (
 	"github.com/tetratelabs/wazero/internal/wazeroir"
 )
 
-func newArchContext() archContext {
-	return archContext{
-		minimum32BitSignedInt: math.MinInt32,
-		minimum64BitSignedInt: math.MinInt64,
-	}
-}
-
-// archContext is embedded in callEngine in order to store architecture-specific data.
-type archContext struct {
-	// jitCallReturnAddress holds the absolute return address for jitcall.
-	// The value is set whenever jitcall is executed and done in jit_arm64.s
-	// Native code can return back to the ce.execWasmFunction's main loop back by
-	// executing "ret" instruction with this value. See arm64Compiler.exit.
-	// Note: this is only used by JIT code so mark this as nolint.
-	jitCallReturnAddress uint64 //nolint
-
-	// Loading large constants in arm64 is a bit costly, so we place the following
-	// consts on callEngine struct so that we can quickly access them during various operations.
-
-	// minimum32BitSignedInt is used for overflow check for 32-bit signed division.
-	// Note: this can be obtained by moving $1 and doing left-shift with 31, but it is
-	// slower than directly loading fron this location.
-	minimum32BitSignedInt int32
-	// Note: this can be obtained by moving $1 and doing left-shift with 63, but it is
-	// slower than directly loading fron this location.
-	// minimum64BitSignedInt is used for overflow check for 64-bit signed division.
-	minimum64BitSignedInt int64
-}
-
-const (
-	// callEngineArchContextJITCallReturnAddressOffset is the offset of archContext.jitCallReturnAddress in callEngine.
-	callEngineArchContextJITCallReturnAddressOffset = 120
-	// callEngineArchContextMinimum32BitSignedIntOffset is the offset of archContext.minimum32BitSignedIntAddress in callEngine.
-	callEngineArchContextMinimum32BitSignedIntOffset = 128
-	// callEngineArchContextMinimum64BitSignedIntOffset is the offset of archContext.minimum64BitSignedIntAddress in callEngine.
-	callEngineArchContextMinimum64BitSignedIntOffset = 136
-)
-
-// jitcall is implemented in jit_arm64.s as a Go Assembler function.
-// This is used by callEngine.execWasmFunction and the entrypoint to enter the JITed native code.
-// codeSegment is the pointer to the initial instruction of the compiled native code.
-// ce is "*callEngine" as uintptr.
-func jitcall(codeSegment, ce uintptr)
-
-// golang-asm is not goroutine-safe so we take lock until we complete the compilation.
-// TODO: delete after https://github.com/tetratelabs/wazero/issues/233
-var assemblerMutex = &sync.Mutex{}
-
-func unlockAssembler() {
-	assemblerMutex.Unlock()
-}
-
-// newCompiler returns a new compiler interface which can be used to compile the given function instance.
-// The function returned must be invoked when finished compiling, so use `defer` to ensure this.
-// Note: ir param can be nil for host functions.
-func newCompiler(f *wasm.FunctionInstance, ir *wazeroir.CompilationResult) (c compiler, done func(), err error) {
-	// golang-asm is not goroutine-safe so we take lock until we complete the compilation.
-	// TODO: delete after https://github.com/tetratelabs/wazero/issues/233
-	assemblerMutex.Lock()
-
-	// We can choose arbitrary number instead of 1024 which indicates the cache size in the compiler.
-	// TODO: optimize the number.
-	b, err := asm.NewBuilder("arm64", 1024)
-	if err != nil {
-		return nil, unlockAssembler, fmt.Errorf("failed to create a new assembly builder: %w", err)
-	}
-
-	compiler := &arm64Compiler{
-		f:             f,
-		builder:       b,
-		locationStack: newValueLocationStack(),
-		ir:            ir,
-		labels:        map[string]*labelInfo{},
-	}
-	return compiler, unlockAssembler, nil
-}
-
 type arm64Compiler struct {
 	builder *asm.Builder
 	f       *wasm.FunctionInstance
@@ -124,6 +46,28 @@ type arm64Compiler struct {
 	staticData compiledFunctionStaticData
 }
 
+const (
+	// reservedRegisterForCallEngine holds the pointer to callEngine instance (i.e. *callEngine as uintptr)
+	reservedRegisterForCallEngine int16 = arm64.REG_R0
+	// reservedRegisterForStackBasePointerAddress holds stack base pointer's address (callEngine.stackBasePointer) in the current function call.
+	reservedRegisterForStackBasePointerAddress int16 = arm64.REG_R1
+	// reservedRegisterForMemory holds the pointer to the memory slice's data (i.e. &memory.Buffer[0] as uintptr).
+	reservedRegisterForMemory    int16 = arm64.REG_R2
+	reservedRegisterForTemporary int16 = arm64.REG_R3
+	// zeroRegister is the alias of the arm64-specific zero register for readability.
+	zeroRegister int16 = arm64.REGZERO
+)
+
+// simdRegisterForScalarFloatRegister returns SIMD register which corresponds to the given scalar float register.
+// In other words, this returns: REG_F0 -> REG_V0, REG_F1 -> REG_V1, ...., REG_F31 -> REG_V31.
+func simdRegisterForScalarFloatRegister(freg int16) int16 {
+	return freg + (arm64.REG_F31 - arm64.REG_F0) + 1
+}
+
+func isZeroRegister(r int16) bool {
+	return r == zeroRegister
+}
+
 func (c *arm64Compiler) addStaticData(d []byte) {
 	c.staticData = append(c.staticData, d)
 }
@@ -197,6 +141,12 @@ func (c *arm64Compiler) setBranchTargetOnNext(progs ...*obj.Prog) {
 	c.setBranchTargetOnNextInstructions = append(c.setBranchTargetOnNextInstructions, progs...)
 }
 
+func (c *arm64Compiler) pushValueLocationOnRegister(reg int16) (ret *valueLocation) {
+	ret = c.locationStack.pushValueLocationOnRegister(reg)
+	c.markRegisterUsed(reg)
+	return
+}
+
 func (c *arm64Compiler) markRegisterUsed(regs ...int16) {
 	for _, reg := range regs {
 		if !isZeroRegister(reg) {
@@ -493,7 +443,7 @@ func (c *arm64Compiler) compileMaybeGrowValueStack() error {
 	// Otherwise, skip calling it.
 	c.setBranchTargetOnNext(brIfValueStackOK)
 
-	c.locationStack.markRegisterUnused(tmpRegs...)
+	c.markRegisterUnused(tmpRegs...)
 	return nil
 }
 
@@ -576,7 +526,7 @@ func (c *arm64Compiler) compileReturnFunction() error {
 		tmpReg)
 	c.compileUnconditionalBranchToAddressOnRegister(tmpReg)
 
-	c.locationStack.markRegisterUnused(tmpRegs...)
+	c.markRegisterUnused(tmpRegs...)
 	return nil
 }
 
@@ -726,7 +676,7 @@ func (c *arm64Compiler) compileGlobalGet(o *wazeroir.OperationGlobalGet) error {
 		c.compileRegisterToRegisterInstruction(floatMov, intReg, resultReg)
 	}
 
-	c.locationStack.pushValueLocationOnRegister(resultReg)
+	c.pushValueLocationOnRegister(resultReg)
 	return nil
 }
 
@@ -1096,7 +1046,7 @@ func (c *arm64Compiler) compileCallImpl(index wasm.Index, compiledFunctionAddres
 	if !found {
 		return fmt.Errorf("BUG: all registers except indexReg should be free at this point")
 	}
-	c.locationStack.markRegisterUsed(freeRegisters...)
+	c.markRegisterUsed(freeRegisters...)
 
 	// Alias for readability.
 	callFrameStackPointerRegister, callFrameStackTopAddressRegister, compiledFunctionRegister, oldStackBasePointer,
@@ -1123,7 +1073,7 @@ func (c *arm64Compiler) compileCallImpl(index wasm.Index, compiledFunctionAddres
 	if !isNilRegister(compiledFunctionAddressRegister) {
 		// If we need to get the target funcaddr from register (call_indirect case), we must save it before growing the
 		// call-frame stack, as the register is not saved across function calls.
-		savedOffsetLocation := c.locationStack.pushValueLocationOnRegister(compiledFunctionAddressRegister)
+		savedOffsetLocation := c.pushValueLocationOnRegister(compiledFunctionAddressRegister)
 		c.compileReleaseRegisterToStack(savedOffsetLocation)
 	}
 
@@ -1517,7 +1467,7 @@ func (c *arm64Compiler) compileSelect() error {
 
 	if isZeroRegister(x1.register) && isZeroRegister(x2.register) {
 		// If both values are zero, the result is always zero.
-		c.locationStack.pushValueLocationOnRegister(zeroRegister)
+		c.pushValueLocationOnRegister(zeroRegister)
 		c.markRegisterUnused(cv.register)
 		return nil
 	}
@@ -1553,7 +1503,7 @@ func (c *arm64Compiler) compileSelect() error {
 	} else {
 		c.compileRegisterToRegisterInstruction(arm64.AFMOVD, x2.register, x1.register)
 	}
-	c.locationStack.pushValueLocationOnRegister(x1.register)
+	c.pushValueLocationOnRegister(x1.register)
 
 	// Otherwise, nothing to do for select.
 	c.setBranchTargetOnNext(brIfNotZero)
@@ -1593,7 +1543,7 @@ func (c *arm64Compiler) compilePick(o *wazeroir.OperationPick) error {
 
 	// Now we have the value of the target on the pickedRegister,
 	// so push the location.
-	c.locationStack.pushValueLocationOnRegister(pickedRegister)
+	c.pushValueLocationOnRegister(pickedRegister)
 	return nil
 }
 
@@ -1606,10 +1556,10 @@ func (c *arm64Compiler) compileAdd(o *wazeroir.OperationAdd) error {
 
 	// Addition can be nop if one of operands is zero.
 	if isZeroRegister(x1.register) {
-		c.locationStack.pushValueLocationOnRegister(x2.register)
+		c.pushValueLocationOnRegister(x2.register)
 		return nil
 	} else if isZeroRegister(x2.register) {
-		c.locationStack.pushValueLocationOnRegister(x1.register)
+		c.pushValueLocationOnRegister(x1.register)
 		return nil
 	}
 
@@ -1627,7 +1577,7 @@ func (c *arm64Compiler) compileAdd(o *wazeroir.OperationAdd) error {
 
 	c.compileRegisterToRegisterInstruction(inst, x2.register, x1.register)
 	// The result is placed on a register for x1, so record it.
-	c.locationStack.pushValueLocationOnRegister(x1.register)
+	c.pushValueLocationOnRegister(x1.register)
 	return nil
 }
 
@@ -1640,7 +1590,7 @@ func (c *arm64Compiler) compileSub(o *wazeroir.OperationSub) error {
 
 	// If both of registers are zeros, this can be nop and push the zero register.
 	if isZeroRegister(x1.register) && isZeroRegister(x2.register) {
-		c.locationStack.pushValueLocationOnRegister(zeroRegister)
+		c.pushValueLocationOnRegister(zeroRegister)
 		return nil
 	}
 
@@ -1664,7 +1614,7 @@ func (c *arm64Compiler) compileSub(o *wazeroir.OperationSub) error {
 	}
 
 	c.compileTwoRegistersToRegisterInstruction(inst, x2.register, x1.register, destinationReg)
-	c.locationStack.pushValueLocationOnRegister(destinationReg)
+	c.pushValueLocationOnRegister(destinationReg)
 	return nil
 }
 
@@ -1677,7 +1627,7 @@ func (c *arm64Compiler) compileMul(o *wazeroir.OperationMul) error {
 
 	// Multiplication can be done by putting a zero register if one of operands is zero.
 	if isZeroRegister(x1.register) || isZeroRegister(x2.register) {
-		c.locationStack.pushValueLocationOnRegister(zeroRegister)
+		c.pushValueLocationOnRegister(zeroRegister)
 		return nil
 	}
 
@@ -1695,7 +1645,7 @@ func (c *arm64Compiler) compileMul(o *wazeroir.OperationMul) error {
 
 	c.compileRegisterToRegisterInstruction(inst, x2.register, x1.register)
 	// The result is placed on a register for x1, so record it.
-	c.locationStack.pushValueLocationOnRegister(x1.register)
+	c.pushValueLocationOnRegister(x1.register)
 	return nil
 }
 
@@ -1718,7 +1668,7 @@ func (c *arm64Compiler) compileClz(o *wazeroir.OperationClz) error {
 		} else {
 			c.compileConstToRegisterInstruction(arm64.AMOVD, 64, reg)
 		}
-		c.locationStack.pushValueLocationOnRegister(reg)
+		c.pushValueLocationOnRegister(reg)
 		return nil
 	}
 
@@ -1728,7 +1678,7 @@ func (c *arm64Compiler) compileClz(o *wazeroir.OperationClz) error {
 	} else {
 		c.compileRegisterToRegisterInstruction(arm64.ACLZ, reg, reg)
 	}
-	c.locationStack.pushValueLocationOnRegister(reg)
+	c.pushValueLocationOnRegister(reg)
 	return nil
 }
 
@@ -1752,7 +1702,7 @@ func (c *arm64Compiler) compileCtz(o *wazeroir.OperationCtz) error {
 		} else {
 			c.compileConstToRegisterInstruction(arm64.AMOVD, 64, reg)
 		}
-		c.locationStack.pushValueLocationOnRegister(reg)
+		c.pushValueLocationOnRegister(reg)
 		return nil
 	}
 
@@ -1766,7 +1716,7 @@ func (c *arm64Compiler) compileCtz(o *wazeroir.OperationCtz) error {
 		c.compileRegisterToRegisterInstruction(arm64.ARBIT, reg, reg)
 		c.compileRegisterToRegisterInstruction(arm64.ACLZ, reg, reg)
 	}
-	c.locationStack.pushValueLocationOnRegister(reg)
+	c.pushValueLocationOnRegister(reg)
 	return nil
 }
 
@@ -1779,7 +1729,7 @@ func (c *arm64Compiler) compilePopcnt(o *wazeroir.OperationPopcnt) error {
 
 	reg := v.register
 	if isZeroRegister(reg) {
-		c.locationStack.pushValueLocationOnRegister(reg)
+		c.pushValueLocationOnRegister(reg)
 		return nil
 	}
 
@@ -1807,7 +1757,7 @@ func (c *arm64Compiler) compilePopcnt(o *wazeroir.OperationPopcnt) error {
 	c.compileRegisterToRegisterInstruction(arm64.AVUADDLV, vreg&31+arm64.REG_ARNG+(arm64.ARNG_8B&15)<<5, vreg)
 	c.compileRegisterToRegisterInstruction(arm64.AFMOVD, freg, reg)
 
-	c.locationStack.pushValueLocationOnRegister(reg)
+	c.pushValueLocationOnRegister(reg)
 	return nil
 }
 
@@ -1856,7 +1806,7 @@ func (c *arm64Compiler) compileDiv(o *wazeroir.OperationDiv) error {
 
 	c.compileRegisterToRegisterInstruction(inst, divisor.register, dividend.register)
 
-	c.locationStack.pushValueLocationOnRegister(dividend.register)
+	c.pushValueLocationOnRegister(dividend.register)
 	return nil
 }
 
@@ -1983,7 +1933,7 @@ func (c *arm64Compiler) compileRem(o *wazeroir.OperationRem) error {
 	c.compileTwoRegistersInstruction(msubInst, divisorReg, dividendReg, resultReg, resultReg)
 
 	c.markRegisterUnused(dividend.register, divisor.register)
-	c.locationStack.pushValueLocationOnRegister(resultReg)
+	c.pushValueLocationOnRegister(resultReg)
 	return nil
 }
 
@@ -1997,7 +1947,7 @@ func (c *arm64Compiler) compileAnd(o *wazeroir.OperationAnd) error {
 	// If either of the registers x1 or x2 is zero,
 	// the result will always be zero.
 	if isZeroRegister(x1.register) || isZeroRegister(x2.register) {
-		c.locationStack.pushValueLocationOnRegister(zeroRegister)
+		c.pushValueLocationOnRegister(zeroRegister)
 		return nil
 	}
 
@@ -2017,7 +1967,7 @@ func (c *arm64Compiler) compileAnd(o *wazeroir.OperationAnd) error {
 	}
 
 	c.compileTwoRegistersToRegisterInstruction(inst, x2.register, x1.register, destinationReg)
-	c.locationStack.pushValueLocationOnRegister(x1.register)
+	c.pushValueLocationOnRegister(x1.register)
 	return nil
 }
 
@@ -2029,11 +1979,11 @@ func (c *arm64Compiler) compileOr(o *wazeroir.OperationOr) error {
 	}
 
 	if isZeroRegister(x1.register) {
-		c.locationStack.pushValueLocationOnRegister(x2.register)
+		c.pushValueLocationOnRegister(x2.register)
 		return nil
 	}
 	if isZeroRegister(x2.register) {
-		c.locationStack.pushValueLocationOnRegister(x1.register)
+		c.pushValueLocationOnRegister(x1.register)
 		return nil
 	}
 
@@ -2046,7 +1996,7 @@ func (c *arm64Compiler) compileOr(o *wazeroir.OperationOr) error {
 	}
 
 	c.compileTwoRegistersToRegisterInstruction(inst, x2.register, x1.register, x1.register)
-	c.locationStack.pushValueLocationOnRegister(x1.register)
+	c.pushValueLocationOnRegister(x1.register)
 	return nil
 }
 
@@ -2073,7 +2023,7 @@ func (c *arm64Compiler) compileXor(o *wazeroir.OperationXor) error {
 	}
 
 	c.compileTwoRegistersToRegisterInstruction(inst, x2.register, x1.register, destinationReg)
-	c.locationStack.pushValueLocationOnRegister(destinationReg)
+	c.pushValueLocationOnRegister(destinationReg)
 	return nil
 }
 
@@ -2085,7 +2035,7 @@ func (c *arm64Compiler) compileShl(o *wazeroir.OperationShl) error {
 	}
 
 	if isZeroRegister(x1.register) || isZeroRegister(x2.register) {
-		c.locationStack.pushValueLocationOnRegister(x1.register)
+		c.pushValueLocationOnRegister(x1.register)
 		return nil
 	}
 
@@ -2098,7 +2048,7 @@ func (c *arm64Compiler) compileShl(o *wazeroir.OperationShl) error {
 	}
 
 	c.compileTwoRegistersToRegisterInstruction(inst, x2.register, x1.register, x1.register)
-	c.locationStack.pushValueLocationOnRegister(x1.register)
+	c.pushValueLocationOnRegister(x1.register)
 	return nil
 }
 
@@ -2110,7 +2060,7 @@ func (c *arm64Compiler) compileShr(o *wazeroir.OperationShr) error {
 	}
 
 	if isZeroRegister(x1.register) || isZeroRegister(x2.register) {
-		c.locationStack.pushValueLocationOnRegister(x1.register)
+		c.pushValueLocationOnRegister(x1.register)
 		return nil
 	}
 
@@ -2127,7 +2077,7 @@ func (c *arm64Compiler) compileShr(o *wazeroir.OperationShr) error {
 	}
 
 	c.compileTwoRegistersToRegisterInstruction(inst, x2.register, x1.register, x1.register)
-	c.locationStack.pushValueLocationOnRegister(x1.register)
+	c.pushValueLocationOnRegister(x1.register)
 	return nil
 }
 
@@ -2139,7 +2089,7 @@ func (c *arm64Compiler) compileRotl(o *wazeroir.OperationRotl) error {
 	}
 
 	if isZeroRegister(x1.register) || isZeroRegister(x2.register) {
-		c.locationStack.pushValueLocationOnRegister(x1.register)
+		c.pushValueLocationOnRegister(x1.register)
 		return nil
 	}
 
@@ -2162,7 +2112,7 @@ func (c *arm64Compiler) compileRotl(o *wazeroir.OperationRotl) error {
 	c.compileRegisterToRegisterInstruction(neginst, x2.register, x2.register)
 
 	c.compileTwoRegistersToRegisterInstruction(inst, x2.register, x1.register, x1.register)
-	c.locationStack.pushValueLocationOnRegister(x1.register)
+	c.pushValueLocationOnRegister(x1.register)
 	return nil
 }
 
@@ -2174,7 +2124,7 @@ func (c *arm64Compiler) compileRotr(o *wazeroir.OperationRotr) error {
 	}
 
 	if isZeroRegister(x1.register) || isZeroRegister(x2.register) {
-		c.locationStack.pushValueLocationOnRegister(x1.register)
+		c.pushValueLocationOnRegister(x1.register)
 		return nil
 	}
 
@@ -2187,7 +2137,7 @@ func (c *arm64Compiler) compileRotr(o *wazeroir.OperationRotr) error {
 	}
 
 	c.compileTwoRegistersToRegisterInstruction(inst, x2.register, x1.register, x1.register)
-	c.locationStack.pushValueLocationOnRegister(x1.register)
+	c.pushValueLocationOnRegister(x1.register)
 	return nil
 }
 
@@ -2278,7 +2228,7 @@ func (c *arm64Compiler) compileSimpleFloatBinop(inst obj.As) error {
 		return err
 	}
 	c.compileRegisterToRegisterInstruction(inst, x2.register, x1.register)
-	c.locationStack.pushValueLocationOnRegister(x1.register)
+	c.pushValueLocationOnRegister(x1.register)
 	return nil
 }
 
@@ -2338,7 +2288,7 @@ func (c *arm64Compiler) compileCopysign(o *wazeroir.OperationCopysign) error {
 	)
 
 	c.markRegisterUnused(x2.register)
-	c.locationStack.pushValueLocationOnRegister(x1.register)
+	c.pushValueLocationOnRegister(x1.register)
 	return nil
 }
 
@@ -2383,7 +2333,7 @@ func (c *arm64Compiler) compileITruncFromF(o *wazeroir.OperationITruncFromF) err
 	}
 
 	c.compileRegisterToRegisterInstruction(convinst, source.register, destinationReg)
-	c.locationStack.pushValueLocationOnRegister(destinationReg)
+	c.pushValueLocationOnRegister(destinationReg)
 
 	// Obtain the floating point status register value into the general purpose register,
 	// so that we can check if the conversion resulted in undefined behavior.
@@ -2505,7 +2455,7 @@ func (c *arm64Compiler) compileSimpleConversion(inst obj.As, destinationRegType
 	}
 
 	c.compileRegisterToRegisterInstruction(inst, source.register, destinationReg)
-	c.locationStack.pushValueLocationOnRegister(destinationReg)
+	c.pushValueLocationOnRegister(destinationReg)
 	return nil
 }
 
@@ -2550,7 +2500,7 @@ func (c *arm64Compiler) compileSimpleUnop(inst obj.As) error {
 	}
 	reg := v.register
 	c.compileRegisterToRegisterInstruction(inst, reg, reg)
-	c.locationStack.pushValueLocationOnRegister(reg)
+	c.pushValueLocationOnRegister(reg)
 	return nil
 }
 
@@ -2854,7 +2804,7 @@ func (c *arm64Compiler) compileLoadImpl(offsetArg uint32, loadInst obj.As, targe
 		resultRegister,
 	)
 
-	c.locationStack.pushValueLocationOnRegister(resultRegister)
+	c.pushValueLocationOnRegister(resultRegister)
 	return nil
 }
 
@@ -3006,7 +2956,7 @@ func (c *arm64Compiler) compileMemorySize() error {
 		reg,
 	)
 
-	c.locationStack.pushValueLocationOnRegister(reg)
+	c.pushValueLocationOnRegister(reg)
 	return nil
 }
 
@@ -3023,7 +2973,7 @@ func (c *arm64Compiler) compileCallGoFunction(jitStatus jitCallStatusCode, built
 	if !found {
 		return fmt.Errorf("BUG: all registers except indexReg should be free at this point")
 	}
-	c.locationStack.markRegisterUsed(freeRegs...)
+	c.markRegisterUsed(freeRegs...)
 
 	// Alias these free tmp registers for readability.
 	tmp, currentCallFrameStackPointerRegister, currentCallFrameTopAddressRegister, returnAddressRegister :=
@@ -3105,7 +3055,7 @@ func (c *arm64Compiler) compileIntConstant(is32bit bool, value uint64) error {
 		}
 		c.compileConstToRegisterInstruction(inst, int64(value), reg)
 
-		c.locationStack.pushValueLocationOnRegister(reg)
+		c.pushValueLocationOnRegister(reg)
 	}
 	return nil
 }
@@ -3153,12 +3103,12 @@ func (c *arm64Compiler) compileFloatConstant(is32bit bool, value uint64) error {
 	}
 	c.compileRegisterToRegisterInstruction(inst, tmpReg, reg)
 
-	c.locationStack.pushValueLocationOnRegister(reg)
+	c.pushValueLocationOnRegister(reg)
 	return nil
 }
 
 func (c *arm64Compiler) pushZeroValue() {
-	c.locationStack.pushValueLocationOnRegister(zeroRegister)
+	c.pushValueLocationOnRegister(zeroRegister)
 }
 
 // popTwoValuesOnRegisters pops two values from the location stacks, ensures
@@ -3214,7 +3164,7 @@ func (c *arm64Compiler) compileEnsureOnGeneralPurposeRegister(loc *valueLocation
 
 		// Record that the value holds the register and the register is marked used.
 		loc.setRegister(reg)
-		c.locationStack.markRegisterUsed(reg)
+		c.markRegisterUsed(reg)
 
 		c.compileLoadValueOnStackToRegister(loc)
 	} else if loc.onConditionalRegister() {
@@ -3363,7 +3313,7 @@ func (c *arm64Compiler) compileModuleContextInitialization() error {
 	if !found {
 		return fmt.Errorf("BUG: all the registers should be free at this point")
 	}
-	c.locationStack.markRegisterUsed(regs...)
+	c.markRegisterUsed(regs...)
 
 	// Alias these free registers for readability.
 	moduleInstanceAddressRegister, tmpX, tmpY := regs[0], regs[1], regs[2]
@@ -3525,6 +3475,6 @@ func (c *arm64Compiler) compileModuleContextInitialization() error {
 	}
 
 	c.setBranchTargetOnNext(brIfModuleUnchanged)
-	c.locationStack.markRegisterUnused(regs...)
+	c.markRegisterUnused(regs...)
 	return nil
 }
diff --git a/internal/wasm/jit/jit_other.go b/internal/wasm/jit/jit_other.go
deleted file mode 100644
index 6591c629..00000000
--- a/internal/wasm/jit/jit_other.go
+++ /dev/null
@@ -1,34 +0,0 @@
-//go:build !amd64 && !arm64
-
-package jit
-
-import (
-	wasm "github.com/tetratelabs/wazero/internal/wasm"
-	"github.com/tetratelabs/wazero/internal/wazeroir"
-)
-
-type archContext struct{}
-
-func newArchContext() (ret archContext) { return }
-
-func jitcall(codeSegment, ce uintptr) {
-	panic("unsupported GOARCH")
-}
-
-func newCompiler(f *wasm.FunctionInstance, ir *wazeroir.CompilationResult) (compiler, func(), error) {
-	panic("unsupported GOARCH")
-}
-
-const (
-	reservedRegisterForCallEngine              = 0
-	reservedRegisterForStackBasePointerAddress = 0
-	reservedRegisterForMemory                  = 0
-	reservedRegisterForTemporary               = 9
-)
-
-var (
-	generalPurposeFloatRegisters         = []int16{}
-	unreservedGeneralPurposeIntRegisters = []int16{}
-)
-
-const zeroRegister int16 = nilRegister
diff --git a/internal/wasm/jit/jit_value_location.go b/internal/wasm/jit/jit_value_location.go
index e60baadc..6df04f2c 100644
--- a/internal/wasm/jit/jit_value_location.go
+++ b/internal/wasm/jit/jit_value_location.go
@@ -3,8 +3,6 @@ package jit
 import (
 	"fmt"
 	"strings"
-
-	"github.com/tetratelabs/wazero/internal/wasm/buildoptions"
 )
 
 // nilRegister is used to indicate a register argument a variable is invalid and not an actual register.
@@ -19,11 +17,7 @@ func isIntRegister(r int16) bool {
 }
 
 func isFloatRegister(r int16) bool {
-	return generalPurposeFloatRegisters[0] <= r && r <= generalPurposeFloatRegisters[len(generalPurposeFloatRegisters)-1]
-}
-
-func isZeroRegister(r int16) bool {
-	return r == zeroRegister
+	return unreservedGeneralPurposeFloatRegisters[0] <= r && r <= unreservedGeneralPurposeFloatRegisters[len(unreservedGeneralPurposeFloatRegisters)-1]
 }
 
 // conditionalRegisterState indicates a state of the conditional flag register.
@@ -145,18 +139,12 @@ func (s *valueLocationStack) clone() *valueLocationStack {
 // the location stack.
 func (s *valueLocationStack) pushValueLocationOnRegister(reg int16) (loc *valueLocation) {
 	loc = &valueLocation{register: reg, conditionalRegister: conditionalRegisterStateUnset}
-	if buildoptions.IsDebugMode {
-		if _, ok := s.usedRegisters[loc.register]; ok {
-			panic("bug in compiler: try pushing a register which is already in use")
-		}
-	}
 
 	if isIntRegister(reg) {
 		loc.setRegisterType(generalPurposeRegisterTypeInt)
 	} else if isFloatRegister(reg) {
 		loc.setRegisterType(generalPurposeRegisterTypeFloat)
 	}
-	s.markRegisterUsed(reg)
 	s.push(loc)
 	return
 }
@@ -217,9 +205,7 @@ func (s *valueLocationStack) markRegisterUnused(regs ...int16) {
 
 func (s *valueLocationStack) markRegisterUsed(regs ...int16) {
 	for _, reg := range regs {
-		if !isZeroRegister(reg) {
-			s.usedRegisters[reg] = struct{}{}
-		}
+		s.usedRegisters[reg] = struct{}{}
 	}
 }
 
@@ -245,7 +231,7 @@ func (s *valueLocationStack) takeFreeRegister(tp generalPurposeRegisterType) (re
 	var targetRegs []int16
 	switch tp {
 	case generalPurposeRegisterTypeFloat:
-		targetRegs = generalPurposeFloatRegisters
+		targetRegs = unreservedGeneralPurposeFloatRegisters
 	case generalPurposeRegisterTypeInt:
 		targetRegs = unreservedGeneralPurposeIntRegisters
 	}
@@ -262,7 +248,7 @@ func (s *valueLocationStack) takeFreeRegisters(tp generalPurposeRegisterType, nu
 	var targetRegs []int16
 	switch tp {
 	case generalPurposeRegisterTypeFloat:
-		targetRegs = generalPurposeFloatRegisters
+		targetRegs = unreservedGeneralPurposeFloatRegisters
 	case generalPurposeRegisterTypeInt:
 		targetRegs = unreservedGeneralPurposeIntRegisters
 	}
diff --git a/internal/wasm/jit/jit_value_location_amd64.go b/internal/wasm/jit/jit_value_location_amd64.go
deleted file mode 100644
index ae8c27ac..00000000
--- a/internal/wasm/jit/jit_value_location_amd64.go
+++ /dev/null
@@ -1,64 +0,0 @@
-package jit
-
-import (
-	"github.com/twitchyliquid64/golang-asm/obj/x86"
-)
-
-// Reserved registers.
-const (
-	// reservedRegisterForCallEngine: pointer to callEngine (i.e. *callEngine as uintptr)
-	reservedRegisterForCallEngine = x86.REG_R13
-	// reservedRegisterForStackBasePointerAddress: stack base pointer's address (callEngine.stackBasePointer) in the current function call.
-	reservedRegisterForStackBasePointerAddress = x86.REG_R14
-	// reservedRegisterForMemory: pointer to the memory slice's data (i.e. &memory.Buffer[0] as uintptr).
-	reservedRegisterForMemory    = x86.REG_R15
-	reservedRegisterForTemporary = nilRegister
-)
-
-const zeroRegister int16 = nilRegister
-
-var (
-	generalPurposeFloatRegisters = []int16{
-		x86.REG_X0, x86.REG_X1, x86.REG_X2, x86.REG_X3,
-		x86.REG_X4, x86.REG_X5, x86.REG_X6, x86.REG_X7,
-		x86.REG_X8, x86.REG_X9, x86.REG_X10, x86.REG_X11,
-		x86.REG_X12, x86.REG_X13, x86.REG_X14, x86.REG_X15,
-	}
-	// Note that we never invoke "call" instruction,
-	// so we don't need to care about the calling convention.
-	// TODO: Maybe it is safe just save rbp, rsp somewhere
-	// in Go-allocated variables, and reuse these registers
-	// in JITed functions and write them back before returns.
-	unreservedGeneralPurposeIntRegisters = []int16{
-		x86.REG_AX, x86.REG_CX, x86.REG_DX, x86.REG_BX,
-		x86.REG_SI, x86.REG_DI, x86.REG_R8, x86.REG_R9,
-		x86.REG_R10, x86.REG_R11, x86.REG_R12,
-	}
-)
-
-const (
-	conditionalRegisterStateE  = conditionalRegisterStateUnset + 1 + iota // ZF equal to zero
-	conditionalRegisterStateNE                                            //˜ZF not equal to zero
-	conditionalRegisterStateS                                             // SF negative
-	conditionalRegisterStateNS                                            // ˜SF non-negative
-	conditionalRegisterStateG                                             // ˜(SF xor OF) & ˜ ZF greater (signed >)
-	conditionalRegisterStateGE                                            // ˜(SF xor OF) greater or equal (signed >=)
-	conditionalRegisterStateL                                             // SF xor OF less (signed <)
-	conditionalRegisterStateLE                                            // (SF xor OF) | ZF less or equal (signed <=)
-	conditionalRegisterStateA                                             // ˜CF & ˜ZF above (unsigned >)
-	conditionalRegisterStateAE                                            // ˜CF above or equal (unsigned >=)
-	conditionalRegisterStateB                                             // CF below (unsigned <)
-	conditionalRegisterStateBE                                            // CF | ZF below or equal (unsigned <=)
-)
-
-// 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
-}
diff --git a/internal/wasm/jit/jit_value_location_arm64.go b/internal/wasm/jit/jit_value_location_arm64.go
deleted file mode 100644
index 27eaac36..00000000
--- a/internal/wasm/jit/jit_value_location_arm64.go
+++ /dev/null
@@ -1,45 +0,0 @@
-package jit
-
-import "github.com/twitchyliquid64/golang-asm/obj/arm64"
-
-// Reserved registers.
-const (
-	// reservedRegisterForCallEngine holds the pointer to callEngine instance (i.e. *callEngine as uintptr)
-	reservedRegisterForCallEngine = arm64.REG_R0
-	// reservedRegisterForStackBasePointerAddress holds stack base pointer's address (callEngine.stackBasePointer) in the current function call.
-	reservedRegisterForStackBasePointerAddress = arm64.REG_R1
-	// reservedRegisterForMemory holds the pointer to the memory slice's data (i.e. &memory.Buffer[0] as uintptr).
-	reservedRegisterForMemory    = arm64.REG_R2
-	reservedRegisterForTemporary = arm64.REG_R3
-)
-
-// zeroRegister is the alias of the arm64-specific zero register for readability.
-const zeroRegister int16 = arm64.REGZERO
-
-var (
-	generalPurposeFloatRegisters = []int16{
-		arm64.REG_F0, arm64.REG_F1, arm64.REG_F2, arm64.REG_F3,
-		arm64.REG_F4, arm64.REG_F5, arm64.REG_F6, arm64.REG_F7, arm64.REG_F8,
-		arm64.REG_F9, arm64.REG_F10, arm64.REG_F11, arm64.REG_F12, arm64.REG_F13,
-		arm64.REG_F14, arm64.REG_F15, arm64.REG_F16, arm64.REG_F17, arm64.REG_F18,
-		arm64.REG_F19, arm64.REG_F20, arm64.REG_F21, arm64.REG_F22, arm64.REG_F23,
-		arm64.REG_F24, arm64.REG_F25, arm64.REG_F26, arm64.REG_F27, arm64.REG_F28,
-		arm64.REG_F29, arm64.REG_F30, arm64.REG_F31,
-	}
-	// Note (see arm64 section in https://go.dev/doc/asm):
-	// * REG_R18 is reserved as a platform register, and we don't use it in JIT.
-	// * REG_R28 is reserved for Goroutine by Go runtime, and we don't use it in JIT.
-	unreservedGeneralPurposeIntRegisters = []int16{
-		arm64.REG_R4, arm64.REG_R5, arm64.REG_R6, arm64.REG_R7, arm64.REG_R8,
-		arm64.REG_R9, arm64.REG_R10, arm64.REG_R11, arm64.REG_R12, arm64.REG_R13,
-		arm64.REG_R14, arm64.REG_R15, arm64.REG_R16, arm64.REG_R17, arm64.REG_R19,
-		arm64.REG_R20, arm64.REG_R21, arm64.REG_R22, arm64.REG_R23, arm64.REG_R24,
-		arm64.REG_R25, arm64.REG_R26, arm64.REG_R27, arm64.REG_R29, arm64.REG_R30,
-	}
-)
-
-// simdRegisterForScalarFloatRegister returns SIMD register which corresponds to the given scalar float register.
-// In other words, this returns: REG_F0 -> REG_V0, REG_F1 -> REG_V1, ...., REG_F31 -> REG_V31.
-func simdRegisterForScalarFloatRegister(freg int16) int16 {
-	return freg + (arm64.REG_F31 - arm64.REG_F0) + 1
-}


Test patch:
diff --git a/internal/wasm/jit/arch_arm64_test.go b/internal/wasm/jit/arch_arm64_test.go
new file mode 100644
index 00000000..25fae6b2
--- /dev/null
+++ b/internal/wasm/jit/arch_arm64_test.go
@@ -0,0 +1,15 @@
+package jit
+
+import (
+	"testing"
+	"unsafe"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestArchContextOffsetInArm64Engine(t *testing.T) {
+	var ctx callEngine
+	require.Equal(t, int(unsafe.Offsetof(ctx.jitCallReturnAddress)), callEngineArchContextJITCallReturnAddressOffset, "fix consts in jit_arm64.s")
+	require.Equal(t, int(unsafe.Offsetof(ctx.minimum32BitSignedInt)), callEngineArchContextMinimum32BitSignedIntOffset)
+	require.Equal(t, int(unsafe.Offsetof(ctx.minimum64BitSignedInt)), callEngineArchContextMinimum64BitSignedIntOffset)
+}
diff --git a/internal/wasm/jit/jit_amd64_test.go b/internal/wasm/jit/jit_amd64_test.go
index 59ed5945..9b94c83c 100644
--- a/internal/wasm/jit/jit_amd64_test.go
+++ b/internal/wasm/jit/jit_amd64_test.go
@@ -100,19 +100,19 @@ func TestAmd64Compiler_compile_Mul_Div_Rem(t *testing.T) {
 						// Here, we put it just before two operands as ["any value used by DX", x1, x2]
 						// but in reality, it can exist in any position of stack.
 						compiler.compileConstToRegisterInstruction(x86.AMOVQ, int64(dxValue), x86.REG_DX)
-						prevOnDX := compiler.valueLocationStack().pushValueLocationOnRegister(x86.REG_DX)
+						prevOnDX := compiler.pushValueLocationOnRegister(x86.REG_DX)
 
 						// Setup values.
 						if tc.x1Reg != nilRegister {
 							compiler.compileConstToRegisterInstruction(x86.AMOVQ, int64(x1Value), tc.x1Reg)
-							compiler.valueLocationStack().pushValueLocationOnRegister(tc.x1Reg)
+							compiler.pushValueLocationOnRegister(tc.x1Reg)
 						} else {
 							loc := compiler.valueLocationStack().pushValueLocationOnStack()
 							env.stack()[loc.stackPointer] = uint64(x1Value)
 						}
 						if tc.x2Reg != nilRegister {
 							compiler.compileConstToRegisterInstruction(x86.AMOVQ, int64(x2Value), tc.x2Reg)
-							compiler.valueLocationStack().pushValueLocationOnRegister(tc.x2Reg)
+							compiler.pushValueLocationOnRegister(tc.x2Reg)
 						} else {
 							loc := compiler.valueLocationStack().pushValueLocationOnStack()
 							env.stack()[loc.stackPointer] = uint64(x2Value)
@@ -222,19 +222,19 @@ func TestAmd64Compiler_compile_Mul_Div_Rem(t *testing.T) {
 						// Here, we put it just before two operands as ["any value used by DX", x1, x2]
 						// but in reality, it can exist in any position of stack.
 						compiler.compileConstToRegisterInstruction(x86.AMOVQ, int64(dxValue), x86.REG_DX)
-						prevOnDX := compiler.valueLocationStack().pushValueLocationOnRegister(x86.REG_DX)
+						prevOnDX := compiler.pushValueLocationOnRegister(x86.REG_DX)
 
 						// Setup values.
 						if tc.x1Reg != nilRegister {
 							compiler.compileConstToRegisterInstruction(x86.AMOVQ, int64(x1Value), tc.x1Reg)
-							compiler.valueLocationStack().pushValueLocationOnRegister(tc.x1Reg)
+							compiler.pushValueLocationOnRegister(tc.x1Reg)
 						} else {
 							loc := compiler.valueLocationStack().pushValueLocationOnStack()
 							env.stack()[loc.stackPointer] = uint64(x1Value)
 						}
 						if tc.x2Reg != nilRegister {
 							compiler.compileConstToRegisterInstruction(x86.AMOVQ, int64(x2Value), tc.x2Reg)
-							compiler.valueLocationStack().pushValueLocationOnRegister(tc.x2Reg)
+							compiler.pushValueLocationOnRegister(tc.x2Reg)
 						} else {
 							loc := compiler.valueLocationStack().pushValueLocationOnStack()
 							env.stack()[loc.stackPointer] = uint64(x2Value)
diff --git a/internal/wasm/jit/jit_arm64_test.go b/internal/wasm/jit/jit_arm64_test.go
index 309b2bae..6df22fc8 100644
--- a/internal/wasm/jit/jit_arm64_test.go
+++ b/internal/wasm/jit/jit_arm64_test.go
@@ -2,10 +2,10 @@ package jit
 
 import (
 	"testing"
-	"unsafe"
 
 	"github.com/stretchr/testify/require"
 	"github.com/twitchyliquid64/golang-asm/obj"
+	"github.com/twitchyliquid64/golang-asm/obj/arm64"
 
 	"github.com/tetratelabs/wazero/internal/wazeroir"
 )
@@ -30,11 +30,9 @@ func (c *arm64Compiler) setValueLocationStack(s *valueLocationStack) {
 	c.locationStack = s
 }
 
-func TestArchContextOffsetInEngine(t *testing.T) {
-	var ctx callEngine
-	require.Equal(t, int(unsafe.Offsetof(ctx.jitCallReturnAddress)), callEngineArchContextJITCallReturnAddressOffset, "fix consts in jit_arm64.s")
-	require.Equal(t, int(unsafe.Offsetof(ctx.minimum32BitSignedInt)), callEngineArchContextMinimum32BitSignedIntOffset)
-	require.Equal(t, int(unsafe.Offsetof(ctx.minimum64BitSignedInt)), callEngineArchContextMinimum64BitSignedIntOffset)
+func Test_simdRegisterForScalarFloatRegister(t *testing.T) {
+	require.Equal(t, int16(arm64.REG_V0), simdRegisterForScalarFloatRegister(arm64.REG_F0))
+	require.Equal(t, int16(arm64.REG_V30), simdRegisterForScalarFloatRegister(arm64.REG_F30))
 }
 
 func TestArm64Compiler_readInstructionAddress(t *testing.T) {
diff --git a/internal/wasm/jit/jit_test.go b/internal/wasm/jit/jit_test.go
index 08a7f838..d4e9fc29 100644
--- a/internal/wasm/jit/jit_test.go
+++ b/internal/wasm/jit/jit_test.go
@@ -138,12 +138,16 @@ func (j *jitEnv) exec(code []byte) {
 }
 
 func (j *jitEnv) requireNewCompiler(t *testing.T, functype *wasm.FunctionType) compilerImpl {
+	// golang-asm is not goroutine-safe so we take lock until we complete the compilation.
+	// TODO: delete after https://github.com/tetratelabs/wazero/issues/233
+	assemblerMutex.Lock()
+
 	requireSupportedOSArch(t)
-	c, release, err := newCompiler(
+	c, err := newCompiler(
 		&wasm.FunctionInstance{Module: j.moduleInstance, Kind: wasm.FunctionKindWasm, Type: functype},
 		&wazeroir.CompilationResult{LabelCallers: map[string]uint32{}},
 	)
-	t.Cleanup(release)
+	t.Cleanup(func() { assemblerMutex.Unlock() })
 	require.NoError(t, err)
 
 	ret, ok := c.(compilerImpl)
diff --git a/internal/wasm/jit/jit_value_location_arm64_test.go b/internal/wasm/jit/jit_value_location_arm64_test.go
deleted file mode 100644
index 7b3fe46c..00000000
--- a/internal/wasm/jit/jit_value_location_arm64_test.go
+++ /dev/null
@@ -1,13 +0,0 @@
-package jit
-
-import (
-	"testing"
-
-	"github.com/stretchr/testify/require"
-	"github.com/twitchyliquid64/golang-asm/obj/arm64"
-)
-
-func Test_simdRegisterForScalarFloatRegister(t *testing.T) {
-	require.Equal(t, int16(arm64.REG_V0), simdRegisterForScalarFloatRegister(arm64.REG_F0))
-	require.Equal(t, int16(arm64.REG_V30), simdRegisterForScalarFloatRegister(arm64.REG_F30))
-}
diff --git a/internal/wasm/jit/jit_value_location_test.go b/internal/wasm/jit/jit_value_location_test.go
index b7c17f0e..c444bc5a 100644
--- a/internal/wasm/jit/jit_value_location_test.go
+++ b/internal/wasm/jit/jit_value_location_test.go
@@ -9,16 +9,6 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-// TestReservedRegisters ensures that reserved registers are not contained in unreservedGeneralPurposeIntRegisters.
-func TestReservedRegisters(t *testing.T) {
-	require.NotContains(t, unreservedGeneralPurposeIntRegisters, reservedRegisterForCallEngine)
-	require.NotContains(t, unreservedGeneralPurposeIntRegisters, reservedRegisterForStackBasePointerAddress)
-	require.NotContains(t, unreservedGeneralPurposeIntRegisters, reservedRegisterForMemory)
-	require.NotContains(t, unreservedGeneralPurposeIntRegisters, reservedRegisterForTemporary)
-	require.NotContains(t, unreservedGeneralPurposeIntRegisters, zeroRegister)
-	require.NotContains(t, unreservedGeneralPurposeIntRegisters, nilRegister)
-}
-
 func Test_isIntRegister(t *testing.T) {
 	for _, r := range unreservedGeneralPurposeIntRegisters {
 		require.True(t, isIntRegister(r))
@@ -26,7 +16,7 @@ func Test_isIntRegister(t *testing.T) {
 }
 
 func Test_isFloatRegister(t *testing.T) {
-	for _, r := range generalPurposeFloatRegisters {
+	for _, r := range unreservedGeneralPurposeFloatRegisters {
 		require.True(t, isFloatRegister(r))
 	}
 }
@@ -43,7 +33,6 @@ func TestValueLocationStack_basic(t *testing.T) {
 	require.Equal(t, uint64(2), s.sp)
 	require.Equal(t, uint64(1), loc.stackPointer)
 	require.Equal(t, tmpReg, loc.register)
-	require.Contains(t, s.usedRegisters, loc.register)
 	// markRegisterUsed.
 	tmpReg2 := unreservedGeneralPurposeIntRegisters[1]
 	s.markRegisterUsed(tmpReg2)
@@ -89,7 +78,7 @@ func TestValueLocationStack_takeFreeRegister(t *testing.T) {
 	require.True(t, ok)
 	require.True(t, isFloatRegister(r))
 	// Mark all the float registers used.
-	for _, r := range generalPurposeFloatRegisters {
+	for _, r := range unreservedGeneralPurposeFloatRegisters {
 		s.markRegisterUsed(r)
 	}
 	// Now we cannot take free ones for floats.
@@ -101,7 +90,7 @@ func TestValueLocationStack_takeStealTargetFromUsedRegister(t *testing.T) {
 	s := newValueLocationStack()
 	intReg := unreservedGeneralPurposeIntRegisters[0]
 	intLocation := &valueLocation{register: intReg}
-	floatReg := generalPurposeFloatRegisters[0]
+	floatReg := unreservedGeneralPurposeFloatRegisters[0]
 	floatLocation := &valueLocation{register: floatReg}
 	s.push(intLocation)
 	s.push(floatLocation)