Repository: tetratelabs/wazero
Pull request #1656 (2023-08-24 01:37:35)

wazevo: pass memory_grow.wast
#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/compiler.go b/internal/engine/wazevo/backend/compiler.go
index 5ed6fcc9..917dc689 100644
--- a/internal/engine/wazevo/backend/compiler.go
+++ b/internal/engine/wazevo/backend/compiler.go
@@ -143,11 +143,17 @@ type compiler struct {
 // Compile implements Compiler.Compile.
 func (c *compiler) Compile() ([]byte, []RelocationInfo, int, error) {
 	c.Lower()
-	if true {
-		fmt.Printf("[[[after lowering]]]%s", c.Format())
+	if false {
+		fmt.Printf("[[[after lowering]]]%s\n", c.Format())
 	}
 	c.RegAlloc()
+	if false {
+		fmt.Printf("[[[after regalloc]]]%s\n", c.Format())
+	}
 	c.Finalize()
+	if false {
+		fmt.Printf("[[[after finalize]]]%s\n", c.Format())
+	}
 	goPreambleSize := c.Encode()
 	return c.buf, c.relocations, goPreambleSize, nil
 }
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine.go b/internal/engine/wazevo/backend/isa/arm64/machine.go
index 24bb6f6d..65fe1f6b 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine.go
@@ -129,6 +129,13 @@ func (m *machine) Reset() {
 		v.begin, v.end = nil, nil
 	}
 	m.clobberedRegs = m.clobberedRegs[:0]
+	for key := range m.spillSlots {
+		m.clobberedRegs = append(m.clobberedRegs, regalloc.VReg(key))
+	}
+	for _, key := range m.clobberedRegs {
+		delete(m.spillSlots, regalloc.VRegID(key))
+	}
+	m.clobberedRegs = m.clobberedRegs[:0]
 	m.orderedLabels = m.orderedLabels[:0]
 	m.regAllocFn.reset()
 	m.spillSlotSize = 0
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
index 1a63031e..16dc99ef 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
@@ -41,6 +41,10 @@ func (m *machine) SetupPrologue() {
 	}
 
 	// Decrement SP if spillSlotSize > 0.
+	if m.spillSlotSize == 0 && len(m.spillSlots) != 0 {
+		panic(fmt.Sprintf("BUG: spillSlotSize=%d, spillSlots=%v\n", m.spillSlotSize, m.spillSlots))
+	}
+
 	if size := m.spillSlotSize; size > 0 {
 		// Check if size is 16-byte aligned.
 		if size&0xf != 0 {
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go b/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
index 03af6c8f..5819de40 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
@@ -178,14 +178,14 @@ func (f *regAllocFunctionImpl) Done() {
 
 // ID implements regalloc.Block ID.
 func (r *regAllocBlockImpl) ID() int {
-	return int(r.l)
+	return int(r.sb.ID())
 }
 
 // Preds implements regalloc.Block Preds.
 func (r *regAllocBlockImpl) Preds() []regalloc.Block {
 	sb := r.sb
 	r.f.predsSlice = r.f.predsSlice[:0]
-	for pred := sb.NextPredIterator(); pred != nil; pred = sb.NextPredIterator() {
+	for pred := sb.BeginPredIterator(); pred != nil; pred = sb.NextPredIterator() {
 		l := r.f.m.ssaBlockIDToLabels[pred.ID()]
 		index := r.f.labelToRegAllocBlockIndex[l]
 		r.f.predsSlice = append(r.f.predsSlice, &r.f.reversePostOrderBlocks[index])
diff --git a/internal/engine/wazevo/backend/regalloc/assign.go b/internal/engine/wazevo/backend/regalloc/assign.go
index 1242142d..d49b469a 100644
--- a/internal/engine/wazevo/backend/regalloc/assign.go
+++ b/internal/engine/wazevo/backend/regalloc/assign.go
@@ -59,7 +59,7 @@ func (a *Allocator) assignRegistersPerInstr(f Function, pc programCounter, instr
 			a.vs = append(a.vs, u)
 			continue
 		}
-		if true {
+		if false {
 			fmt.Printf("%s uses %d\n", instr, u.ID())
 		}
 		n := vRegIDToNode[u.ID()]
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc.go b/internal/engine/wazevo/backend/regalloc/regalloc.go
index 207d1d43..943b38cb 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc.go
@@ -267,8 +267,16 @@ func (a *Allocator) upAndMarkStack(b Block, v VReg, depth int) {
 
 	// Now we can safely mark v as a part of live-in
 	info.liveIns[v] = struct{}{}
+	preds := b.Preds()
+	if len(preds) == 0 {
+		panic("BUG: block with no predecessors while requiring live-in")
+	}
+
 	// and climb up the CFG.
-	for _, pred := range b.Preds() {
+	for _, pred := range preds {
+		if debug {
+			fmt.Printf("%sadding %v live-out at block[%d]\n", strings.Repeat("\t", depth+1), v, pred.ID())
+		}
 		a.blockInfoAt(pred.ID()).liveOuts[v] = struct{}{}
 		a.upAndMarkStack(pred, v, depth+1)
 	}
diff --git a/internal/engine/wazevo/engine.go b/internal/engine/wazevo/engine.go
index 025b94c4..12c17a4b 100644
--- a/internal/engine/wazevo/engine.go
+++ b/internal/engine/wazevo/engine.go
@@ -100,7 +100,10 @@ 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)
+		const debug = false
+		if debug {
+			fmt.Printf("------------------------------------------ %d/%d ------------------------------------------\n", i, len(module.CodeSection)-1)
+		}
 		fidx := wasm.Index(i + importedFns)
 		fref := frontend.FunctionIndexToFuncRef(fidx)
 
@@ -125,8 +128,6 @@ 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())
 		}
diff --git a/internal/engine/wazevo/ssa/basic_block.go b/internal/engine/wazevo/ssa/basic_block.go
index d1281d98..a916db84 100644
--- a/internal/engine/wazevo/ssa/basic_block.go
+++ b/internal/engine/wazevo/ssa/basic_block.go
@@ -248,6 +248,7 @@ func (bb *basicBlock) reset() {
 	// TODO: reuse the map!
 	bb.unknownValues = make(map[Variable]Value)
 	bb.lastDefinitions = make(map[Variable]Value)
+	bb.reversePostOrder = -1
 }
 
 // addPred adds a predecessor to this block specified by the branch instruction.
@@ -285,8 +286,10 @@ 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: %s",
-					len(pred.branch.vs), len(bb.params), pred.branch.Format(b),
+				panic(fmt.Sprintf(
+					"BUG: len(argument at %s) != len(params at %s): %d != %d: %s",
+					pred.blk.Name(), bb.Name(),
+					len(bb.params), len(pred.branch.vs), pred.branch.Format(b),
 				))
 			}
 			if pred.blk.invalid {
@@ -302,6 +305,30 @@ func (bb *basicBlock) FormatHeader(b Builder) string {
 	}
 }
 
+// validates validates the basicBlock for debugging purpose.
+func (bb *basicBlock) validate(b *builder) {
+	if bb.invalid {
+		panic("BUG: trying to validate an invalid block: " + bb.Name())
+	}
+	if len(bb.preds) > 0 {
+		for _, pred := range bb.preds {
+			if pred.branch.opcode != OpcodeBrTable {
+				if target := pred.branch.blk; target != bb {
+					panic(fmt.Sprintf("BUG: '%s' is not branch to %s, but to %s",
+						pred.branch.Format(b), bb.Name(), target.Name()))
+				}
+			}
+			if len(pred.branch.vs) != len(bb.params) {
+				panic(fmt.Sprintf(
+					"BUG: len(argument at %s) != len(params at %s): %d != %d: %s",
+					pred.blk.Name(), bb.Name(),
+					len(bb.params), len(pred.branch.vs), pred.branch.Format(b),
+				))
+			}
+		}
+	}
+}
+
 // String implements fmt.Stringer for debugging purpose only.
 func (bb *basicBlock) String() string {
 	return strconv.Itoa(int(bb.id))
diff --git a/internal/engine/wazevo/ssa/builder.go b/internal/engine/wazevo/ssa/builder.go
index 500587b9..45b5c294 100644
--- a/internal/engine/wazevo/ssa/builder.go
+++ b/internal/engine/wazevo/ssa/builder.go
@@ -664,6 +664,8 @@ func (b *builder) LayoutBlocks() {
 		panic("LayoutBlocks must be called after all passes are done")
 	}
 
+	b.clearBlkVisited()
+
 	// We might end up splitting critical edges which adds more basic blocks,
 	// so we store the currently existing basic blocks in nonSplitBlocks temporarily.
 	// That way we can iterate over the original basic blocks while appending new ones into reversePostOrderedBasicBlocks.
@@ -678,8 +680,7 @@ func (b *builder) LayoutBlocks() {
 		}
 	}
 
-	b.clearBlkVisited()
-	inserted := b.blkVisited
+	trampolines := []*basicBlock{}
 
 	// Reset the order slice since we update on the fly by splitting critical edges.
 	b.reversePostOrderedBasicBlocks = b.reversePostOrderedBasicBlocks[:0]
@@ -687,13 +688,13 @@ func (b *builder) LayoutBlocks() {
 	for _, blk := range nonSplitBlocks {
 		for i := range blk.preds {
 			pred := blk.preds[i].blk
-			if _, ok := inserted[pred]; ok || !pred.Valid() {
+			if _, ok := b.blkVisited[pred]; ok || !pred.Valid() {
 				continue
 			} else if pred.reversePostOrder < blk.reversePostOrder {
 				// This means the edge is critical, and this pred is the trampoline and yet to be inserted.
 				// Split edge trampolines must come before the destination in reverse post-order.
 				b.reversePostOrderedBasicBlocks = append(b.reversePostOrderedBasicBlocks, pred)
-				b.blkVisited[blk] = 0 // mark as inserted, the value is not used.
+				b.blkVisited[pred] = 0 // mark as inserted, the value is not used.
 			}
 		}
 
@@ -733,12 +734,19 @@ func (b *builder) LayoutBlocks() {
 				panic("BUG: predecessor info not found while the successor exists in successors list")
 			}
 
+			if debug {
+				fmt.Printf("trying to split edge from %d->%d at %s\n",
+					blk.ID(), succ.ID(), predInfo.branch.Format(b))
+			}
+
 			trampoline := b.splitCriticalEdge(blk, succ, predInfo)
 			// Update the successors slice because the target is no longer the original `succ`.
 			blk.success[sidx] = trampoline
 
+			trampolines = append(trampolines, trampoline)
+
 			if debug {
-				fmt.Printf("splitting edge from %d->%d at %s as %d->%d->%d \n",
+				fmt.Printf("edge split from %d->%d at %s as %d->%d->%d \n",
 					blk.ID(), succ.ID(), predInfo.branch.Format(b),
 					blk.ID(), trampoline.ID(), succ.ID())
 			}
@@ -747,7 +755,7 @@ func (b *builder) LayoutBlocks() {
 			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.
+				b.blkVisited[trampoline] = 0 // mark as inserted, the value is not used.
 			} else {
 				uninsertedTrampolines = append(uninsertedTrampolines, trampoline)
 			}
@@ -757,12 +765,32 @@ func (b *builder) LayoutBlocks() {
 			if trampoline.success[0].reversePostOrder < trampoline.reversePostOrder {
 				// This means the critical edge was backward, so we insert after the current block immediately.
 				b.reversePostOrderedBasicBlocks = append(b.reversePostOrderedBasicBlocks, trampoline)
-				inserted[trampoline] = 0 // mark as inserted, the value is not used.
+				b.blkVisited[trampoline] = 0 // mark as inserted, the value is not used.
 			} // If the target is forward, we can wait to insert until the target is inserted.
 		}
 		uninsertedTrampolines = uninsertedTrampolines[:0] // Reuse the stack for the next block.
 	}
 
+	if debug {
+		var bs []string
+		for _, blk := range b.reversePostOrderedBasicBlocks {
+			bs = append(bs, blk.Name())
+		}
+		fmt.Println("ordered blocks: ", strings.Join(bs, ", "))
+		bs = bs[:0]
+		for visited := range b.blkVisited {
+			bs = append(bs, visited.Name())
+		}
+		sort.Slice(bs, func(i, j int) bool { return bs[i] < bs[j] })
+		fmt.Println("visited blocks: ", strings.Join(bs, ", "))
+		for _, trampoline := range trampolines {
+			if _, ok := b.blkVisited[trampoline]; !ok {
+				panic("BUG: trampoline block not inserted: " + trampoline.FormatHeader(b))
+			}
+			trampoline.validate(b)
+		}
+	}
+
 	// Reuse the stack for the next iteration.
 	b.blkStack2 = uninsertedTrampolines[:0]
 
@@ -914,6 +942,15 @@ func (b *builder) splitCriticalEdge(pred, succ *basicBlock, predInfo *basicBlock
 
 	// Update the original branch to point to the trampoline.
 	predInfo.blk = trampoline
+	predInfo.branch = originalBranch
+
+	if debug {
+		trampoline.validate(b)
+	}
+
+	if len(trampoline.params) > 0 {
+		panic("trampoline should not have params")
+	}
 
 	// Assign the same order as the original block so that this will be placed before the actual destination.
 	trampoline.reversePostOrder = pred.reversePostOrder
diff --git a/internal/engine/wazevo/ssa/instructions.go b/internal/engine/wazevo/ssa/instructions.go
index 5eb87b60..8c6e1cf0 100644
--- a/internal/engine/wazevo/ssa/instructions.go
+++ b/internal/engine/wazevo/ssa/instructions.go
@@ -51,6 +51,7 @@ func (i *Instruction) reset() {
 	i.v3 = ValueInvalid
 	i.rValue = ValueInvalid
 	i.typ = typeInvalid
+	i.vs = nil
 }
 
 // InstructionGroupID is assigned to each instruction and represents a group of instructions
@@ -1177,9 +1178,10 @@ func (i *Instruction) AsF64const(f float64) {
 }
 
 // AsReturn initializes this instruction as a return instruction with OpcodeReturn.
-func (i *Instruction) AsReturn(vs []Value) {
+func (i *Instruction) AsReturn(vs []Value) *Instruction {
 	i.opcode = OpcodeReturn
 	i.vs = vs
+	return i
 }
 
 // ReturnVals returns the return values of OpcodeReturn.
diff --git a/internal/engine/wazevo/ssa/pass.go b/internal/engine/wazevo/ssa/pass.go
index 7eae17be..07b23951 100644
--- a/internal/engine/wazevo/ssa/pass.go
+++ b/internal/engine/wazevo/ssa/pass.go
@@ -275,8 +275,12 @@ func passDeadCodeEliminationOpt(b *builder) {
 
 // clearBlkVisited clears the b.blkVisited map so that we can reuse it for multiple places.
 func (b *builder) clearBlkVisited() {
-	for i := 0; i < b.basicBlocksPool.Allocated(); i++ {
-		blk := b.basicBlocksPool.View(i)
+	b.blkStack2 = b.blkStack2[:0]
+	for key := range b.blkVisited {
+		b.blkStack2 = append(b.blkStack2, key)
+	}
+	for _, blk := range b.blkStack2 {
 		delete(b.blkVisited, blk)
 	}
+	b.blkStack2 = b.blkStack2[:0]
 }


Test patch:
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index b3c1d468..ae6c591c 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -37,18 +37,22 @@ func TestSpectestV1(t *testing.T) {
 	for _, tc := range []struct {
 		name string
 	}{
-		{name: "address.json"},
-		{name: "binary.json"},
-		{name: "binary-leb128.json"},
-		{name: "comments.json"},
-		{name: "custom.json"},
-		{name: "const.json"},
-		{name: "data.json"},
-		{name: "local_get.json"},
-		{name: "memory_size.json"},
-		{name: "nop.json"},
+		{name: "address"},
+		{name: "binary"},
+		{name: "binary-leb128"},
+		{name: "comments"},
+		{name: "custom"},
+		{name: "const"},
+		{name: "data"},
+		{name: "local_get"},
+		{name: "memory_size"},
+		{name: "memory_grow"},
+		{name: "nop"},
 	} {
-		t.Run(tc.name, func(t *testing.T) { spectest.RunJson(t, v1.Testcases, tc.name, context.Background(), config) })
+		t.Run(tc.name, func(t *testing.T) {
+			spectest.RunCase(t, v1.Testcases, tc.name, context.Background(), config,
+				-1, 0, math.MaxInt)
+		})
 	}
 }
 
diff --git a/internal/engine/wazevo/ssa/builder_test.go b/internal/engine/wazevo/ssa/builder_test.go
index 52be0e24..38b223d5 100644
--- a/internal/engine/wazevo/ssa/builder_test.go
+++ b/internal/engine/wazevo/ssa/builder_test.go
@@ -522,6 +522,35 @@ func TestBuilder_LayoutBlocks(t *testing.T) {
 				4,
 			},
 		},
+		{
+			name: "brz with arg",
+			setup: func(b *builder) {
+				b0, b1, b2 := b.allocateBasicBlock(), b.allocateBasicBlock(), b.allocateBasicBlock()
+				p := b0.AddParam(b, TypeI32)
+				retval := b1.AddParam(b, TypeI32)
+
+				b.SetCurrentBlock(b0)
+				{
+					arg := b.AllocateInstruction().AsIconst32(1000).Insert(b).Return()
+					insertBrz(b, b0, b1, p, arg)
+					insertJump(b, b0, b2)
+				}
+				b.SetCurrentBlock(b1)
+				{
+					b.AllocateInstruction().AsReturn([]Value{retval}).Insert(b)
+				}
+				b.SetCurrentBlock(b2)
+				{
+					arg := b.AllocateInstruction().AsIconst32(1).Insert(b).Return()
+					insertJump(b, b2, b1, arg)
+				}
+
+				b.Seal(b0)
+				b.Seal(b1)
+				b.Seal(b2)
+			},
+			exp: []BasicBlockID{0x0, 0x3, 0x1, 0x2},
+		},
 		{
 			name: "loop with output",
 			exp:  []BasicBlockID{0x0, 0x2, 0x4, 0x1, 0x3, 0x6, 0x5},
diff --git a/internal/integration_test/spectest/spectest.go b/internal/integration_test/spectest/spectest.go
index f7b731c7..77c16d59 100644
--- a/internal/integration_test/spectest/spectest.go
+++ b/internal/integration_test/spectest/spectest.go
@@ -326,26 +326,35 @@ func Run(t *testing.T, testDataFS embed.FS, ctx context.Context, config wazero.R
 	files, err := testDataFS.ReadDir("testdata")
 	require.NoError(t, err)
 
-	jsonfiles := make([]string, 0, len(files))
+	caseNames := make([]string, 0, len(files))
 	for _, f := range files {
 		filename := f.Name()
 		if strings.HasSuffix(filename, ".json") {
-			jsonfiles = append(jsonfiles, filename)
+			caseNames = append(caseNames, strings.TrimSuffix(filename, ".json"))
 		}
 	}
 
 	// If the go:embed path resolution was wrong, this fails.
 	// https://github.com/tetratelabs/wazero/issues/247
-	require.True(t, len(jsonfiles) > 1, "len(jsonfiles)=%d (not greater than one)", len(jsonfiles))
+	require.True(t, len(caseNames) > 1, "len(caseNames)=%d (not greater than one)", len(caseNames))
 
-	for _, f := range jsonfiles {
-		RunJson(t, testDataFS, f, ctx, config)
+	for _, f := range caseNames {
+		RunCase(t, testDataFS, f, ctx, config, -1, 0, math.MaxInt)
 	}
 }
 
-// RunJson runs the test case described by the given spectest JSON file name in the testDataFS file system.
-func RunJson(t *testing.T, testDataFS embed.FS, f string, ctx context.Context, config wazero.RuntimeConfig) {
-	raw, err := testDataFS.ReadFile(testdataPath(f))
+// RunCase runs the test case described by the given spectest file name (without .wast!) in the testDataFS file system.
+// lineBegin and lineEnd are the line numbers to run. If lineBegin == 0 and lineEnd == math.MaxInt, all the lines are run.
+//
+// For example, if you want to run memory_grow.wast:66 to 70, you can do:
+//
+//	RunCase(t, testDataFS, "memory_grow", ctx, config, mandatoryLine, 66, 70)
+//
+// where mandatoryLine is the line number which can be run regardless of the lineBegin and lineEnd. It is useful when
+// we only want to run specific command while running "module" command to instantiate a module. If you don't need it,
+// just pass -1.
+func RunCase(t *testing.T, testDataFS embed.FS, f string, ctx context.Context, config wazero.RuntimeConfig, mandatoryLine, lineBegin, lineEnd int) {
+	raw, err := testDataFS.ReadFile(testdataPath(f + ".json"))
 	require.NoError(t, err)
 
 	var base testbase
@@ -366,6 +375,11 @@ func RunJson(t *testing.T, testDataFS embed.FS, f string, ctx context.Context, c
 		var lastInstantiatedModule api.Module
 		for i := 0; i < len(base.Commands); i++ {
 			c := &base.Commands[i]
+			line := c.Line
+			if mandatoryLine > -1 && c.Line == mandatoryLine {
+			} else if line < lineBegin || line > lineEnd {
+				continue
+			}
 			t.Run(fmt.Sprintf("%s/line:%d", c.CommandType, c.Line), func(t *testing.T) {
 				msg := fmt.Sprintf("%s:%d %s", wastName, c.Line, c.CommandType)
 				switch c.CommandType {