Repository: tetratelabs/wazero
Pull request #217 (2022-02-09 03:53:47)

jit(arm64): support for br, br_if and labels
part of #187 

there's some of overlaps with amd64, but I intentionally not share the code,
as they are slightly different, and I'm not sure if the new JIT target (e.g. RISCV)
can reuse the code.



mathetake: ok done with label and br including tests.. left is notorious br_if
mathetake: cool done with all
mathetake: After this, the only thing is to implement `call` to support fibonacci!
codefromthecrypt: >
>
> uint32 is enough as that means at least 4GB instructions for a function :D
>
I was actually wondering smaller, but we can do that by limit vs via type
when/if needed

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

codefromthecrypt: can you add the clarification on zero to the docs on the field name? If you tab indent after '//" at least goland formats properly. Edge case notes like this are worth keeping in the doc vs PR comments.

```
// Note: zero possible and allowed in wasm. Ex.
//
//	(block
//	  (br 0)
//	  (block i32.const 1111)
//	)
//
// This example the label corresponding to `(block i32.const 1111)` is never be reached at runtime because `br 0` exits the function before we reach there
```
mathetake: added on the wazeroir side!
mathetake: I think I should reuse the compilationResult to reduce the drifts.. on it


Patch:
diff --git a/RATIONALE.md b/RATIONALE.md
index 2858365d..5db8bb9b 100644
--- a/RATIONALE.md
+++ b/RATIONALE.md
@@ -44,7 +44,7 @@ All the functions are statically analyzed at module instantiation phase, and if
 ### Number of globals in a module
 
 Theoretically, a module can declare globals (including imports) up to 2^32 times. However, wazero limits this to  2^27(134,217,728) per module.
-That is because internally we store globals in a slice with pointer types (meaning 8 bytes on 64-bit platforms), and thefore 2^27 globals
+That is because internally we store globals in a slice with pointer types (meaning 8 bytes on 64-bit platforms), and therefore 2^27 globals
 means that we have 1 GiB size of slice which seems large enough for most applications.
 
 ## JIT engine implementation
