Repository: tetratelabs/wazero
Pull request #2285 (2024-07-09 23:45:23)

ssa: removes map use for Value aliasing
This improves the compilation perf by up to ~30% without 
runtime perf impacts.

```
goos: darwin
goarch: arm64
pkg: github.com/tetratelabs/wazero/internal/integration_test/stdlibs
                                            │   old.txt    │              new.txt               │
                                            │    sec/op    │   sec/op     vs base               │
Zig/Compile/test-opt.wasm-10                    2.683 ± 1%    2.591 ± 1%   -3.44% (p=0.002 n=6)
Zig/Run/test-opt.wasm-10                        18.75 ± 0%    19.03 ± 1%   +1.48% (p=0.004 n=6)
Zig/Compile/test.wasm-10                        3.529 ± 0%    3.516 ± 1%        ~ (p=0.180 n=6)
Zig/Run/test.wasm-10                            19.55 ± 2%    19.48 ± 1%        ~ (p=0.310 n=6)
TinyGo/Compile/container_heap.test-10          298.1m ± 1%   210.1m ± 0%  -29.51% (p=0.002 n=6)
TinyGo/Run/container_heap.test-10              14.59m ± 2%   14.55m ± 1%        ~ (p=0.310 n=6)
TinyGo/Compile/container_list.test-10          290.6m ± 0%   208.6m ± 0%  -28.20% (p=0.002 n=6)
TinyGo/Run/container_list.test-10              14.47m ± 1%   14.40m ± 0%        ~ (p=0.132 n=6)
TinyGo/Compile/container_ring.test-10          286.4m ± 0%   205.1m ± 0%  -28.42% (p=0.002 n=6)
TinyGo/Run/container_ring.test-10              14.36m ± 1%   14.42m ± 1%        ~ (p=0.180 n=6)
TinyGo/Compile/crypto_des.test-10              296.6m ± 0%   214.2m ± 2%  -27.77% (p=0.002 n=6)
TinyGo/Run/crypto_des.test-10                  18.70m ± 4%   18.64m ± 0%        ~ (p=0.132 n=6)
TinyGo/Compile/crypto_md5.test-10              295.8m ± 5%   211.4m ± 0%  -28.53% (p=0.002 n=6)
TinyGo/Run/crypto_md5.test-10                  21.72m ± 1%   20.67m ± 1%   -4.85% (p=0.002 n=6)
TinyGo/Compile/crypto_rc4.test-10              291.4m ± 2%   203.4m ± 0%  -30.20% (p=0.002 n=6)
TinyGo/Run/crypto_rc4.test-10                  161.0m ± 1%   160.7m ± 1%        ~ (p=0.818 n=6)
TinyGo/Compile/crypto_sha1.test-10             296.0m ± 0%   211.9m ± 1%  -28.43% (p=0.002 n=6)
TinyGo/Run/crypto_sha1.test-10                 16.14m ± 1%   16.17m ± 1%        ~ (p=0.485 n=6)
TinyGo/Compile/crypto_sha256.test-10           301.0m ± 1%   216.1m ± 1%  -28.21% (p=0.002 n=6)
TinyGo/Run/crypto_sha256.test-10               16.21m ± 1%   16.27m ± 1%        ~ (p=0.310 n=6)
Wasip1/Compile/src_archive_tar.test-10          1.596 ± 0%    1.490 ± 0%   -6.63% (p=0.002 n=6)
Wasip1/Run/src_archive_tar.test-10             396.5m ± 0%   396.5m ± 1%        ~ (p=0.818 n=6)
Wasip1/Compile/src_bufio.test-10                1.089 ± 0%    1.008 ± 0%   -7.44% (p=0.002 n=6)
Wasip1/Run/src_bufio.test-10                   121.4m ± 0%   121.5m ± 0%        ~ (p=0.240 n=6)
Wasip1/Compile/src_bytes.test-10                1.127 ± 0%    1.038 ± 0%   -7.91% (p=0.002 n=6)
Wasip1/Run/src_bytes.test-10                   468.1m ± 1%   466.8m ± 1%        ~ (p=0.132 n=6)
Wasip1/Compile/src_context.test-10              1.213 ± 1%    1.119 ± 0%   -7.76% (p=0.002 n=6)
Wasip1/Run/src_context.test-10                 33.69m ± 1%   33.57m ± 0%   -0.33% (p=0.026 n=6)
Wasip1/Compile/src_encoding_ascii85.test-10    979.9m ± 0%   905.5m ± 0%   -7.59% (p=0.002 n=6)
Wasip1/Run/src_encoding_ascii85.test-10        7.320m ± 1%   7.324m ± 3%        ~ (p=0.589 n=6)
Wasip1/Compile/src_encoding_asn1.test-10        1.128 ± 1%    1.043 ± 0%   -7.48% (p=0.002 n=6)
Wasip1/Run/src_encoding_asn1.test-10           8.221m ± 2%   8.239m ± 1%        ~ (p=0.394 n=6)
Wasip1/Compile/src_encoding_base32.test-10    1019.3m ± 1%   940.4m ± 0%   -7.74% (p=0.002 n=6)
Wasip1/Run/src_encoding_base32.test-10         28.57m ± 1%   28.78m ± 0%   +0.74% (p=0.041 n=6)
Wasip1/Compile/src_encoding_base64.test-10    1024.8m ± 1%   945.4m ± 0%   -7.75% (p=0.002 n=6)
Wasip1/Run/src_encoding_base64.test-10         11.78m ± 1%   11.80m ± 4%        ~ (p=0.589 n=6)
```

