Repository: tetratelabs/wazero
Pull request #1653 (2023-08-23 04:49:58)

wazevo: pass memory_size.wast
#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/compiler.go b/internal/engine/wazevo/backend/compiler.go
index 0907d44b..5ed6fcc9 100644
--- a/internal/engine/wazevo/backend/compiler.go
+++ b/internal/engine/wazevo/backend/compiler.go
@@ -143,7 +143,7 @@ type compiler struct {
 // Compile implements Compiler.Compile.
 func (c *compiler) Compile() ([]byte, []RelocationInfo, int, error) {
 	c.Lower()
-	if false {
+	if true {
 		fmt.Printf("[[[after lowering]]]%s", c.Format())
 	}
 	c.RegAlloc()
diff --git a/internal/engine/wazevo/backend/compiler_lower.go b/internal/engine/wazevo/backend/compiler_lower.go
index 6669dbd8..c5f58915 100644
--- a/internal/engine/wazevo/backend/compiler_lower.go
+++ b/internal/engine/wazevo/backend/compiler_lower.go
@@ -1,6 +1,8 @@
 package backend
 
 import (
+	"fmt"
+
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
 )
@@ -14,10 +16,15 @@ func (c *compiler) Lower() {
 	c.mach.EndLoweringFunction()
 }
 
+const debug = false
+
 // lowerBlocks lowers each block in the ssa.Builder.
 func (c *compiler) lowerBlocks() {
 	builder := c.ssaBuilder
 	for blk := builder.BlockIteratorReversePostOrderBegin(); blk != nil; blk = builder.BlockIteratorReversePostOrderNext() {
+		if debug {
+			fmt.Printf("lowering block %s\n", blk.Name())
+		}
 		c.lowerBlock(blk)
 	}
 	// After lowering all blocks, we need to link adjacent blocks to layout one single instruction list.
@@ -64,6 +71,10 @@ func (c *compiler) lowerBlock(blk ssa.BasicBlock) {
 			continue
 		}
 
+		if debug {
+			fmt.Printf("\tlowering instr %s\n", cur.Format(c.ssaBuilder))
+		}
+
 		switch cur.Opcode() {
 		case ssa.OpcodeReturn:
 			c.lowerFunctionReturns(cur.ReturnVals())
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
index b250d556..0022c2cd 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
@@ -7,6 +7,8 @@ package arm64
 // and merge the multiple instructions if possible. It can be considered as "N:1" instruction selection.
 
 import (
+	"fmt"
+
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/wazevoapi"
@@ -71,7 +73,11 @@ func (m *machine) lowerBrTable(i *ssa.Instruction) {
 func (m *machine) LowerConditionalBranch(b *ssa.Instruction) {
 	cval, args, targetBlk := b.BranchData()
 	if len(args) > 0 {
-		panic("conditional branch shouldn't have args; likely a bug in critical edge splitting")
+		panic(fmt.Sprintf(
+			"conditional branch shouldn't have args; likely a bug in critical edge splitting: from %s to %s",
+			m.currentSSABlk,
+			targetBlk,
+		))
 	}
 
 	target := m.getOrAllocateSSABlockLabel(targetBlk)
diff --git a/internal/engine/wazevo/backend/regalloc/assign.go b/internal/engine/wazevo/backend/regalloc/assign.go
index bfca69da..1242142d 100644
--- a/internal/engine/wazevo/backend/regalloc/assign.go
+++ b/internal/engine/wazevo/backend/regalloc/assign.go
@@ -55,13 +55,13 @@ func (a *Allocator) assignRegistersPerInstr(f Function, pc programCounter, instr
 	a.vs = a.vs[:0]
 	uses := instr.Uses()
 	for _, u := range uses {
-		if false {
-			fmt.Printf("%s uses %d\n", instr, u.ID())
-		}
 		if u.IsRealReg() {
 			a.vs = append(a.vs, u)
 			continue
 		}
+		if true {
+			fmt.Printf("%s uses %d\n", instr, u.ID())
+		}
 		n := vRegIDToNode[u.ID()]
 		if !n.spill() {
 			a.vs = append(a.vs, u.SetRealReg(n.r))
diff --git a/internal/engine/wazevo/engine.go b/internal/engine/wazevo/engine.go
index 90c0a4b0..025b94c4 100644
--- a/internal/engine/wazevo/engine.go
+++ b/internal/engine/wazevo/engine.go
@@ -100,6 +100,7 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []exper
 	cm.functionOffsets = make([]compiledFunctionOffset, localFns)
 	bodies := make([][]byte, localFns)
 	for i := range module.CodeSection {
+		fmt.Printf("------------------------------------------ %d/%d ------------------------------------------\n", i, len(module.CodeSection)-1)
 		fidx := wasm.Index(i + importedFns)
 		fref := frontend.FunctionIndexToFuncRef(fidx)
 
@@ -124,16 +125,26 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []exper
 			return fmt.Errorf("wasm->ssa: %v", err)
 		}
 
+		const debug = true
+
+		if debug {
+			fmt.Printf("[[[SSA]]]%s\n", ssaBuilder.Format())
+		}
+
 		// Run SSA-level optimization passes.
 		ssaBuilder.RunPasses()
 
-		if false {
-			fmt.Printf("[[[SSA for %d/%d]]]%s", i, len(module.CodeSection)-1, ssaBuilder.Format())
+		if debug {
+			fmt.Printf("[[[optimized SSA]]]%s", ssaBuilder.Format())
 		}
 
 		// Finalize the layout of SSA blocks which might use the optimization results.
 		ssaBuilder.LayoutBlocks()
 
+		if debug {
+			fmt.Printf("[[[laidout SSA]]]%s", ssaBuilder.Format())
+		}
+
 		// Now our ssaBuilder contains the necessary information to further lower them to
 		// machine code.
 		body, rels, goPreambleSize, err := be.Compile()
diff --git a/internal/engine/wazevo/frontend/lower.go b/internal/engine/wazevo/frontend/lower.go
index 5e2ddf54..93487ada 100644
--- a/internal/engine/wazevo/frontend/lower.go
+++ b/internal/engine/wazevo/frontend/lower.go
@@ -164,9 +164,7 @@ func (c *Compiler) lowerOpcode(op wasm.Opcode) {
 			return
 		}
 
-		iconst := builder.AllocateInstruction()
-		iconst.AsIconst32(uint32(c))
-		builder.InsertInstruction(iconst)
+		iconst := builder.AllocateInstruction().AsIconst32(uint32(c)).Insert(builder)
 		value := iconst.Return()
 		state.push(value)
 	case wasm.OpcodeI64Const:
@@ -174,9 +172,7 @@ func (c *Compiler) lowerOpcode(op wasm.Opcode) {
 		if state.unreachable {
 			return
 		}
-		iconst := builder.AllocateInstruction()
-		iconst.AsIconst64(uint64(c))
-		builder.InsertInstruction(iconst)
+		iconst := builder.AllocateInstruction().AsIconst64(uint64(c)).Insert(builder)
 		value := iconst.Return()
 		state.push(value)
 	case wasm.OpcodeF32Const:
@@ -921,6 +917,7 @@ func (c *Compiler) lowerOpcode(op wasm.Opcode) {
 		c.insertJumpToBlock(nil, elseBlk)
 
 		// Now start translating the instructions after br_if.
+		builder.Seal(elseBlk) // Else of br_if has the current block as the only one successor.
 		builder.SetCurrentBlock(elseBlk)
 
 	case wasm.OpcodeBrTable:
diff --git a/internal/engine/wazevo/ssa/basic_block.go b/internal/engine/wazevo/ssa/basic_block.go
index 64119564..d1281d98 100644
--- a/internal/engine/wazevo/ssa/basic_block.go
+++ b/internal/engine/wazevo/ssa/basic_block.go
@@ -285,8 +285,9 @@ func (bb *basicBlock) FormatHeader(b Builder) string {
 		preds := make([]string, 0, len(bb.preds))
 		for _, pred := range bb.preds {
 			if len(pred.branch.vs) != len(bb.params) {
-				panic(fmt.Sprintf("BUG: len(argument) != len(params): %d != %d",
-					len(pred.branch.vs), len(bb.params)))
+				panic(fmt.Sprintf("BUG: len(argument) != len(params): %d != %d: %s",
+					len(pred.branch.vs), len(bb.params), pred.branch.Format(b),
+				))
 			}
 			if pred.blk.invalid {
 				continue
diff --git a/internal/engine/wazevo/ssa/builder.go b/internal/engine/wazevo/ssa/builder.go
index 9a1f5013..500587b9 100644
--- a/internal/engine/wazevo/ssa/builder.go
+++ b/internal/engine/wazevo/ssa/builder.go
@@ -222,6 +222,8 @@ func (b *builder) Init(s *Signature) {
 	}
 	b.nextValueID = 0
 	b.reversePostOrderedBasicBlocks = b.reversePostOrderedBasicBlocks[:0]
+	b.donePasses = false
+	b.doneBlockLayout = false
 }
 
 // Signature implements Builder.Signature.
@@ -316,6 +318,9 @@ func (b *builder) DefineVariable(variable Variable, value Value, block BasicBloc
 		panic("BUG: trying to define variable " + variable.String() + " but is not declared yet")
 	}
 
+	if b.variables[variable] != value.Type() {
+		panic(fmt.Sprintf("BUG: inconsistent type for variable %d: expected %s but got %s", variable, b.variables[variable], value.Type()))
+	}
 	bb := block.(*basicBlock)
 	bb.lastDefinitions[variable] = value
 }
@@ -388,6 +393,9 @@ func (b *builder) findValue(typ Type, variable Variable, blk *basicBlock, must b
 			// and record it as unknown.
 			// The unknown values are resolved when we call seal this block via BasicBlock.Seal().
 			value := b.allocateValue(typ)
+			if debug {
+				fmt.Printf("adding unknown value placeholder for %s at %d\n", variable, blk.id)
+			}
 			blk.lastDefinitions[variable] = value
 			blk.unknownValues[variable] = value
 			return value
@@ -729,8 +737,14 @@ func (b *builder) LayoutBlocks() {
 			// Update the successors slice because the target is no longer the original `succ`.
 			blk.success[sidx] = trampoline
 
+			if debug {
+				fmt.Printf("splitting edge from %d->%d at %s as %d->%d->%d \n",
+					blk.ID(), succ.ID(), predInfo.branch.Format(b),
+					blk.ID(), trampoline.ID(), succ.ID())
+			}
+
 			fallthroughBranch := blk.currentInstr
-			if fallthroughBranch.opcode != OpcodeBrTable && fallthroughBranch.blk == trampoline {
+			if fallthroughBranch.opcode == OpcodeJump && fallthroughBranch.blk == trampoline {
 				// This can be lowered as fallthrough at the end of the block.
 				b.reversePostOrderedBasicBlocks = append(b.reversePostOrderedBasicBlocks, trampoline)
 				inserted[trampoline] = 0 // mark as inserted, the value is not used.
@@ -785,6 +799,11 @@ func maybeInvertBranches(now *basicBlock, nextInRPO *basicBlock) bool {
 		return false
 	}
 
+	if len(fallthroughBranch.vs) != 0 || len(condBranch.vs) != 0 {
+		// If either one of them has arguments, we don't invert the branches.
+		return false
+	}
+
 	// So this block has two branches (a conditional branch followed by an unconditional branch) at the end.
 	// We can invert the condition of the branch if it makes the fallthrough more likely.
 
@@ -834,9 +853,16 @@ invert:
 	condBranch.InvertBrx()
 	condBranch.blk = fallthroughTarget
 	fallthroughBranch.blk = condTarget
+	if debug {
+		fmt.Printf("inverting branches at %d->%d and %d->%d\n",
+			now.ID(), fallthroughTarget.ID(), now.ID(), condTarget.ID())
+	}
+
 	return true
 }
 
+const debug = false
+
 // splitCriticalEdge splits the critical edge between the given predecessor (`pred`) and successor (owning `predInfo`).
 //
 // - `pred` is the source of the critical edge,
diff --git a/internal/engine/wazevo/ssa/instructions.go b/internal/engine/wazevo/ssa/instructions.go
index 9b27e454..5eb87b60 100644
--- a/internal/engine/wazevo/ssa/instructions.go
+++ b/internal/engine/wazevo/ssa/instructions.go
@@ -20,6 +20,7 @@ type Instruction struct {
 	u64        uint64
 	v          Value
 	v2         Value
+	v3         Value
 	vs         []Value
 	typ        Type
 	blk        BasicBlock
@@ -47,6 +48,7 @@ func (i *Instruction) reset() {
 	*i = Instruction{}
 	i.v = ValueInvalid
 	i.v2 = ValueInvalid
+	i.v3 = ValueInvalid
 	i.rValue = ValueInvalid
 	i.typ = typeInvalid
 }
@@ -78,8 +80,8 @@ func (i *Instruction) Return() (first Value) {
 }
 
 // Args returns the arguments to this instruction.
-func (i *Instruction) Args() (v1, v2 Value, vs []Value) {
-	return i.v, i.v2, i.vs
+func (i *Instruction) Args() (v1, v2, v3 Value, vs []Value) {
+	return i.v, i.v2, i.v3, i.vs
 }
 
 // Arg returns the first argument to this instruction.
@@ -1008,17 +1010,19 @@ func (i *Instruction) StoreData() (value, ptr Value, offset uint32, storeSizeInB
 }
 
 // AsIconst64 initializes this instruction as a 64-bit integer constant instruction with OpcodeIconst.
-func (i *Instruction) AsIconst64(v uint64) {
+func (i *Instruction) AsIconst64(v uint64) *Instruction {
 	i.opcode = OpcodeIconst
 	i.typ = TypeI64
 	i.u64 = v
+	return i
 }
 
 // AsIconst32 initializes this instruction as a 32-bit integer constant instruction with OpcodeIconst.
-func (i *Instruction) AsIconst32(v uint32) {
+func (i *Instruction) AsIconst32(v uint32) *Instruction {
 	i.opcode = OpcodeIconst
 	i.typ = TypeI32
 	i.u64 = uint64(v)
+	return i
 }
 
 // BinaryData return the operands for a binary instruction.
@@ -1027,11 +1031,12 @@ func (i *Instruction) BinaryData() (x, y Value) {
 }
 
 // AsIadd initializes this instruction as an integer addition instruction with OpcodeIadd.
-func (i *Instruction) AsIadd(x, y Value) {
+func (i *Instruction) AsIadd(x, y Value) *Instruction {
 	i.opcode = OpcodeIadd
 	i.v = x
 	i.v2 = y
 	i.typ = x.Type()
+	return i
 }
 
 // AsImul initializes this instruction as an integer addition instruction with OpcodeImul.
@@ -1042,21 +1047,28 @@ func (i *Instruction) AsImul(x, y Value) {
 	i.typ = x.Type()
 }
 
+func (i *Instruction) Insert(b Builder) *Instruction {
+	b.InsertInstruction(i)
+	return i
+}
+
 // AsIsub initializes this instruction as an integer subtraction instruction with OpcodeIsub.
-func (i *Instruction) AsIsub(x, y Value) {
+func (i *Instruction) AsIsub(x, y Value) *Instruction {
 	i.opcode = OpcodeIsub
 	i.v = x
 	i.v2 = y
 	i.typ = x.Type()
+	return i
 }
 
 // AsIcmp initializes this instruction as an integer comparison instruction with OpcodeIcmp.
-func (i *Instruction) AsIcmp(x, y Value, c IntegerCmpCond) {
+func (i *Instruction) AsIcmp(x, y Value, c IntegerCmpCond) *Instruction {
 	i.opcode = OpcodeIcmp
 	i.v = x
 	i.v2 = y
 	i.u64 = uint64(c)
 	i.typ = TypeI32
+	return i
 }
 
 // AsFcmp initializes this instruction as an integer comparison instruction with OpcodeFcmp.
@@ -1418,7 +1430,7 @@ func (i *Instruction) AsSelect(c, x, y Value) {
 	i.opcode = OpcodeSelect
 	i.v = c
 	i.v2 = x
-	i.u64 = uint64(y)
+	i.v3 = y
 	i.typ = x.Type()
 }
 
@@ -1426,7 +1438,7 @@ func (i *Instruction) AsSelect(c, x, y Value) {
 func (i *Instruction) SelectData() (c, x, y Value) {
 	c = i.v
 	x = i.v2
-	y = Value(i.u64)
+	y = i.v3
 	return
 }
 
@@ -1471,7 +1483,7 @@ func (i *Instruction) Format(b Builder) string {
 	case OpcodeUload8, OpcodeUload16, OpcodeUload32, OpcodeSload8, OpcodeSload16, OpcodeSload32:
 		instSuffix = fmt.Sprintf(" %s, %#x", i.v.Format(b), int32(i.u64))
 	case OpcodeSelect:
-		instSuffix = fmt.Sprintf(" %s, %s, %s", i.v.Format(b), i.v2.Format(b), Value(i.u64).Format(b))
+		instSuffix = fmt.Sprintf(" %s, %s, %s", i.v.Format(b), i.v2.Format(b), i.v3.Format(b))
 	case OpcodeIconst:
 		switch i.typ {
 		case TypeI32:
diff --git a/internal/engine/wazevo/ssa/pass.go b/internal/engine/wazevo/ssa/pass.go
index 5a24cb6f..7eae17be 100644
--- a/internal/engine/wazevo/ssa/pass.go
+++ b/internal/engine/wazevo/ssa/pass.go
@@ -1,5 +1,7 @@
 package ssa
 
+import "fmt"
+
 // RunPasses implements Builder.RunPasses.
 //
 // The order here matters; some pass depends on the previous ones.
@@ -39,6 +41,10 @@ func passDeadBlockEliminationOpt(b *builder) {
 		b.blkStack = b.blkStack[:len(b.blkStack)-1]
 		b.blkVisited[reachableBlk] = 0 // the value won't be used in this pass.
 
+		if !reachableBlk.sealed && !reachableBlk.ReturnBlock() {
+			panic(fmt.Sprintf("%s is not sealed", reachableBlk))
+		}
+
 		for _, succ := range reachableBlk.success {
 			if _, ok := b.blkVisited[succ]; ok {
 				continue
@@ -200,7 +206,7 @@ func passDeadCodeEliminationOpt(b *builder) {
 		// Before we walk, we need to resolve the alias first.
 		b.resolveArgumentAlias(live)
 
-		v1, v2, vs := live.Args()
+		v1, v2, v3, vs := live.Args()
 		if v1.Valid() {
 			producingInst := b.valueIDToInstruction[v1.ID()]
 			if producingInst != nil {
@@ -215,6 +221,13 @@ func passDeadCodeEliminationOpt(b *builder) {
 			}
 		}
 
+		if v3.Valid() {
+			producingInst := b.valueIDToInstruction[v3.ID()]
+			if producingInst != nil {
+				liveInstructions = append(liveInstructions, producingInst)
+			}
+		}
+
 		for _, v := range vs {
 			producingInst := b.valueIDToInstruction[v.ID()]
 			if producingInst != nil {
@@ -241,13 +254,16 @@ func passDeadCodeEliminationOpt(b *builder) {
 
 			// If the value alive, we can be sure that arguments are used definitely.
 			// Hence, we can increment the value reference counts.
-			v1, v2, vs := cur.Args()
+			v1, v2, v3, vs := cur.Args()
 			if v1.Valid() {
 				b.valueRefCounts[v1.ID()]++
 			}
 			if v2.Valid() {
 				b.valueRefCounts[v2.ID()]++
 			}
+			if v3.Valid() {
+				b.valueRefCounts[v3.ID()]++
+			}
 			for _, v := range vs {
 				b.valueRefCounts[v.ID()]++
 			}


Test patch:
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index 77446e30..b3c1d468 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -45,6 +45,7 @@ func TestSpectestV1(t *testing.T) {
 		{name: "const.json"},
 		{name: "data.json"},
 		{name: "local_get.json"},
+		{name: "memory_size.json"},
 		{name: "nop.json"},
 	} {
 		t.Run(tc.name, func(t *testing.T) { spectest.RunJson(t, v1.Testcases, tc.name, context.Background(), config) })
diff --git a/internal/engine/wazevo/ssa/builder_test.go b/internal/engine/wazevo/ssa/builder_test.go
index 2a4f808c..52be0e24 100644
--- a/internal/engine/wazevo/ssa/builder_test.go
+++ b/internal/engine/wazevo/ssa/builder_test.go
@@ -274,21 +274,19 @@ func Test_swapInstruction(t *testing.T) {
 }
 
 func TestBuilder_LayoutBlocks(t *testing.T) {
-	insertJump := func(b *builder, src, dst *basicBlock) {
+	insertJump := func(b *builder, src, dst *basicBlock, vs ...Value) {
 		b.SetCurrentBlock(src)
 		jump := b.AllocateInstruction()
-		jump.AsJump(nil, dst)
+		jump.AsJump(vs, dst)
 		b.InsertInstruction(jump)
 	}
 
-	insertBrz := func(b *builder, src, dst *basicBlock) {
+	insertBrz := func(b *builder, src, dst *basicBlock, condVal Value, vs ...Value) {
 		b.SetCurrentBlock(src)
-		vinst := b.AllocateInstruction()
-		vinst.AsIconst32(0)
+		vinst := b.AllocateInstruction().AsIconst32(0)
 		b.InsertInstruction(vinst)
-		v := vinst.Return()
 		brz := b.AllocateInstruction()
-		brz.AsBrz(v, nil, dst)
+		brz.AsBrz(condVal, vs, dst)
 		b.InsertInstruction(brz)
 	}
 
@@ -304,6 +302,10 @@ func TestBuilder_LayoutBlocks(t *testing.T) {
 				insertJump(b, b1, b2)
 				insertJump(b, b2, b3)
 				insertJump(b, b3, b4)
+				b.Seal(b1)
+				b.Seal(b2)
+				b.Seal(b3)
+				b.Seal(b4)
 			},
 			exp: []BasicBlockID{0, 1, 2, 3},
 		},
@@ -314,6 +316,9 @@ func TestBuilder_LayoutBlocks(t *testing.T) {
 				insertJump(b, b0, b2)
 				insertJump(b, unreachable, b2)
 				unreachable.invalid = true
+				b.Seal(b0)
+				b.Seal(unreachable)
+				b.Seal(b2)
 			},
 			exp: []BasicBlockID{0, 2},
 		},
@@ -325,10 +330,17 @@ func TestBuilder_LayoutBlocks(t *testing.T) {
 			// 2 ---------
 			setup: func(b *builder) {
 				b0, b1, b2, b3 := b.allocateBasicBlock(), b.allocateBasicBlock(), b.allocateBasicBlock(), b.allocateBasicBlock()
-				insertBrz(b, b0, b2)
+				b.SetCurrentBlock(b0)
+				c := b.AllocateInstruction().AsIconst32(0)
+				b.InsertInstruction(c)
+				insertBrz(b, b0, b2, c.Return())
 				insertJump(b, b0, b1)
 				insertJump(b, b1, b3)
 				insertJump(b, b2, b3)
+				b.Seal(b0)
+				b.Seal(b1)
+				b.Seal(b2)
+				b.Seal(b3)
 			},
 			exp: []BasicBlockID{0, 2, 1, 3},
 		},
@@ -355,8 +367,15 @@ func TestBuilder_LayoutBlocks(t *testing.T) {
 				b0, b1, b2, b3 := b.allocateBasicBlock(), b.allocateBasicBlock(), b.allocateBasicBlock(), b.allocateBasicBlock()
 				insertJump(b, b0, b1)
 				insertJump(b, b1, b2)
-				insertBrz(b, b2, b1)
+				b.SetCurrentBlock(b2)
+				c := b.AllocateInstruction().AsIconst32(0)
+				b.InsertInstruction(c)
+				insertBrz(b, b2, b1, c.Return())
 				insertJump(b, b2, b3)
+				b.Seal(b0)
+				b.Seal(b1)
+				b.Seal(b2)
+				b.Seal(b3)
 			},
 			// The trampoline 4 is placed right after 2, which is the hot path of the loop.
 			exp: []BasicBlockID{0, 1, 2, 4, 3},
@@ -384,8 +403,15 @@ func TestBuilder_LayoutBlocks(t *testing.T) {
 				b0, b1, b2, b3 := b.allocateBasicBlock(), b.allocateBasicBlock(), b.allocateBasicBlock(), b.allocateBasicBlock()
 				insertJump(b, b0, b1)
 				insertJump(b, b1, b2)
-				insertBrz(b, b2, b3)
+				b.SetCurrentBlock(b2)
+				c := b.AllocateInstruction().AsIconst32(0)
+				b.InsertInstruction(c)
+				insertBrz(b, b2, b3, c.Return())
 				insertJump(b, b2, b1)
+				b.Seal(b0)
+				b.Seal(b1)
+				b.Seal(b2)
+				b.Seal(b3)
 			},
 			// The trampoline 4 is placed right after 2, which is the hot path of the loop.
 			exp: []BasicBlockID{0, 1, 2, 4, 3},
@@ -417,12 +443,24 @@ func TestBuilder_LayoutBlocks(t *testing.T) {
 				b0, b1, b2, b3, b4, b5 := b.allocateBasicBlock(), b.allocateBasicBlock(), b.allocateBasicBlock(),
 					b.allocateBasicBlock(), b.allocateBasicBlock(), b.allocateBasicBlock()
 				insertJump(b, b0, b1)
-				insertBrz(b, b1, b2)
+				b.SetCurrentBlock(b0)
+				c1 := b.AllocateInstruction().AsIconst32(0)
+				b.InsertInstruction(c1)
+				insertBrz(b, b1, b2, c1.Return())
 				insertJump(b, b1, b3)
 				insertJump(b, b3, b4)
 				insertJump(b, b2, b4)
-				insertBrz(b, b4, b1)
+				b.SetCurrentBlock(b4)
+				c2 := b.AllocateInstruction().AsIconst32(0)
+				b.InsertInstruction(c2)
+				insertBrz(b, b4, b1, c2.Return())
 				insertJump(b, b4, b5)
+				b.Seal(b0)
+				b.Seal(b1)
+				b.Seal(b2)
+				b.Seal(b3)
+				b.Seal(b4)
+				b.Seal(b5)
 			},
 			// The trampoline 6 is placed right after 4, which is the hot path of the loop.
 			exp: []BasicBlockID{0, 1, 2, 3, 4, 6, 5},
@@ -453,12 +491,24 @@ func TestBuilder_LayoutBlocks(t *testing.T) {
 				b0, b1, b2, b3, b4 := b.allocateBasicBlock(), b.allocateBasicBlock(), b.allocateBasicBlock(),
 					b.allocateBasicBlock(), b.allocateBasicBlock()
 				insertJump(b, b0, b1)
-				insertBrz(b, b1, b2)
+				b.SetCurrentBlock(b1)
+				c1 := b.AllocateInstruction().AsIconst32(0)
+				b.InsertInstruction(c1)
+				insertBrz(b, b1, b2, c1.Return())
 				insertJump(b, b1, b3)
 
-				insertBrz(b, b2, b1)
+				b.SetCurrentBlock(b2)
+				c2 := b.AllocateInstruction().AsIconst32(0)
+				b.InsertInstruction(c2)
+				insertBrz(b, b2, b1, c2.Return())
 				insertJump(b, b2, b3)
 				insertJump(b, b3, b4)
+
+				b.Seal(b0)
+				b.Seal(b1)
+				b.Seal(b2)
+				b.Seal(b3)
+				b.Seal(b4)
 			},
 			exp: []BasicBlockID{
 				0, 1,
@@ -472,6 +522,51 @@ func TestBuilder_LayoutBlocks(t *testing.T) {
 				4,
 			},
 		},
+		{
+			name: "loop with output",
+			exp:  []BasicBlockID{0x0, 0x2, 0x4, 0x1, 0x3, 0x6, 0x5},
+			setup: func(b *builder) {
+				b0, b1, b2, b3 := b.allocateBasicBlock(), b.allocateBasicBlock(), b.allocateBasicBlock(), b.allocateBasicBlock()
+
+				b.SetCurrentBlock(b0)
+				funcParam := b0.AddParam(b, TypeI32)
+				b2Param := b2.AddParam(b, TypeI32)
+				insertJump(b, b0, b2, funcParam)
+
+				b.SetCurrentBlock(b1)
+				{
+					returnParam := b1.AddParam(b, TypeI32)
+					insertJump(b, b1, b.returnBlk, returnParam)
+				}
+
+				b.SetCurrentBlock(b2)
+				{
+					c := b.AllocateInstruction().AsIconst32(100).Insert(b)
+					cmp := b.AllocateInstruction().
+						AsIcmp(b2Param, c.Return(), IntegerCmpCondUnsignedLessThan).
+						Insert(b)
+					insertBrz(b, b2, b1, cmp.Return(), b2Param)
+					insertJump(b, b2, b3)
+				}
+
+				b.SetCurrentBlock(b3)
+				{
+					one := b.AllocateInstruction().AsIconst32(1).Insert(b)
+					minusOned := b.AllocateInstruction().AsIsub(b2Param, one.Return()).Insert(b)
+					c := b.AllocateInstruction().AsIconst32(150).Insert(b)
+					cmp := b.AllocateInstruction().
+						AsIcmp(b2Param, c.Return(), IntegerCmpCondEqual).
+						Insert(b)
+					insertBrz(b, b3, b1, cmp.Return(), minusOned.Return())
+					insertJump(b, b3, b2, minusOned.Return())
+				}
+
+				b.Seal(b0)
+				b.Seal(b1)
+				b.Seal(b2)
+				b.Seal(b3)
+			},
+		},
 	} {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
@@ -479,6 +574,8 @@ func TestBuilder_LayoutBlocks(t *testing.T) {
 			tc.setup(b)
 
 			b.RunPasses() // LayoutBlocks() must be called after RunPasses().
+			fmt.Println("============ SSA before block layout ============")
+			fmt.Println(b.Format())
 			b.LayoutBlocks()
 			fmt.Println("============ SSA after block layout ============")
 			fmt.Println(b.Format())
diff --git a/internal/engine/wazevo/ssa/pass_test.go b/internal/engine/wazevo/ssa/pass_test.go
index f05aff66..7d22debb 100644
--- a/internal/engine/wazevo/ssa/pass_test.go
+++ b/internal/engine/wazevo/ssa/pass_test.go
@@ -73,6 +73,11 @@ func TestBuilder_passes(t *testing.T) {
 					jmp.AsJump(nil, middle1)
 					b.InsertInstruction(jmp)
 				}
+
+				b.Seal(entry)
+				b.Seal(middle1)
+				b.Seal(middle2)
+				b.Seal(end)
 				return nil
 			},
 			before: `