diff --git a/wasm/internal/wazeroir/compiler.go b/wasm/internal/wazeroir/compiler.go
index 4ea755c6..20e67cd2 100644
--- a/wasm/internal/wazeroir/compiler.go
+++ b/wasm/internal/wazeroir/compiler.go
@@ -138,14 +138,23 @@ type CompilationResult struct {
 	// LabelCallers maps Label.String() to the number of callers to that label.
 	// Here "callers" means that the callsites which jumps to the label with br, br_if or br_table
 	// instructions.
-	LabelCallers map[string]int
+	//
+	// Note: zero possible and allowed in wasm. Ex.
+	//
+	//	(block
+	//	  (br 0)
+	//	  (block i32.const 1111)
+	//	)
+	//
+	// This example the label corresponding to `(block i32.const 1111)` is never be reached at runtime because `br 0` exits the function before we reach there
+	LabelCallers map[string]uint32
 }
 
 // Compile lowers given function instance into wazeroir operations
 // so that the resulting operations can be consumed by the interpreter
 // or the JIT compilation engine.
 func Compile(f *wasm.FunctionInstance) (*CompilationResult, error) {
-	c := compiler{controlFrames: &controlFrames{}, f: f, result: CompilationResult{LabelCallers: map[string]int{}}}
+	c := compiler{controlFrames: &controlFrames{}, f: f, result: CompilationResult{LabelCallers: map[string]uint32{}}}
 
 	// Push function arguments.
 	for _, t := range f.FunctionType.Type.Params {
diff --git a/wasm/jit/compiler.go b/wasm/jit/compiler.go
index 71e35f55..710bccbc 100644
--- a/wasm/jit/compiler.go
+++ b/wasm/jit/compiler.go
@@ -14,9 +14,9 @@ type compiler interface {
 	// This is used, for example, to initilize the reserved registers, etc.
 	emitPreamble() error
 	// compile generates the byte slice of native code.
-	// maxStackPointer is the max stack pointer that the target function would reach.
+	// stackPointerCeil is the max stack pointer that the target function would reach.
 	// staticData is compiledFunctionStaticData for the resutling native code.
-	compile() (code []byte, staticData compiledFunctionStaticData, maxStackPointer uint64, err error)
+	compile() (code []byte, staticData compiledFunctionStaticData, stackPointerCeil uint64, err error)
 	// compileHostFunction emits the trampoline code from which native code can jump into the host function.
 	// TODO: maybe we wouldn't need to have trampoline for host functions.
 	compileHostFunction(address wasm.FunctionAddress) error
diff --git a/wasm/jit/engine.go b/wasm/jit/engine.go
index 0213cb8f..10475fde 100644
--- a/wasm/jit/engine.go
+++ b/wasm/jit/engine.go
@@ -158,7 +158,7 @@ type (
 		// 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.
-		maxStackPointer uint64
+		stackPointerCeil uint64
 
 		// Followings are not accessed by JITed code.
 
@@ -214,7 +214,7 @@ const (
 
 	// Offsets for compiledFunction.
 	compiledFunctionCodeInitialAddressOffset = 0
-	compiledFunctionMaxStackPointerOffset    = 8
+	compiledFunctionStackPointerCeilOffset   = 8
 
 	// Offsets for wasm.TableElement
 	tableElementFunctionAddressOffset = 0
@@ -551,7 +551,7 @@ jitentry:
 				e.builtinFunctionMemoryGrow(callerCompiledFunction.source.ModuleInstance.Memory)
 			case builtinFunctionAddressGrowValueStack:
 				callerCompiledFunction := e.callFrameTop().compiledFunction
-				e.builtinFunctionGrowValueStack(callerCompiledFunction.maxStackPointer)
+				e.builtinFunctionGrowValueStack(callerCompiledFunction.stackPointerCeil)
 			case builtinFunctionAddressGrowCallFrameStack:
 				e.builtinFunctionGrowCallFrameStack()
 			}
@@ -590,9 +590,9 @@ func (e *engine) pushCallFrame(f *compiledFunction) {
 		e.valueStackContext.stackBasePointer + e.valueStackContext.stackPointer - uint64(len(f.source.FunctionType.Type.Params))
 }
 
-func (e *engine) builtinFunctionGrowValueStack(maxStackPointer uint64) {
-	// Extends the valueStack's length to currentLen*2+maxStackPointer.
-	newLen := e.globalContext.valueStackLen*2 + (maxStackPointer)
+func (e *engine) builtinFunctionGrowValueStack(stackPointerCeil uint64) {
+	// Extends the valueStack's length to currentLen*2+stackPointerCeil.
+	newLen := e.globalContext.valueStackLen*2 + (stackPointerCeil)
 	newStack := make([]uint64, newLen)
 	top := e.valueStackContext.stackBasePointer + e.valueStackContext.stackPointer
 	copy(newStack[:top], e.valueStack[:top])
@@ -675,16 +675,16 @@ func compileHostFunction(f *wasm.FunctionInstance) (*compiledFunction, error) {
 		return nil, err
 	}
 
-	maxStackPointer := uint64(len(f.FunctionType.Type.Params))
-	if res := uint64(len(f.FunctionType.Type.Results)); maxStackPointer < res {
-		maxStackPointer = res
+	stackPointerCeil := uint64(len(f.FunctionType.Type.Params))
+	if res := uint64(len(f.FunctionType.Type.Results)); stackPointerCeil < res {
+		stackPointerCeil = res
 	}
 
 	return &compiledFunction{
 		source:             f,
 		codeSegment:        code,
 		codeInitialAddress: uintptr(unsafe.Pointer(&code[0])),
-		maxStackPointer:    maxStackPointer,
+		stackPointerCeil:   stackPointerCeil,
 	}, nil
 }
 
@@ -866,7 +866,7 @@ func compileWasmFunction(f *wasm.FunctionInstance) (*compiledFunction, error) {
 		}
 	}
 
-	code, staticData, maxStackPointer, err := compiler.compile()
+	code, staticData, stackPointerCeil, err := compiler.compile()
 	if err != nil {
 		return nil, fmt.Errorf("failed to compile: %w", err)
 	}
@@ -879,7 +879,7 @@ func compileWasmFunction(f *wasm.FunctionInstance) (*compiledFunction, error) {
 		source:             f,
 		codeSegment:        code,
 		codeInitialAddress: uintptr(unsafe.Pointer(&code[0])),
-		maxStackPointer:    maxStackPointer,
+		stackPointerCeil:   stackPointerCeil,
 		staticData:         staticData,
 	}, nil
 }
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index 16b836f5..308aa8bb 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -101,13 +101,8 @@ func newCompiler(f *wasm.FunctionInstance, ir *wazeroir.CompilationResult) (comp
 		builder:       b,
 		locationStack: newValueLocationStack(),
 		currentLabel:  wazeroir.EntrypointLabel,
-	}
-
-	if ir != nil {
-		compiler.labels = make(map[string]*labelInfo, len(ir.LabelCallers))
-		for key, callers := range ir.LabelCallers {
-			compiler.labels[key] = &labelInfo{callers: callers}
-		}
+		ir:            ir,
+		labels:        map[string]*labelInfo{},
 	}
 	return compiler, nil
 }
@@ -119,31 +114,32 @@ func (c *amd64Compiler) String() string {
 type amd64Compiler struct {
 	builder *asm.Builder
 	f       *wasm.FunctionInstance
+	ir      *wazeroir.CompilationResult
 	// setJmpOrigins sets jmp kind instructions where you want to set the next coming
 	// instruction as the destination of the jmp instruction.
 	setJmpOrigins []*obj.Prog
 	// locationStack holds the state of wazeroir virtual stack.
 	// and each item is either placed in register or the actual memory stack.
 	locationStack *valueLocationStack
-	// labels holds per wazeroir label specific informationsin this function.
+	// labels hold per wazeroir label specific information in this function.
 	labels map[string]*labelInfo
-	// maxStackPointer tracks the maximum value of stack pointer (from valueLocationStack).
-	maxStackPointer uint64
+	// stackPointerCeil is the greatest stack pointer value (from valueLocationStack) seen during compilation.
+	stackPointerCeil uint64
 	// currentLabel holds a currently compiled wazeroir label key. For debugging only.
 	currentLabel string
-	// onMaxStackPointerDeterminedCallBack hold a callback which are called when the max stack pointer is determined BEFORE generating native code.
-	onMaxStackPointerDeterminedCallBack func(maxStackPointer uint64)
+	// onStackPointerCeilDeterminedCallBack hold a callback which are called when the max stack pointer is determined BEFORE generating native code.
+	onStackPointerCeilDeterminedCallBack func(stackPointerCeil uint64)
 	// onGenerateCallbacks holds the callbacks which are called AFTER generating native code.
 	onGenerateCallbacks []func(code []byte) error
 	staticData          compiledFunctionStaticData
 }
 
-// replaceLocationStack sets the given valueLocationStack to .locationStack field,
-// while allowing us to track valueLocationStack.maxStackPointer across multiple stacks.
+// setLocationStack sets the given valueLocationStack to .locationStack field,
+// while allowing us to track valueLocationStack.stackPointerCeil across multiple stacks.
 // This is called when we branch into different block.
-func (c *amd64Compiler) replaceLocationStack(newStack *valueLocationStack) {
-	if c.maxStackPointer < c.locationStack.maxStackPointer {
-		c.maxStackPointer = c.locationStack.maxStackPointer
+func (c *amd64Compiler) setLocationStack(newStack *valueLocationStack) {
+	if c.stackPointerCeil < c.locationStack.stackPointerCeil {
+		c.stackPointerCeil = c.locationStack.stackPointerCeil
 	}
 	c.locationStack = newStack
 }
@@ -153,8 +149,6 @@ func (c *amd64Compiler) addStaticData(d []byte) {
 }
 
 type labelInfo struct {
-	// callers is the number of call sites which may jump into this label.
-	callers int
 	// initialInstruction is the initial instruction for this label so other block can jump into it.
 	initialInstruction *obj.Prog
 	// initialStack is the initial value location stack from which we start compiling this label.
@@ -202,19 +196,19 @@ func (c *amd64Compiler) compileHostFunction(address wasm.FunctionAddress) error
 }
 
 // compile implements compiler.compile for the amd64 architecture.
-func (c *amd64Compiler) compile() (code []byte, staticData compiledFunctionStaticData, maxStackPointer uint64, err error) {
-	// c.maxStackPointer tracks the maximum stack pointer across all valueLocationStack(s)
-	// used for all labels (via replaceLocationStack), excluding the current one.
-	// Hence, we check here if the final block's max one exceeds the current c.maxStackPointer.
-	maxStackPointer = c.maxStackPointer
-	if maxStackPointer < c.locationStack.maxStackPointer {
-		maxStackPointer = c.locationStack.maxStackPointer
+func (c *amd64Compiler) compile() (code []byte, staticData compiledFunctionStaticData, stackPointerCeil uint64, err error) {
+	// c.stackPointerCeil tracks the stack pointer ceiling (max seen) value across all valueLocationStack(s)
+	// used for all labels (via setLocationStack), excluding the current one.
+	// Hence, we check here if the final block's max one exceeds the current c.stackPointerCeil.
+	stackPointerCeil = c.stackPointerCeil
+	if stackPointerCeil < c.locationStack.stackPointerCeil {
+		stackPointerCeil = c.locationStack.stackPointerCeil
 	}
 
 	// Now that the max stack pointer is determined, we are invoking the callback.
 	// Note this MUST be called before Assemble() befolow.
-	if c.onMaxStackPointerDeterminedCallBack != nil {
-		c.onMaxStackPointerDeterminedCallBack(maxStackPointer)
+	if c.onStackPointerCeilDeterminedCallBack != nil {
+		c.onStackPointerCeilDeterminedCallBack(stackPointerCeil)
 	}
 
 	code, err = mmapCodeSegment(c.builder.Assemble())
@@ -524,14 +518,11 @@ func (c *amd64Compiler) compileBr(o *wazeroir.OperationBr) error {
 // branchInto adds instruction necessary to jump into the given branch target.
 func (c *amd64Compiler) branchInto(target *wazeroir.BranchTarget) error {
 	if target.IsReturnTarget() {
-		if err := c.returnFunction(); err != nil {
-			return err
-		}
+		return c.returnFunction()
 	} else {
 		labelKey := target.String()
-		targetLabel := c.label(labelKey)
-		if targetLabel.callers > 1 {
-			// If the number of callers to the target label is larget than one,
+		if c.ir.LabelCallers[labelKey] > 1 {
+			// If the number of callers to the target label is larger than one,
 			// we have multiple origins to the target branch. In that case,
 			// we must have unique register state.
 			if err := c.preLabelJumpRegisterAdjustment(); err != nil {
@@ -541,6 +532,7 @@ func (c *amd64Compiler) branchInto(target *wazeroir.BranchTarget) error {
 		// Set the initial stack of the target label, so we can start compiling the label
 		// with the appropriate value locations. Note we clone the stack here as we maybe
 		// manipulate the stack before compiler reaches the label.
+		targetLabel := c.label(labelKey)
 		if targetLabel.initialStack == nil {
 			// It seems unnecessary to clone as branchInto is always the tail of the current block.
 			// TODO: verify ^^.
@@ -635,15 +627,14 @@ func (c *amd64Compiler) compileBrIf(o *wazeroir.OperationBrIf) error {
 
 	// Emit for else branches
 	saved := c.locationStack
-	c.replaceLocationStack(saved.clone())
+	c.setLocationStack(saved.clone())
 	if elseTarget.Target.IsReturnTarget() {
 		if err := c.returnFunction(); err != nil {
 			return err
 		}
 	} else {
 		elseLabelKey := elseTarget.Target.Label.String()
-		labelInfo := c.label(elseLabelKey)
-		if labelInfo.callers > 1 {
+		if c.ir.LabelCallers[elseLabelKey] > 1 {
 			if err := c.preLabelJumpRegisterAdjustment(); err != nil {
 				return err
 			}
@@ -651,6 +642,7 @@ func (c *amd64Compiler) compileBrIf(o *wazeroir.OperationBrIf) error {
 		// Set the initial stack of the target label, so we can start compiling the label
 		// with the appropriate value locations. Note we clone the stack here as we maybe
 		// manipulate the stack before compiler reaches the label.
+		labelInfo := c.label(elseLabelKey)
 		if labelInfo.initialStack == nil {
 			labelInfo.initialStack = c.locationStack
 		}
@@ -663,18 +655,15 @@ func (c *amd64Compiler) compileBrIf(o *wazeroir.OperationBrIf) error {
 
 	// Handle then branch.
 	c.addSetJmpOrigins(jmpWithCond)
-	c.replaceLocationStack(saved)
+	c.setLocationStack(saved)
 	if err := c.emitDropRange(thenTarget.ToDrop); err != nil {
 		return err
 	}
 	if thenTarget.Target.IsReturnTarget() {
-		if err := c.returnFunction(); err != nil {
-			return err
-		}
+		return c.returnFunction()
 	} else {
 		thenLabelKey := thenTarget.Target.Label.String()
-		labelInfo := c.label(thenLabelKey)
-		if c.label(thenLabelKey).callers > 1 {
+		if c.ir.LabelCallers[thenLabelKey] > 1 {
 			if err := c.preLabelJumpRegisterAdjustment(); err != nil {
 				return err
 			}
@@ -682,6 +671,7 @@ func (c *amd64Compiler) compileBrIf(o *wazeroir.OperationBrIf) error {
 		// Set the initial stack of the target label, so we can start compiling the label
 		// with the appropriate value locations. Note we clone the stack here as we maybe
 		// manipulate the stack before compiler reaches the label.
+		labelInfo := c.label(thenLabelKey)
 		if labelInfo.initialStack == nil {
 			labelInfo.initialStack = c.locationStack
 		}
@@ -690,8 +680,8 @@ func (c *amd64Compiler) compileBrIf(o *wazeroir.OperationBrIf) error {
 		thenJmp.To.Type = obj.TYPE_BRANCH
 		c.addInstruction(thenJmp)
 		c.assignJumpTarget(thenLabelKey, thenJmp)
+		return nil
 	}
-	return nil
 }
 
 // compileBrTable implements compiler.compileBrTable for the amd64 architecture.
@@ -846,7 +836,7 @@ func (c *amd64Compiler) compileBrTable(o *wazeroir.OperationBrTable) error {
 			// as this is the last code in this block.
 			locationStack = saved
 		}
-		c.replaceLocationStack(locationStack)
+		c.setLocationStack(locationStack)
 		if err := c.emitDropRange(target.ToDrop); err != nil {
 			return err
 		}
@@ -928,7 +918,7 @@ func (c *amd64Compiler) compileLabel(o *wazeroir.OperationLabel) (skipLabel bool
 	labelInfo.initialInstruction = labelBegin
 
 	// Set the initial stack.
-	c.replaceLocationStack(labelInfo.initialStack)
+	c.setLocationStack(labelInfo.initialStack)
 
 	// Invoke callbacks to notify the forward branching
 	// instructions can properly jump to this label.
@@ -940,7 +930,7 @@ func (c *amd64Compiler) compileLabel(o *wazeroir.OperationLabel) (skipLabel bool
 	labelInfo.labelBeginningCallbacks = nil
 
 	if buildoptions.IsDebugMode {
-		fmt.Printf("[label %s (num callers=%d)]\n%s\n", labelKey, labelInfo.callers, c.locationStack)
+		fmt.Printf("[label %s (num callers=%d)]\n%s\n", labelKey, c.ir.LabelCallers[labelKey], c.locationStack)
 	}
 	c.currentLabel = labelKey
 	return
@@ -5336,15 +5326,15 @@ func (c *amd64Compiler) maybeGrowValueStack() error {
 	c.addInstruction(subStackBasePointer)
 
 	// If stack base pointer + max stack poitner > valueStackLen, we need to grow the stack.
-	cmpWithMaxStackPointer := c.newProg()
-	cmpWithMaxStackPointer.As = x86.ACMPQ
-	cmpWithMaxStackPointer.From.Type = obj.TYPE_REG
-	cmpWithMaxStackPointer.From.Reg = tmpRegister
-	cmpWithMaxStackPointer.To.Type = obj.TYPE_CONST
+	cmpWithStackPointerCeil := c.newProg()
+	cmpWithStackPointerCeil.As = x86.ACMPQ
+	cmpWithStackPointerCeil.From.Type = obj.TYPE_REG
+	cmpWithStackPointerCeil.From.Reg = tmpRegister
+	cmpWithStackPointerCeil.To.Type = obj.TYPE_CONST
 	// We don't yet know the max stack poitner at this point.
 	// The max stack pointer is determined after emitting all the instructions.
-	c.onMaxStackPointerDeterminedCallBack = func(maxStackPointer uint64) { cmpWithMaxStackPointer.To.Offset = int64(maxStackPointer) }
-	c.addInstruction(cmpWithMaxStackPointer)
+	c.onStackPointerCeilDeterminedCallBack = func(stackPointerCeil uint64) { cmpWithStackPointerCeil.To.Offset = int64(stackPointerCeil) }
+	c.addInstruction(cmpWithStackPointerCeil)
 
 	// Jump if we have no need to grow.
 	jmpIfNoNeedToGrowStack := c.newProg()
diff --git a/wasm/jit/jit_arm64.go b/wasm/jit/jit_arm64.go
index 32012bb8..2985d3f6 100644
--- a/wasm/jit/jit_arm64.go
+++ b/wasm/jit/jit_arm64.go
@@ -57,21 +57,38 @@ func newCompiler(f *wasm.FunctionInstance, ir *wazeroir.CompilationResult) (comp
 		f:             f,
 		builder:       b,
 		locationStack: newValueLocationStack(),
+		ir:            ir,
+		labels:        map[string]*labelInfo{},
 	}
-
 	return compiler, nil
 }
 
 type arm64Compiler struct {
 	builder *asm.Builder
 	f       *wasm.FunctionInstance
+	ir      *wazeroir.CompilationResult
+	// setBRTargetOnNextInstructions holds branch kind instructions (BR, conditional BR, etc)
+	// where we want to set the next coming instruction as the destination of these BR instructions.
+	setBRTargetOnNextInstructions []*obj.Prog
 	// locationStack holds the state of wazeroir virtual stack.
 	// and each item is either placed in register or the actual memory stack.
 	locationStack *valueLocationStack
+	// labels maps a label (Ex. ".L1_then") to *labelInfo.
+	labels map[string]*labelInfo
+	// stackPointerCeil is the greatest stack pointer value (from valueLocationStack) seen during compilation.
+	stackPointerCeil uint64
 }
 
 // compile implements compiler.compile for the arm64 architecture.
-func (c *arm64Compiler) compile() (code []byte, staticData compiledFunctionStaticData, maxStackPointer uint64, err error) {
+func (c *arm64Compiler) compile() (code []byte, staticData compiledFunctionStaticData, stackPointerCeil uint64, err error) {
+	// c.stackPointerCeil tracks the stack pointer ceiling (max seen) value across all valueLocationStack(s)
+	// used for all labels (via setLocationStack), excluding the current one.
+	// Hence, we check here if the final block's max one exceeds the current c.stackPointerCeil.
+	stackPointerCeil = c.stackPointerCeil
+	if stackPointerCeil < c.locationStack.stackPointerCeil {
+		stackPointerCeil = c.locationStack.stackPointerCeil
+	}
+
 	code, err = mmapCodeSegment(c.builder.Assemble())
 	if err != nil {
 		return
@@ -79,8 +96,31 @@ func (c *arm64Compiler) compile() (code []byte, staticData compiledFunctionStati
 	return
 }
 
+// labelInfo holds a wazeroir label specific information in this function.
+type labelInfo struct {
+	// initialInstruction is the initial instruction for this label so other block can branch into it.
+	initialInstruction *obj.Prog
+	// initialStack is the initial value location stack from which we start compiling this label.
+	initialStack *valueLocationStack
+	// labelBeginningCallbacks holds callbacks should to be called with initialInstruction
+	labelBeginningCallbacks []func(*obj.Prog)
+}
+
+func (c *arm64Compiler) label(labelKey string) *labelInfo {
+	ret, ok := c.labels[labelKey]
+	if ok {
+		return ret
+	}
+	c.labels[labelKey] = &labelInfo{}
+	return c.labels[labelKey]
+}
+
 func (c *arm64Compiler) newProg() (inst *obj.Prog) {
 	inst = c.builder.NewProg()
+	for _, origin := range c.setBRTargetOnNextInstructions {
+		origin.To.SetTarget(inst)
+	}
+	c.setBRTargetOnNextInstructions = nil
 	return
 }
 
@@ -88,6 +128,10 @@ func (c *arm64Compiler) addInstruction(inst *obj.Prog) {
 	c.builder.AddInstruction(inst)
 }
 
+func (c *arm64Compiler) setBRTargetOnNext(progs ...*obj.Prog) {
+	c.setBRTargetOnNextInstructions = append(c.setBRTargetOnNextInstructions, progs...)
+}
+
 func (c *arm64Compiler) markRegisterUsed(reg int16) {
 	c.locationStack.markRegisterUsed(reg)
 }
@@ -217,6 +261,14 @@ func (c *arm64Compiler) applyTwoRegistersToNoneInstruction(instruction obj.As, s
 	c.addInstruction(inst)
 }
 
+func (c *arm64Compiler) emitUnconditionalBRInstruction(targetType obj.AddrType) (jmp *obj.Prog) {
+	jmp = c.newProg()
+	jmp.As = obj.AJMP
+	jmp.To.Type = targetType
+	c.addInstruction(jmp)
+	return
+}
+
 func (c *arm64Compiler) String() (ret string) { return }
 
 // pushFunctionParams pushes any function parameters onto the stack, setting appropriate register types.
@@ -251,7 +303,7 @@ func (c *arm64Compiler) emitPreamble() error {
 
 // returnFunction emits instructions to return from the current function frame.
 // If the current frame is the bottom, the code goes back to the Go code with jitCallStatusCodeReturned status.
-// Otherwise, we jump into the caller's return address (TODO).
+// Otherwise, we branch into the caller's return address (TODO).
 func (c *arm64Compiler) returnFunction() error {
 	// TODO: we don't support function calls yet.
 	// For now the following code just returns to Go code.
@@ -314,7 +366,45 @@ func (c *arm64Compiler) compileHostFunction(address wasm.FunctionAddress) error
 	return errors.New("TODO: implement compileHostFunction on arm64")
 }
 
+// setLocationStack sets the given valueLocationStack to .locationStack field,
+// while allowing us to track valueLocationStack.stackPointerCeil across multiple stacks.
+// This is called when we branch into different block.
+func (c *arm64Compiler) setLocationStack(newStack *valueLocationStack) {
+	if c.stackPointerCeil < c.locationStack.stackPointerCeil {
+		c.stackPointerCeil = c.locationStack.stackPointerCeil
+	}
+	c.locationStack = newStack
+}
+
+// arm64Compiler implements compiler.arm64Compiler for the arm64 architecture.
 func (c *arm64Compiler) compileLabel(o *wazeroir.OperationLabel) (skipThisLabel bool) {
+	labelKey := o.Label.String()
+	labelInfo := c.label(labelKey)
+
+	// If initialStack is not set, that means this label has never been reached.
+	if labelInfo.initialStack == nil {
+		skipThisLabel = true
+		return
+	}
+
+	// We use NOP as a beginning of instructions in a label.
+	// This should be eventually optimized out by assembler.
+	labelBegin := c.newProg()
+	labelBegin.As = obj.ANOP
+	c.addInstruction(labelBegin)
+
+	// Save the instructions so that backward branching
+	// instructions can branch to this label.
+	labelInfo.initialInstruction = labelBegin
+
+	// Set the initial stack.
+	c.setLocationStack(labelInfo.initialStack)
+
+	// Invoke callbacks to notify the forward branching
+	// instructions can properly branch to this label.
+	for _, cb := range labelInfo.labelBeginningCallbacks {
+		cb(labelBegin)
+	}
 	return false
 }
 
@@ -336,15 +426,128 @@ func (c *arm64Compiler) compileGlobalSet(o *wazeroir.OperationGlobalSet) error {
 
 // compileBr implements compiler.compileBr for the arm64 architecture.
 func (c *arm64Compiler) compileBr(o *wazeroir.OperationBr) error {
-	if o.Target.IsReturnTarget() {
+	c.maybeMoveTopConditionalToFreeGeneralPurposeRegister()
+	return c.branchInto(o.Target)
+}
+
+// compileBrIf implements compiler.compileBrIf for the arm64 architecture.
+func (c *arm64Compiler) compileBrIf(o *wazeroir.OperationBrIf) error {
+	cond := c.locationStack.pop()
+
+	conditionalBR := c.newProg()
+	conditionalBR.To.Type = obj.TYPE_BRANCH
+	if cond.onConditionalRegister() {
+		// If the cond is on a conditional register, it corresponds to one of "conditonal codes"
+		// https://developer.arm.com/documentation/dui0801/a/Condition-Codes/Condition-code-suffixes
+		// Here we represent the conditional codes by using arm64.COND_** registers, and that means the
+		// conditional jump can be performed if we use arm64.AB**.
+		// For example, if we have arm64.COND_EQ on cond, that means we performed compileEq right before
+		// this compileBrIf and BrIf can be achieved by arm64.ABEQ.
+		switch cond.conditionalRegister {
+		case arm64.COND_EQ:
+			conditionalBR.As = arm64.ABEQ
+		case arm64.COND_NE:
+			conditionalBR.As = arm64.ABNE
+		case arm64.COND_HS:
+			conditionalBR.As = arm64.ABHS
+		case arm64.COND_LO:
+			conditionalBR.As = arm64.ABLO
+		case arm64.COND_MI:
+			conditionalBR.As = arm64.ABMI
+		case arm64.COND_HI:
+			conditionalBR.As = arm64.ABHI
+		case arm64.COND_LS:
+			conditionalBR.As = arm64.ABLS
+		case arm64.COND_GE:
+			conditionalBR.As = arm64.ABGE
+		case arm64.COND_LT:
+			conditionalBR.As = arm64.ABLT
+		case arm64.COND_GT:
+			conditionalBR.As = arm64.ABGT
+		case arm64.COND_LE:
+			conditionalBR.As = arm64.ABLE
+		default:
+			// BUG: This means that we use the cond.conditionalRegister somewhere in this file,
+			// but not covered in switch ^. That shouldn't happen.
+			return fmt.Errorf("unsupported condition for br_if: %v", cond.conditionalRegister)
+		}
+	} else {
+		// If the value is not on the conditional register, we compare the value with the zero register,
+		// and then do the conditional BR if the value does't equal zero.
+		if err := c.ensureOnGeneralPurposeRegister(cond); err != nil {
+			return err
+		}
+		// Compare the value with zero register. Note that the value is ensured to be i32 by function validation phase,
+		// so we use CMPW (32-bit compare) here.
+		c.applyTwoRegistersToNoneInstruction(arm64.ACMPW, cond.register, zeroRegister)
+		conditionalBR.As = arm64.ABNE
+
+		c.markRegisterUnused(cond.register)
+	}
+
+	c.addInstruction(conditionalBR)
+
+	// Emit the code for branching into else branch.
+	// We save and clone the location stack because we might end up modifying it inside of branchInto,
+	// and we have to avoid affecting the code generation for Then branch afterwards.
+	saved := c.locationStack
+	c.setLocationStack(saved.clone())
+	if err := c.emitDropRange(o.Else.ToDrop); err != nil {
+		return err
+	}
+	c.branchInto(o.Else.Target)
+
+	// Now ready to emit the code for branching into then branch.
+	// Retrieve the original value location stack so that the code below wont'be affected by the Else branch ^^.
+	c.setLocationStack(saved)
+	// We branch into here from the original conditional BR (conditionalBR).
+	c.setBRTargetOnNext(conditionalBR)
+	if err := c.emitDropRange(o.Then.ToDrop); err != nil {
+		return err
+	}
+	c.branchInto(o.Then.Target)
+	return nil
+}
+
+func (c *arm64Compiler) branchInto(target *wazeroir.BranchTarget) error {
+	if target.IsReturnTarget() {
 		return c.returnFunction()
 	} else {
-		return fmt.Errorf("TODO: only return target is available on arm64")
+		labelKey := target.String()
+		if c.ir.LabelCallers[labelKey] > 1 {
+			// We can only re-use register state if when there's a single call-site.
+			// Release existing values on registers to the stack if there's multiple ones to have
+			// the consistent value location state at the beginning of label.
+			if err := c.releaseAllRegistersToStack(); err != nil {
+				return err
+			}
+		}
+		// Set the initial stack of the target label, so we can start compiling the label
+		// with the appropriate value locations. Note we clone the stack here as we maybe
+		// manipulate the stack before compiler reaches the label.
+		targetLabel := c.label(labelKey)
+		if targetLabel.initialStack == nil {
+			targetLabel.initialStack = c.locationStack.clone()
+		}
+
+		jmp := c.emitUnconditionalBRInstruction(obj.TYPE_BRANCH)
+		c.assignBranchTarget(labelKey, jmp)
+		return nil
 	}
 }
 
-func (c *arm64Compiler) compileBrIf(o *wazeroir.OperationBrIf) error {
-	return fmt.Errorf("TODO: unsupported on arm64")
+// assignBranchTarget assigns the given label's initial instruction to the destination of br.
+func (c *arm64Compiler) assignBranchTarget(labelKey string, br *obj.Prog) {
+	target := c.label(labelKey)
+	if target.initialInstruction != nil {
+		br.To.SetTarget(target.initialInstruction)
+	} else {
+		// This case, the target label hasn't been compiled yet, so we append the callback and assign
+		// the target instruction when compileLabel is called for the label.
+		target.labelBeginningCallbacks = append(target.labelBeginningCallbacks, func(labelInitialInstruction *obj.Prog) {
+			br.To.SetTarget(labelInitialInstruction)
+		})
+	}
 }
 
 func (c *arm64Compiler) compileBrTable(o *wazeroir.OperationBrTable) error {
@@ -361,7 +564,11 @@ func (c *arm64Compiler) compileCallIndirect(o *wazeroir.OperationCallIndirect) e
 
 // compileDrop implements compiler.compileDrop for the arm64 architecture.
 func (c *arm64Compiler) compileDrop(o *wazeroir.OperationDrop) error {
-	r := o.Range
+	return c.emitDropRange(o.Range)
+}
+
+// emitDropRange is the implementation of compileDrop. See compiler.compileDrop.
+func (c *arm64Compiler) emitDropRange(r *wazeroir.InclusiveRange) error {
 	if r == nil {
 		return nil
 	} else if r.Start == 0 {
diff --git a/wasm/jit/jit_value_location.go b/wasm/jit/jit_value_location.go
index 4d3f0b25..c3753ffb 100644
--- a/wasm/jit/jit_value_location.go
+++ b/wasm/jit/jit_value_location.go
@@ -90,14 +90,14 @@ func newValueLocationStack() *valueLocationStack {
 // moving the variable to registers to perform actual CPU instruction
 // to achieve wazeroir's add operation.
 type valueLocationStack struct {
-	// Holds all the variables.
+	// stack holds all the variables.
 	stack []*valueLocation
-	// The current stack pointer.
+	// sp is the current stack pointer.
 	sp uint64
-	// Stores the used registers.
+	// usedRegisters stores the used registers.
 	usedRegisters map[int16]struct{}
-	// Records max(.sp) across the lifespan of this struct.
-	maxStackPointer uint64
+	// stackPointerCeil tracks max(.sp) across the lifespan of this struct.
+	stackPointerCeil uint64
 }
 
 func (v *valueLocationStack) String() string {
@@ -128,7 +128,7 @@ func (s *valueLocationStack) clone() *valueLocationStack {
 			register:            v.register,
 		}
 	}
-	ret.maxStackPointer = s.maxStackPointer
+	ret.stackPointerCeil = s.stackPointerCeil
 	return ret
 }
 
@@ -177,8 +177,8 @@ func (s *valueLocationStack) push(loc *valueLocation) {
 	} else {
 		s.stack[s.sp] = loc
 	}
-	if s.sp > s.maxStackPointer {
-		s.maxStackPointer = s.sp
+	if s.sp > s.stackPointerCeil {
+		s.stackPointerCeil = s.sp
 	}
 	s.sp++
 }


Test patch:
diff --git a/wasm/jit/engine_test.go b/wasm/jit/engine_test.go
index abcd2949..20d340a6 100644
--- a/wasm/jit/engine_test.go
+++ b/wasm/jit/engine_test.go
@@ -51,7 +51,7 @@ func TestVerifyOffsetValue(t *testing.T) {
 	// Offsets for compiledFunction.
 	var compiledFunc compiledFunction
 	require.Equal(t, int(unsafe.Offsetof(compiledFunc.codeInitialAddress)), compiledFunctionCodeInitialAddressOffset)
-	require.Equal(t, int(unsafe.Offsetof(compiledFunc.maxStackPointer)), compiledFunctionMaxStackPointerOffset)
+	require.Equal(t, int(unsafe.Offsetof(compiledFunc.stackPointerCeil)), compiledFunctionStackPointerCeilOffset)
 
 	// Offsets for wasm.TableElement.
 	var tableElement wasm.TableElement
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index d260f2bb..7542fd5c 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -54,13 +54,13 @@ func TestAmd64Compiler_maybeGrowValueStack(t *testing.T) {
 			compiler.initializeReservedStackBasePointer()
 			err := compiler.maybeGrowValueStack()
 			require.NoError(t, err)
-			require.NotNil(t, compiler.onMaxStackPointerDeterminedCallBack)
+			require.NotNil(t, compiler.onStackPointerCeilDeterminedCallBack)
 
 			valueStackLen := uint64(len(env.stack()))
-			maxStackPointer := uint64(5)
+			stackPointerCeil := uint64(5)
 			stackBasePointer := valueStackLen - baseOffset // Base + Max <= valueStackLen = no need to grow!
-			compiler.onMaxStackPointerDeterminedCallBack(maxStackPointer)
-			compiler.onMaxStackPointerDeterminedCallBack = nil
+			compiler.onStackPointerCeilDeterminedCallBack(stackPointerCeil)
+			compiler.onStackPointerCeilDeterminedCallBack = nil
 			env.setValueStackBasePointer(stackBasePointer)
 
 			compiler.exit(jitCallStatusCodeReturned)
@@ -87,8 +87,8 @@ func TestAmd64Compiler_maybeGrowValueStack(t *testing.T) {
 		// On the return from grow value stack, we just exit with "Returned" status.
 		compiler.exit(jitCallStatusCodeReturned)
 
-		maxStackPointer := uint64(6)
-		compiler.maxStackPointer = maxStackPointer
+		stackPointerCeil := uint64(6)
+		compiler.stackPointerCeil = stackPointerCeil
 		valueStackLen := uint64(len(env.stack()))
 		stackBasePointer := valueStackLen - 5 // Base + Max > valueStackLen = need to grow!
 		env.setValueStackBasePointer(stackBasePointer)
@@ -326,7 +326,7 @@ func TestAmd64Compiler_compileBrTable(t *testing.T) {
 		// Emit code for each label which returns the frame ID.
 		for returnValue := uint32(0); returnValue < 10; returnValue++ {
 			label := &wazeroir.Label{Kind: wazeroir.LabelKindHeader, FrameID: returnValue}
-			c.label(label.String()).callers = 1
+			c.ir.LabelCallers[label.String()] = 1
 			_ = c.compileLabel(&wazeroir.OperationLabel{Label: label})
 			_ = c.compileConstI32(&wazeroir.OperationConstI32{Value: label.FrameID})
 			err := c.releaseAllRegistersToStack()
@@ -477,6 +477,8 @@ func TestAmd64Compiler_compileBrTable(t *testing.T) {
 					t.Run(tc.name, func(t *testing.T) {
 						env := newJITEnvironment()
 						compiler := env.requireNewCompiler(t)
+						compiler.ir = &wazeroir.CompilationResult{LabelCallers: map[string]uint32{}}
+
 						err := compiler.emitPreamble()
 						require.NoError(t, err)
 
@@ -5572,31 +5574,31 @@ func TestAmd64Compiler_generate(t *testing.T) {
 			compiler.addInstruction(ret)
 			return
 		}
-		verify := func(t *testing.T, compiler *amd64Compiler, expectedMaxStackPointer uint64) {
+		verify := func(t *testing.T, compiler *amd64Compiler, expectedStackPointerCeil uint64) {
 			var called bool
-			compiler.onMaxStackPointerDeterminedCallBack = func(acutalMaxStackPointerInCallBack uint64) {
+			compiler.onStackPointerCeilDeterminedCallBack = func(acutalStackPointerCeilInCallBack uint64) {
 				called = true
-				require.Equal(t, expectedMaxStackPointer, acutalMaxStackPointerInCallBack)
+				require.Equal(t, expectedStackPointerCeil, acutalStackPointerCeilInCallBack)
 			}
 
-			_, _, acutalMaxStackPointer, err := compiler.compile()
+			_, _, acutalStackPointerCeil, err := compiler.compile()
 			require.NoError(t, err)
 			require.True(t, called)
-			require.Equal(t, expectedMaxStackPointer, acutalMaxStackPointer)
+			require.Equal(t, expectedStackPointerCeil, acutalStackPointerCeil)
 		}
 		t.Run("current one win", func(t *testing.T) {
 			compiler := getCompiler(t)
-			const expectedMaxStackPointer uint64 = 100
-			compiler.maxStackPointer = expectedMaxStackPointer
-			compiler.locationStack.maxStackPointer = expectedMaxStackPointer - 1
-			verify(t, compiler, expectedMaxStackPointer)
+			const expectedStackPointerCeil uint64 = 100
+			compiler.stackPointerCeil = expectedStackPointerCeil
+			compiler.locationStack.stackPointerCeil = expectedStackPointerCeil - 1
+			verify(t, compiler, expectedStackPointerCeil)
 		})
 		t.Run("previous one win", func(t *testing.T) {
 			compiler := getCompiler(t)
-			const expectedMaxStackPointer uint64 = 100
-			compiler.locationStack.maxStackPointer = expectedMaxStackPointer
-			compiler.maxStackPointer = expectedMaxStackPointer - 1
-			verify(t, compiler, expectedMaxStackPointer)
+			const expectedStackPointerCeil uint64 = 100
+			compiler.locationStack.stackPointerCeil = expectedStackPointerCeil
+			compiler.stackPointerCeil = expectedStackPointerCeil - 1
+			verify(t, compiler, expectedStackPointerCeil)
 		})
 	})
 
@@ -6389,7 +6391,7 @@ func TestAmd64Compiler_readInstructionAddress(t *testing.T) {
 		compiler.addInstruction(ret)
 
 		// This could be the read instruction target as this is the
-		// right after RET. Thefore, the jmp instruction above
+		// right after RET. Therefore, the jmp instruction above
 		// must target here.
 		const expectedReturnValue uint32 = 10000
 		err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: expectedReturnValue})
diff --git a/wasm/jit/jit_arm64_test.go b/wasm/jit/jit_arm64_test.go
index 1a25ae9d..51ff7159 100644
--- a/wasm/jit/jit_arm64_test.go
+++ b/wasm/jit/jit_arm64_test.go
@@ -11,17 +11,41 @@ import (
 	"unsafe"
 
 	"github.com/stretchr/testify/require"
+	"github.com/twitchyliquid64/golang-asm/obj"
 	"github.com/twitchyliquid64/golang-asm/obj/arm64"
 
 	"github.com/tetratelabs/wazero/wasm"
 	"github.com/tetratelabs/wazero/wasm/internal/wazeroir"
 )
 
+func requireAddLabel(t *testing.T, compiler *arm64Compiler, label *wazeroir.Label) {
+	// We set a value location stack so that the label must not be skipped.
+	compiler.label(label.String()).initialStack = newValueLocationStack()
+	skip := compiler.compileLabel(&wazeroir.OperationLabel{Label: label})
+	require.False(t, skip)
+}
+
+func requirePushTwoInt32Consts(t *testing.T, x1, x2 uint32, compiler *arm64Compiler) {
+	err := compiler.compileConstI32(&wazeroir.OperationConstI32{Value: x1})
+	require.NoError(t, err)
+	err = compiler.compileConstI32(&wazeroir.OperationConstI32{Value: x2})
+	require.NoError(t, err)
+}
+
+func requirePushTwoFloat32Consts(t *testing.T, x1, x2 float32, compiler *arm64Compiler) {
+	err := compiler.compileConstF32(&wazeroir.OperationConstF32{Value: x1})
+	require.NoError(t, err)
+	err = compiler.compileConstF32(&wazeroir.OperationConstF32{Value: x2})
+	require.NoError(t, err)
+}
+
 func (j *jitEnv) requireNewCompiler(t *testing.T) *arm64Compiler {
 	cmp, err := newCompiler(&wasm.FunctionInstance{ModuleInstance: j.moduleInstance}, nil)
 	require.NoError(t, err)
 	ret, ok := cmp.(*arm64Compiler)
 	require.True(t, ok)
+	ret.labels = make(map[string]*labelInfo)
+	ret.ir = &wazeroir.CompilationResult{}
 	return ret
 }
 
@@ -1008,3 +1032,322 @@ func TestArm64Compiler_compieleDrop(t *testing.T) {
 		}
 	})
 }
+
+func TestArm64Compiler_compileLabel(t *testing.T) {
+	label := &wazeroir.Label{FrameID: 100, Kind: wazeroir.LabelKindContinuation}
+	labelKey := label.String()
+	for _, expectSkip := range []bool{false, true} {
+		expectSkip := expectSkip
+		t.Run(fmt.Sprintf("expect skip=%v", expectSkip), func(t *testing.T) {
+			env := newJITEnvironment()
+			compiler := env.requireNewCompiler(t)
+
+			var callBackCalled bool
+			compiler.labels[labelKey] = &labelInfo{
+				labelBeginningCallbacks: []func(*obj.Prog){func(p *obj.Prog) { callBackCalled = true }},
+			}
+
+			if expectSkip {
+				// If the initial stack is not set, compileLabel must return skip=true.
+				compiler.labels[labelKey].initialStack = nil
+				actual := compiler.compileLabel(&wazeroir.OperationLabel{Label: label})
+				require.True(t, actual)
+				// Also, callback must not be called.
+				require.False(t, callBackCalled)
+			} else {
+				// If the initial stack is not set, compileLabel must return skip=false.
+				compiler.labels[labelKey].initialStack = newValueLocationStack()
+				actual := compiler.compileLabel(&wazeroir.OperationLabel{Label: label})
+				require.False(t, actual)
+				// Also, callback must be called.
+				require.True(t, callBackCalled)
+			}
+		})
+	}
+}
+
+func TestArm64Compiler_compileBr(t *testing.T) {
+	t.Run("return", func(t *testing.T) {
+		env := newJITEnvironment()
+		compiler := env.requireNewCompiler(t)
+		err := compiler.emitPreamble()
+		require.NoError(t, err)
+
+		// Branch into nil label is interpreted as return. See BranchTarget.IsReturnTarget
+		err = compiler.compileBr(&wazeroir.OperationBr{Target: &wazeroir.BranchTarget{Label: nil}})
+		require.NoError(t, err)
+
+		// Compile and execute the code under test.
+		// Note: we don't invoke "compiler.return()" as the code emitted by compilerBr is enough to exit.
+		code, _, _, err := compiler.compile()
+		require.NoError(t, err)
+		env.exec(code)
+
+		require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+	})
+	t.Run("backward br", func(t *testing.T) {
+		env := newJITEnvironment()
+		compiler := env.requireNewCompiler(t)
+
+		// Emit code for the backward label.
+		backwardLabel := &wazeroir.Label{Kind: wazeroir.LabelKindHeader, FrameID: 0}
+		requireAddLabel(t, compiler, backwardLabel)
+		compiler.exit(jitCallStatusCodeReturned)
+
+		// Now emit the body. First we add NOP so that we can execute code after the target label.
+		nop := compiler.newProg()
+		nop.As = obj.ANOP
+		compiler.addInstruction(nop)
+
+		err := compiler.emitPreamble()
+		require.NoError(t, err)
+
+		err = compiler.compileBr(&wazeroir.OperationBr{Target: &wazeroir.BranchTarget{Label: backwardLabel}})
+		require.NoError(t, err)
+
+		// We must not reach the code after Br, so emit the code exiting with unreachable status.
+		compiler.exit(jitCallStatusCodeUnreachable)
+
+		code, _, _, err := compiler.compile()
+		require.NoError(t, err)
+
+		// The generated code looks like this:
+		//
+		// .backwardLabel:
+		//    exit jitCallStatusCodeReturned
+		//    nop
+		//    ... code from emitPreamble()
+		//    br .backwardLabel
+		//    exit jitCallStatusCodeUnreachable
+		//
+		// Therefore, if we start executing from nop, we must end up exiting jitCallStatusCodeReturned.
+		env.exec(code[nop.Pc:])
+		require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+	})
+	t.Run("forward br", func(t *testing.T) {
+		env := newJITEnvironment()
+		compiler := env.requireNewCompiler(t)
+		err := compiler.emitPreamble()
+		require.NoError(t, err)
+
+		// Emit the forward br, meaning that handle Br instruction where the target label hasn't been compiled yet.
+		forwardLabel := &wazeroir.Label{Kind: wazeroir.LabelKindHeader, FrameID: 0}
+		err = compiler.compileBr(&wazeroir.OperationBr{Target: &wazeroir.BranchTarget{Label: forwardLabel}})
+		require.NoError(t, err)
+
+		// We must not reach the code after Br, so emit the code exiting with Unreachable status.
+		compiler.exit(jitCallStatusCodeUnreachable)
+
+		// Emit code for the forward label where we emit the expectedValue and then exit.
+		requireAddLabel(t, compiler, forwardLabel)
+		compiler.exit(jitCallStatusCodeReturned)
+
+		code, _, _, err := compiler.compile()
+		require.NoError(t, err)
+
+		// The generated code looks like this:
+		//
+		//    ... code from emitPreamble()
+		//    br .forwardLabel
+		//    exit jitCallStatusCodeUnreachable
+		// .forwardLabel:
+		//    exit jitCallStatusCodeReturned
+		//
+		// Therefore, if we start executing from the top, we must end up exiting jitCallStatusCodeReturned.
+		env.exec(code)
+		require.Equal(t, jitCallStatusCodeReturned, env.jitStatus())
+	})
+}
+
+func TestArm64Compiler_compileBrIf(t *testing.T) {
+	unreachableStatus, thenLabelExitStatus, elseLabelExitStatus :=
+		jitCallStatusCodeUnreachable, jitCallStatusCodeUnreachable+1, jitCallStatusCodeUnreachable+2
+	thenBranchTarget := &wazeroir.BranchTargetDrop{Target: &wazeroir.BranchTarget{Label: &wazeroir.Label{Kind: wazeroir.LabelKindHeader, FrameID: 1}}}
+	elseBranchTarget := &wazeroir.BranchTargetDrop{Target: &wazeroir.BranchTarget{Label: &wazeroir.Label{Kind: wazeroir.LabelKindHeader, FrameID: 2}}}
+
+	for _, tc := range []struct {
+		name      string
+		setupFunc func(t *testing.T, compiler *arm64Compiler, shoulGoElse bool)
+	}{
+		{
+			name: "cond on register",
+			setupFunc: func(t *testing.T, compiler *arm64Compiler, shoulGoElse bool) {
+				val := uint32(1)
+				if shoulGoElse {
+					val = 0
+				}
+				err := compiler.compileConstI32(&wazeroir.OperationConstI32{Value: val})
+				require.NoError(t, err)
+			},
+		},
+		{
+			name: "LS",
+			setupFunc: func(t *testing.T, compiler *arm64Compiler, shoulGoElse bool) {
+				x1, x2 := uint32(1), uint32(2)
+				if shoulGoElse {
+					x2, x1 = x1, x2
+				}
+				requirePushTwoInt32Consts(t, x1, x2, compiler)
+				// Le on unsigned integer produces the value on COND_LS register.
+				err := compiler.compileLe(&wazeroir.OperationLe{Type: wazeroir.SignedTypeUint32})
+				require.NoError(t, err)
+			},
+		},
+		{
+			name: "LE",
+			setupFunc: func(t *testing.T, compiler *arm64Compiler, shoulGoElse bool) {
+				x1, x2 := uint32(1), uint32(2)
+				if shoulGoElse {
+					x2, x1 = x1, x2
+				}
+				requirePushTwoInt32Consts(t, x1, x2, compiler)
+				// Le on signed integer produces the value on COND_LE register.
+				err := compiler.compileLe(&wazeroir.OperationLe{Type: wazeroir.SignedTypeInt32})
+				require.NoError(t, err)
+			},
+		},
+		// {name: "EQ"} TODO: after compileEq support
+		// {name: "NE"} TODO: after compileNe support
+		{
+			name: "HS",
+			setupFunc: func(t *testing.T, compiler *arm64Compiler, shoulGoElse bool) {
+				x1, x2 := uint32(2), uint32(1)
+				if shoulGoElse {
+					x2, x1 = x1, x2
+				}
+				requirePushTwoInt32Consts(t, x1, x2, compiler)
+				// Ge on unsigned integer produces the value on COND_HS register.
+				err := compiler.compileGe(&wazeroir.OperationGe{Type: wazeroir.SignedTypeUint32})
+				require.NoError(t, err)
+			},
+		},
+		{
+			name: "GE",
+			setupFunc: func(t *testing.T, compiler *arm64Compiler, shoulGoElse bool) {
+				x1, x2 := uint32(2), uint32(1)
+				if shoulGoElse {
+					x2, x1 = x1, x2
+				}
+				requirePushTwoInt32Consts(t, x1, x2, compiler)
+				// Ge on signed integer produces the value on COND_GE register.
+				err := compiler.compileGe(&wazeroir.OperationGe{Type: wazeroir.SignedTypeInt32})
+				require.NoError(t, err)
+			},
+		},
+		{
+			name: "HI",
+			setupFunc: func(t *testing.T, compiler *arm64Compiler, shoulGoElse bool) {
+				x1, x2 := uint32(2), uint32(1)
+				if shoulGoElse {
+					x2, x1 = x1, x2
+				}
+				requirePushTwoInt32Consts(t, x1, x2, compiler)
+				// Gt on unsigned integer produces the value on COND_HI register.
+				err := compiler.compileGt(&wazeroir.OperationGt{Type: wazeroir.SignedTypeUint32})
+				require.NoError(t, err)
+			},
+		},
+		{
+			name: "GT",
+			setupFunc: func(t *testing.T, compiler *arm64Compiler, shoulGoElse bool) {
+				x1, x2 := uint32(2), uint32(1)
+				if shoulGoElse {
+					x2, x1 = x1, x2
+				}
+				requirePushTwoInt32Consts(t, x1, x2, compiler)
+				// Gt on signed integer produces the value on COND_GT register.
+				err := compiler.compileGt(&wazeroir.OperationGt{Type: wazeroir.SignedTypeInt32})
+				require.NoError(t, err)
+			},
+		},
+		{
+			name: "LO",
+			setupFunc: func(t *testing.T, compiler *arm64Compiler, shoulGoElse bool) {
+				x1, x2 := uint32(1), uint32(2)
+				if shoulGoElse {
+					x2, x1 = x1, x2
+				}
+				requirePushTwoInt32Consts(t, x1, x2, compiler)
+				// Lt on unsigned integer produces the value on COND_LO register.
+				err := compiler.compileLt(&wazeroir.OperationLt{Type: wazeroir.SignedTypeUint32})
+				require.NoError(t, err)
+			},
+		},
+		{
+			name: "LT",
+			setupFunc: func(t *testing.T, compiler *arm64Compiler, shoulGoElse bool) {
+				x1, x2 := uint32(1), uint32(2)
+				if shoulGoElse {
+					x2, x1 = x1, x2
+				}
+				requirePushTwoInt32Consts(t, x1, x2, compiler)
+				// Lt on signed integer produces the value on COND_LT register.
+				err := compiler.compileLt(&wazeroir.OperationLt{Type: wazeroir.SignedTypeInt32})
+				require.NoError(t, err)
+			},
+		},
+		{
+			name: "MI",
+			setupFunc: func(t *testing.T, compiler *arm64Compiler, shoulGoElse bool) {
+				x1, x2 := float32(1), float32(2)
+				if shoulGoElse {
+					x2, x1 = x1, x2
+				}
+				requirePushTwoFloat32Consts(t, x1, x2, compiler)
+				// Lt on floats produces the value on COND_MI register.
+				err := compiler.compileLt(&wazeroir.OperationLt{Type: wazeroir.SignedTypeFloat32})
+				require.NoError(t, err)
+			},
+		},
+	} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			for _, shouldGoToElse := range []bool{false, true} {
+				shouldGoToElse := shouldGoToElse
+				t.Run(fmt.Sprintf("should_goto_else=%v", shouldGoToElse), func(t *testing.T) {
+					env := newJITEnvironment()
+					compiler := env.requireNewCompiler(t)
+					err := compiler.emitPreamble()
+					require.NoError(t, err)
+
+					tc.setupFunc(t, compiler, shouldGoToElse)
+					require.Equal(t, uint64(1), compiler.locationStack.sp)
+
+					err = compiler.compileBrIf(&wazeroir.OperationBrIf{Then: thenBranchTarget, Else: elseBranchTarget})
+					require.NoError(t, err)
+					compiler.exit(unreachableStatus)
+
+					// Emit code for .then label.
+					requireAddLabel(t, compiler, thenBranchTarget.Target.Label)
+					compiler.exit(thenLabelExitStatus)
+
+					// Emit code for .else label.
+					requireAddLabel(t, compiler, elseBranchTarget.Target.Label)
+					compiler.exit(elseLabelExitStatus)
+
+					code, _, _, err := compiler.compile()
+					require.NoError(t, err)
+
+					// The generated code looks like this:
+					//
+					//    ... code from emitPreamble()
+					//    br_if .then, .else
+					//    exit $unreachableStatus
+					// .then:
+					//    exit $thenLabelExitStatus
+					// .else:
+					//    exit $elseLabelExitStatus
+					//
+					// Therefore, if we start executing from the top, we must end up exiting with an appropriate status.
+					env.exec(code)
+					require.NotEqual(t, unreachableStatus, env.jitStatus())
+					if shouldGoToElse {
+						require.Equal(t, elseLabelExitStatus, env.jitStatus())
+					} else {
+						require.Equal(t, thenLabelExitStatus, env.jitStatus())
+					}
+				})
+			}
+		})
+	}
+}
diff --git a/wasm/jit/jit_value_location_test.go b/wasm/jit/jit_value_location_test.go
index 3f47ea0d..80853324 100644
--- a/wasm/jit/jit_value_location_test.go
+++ b/wasm/jit/jit_value_location_test.go
@@ -67,7 +67,7 @@ func TestValueLocationStack_basic(t *testing.T) {
 	for i := 0; i < 1000; i++ {
 		s.pop()
 	}
-	require.Equal(t, uint64(1001), s.maxStackPointer)
+	require.Equal(t, uint64(1001), s.stackPointerCeil)
 }
 
 func TestValueLocationStack_takeFreeRegister(t *testing.T) {