#2182 


Patch:
diff --git a/internal/engine/wazevo/backend/compiler.go b/internal/engine/wazevo/backend/compiler.go
index 59bbfe02..0c3d4e58 100644
--- a/internal/engine/wazevo/backend/compiler.go
+++ b/internal/engine/wazevo/backend/compiler.go
@@ -128,8 +128,6 @@ type compiler struct {
 	ssaValueToVRegs [] /* VRegID to */ regalloc.VReg
 	// ssaValueDefinitions maps ssa.ValueID to its definition.
 	ssaValueDefinitions []SSAValueDefinition
-	// ssaValueRefCounts is a cached list obtained by ssa.Builder.ValueRefCounts().
-	ssaValueRefCounts []int
 	// returnVRegs is the list of virtual registers that store the return values.
 	returnVRegs  []regalloc.VReg
 	varEdges     [][2]regalloc.VReg
@@ -206,8 +204,7 @@ func (c *compiler) setCurrentGroupID(gid ssa.InstructionGroupID) {
 // assignVirtualRegisters assigns a virtual register to each ssa.ValueID Valid in the ssa.Builder.
 func (c *compiler) assignVirtualRegisters() {
 	builder := c.ssaBuilder
-	refCounts := builder.ValueRefCounts()
-	c.ssaValueRefCounts = refCounts
+	refCounts := builder.ValuesInfo()
 
 	need := len(refCounts)
 	if need >= len(c.ssaValueToVRegs) {
@@ -242,7 +239,7 @@ func (c *compiler) assignVirtualRegisters() {
 				c.ssaValueDefinitions[id] = SSAValueDefinition{
 					Instr:    cur,
 					N:        0,
-					RefCount: refCounts[id],
+					RefCount: refCounts[id].RefCount,
 				}
 				c.ssaTypeOfVRegID[vReg.ID()] = ssaTyp
 				N++
@@ -255,7 +252,7 @@ func (c *compiler) assignVirtualRegisters() {
 				c.ssaValueDefinitions[id] = SSAValueDefinition{
 					Instr:    cur,
 					N:        N,
-					RefCount: refCounts[id],
+					RefCount: refCounts[id].RefCount,
 				}
 				c.ssaTypeOfVRegID[vReg.ID()] = ssaTyp
 				N++
diff --git a/internal/engine/wazevo/backend/compiler_lower.go b/internal/engine/wazevo/backend/compiler_lower.go
index 9a9414ae..669364c3 100644
--- a/internal/engine/wazevo/backend/compiler_lower.go
+++ b/internal/engine/wazevo/backend/compiler_lower.go
@@ -124,9 +124,10 @@ func (c *compiler) lowerFunctionArguments(entry ssa.BasicBlock) {
 	mach := c.mach
 
 	c.tmpVals = c.tmpVals[:0]
+	data := c.ssaBuilder.ValuesInfo()
 	for i := 0; i < entry.Params(); i++ {
 		p := entry.Param(i)
-		if c.ssaValueRefCounts[p.ID()] > 0 {
+		if data[p.ID()].RefCount > 0 {
 			c.tmpVals = append(c.tmpVals, p)
 		} else {
 			// If the argument is not used, we can just pass an invalid value.
diff --git a/internal/engine/wazevo/backend/vdef.go b/internal/engine/wazevo/backend/vdef.go
index edfa962b..0951f9ab 100644
--- a/internal/engine/wazevo/backend/vdef.go
+++ b/internal/engine/wazevo/backend/vdef.go
@@ -18,7 +18,7 @@ type SSAValueDefinition struct {
 	// N is the index of the return value in the instr's return values list.
 	N int
 	// RefCount is the number of references to the result.
-	RefCount int
+	RefCount uint32
 }
 
 func (d *SSAValueDefinition) IsFromInstr() bool {
diff --git a/internal/engine/wazevo/ssa/builder.go b/internal/engine/wazevo/ssa/builder.go
index 3cc5dbee..60deb535 100644
--- a/internal/engine/wazevo/ssa/builder.go
+++ b/internal/engine/wazevo/ssa/builder.go
@@ -94,9 +94,9 @@ type Builder interface {
 	// Returns nil if there's no unseen BasicBlock.
 	BlockIteratorNext() BasicBlock
 
-	// ValueRefCounts returns the map of ValueID to its reference count.
-	// The returned slice must not be modified.
-	ValueRefCounts() []int
+	// ValuesInfo returns the data per Value used to lower the SSA in backend.
+	// This is indexed by ValueID.
+	ValuesInfo() []ValueInfo
 
 	// BlockIteratorReversePostOrderBegin is almost the same as BlockIteratorBegin except it returns the BasicBlock in the reverse post-order.
 	// This is available after RunPasses is run.
@@ -143,7 +143,6 @@ func NewBuilder() Builder {
 		varLengthPool:           wazevoapi.NewVarLengthPool[Value](),
 		valueAnnotations:        make(map[ValueID]string),
 		signatures:              make(map[SignatureID]*Signature),
-		valueIDAliases:          make(map[ValueID]Value),
 		returnBlk:               &basicBlock{id: basicBlockIDReturnBlock},
 	}
 }
@@ -166,12 +165,11 @@ type builder struct {
 	// nextVariable is used by builder.AllocateVariable.
 	nextVariable Variable
 
-	valueIDAliases   map[ValueID]Value
+	// valueAnnotations contains the annotations for each Value, only used for debugging.
 	valueAnnotations map[ValueID]string
 
-	// valueRefCounts is used to lower the SSA in backend, and will be calculated
-	// by the last SSA-level optimization pass.
-	valueRefCounts []int
+	// valuesInfo contains the data per Value used to lower the SSA in backend. This is indexed by ValueID.
+	valuesInfo []ValueInfo
 
 	// dominators stores the immediate dominator of each BasicBlock.
 	// The index is blockID of the BasicBlock.
@@ -206,6 +204,13 @@ type builder struct {
 	zeros [typeEnd]Value
 }
 
+// ValueInfo contains the data per Value used to lower the SSA in backend.
+type ValueInfo struct {
+	// RefCount is the reference count of the Value.
+	RefCount uint32
+	alias    Value
+}
+
 // redundantParam is a pair of the index of the redundant parameter and the Value.
 // This is used to eliminate the redundant parameters in the optimization pass.
 type redundantParam struct {
@@ -285,8 +290,7 @@ func (b *builder) Init(s *Signature) {
 
 	for v := ValueID(0); v < b.nextValueID; v++ {
 		delete(b.valueAnnotations, v)
-		delete(b.valueIDAliases, v)
-		b.valueRefCounts[v] = 0
+		b.valuesInfo[v] = ValueInfo{alias: ValueInvalid}
 		b.valueIDToInstruction[v] = nil
 	}
 	b.nextValueID = 0
@@ -676,15 +680,24 @@ func (b *builder) blockIteratorReversePostOrderNext() *basicBlock {
 	}
 }
 
-// ValueRefCounts implements Builder.ValueRefCounts.
-func (b *builder) ValueRefCounts() []int {
-	return b.valueRefCounts
+// ValuesInfo implements Builder.ValuesInfo.
+func (b *builder) ValuesInfo() []ValueInfo {
+	return b.valuesInfo
 }
 
 // alias records the alias of the given values. The alias(es) will be
 // eliminated in the optimization pass via resolveArgumentAlias.
 func (b *builder) alias(dst, src Value) {
-	b.valueIDAliases[dst.ID()] = src
+	did := int(dst.ID())
+	if did >= len(b.valuesInfo) {
+		l := did + 1 - len(b.valuesInfo)
+		b.valuesInfo = append(b.valuesInfo, make([]ValueInfo, l)...)
+		view := b.valuesInfo[len(b.valuesInfo)-l:]
+		for i := range view {
+			view[i].alias = ValueInvalid
+		}
+	}
+	b.valuesInfo[did].alias = src
 }
 
 // resolveArgumentAlias resolves the alias of the arguments of the given instruction.
@@ -709,10 +722,13 @@ func (b *builder) resolveArgumentAlias(instr *Instruction) {
 
 // resolveAlias resolves the alias of the given value.
 func (b *builder) resolveAlias(v Value) Value {
+	info := b.valuesInfo
+	l := ValueID(len(info))
 	// Some aliases are chained, so we need to resolve them recursively.
 	for {
-		if src, ok := b.valueIDAliases[v.ID()]; ok {
-			v = src
+		vid := v.ID()
+		if vid < l && info[vid].alias.Valid() {
+			v = info[vid].alias
 		} else {
 			break
 		}
diff --git a/internal/engine/wazevo/ssa/pass.go b/internal/engine/wazevo/ssa/pass.go
index 9f0643cc..64dc5f5c 100644
--- a/internal/engine/wazevo/ssa/pass.go
+++ b/internal/engine/wazevo/ssa/pass.go
@@ -235,8 +235,13 @@ func passRedundantPhiEliminationOpt(b *builder) {
 // TODO: the algorithm here might not be efficient. Get back to this later.
 func passDeadCodeEliminationOpt(b *builder) {
 	nvid := int(b.nextValueID)
-	if nvid >= len(b.valueRefCounts) {
-		b.valueRefCounts = append(b.valueRefCounts, make([]int, nvid-len(b.valueRefCounts)+1)...)
+	if nvid >= len(b.valuesInfo) {
+		l := nvid - len(b.valuesInfo) + 1
+		b.valuesInfo = append(b.valuesInfo, make([]ValueInfo, l)...)
+		view := b.valuesInfo[len(b.valuesInfo)-l:]
+		for i := range view {
+			view[i].alias = ValueInvalid
+		}
 	}
 	if nvid >= len(b.valueIDToInstruction) {
 		b.valueIDToInstruction = append(b.valueIDToInstruction, make([]*Instruction, nvid-len(b.valueIDToInstruction)+1)...)
@@ -356,7 +361,8 @@ func (b *builder) incRefCount(id ValueID, from *Instruction) {
 	if wazevoapi.SSALoggingEnabled {
 		fmt.Printf("v%d referenced from %v\n", id, from.Format(b))
 	}
-	b.valueRefCounts[id]++
+	info := &b.valuesInfo[id]
+	info.RefCount++
 }
 
 // passNopInstElimination eliminates the instructions which is essentially a no-op.
diff --git a/internal/engine/wazevo/ssa/vs.go b/internal/engine/wazevo/ssa/vs.go
index 6e6cce47..2fd12cc6 100644
--- a/internal/engine/wazevo/ssa/vs.go
+++ b/internal/engine/wazevo/ssa/vs.go
@@ -67,7 +67,7 @@ func (v Value) formatWithType(b Builder) (ret string) {
 	if wazevoapi.SSALoggingEnabled { // This is useful to check live value analysis bugs.
 		if bd := b.(*builder); bd.donePostBlockLayoutPasses {
 			id := v.ID()
-			ret += fmt.Sprintf("(ref=%d)", bd.valueRefCounts[id])
+			ret += fmt.Sprintf("(ref=%d)", bd.valuesInfo[id].RefCount)
 		}
 	}
 	return ret


Test patch:
diff --git a/internal/engine/wazevo/ssa/builder_test.go b/internal/engine/wazevo/ssa/builder_test.go
new file mode 100644
index 00000000..94ee7e5e
--- /dev/null
+++ b/internal/engine/wazevo/ssa/builder_test.go
@@ -0,0 +1,26 @@
+package ssa
+
+import (
+	"testing"
+
+	"github.com/tetratelabs/wazero/internal/testing/require"
+)
+
+func TestBuilder_resolveAlias(t *testing.T) {
+	b := NewBuilder().(*builder)
+	v1 := b.allocateValue(TypeI32)
+	v2 := b.allocateValue(TypeI32)
+	v3 := b.allocateValue(TypeI32)
+	v4 := b.allocateValue(TypeI32)
+	v5 := b.allocateValue(TypeI32)
+
+	b.alias(v1, v2)
+	b.alias(v2, v3)
+	b.alias(v3, v4)
+	b.alias(v4, v5)
+	require.Equal(t, v5, b.resolveAlias(v1))
+	require.Equal(t, v5, b.resolveAlias(v2))
+	require.Equal(t, v5, b.resolveAlias(v3))
+	require.Equal(t, v5, b.resolveAlias(v4))
+	require.Equal(t, v5, b.resolveAlias(v5))
+}
diff --git a/internal/engine/wazevo/ssa/pass_test.go b/internal/engine/wazevo/ssa/pass_test.go
index 015a10d0..4b7919b2 100644
--- a/internal/engine/wazevo/ssa/pass_test.go
+++ b/internal/engine/wazevo/ssa/pass_test.go
@@ -264,9 +264,9 @@ blk2: () <-- (blk1)
 					require.True(t, jmp.live)
 					require.True(t, ret.live)
 
-					require.Equal(t, 1, b.valueRefCounts[refOnceVal.ID()])
-					require.Equal(t, 1, b.valueRefCounts[addRes.ID()])
-					require.Equal(t, 3, b.valueRefCounts[refThriceVal.ID()])
+					require.Equal(t, uint32(1), b.valuesInfo[refOnceVal.ID()].RefCount)
+					require.Equal(t, uint32(1), b.valuesInfo[addRes.ID()].RefCount)
+					require.Equal(t, uint32(3), b.valuesInfo[refThriceVal.ID()].RefCount)
 				}
 			},
 			before: `