Repository: tetratelabs/wazero
Pull request #2232 (2024-06-05 21:14:36)

ssa: reuses slice on basicBlock.loopNestingForestChildren
loopNestingForestChildren is non-empty only when the basic block
is a loop header, and hence previously the slice hasn't been reused much.

This resolves it, and as a result, the compilation uses slightly less memory.

```
goos: darwin
goarch: arm64
pkg: github.com/tetratelabs/wazero
               │  old.txt   │           new.txt           │
               │   sec/op   │   sec/op    vs base         │
Compilation-10   2.423 ± 0%   2.436 ± 2%  ~ (p=0.097 n=7)

               │   old.txt    │              new.txt               │
               │     B/op     │     B/op      vs base              │
Compilation-10   340.3Mi ± 0%   339.9Mi ± 0%  -0.10% (p=0.001 n=7)

               │   old.txt   │              new.txt              │
               │  allocs/op  │  allocs/op   vs base              │
Compilation-10   605.5k ± 0%   603.9k ± 0%  -0.27% (p=0.001 n=7)
```

#2182 


Patch:
diff --git a/internal/engine/wazevo/ssa/basic_block.go b/internal/engine/wazevo/ssa/basic_block.go
index 0fc592a0..2e7069dc 100644
--- a/internal/engine/wazevo/ssa/basic_block.go
+++ b/internal/engine/wazevo/ssa/basic_block.go
@@ -49,21 +49,12 @@ type BasicBlock interface {
 	// ReturnBlock returns ture if this block represents the function return.
 	ReturnBlock() bool
 
-	// FormatHeader returns the debug string of this block, not including instruction.
-	FormatHeader(b Builder) string
-
 	// Valid is true if this block is still valid even after optimizations.
 	Valid() bool
 
 	// Sealed is true if this block has been sealed.
 	Sealed() bool
 
-	// BeginPredIterator returns the first predecessor of this block.
-	BeginPredIterator() BasicBlock
-
-	// NextPredIterator returns the next predecessor of this block.
-	NextPredIterator() BasicBlock
-
 	// Preds returns the number of predecessors of this block.
 	Preds() int
 
@@ -90,10 +81,9 @@ type (
 		rootInstr, currentInstr *Instruction
 		// params are Values that represent parameters to a basicBlock.
 		// Each parameter can be considered as an output of PHI instruction in traditional SSA.
-		params   []Value
-		predIter int
-		preds    []basicBlockPredecessorInfo
-		success  []*basicBlock
+		params  []Value
+		preds   []basicBlockPredecessorInfo
+		success []*basicBlock
 		// singlePred is the alias to preds[0] for fast lookup, and only set after Seal is called.
 		singlePred *basicBlock
 		// lastDefinitions maps Variable to its last definition in this block.
@@ -118,7 +108,7 @@ type (
 
 		// loopNestingForestChildren holds the children of this block in the loop nesting forest.
 		// Non-empty if and only if this block is a loop header (i.e. loopHeader=true)
-		loopNestingForestChildren []BasicBlock
+		loopNestingForestChildren wazevoapi.VarLength[BasicBlock]
 
 		// reversePostOrder is used to sort all the blocks in the function in reverse post order.
 		// This is used in builder.LayoutBlocks.
@@ -138,6 +128,9 @@ type (
 	}
 )
 
+// basicBlockVarLengthNil is the default nil value for basicBlock.loopNestingForestChildren.
+var basicBlockVarLengthNil = wazevoapi.NewNilVarLength[BasicBlock]()
+
 const basicBlockIDReturnBlock = 0xffffffff
 
 // Name implements BasicBlock.Name.
@@ -240,22 +233,6 @@ func (bb *basicBlock) NumPreds() int {
 	return len(bb.preds)
 }
 
-// BeginPredIterator implements BasicBlock.BeginPredIterator.
-func (bb *basicBlock) BeginPredIterator() BasicBlock {
-	bb.predIter = 0
-	return bb.NextPredIterator()
-}
-
-// NextPredIterator implements BasicBlock.NextPredIterator.
-func (bb *basicBlock) NextPredIterator() BasicBlock {
-	if bb.predIter >= len(bb.preds) {
-		return nil
-	}
-	pred := bb.preds[bb.predIter].blk
-	bb.predIter++
-	return pred
-}
-
 // Preds implements BasicBlock.Preds.
 func (bb *basicBlock) Preds() int {
 	return len(bb.preds)
@@ -297,7 +274,7 @@ func resetBasicBlock(bb *basicBlock) {
 	bb.unknownValues = bb.unknownValues[:0]
 	bb.lastDefinitions = wazevoapi.ResetMap(bb.lastDefinitions)
 	bb.reversePostOrder = -1
-	bb.loopNestingForestChildren = bb.loopNestingForestChildren[:0]
+	bb.loopNestingForestChildren = basicBlockVarLengthNil
 	bb.loopHeader = false
 	bb.sibling = nil
 	bb.child = nil
@@ -327,8 +304,8 @@ func (bb *basicBlock) addPred(blk BasicBlock, branch *Instruction) {
 	pred.success = append(pred.success, bb)
 }
 
-// FormatHeader implements BasicBlock.FormatHeader.
-func (bb *basicBlock) FormatHeader(b Builder) string {
+// formatHeader returns the string representation of the header of the basicBlock.
+func (bb *basicBlock) formatHeader(b Builder) string {
 	ps := make([]string, len(bb.params))
 	for i, p := range bb.params {
 		ps[i] = p.formatWithType(b)
@@ -390,7 +367,7 @@ func (bb *basicBlock) String() string {
 
 // LoopNestingForestChildren implements BasicBlock.LoopNestingForestChildren.
 func (bb *basicBlock) LoopNestingForestChildren() []BasicBlock {
-	return bb.loopNestingForestChildren
+	return bb.loopNestingForestChildren.View()
 }
 
 // LoopHeader implements BasicBlock.LoopHeader.
diff --git a/internal/engine/wazevo/ssa/builder.go b/internal/engine/wazevo/ssa/builder.go
index f5414079..0671b2a7 100644
--- a/internal/engine/wazevo/ssa/builder.go
+++ b/internal/engine/wazevo/ssa/builder.go
@@ -135,6 +135,7 @@ func NewBuilder() Builder {
 	return &builder{
 		instructionsPool:               wazevoapi.NewPool[Instruction](resetInstruction),
 		basicBlocksPool:                wazevoapi.NewPool[basicBlock](resetBasicBlock),
+		varLengthBasicBlockPool:        wazevoapi.NewVarLengthPool[BasicBlock](),
 		varLengthPool:                  wazevoapi.NewVarLengthPool[Value](),
 		valueAnnotations:               make(map[ValueID]string),
 		signatures:                     make(map[SignatureID]*Signature),
@@ -177,6 +178,8 @@ type builder struct {
 	dominators []*basicBlock
 	sparseTree dominatorSparseTree
 
+	varLengthBasicBlockPool wazevoapi.VarLengthPool[BasicBlock]
+
 	// loopNestingForestRoots are the roots of the loop nesting forest.
 	loopNestingForestRoots []BasicBlock
 
@@ -219,6 +222,7 @@ func (b *builder) Init(s *Signature) {
 	b.instructionsPool.Reset()
 	b.basicBlocksPool.Reset()
 	b.varLengthPool.Reset()
+	b.varLengthBasicBlockPool.Reset()
 	b.donePreBlockLayoutPasses = false
 	b.doneBlockLayout = false
 	b.donePostBlockLayoutPasses = false
@@ -566,7 +570,7 @@ func (b *builder) Format() string {
 	}
 	for bb := iterBegin(); bb != nil; bb = iterNext() {
 		str.WriteByte('\n')
-		str.WriteString(bb.FormatHeader(b))
+		str.WriteString(bb.formatHeader(b))
 		str.WriteByte('\n')
 
 		for cur := bb.Root(); cur != nil; cur = cur.Next() {
diff --git a/internal/engine/wazevo/ssa/pass_blk_layouts.go b/internal/engine/wazevo/ssa/pass_blk_layouts.go
index 9068180a..bf906391 100644
--- a/internal/engine/wazevo/ssa/pass_blk_layouts.go
+++ b/internal/engine/wazevo/ssa/pass_blk_layouts.go
@@ -143,7 +143,7 @@ func passLayoutBlocks(b *builder) {
 	if wazevoapi.SSAValidationEnabled {
 		for _, trampoline := range trampolines {
 			if _, ok := b.blkVisited[trampoline]; !ok {
-				panic("BUG: trampoline block not inserted: " + trampoline.FormatHeader(b))
+				panic("BUG: trampoline block not inserted: " + trampoline.formatHeader(b))
 			}
 			trampoline.validate(b)
 		}
diff --git a/internal/engine/wazevo/ssa/pass_cfg.go b/internal/engine/wazevo/ssa/pass_cfg.go
index 50cb9c47..1bec185e 100644
--- a/internal/engine/wazevo/ssa/pass_cfg.go
+++ b/internal/engine/wazevo/ssa/pass_cfg.go
@@ -180,7 +180,7 @@ func passBuildLoopNestingForest(b *builder) {
 			b.loopNestingForestRoots = append(b.loopNestingForestRoots, blk)
 		} else if n == ent {
 		} else if n.loopHeader {
-			n.loopNestingForestChildren = append(n.loopNestingForestChildren, blk)
+			n.loopNestingForestChildren = n.loopNestingForestChildren.Append(&b.varLengthBasicBlockPool, blk)
 		}
 	}
 
@@ -193,7 +193,7 @@ func passBuildLoopNestingForest(b *builder) {
 
 func printLoopNestingForest(root *basicBlock, depth int) {
 	fmt.Println(strings.Repeat("\t", depth), "loop nesting forest root:", root.ID())
-	for _, child := range root.loopNestingForestChildren {
+	for _, child := range root.loopNestingForestChildren.View() {
 		fmt.Println(strings.Repeat("\t", depth+1), "child:", child.ID())
 		if child.LoopHeader() {
 			printLoopNestingForest(child.(*basicBlock), depth+2)


Test patch:
diff --git a/internal/engine/wazevo/ssa/pass_cfg_test.go b/internal/engine/wazevo/ssa/pass_cfg_test.go
index f6c48e6e..ade6eeb9 100644
--- a/internal/engine/wazevo/ssa/pass_cfg_test.go
+++ b/internal/engine/wazevo/ssa/pass_cfg_test.go
@@ -669,7 +669,7 @@ func TestBuildLoopNestingForest(t *testing.T) {
 			for expBlkID, blk := range blocks {
 				expChildren := tc.expLNF.children[expBlkID]
 				var actualChildren []BasicBlockID
-				for _, child := range blk.loopNestingForestChildren {
+				for _, child := range blk.loopNestingForestChildren.View() {
 					actualChildren = append(actualChildren, child.(*basicBlock).id)
 				}
 				sort.Slice(actualChildren, func(i, j int) bool {