Repository: tetratelabs/wazero
Pull request #1785 (2023-10-16 03:58:28)

wazevo(regalloc): reduces the complexity
#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/regalloc/coloring.go b/internal/engine/wazevo/backend/regalloc/coloring.go
index 2c8cc7de..67d67d8d 100644
--- a/internal/engine/wazevo/backend/regalloc/coloring.go
+++ b/internal/engine/wazevo/backend/regalloc/coloring.go
@@ -3,7 +3,6 @@ package regalloc
 import (
 	"fmt"
 	"sort"
-	"strings"
 
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/wazevoapi"
 )
@@ -31,8 +30,8 @@ func (a *Allocator) buildNeighborsByLiveNodes(lives []liveNodeInBlock) {
 			dstRange := &dst.n.ranges[dst.rangeIndex]
 			if src.n.v.RegType() == dst.n.v.RegType() && // Interfere only if they are the same type.
 				srcRange.intersects(dstRange) {
-				src.n.neighbors[dst.n] = struct{}{}
-				dst.n.neighbors[src.n] = struct{}{}
+				src.n.neighbors = append(src.n.neighbors, dst.n)
+				dst.n.neighbors = append(dst.n.neighbors, src.n)
 			}
 		}
 	}
@@ -70,13 +69,12 @@ func (a *Allocator) coloringFor(allocatable []RealReg) {
 	degreeSortedNodes := a.nodes1 // We assume nodes1 holds all the nodes of the given register type.
 	// Reuses the nodes2 slice and the degrees map from the previous iteration.
 	coloringStack := a.nodes2[:0]
-	currentDegrees := a.nodeSet
 
 	numAllocatable := len(allocatable)
 
 	// Initialize the degree for each node which is defined as the number of neighbors.
 	for _, n := range degreeSortedNodes {
-		currentDegrees[n] = len(n.neighbors)
+		n.degree = len(n.neighbors)
 	}
 
 	// First step of the algorithm:
@@ -87,7 +85,7 @@ func (a *Allocator) coloringFor(allocatable []RealReg) {
 	for len(coloringStack) != total {
 		// Sort the nodes by the current degree.
 		sort.SliceStable(degreeSortedNodes, func(i, j int) bool {
-			return currentDegrees[degreeSortedNodes[i]] < currentDegrees[degreeSortedNodes[j]]
+			return degreeSortedNodes[i].degree < degreeSortedNodes[j].degree
 		})
 		if wazevoapi.RegAllocLoggingEnabled {
 			fmt.Println("-------------------------------")
@@ -108,8 +106,8 @@ func (a *Allocator) coloringFor(allocatable []RealReg) {
 				}
 			}
 			fmt.Printf("\ncurrentDegrees: ")
-			for n, degree := range currentDegrees {
-				fmt.Printf("v%d:%d ", n.v.ID(), degree)
+			for _, n := range degreeSortedNodes {
+				fmt.Printf("v%d:%d ", n.v.ID(), n.degree)
 			}
 			fmt.Println("")
 		}
@@ -117,7 +115,7 @@ func (a *Allocator) coloringFor(allocatable []RealReg) {
 		var popNum int
 		for i := 0; i < len(degreeSortedNodes); i++ {
 			n := degreeSortedNodes[i]
-			if currentDegrees[n] < numAllocatable {
+			if n.degree < numAllocatable {
 				popNum++
 			} else {
 				break
@@ -144,8 +142,8 @@ func (a *Allocator) coloringFor(allocatable []RealReg) {
 
 		// Update the degrees of the affected nodes.
 		for _, popped := range poppoedNodes {
-			for neighbor := range popped.neighbors {
-				currentDegrees[neighbor]--
+			for _, neighbor := range popped.neighbors {
+				neighbor.degree--
 			}
 		}
 
@@ -161,8 +159,8 @@ func (a *Allocator) coloringFor(allocatable []RealReg) {
 					fmt.Printf("v%d ", n.v.ID())
 				}
 				fmt.Printf("\ncurrentDegrees: ")
-				for n, degree := range currentDegrees {
-					fmt.Printf("v%d:%d ", n.v.ID(), degree)
+				for _, n := range degreeSortedNodes {
+					fmt.Printf("v%d:%d ", n.v.ID(), n.degree)
 				}
 				fmt.Println("")
 			}
@@ -180,8 +178,7 @@ func (a *Allocator) coloringFor(allocatable []RealReg) {
 	}
 
 	// Assign colors.
-	neighborColorsSet := a.realRegSet
-	neighborColors := a.realRegs[:0]
+	neighborColorsSet := &a.realRegSet
 	tail := len(coloringStack) - 1
 	for i := range coloringStack {
 		n := coloringStack[tail-i]
@@ -195,19 +192,10 @@ func (a *Allocator) coloringFor(allocatable []RealReg) {
 		}
 
 		// Gather already used colors.
-		for neighbor := range n.neighbors {
+		for _, neighbor := range n.neighbors {
 			if neighborColor := neighbor.r; neighborColor != RealRegInvalid {
-				neighborColorsSet[neighborColor] = struct{}{}
-				neighborColors = append(neighborColors, neighborColor)
-			}
-		}
-
-		if wazevoapi.RegAllocLoggingEnabled {
-			var s []string
-			for _, r := range neighborColors {
-				s = append(s, a.regInfo.RealRegName(r))
+				neighborColorsSet[neighborColor] = true
 			}
-			fmt.Printf("\tneighborColors: %v\n", strings.Join(s, ","))
 		}
 
 		a.assignColor(n, neighborColorsSet, allocatable)
@@ -217,10 +205,9 @@ func (a *Allocator) coloringFor(allocatable []RealReg) {
 		}
 
 		// Reset the map for the next iteration.
-		for _, c := range neighborColors {
-			delete(neighborColorsSet, c)
+		for j := range neighborColorsSet {
+			neighborColorsSet[j] = false
 		}
-		neighborColors = neighborColors[:0]
 	}
 
 	if wazevoapi.RegAllocValidationEnabled {
@@ -228,7 +215,7 @@ func (a *Allocator) coloringFor(allocatable []RealReg) {
 			if n.r == RealRegInvalid {
 				continue
 			}
-			for neighbor := range n.neighbors {
+			for _, neighbor := range n.neighbors {
 				if n.r == neighbor.r {
 					panic(fmt.Sprintf("BUG color conflict: %s vs %s", n.v, neighbor.v))
 				}
@@ -239,16 +226,13 @@ func (a *Allocator) coloringFor(allocatable []RealReg) {
 	// Reuses the slices for the next coloring.
 	a.nodes1 = degreeSortedNodes[:0]
 	a.nodes2 = coloringStack[:0]
-	a.nodeSet = currentDegrees
-	a.realRegSet = neighborColorsSet
-	a.realRegs = neighborColors[:0]
 }
 
-func (a *Allocator) assignColor(n *node, neighborColorsSet map[RealReg]struct{}, allocatable []RealReg) {
+func (a *Allocator) assignColor(n *node, neighborColorsSet *[256]bool, allocatable []RealReg) {
 	if cfv := n.copyFromVReg; cfv != nil && cfv.r != RealRegInvalid {
 		r := cfv.r
 		if _, ok := a.allocatableSet[r]; ok {
-			if _, ok = neighborColorsSet[r]; !ok {
+			if !neighborColorsSet[r] {
 				n.r = r
 				a.allocatedRegSet[r] = struct{}{}
 				return
@@ -259,7 +243,7 @@ func (a *Allocator) assignColor(n *node, neighborColorsSet map[RealReg]struct{},
 	if ctv := n.copyToVReg; ctv != nil && ctv.r != RealRegInvalid {
 		r := ctv.r
 		if _, ok := a.allocatableSet[r]; ok {
-			if _, ok = neighborColorsSet[r]; !ok {
+			if !neighborColorsSet[r] {
 				n.r = r
 				a.allocatedRegSet[r] = struct{}{}
 				return
@@ -269,7 +253,7 @@ func (a *Allocator) assignColor(n *node, neighborColorsSet map[RealReg]struct{},
 
 	if r := n.copyFromReal; r != RealRegInvalid {
 		if _, ok := a.allocatableSet[r]; ok {
-			if _, ok = neighborColorsSet[r]; !ok {
+			if !neighborColorsSet[r] {
 				n.r = r
 				a.allocatedRegSet[r] = struct{}{}
 				return
@@ -279,7 +263,7 @@ func (a *Allocator) assignColor(n *node, neighborColorsSet map[RealReg]struct{},
 
 	if r := n.copyToReal; r != RealRegInvalid {
 		if _, ok := a.allocatableSet[r]; ok {
-			if _, ok := neighborColorsSet[r]; !ok {
+			if !neighborColorsSet[r] {
 				n.r = r
 				a.allocatedRegSet[r] = struct{}{}
 				return
@@ -289,7 +273,7 @@ func (a *Allocator) assignColor(n *node, neighborColorsSet map[RealReg]struct{},
 
 	if n.r == RealRegInvalid {
 		for _, color := range allocatable {
-			if _, ok := neighborColorsSet[color]; !ok {
+			if !neighborColorsSet[color] {
 				n.r = color
 				a.allocatedRegSet[color] = struct{}{}
 				break
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc.go b/internal/engine/wazevo/backend/regalloc/regalloc.go
index af01d593..a134a3bc 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc.go
@@ -22,7 +22,6 @@ func NewAllocator(allocatableRegs *RegisterInfo) Allocator {
 	a := Allocator{
 		regInfo:         allocatableRegs,
 		nodePool:        wazevoapi.NewPool[node](),
-		realRegSet:      make(map[RealReg]struct{}),
 		nodeSet:         make(map[*node]int),
 		allocatedRegSet: make(map[RealReg]struct{}),
 		phis:            make(map[VReg]struct{}),
@@ -71,7 +70,7 @@ type (
 		phis map[VReg]struct{}
 
 		// Followings are re-used during various places e.g. coloring.
-		realRegSet map[RealReg]struct{}
+		realRegSet [256]bool
 		realRegs   []RealReg
 		nodeSet    map[*node]int
 		nodes1     []*node
@@ -110,12 +109,13 @@ type (
 		// r is the real register assigned to this node. It is either a pre-colored register or a register assigned during allocation.
 		r RealReg
 		// neighbors are the nodes that this node interferes with. Such neighbors have the same RegType as this node.
-		neighbors map[*node]struct{}
+		neighbors []*node
 		// copyFromReal and copyToReal are the real registers that this node copies from/to. During the allocation phase,
 		// we try to assign the same RealReg to copyFromReal and copyToReal so that we can remove the redundant copy.
 		copyFromReal, copyToReal RealReg
 		// copyFromVReg and copyToVReg are the same as above, but for VReg not backed by real registers.
 		copyFromVReg, copyToVReg *node
+		degree                   int
 	}
 
 	// liveRange represents a lifetime of a VReg. Both begin (LiveInterval[0]) and end (LiveInterval[1]) are inclusive.
@@ -535,8 +535,7 @@ func (a *Allocator) allocateNode() (n *node) {
 	n.copyToVReg = nil
 	n.copyFromReal = RealRegInvalid
 	n.copyToReal = RealRegInvalid
-	// TODO: reuse!!
-	n.neighbors = make(map[*node]struct{})
+	n.neighbors = n.neighbors[:0]
 	return
 }
 
@@ -657,7 +656,7 @@ func (n *node) String() string {
 	buf.WriteString("]")
 	// Add neighbors
 	buf.WriteString(" neighbors[")
-	for n := range n.neighbors {
+	for _, n := range n.neighbors {
 		buf.WriteString(fmt.Sprintf("v%v ", n.v.ID()))
 	}
 	buf.WriteString("]")


Test patch:
diff --git a/internal/engine/wazevo/backend/regalloc/coloring_test.go b/internal/engine/wazevo/backend/regalloc/coloring_test.go
index 88b02af8..6393b60c 100644
--- a/internal/engine/wazevo/backend/regalloc/coloring_test.go
+++ b/internal/engine/wazevo/backend/regalloc/coloring_test.go
@@ -58,9 +58,6 @@ func TestAllocator_buildNeighborsByLiveNodes(t *testing.T) {
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			a := NewAllocator(&RegisterInfo{})
-			for i := range tc.lives {
-				tc.lives[i].n.neighbors = map[*node]struct{}{}
-			}
 
 			a.buildNeighborsByLiveNodes(tc.lives)
 
@@ -68,12 +65,25 @@ func TestAllocator_buildNeighborsByLiveNodes(t *testing.T) {
 			for _, edge := range tc.expectedEdges {
 				i1, i2 := edge[0], edge[1]
 				n1, n2 := tc.lives[i1].n, tc.lives[i2].n
-				_, n1n2 := n2.neighbors[n1]
-				_, n2n1 := n1.neighbors[n2]
+
+				var found bool
+				for _, neighbor := range n2.neighbors {
+					if neighbor == n1 {
+						found = true
+						break
+					}
+				}
+				require.True(t, found)
+				found = false
+				for _, neighbor := range n1.neighbors {
+					if neighbor == n2 {
+						found = true
+						break
+					}
+				}
+				require.True(t, found)
 				expectedNeighborCounts[n1]++
 				expectedNeighborCounts[n2]++
-				require.True(t, n1n2)
-				require.True(t, n2n1)
 			}
 			for _, n := range tc.lives {
 				require.Equal(t, expectedNeighborCounts[n.n], len(n.n.neighbors))
@@ -101,8 +111,8 @@ func TestAllocator_collectNodesByRegType(t *testing.T) {
 
 func TestAllocator_coloringFor(t *testing.T) {
 	addEdge := func(n1, n2 *node) {
-		n1.neighbors[n2] = struct{}{}
-		n2.neighbors[n1] = struct{}{}
+		n1.neighbors = append(n1.neighbors, n2)
+		n2.neighbors = append(n2.neighbors, n1)
 	}
 
 	for _, tc := range []struct {
@@ -130,7 +140,7 @@ func TestAllocator_coloringFor(t *testing.T) {
 			allocatable: []RealReg{1, 2},
 			links:       [][]int{{1}, {0}},
 			// Interference, so only one can be assigned a register.
-			expRegs: []RealReg{2, 1},
+			expRegs: []RealReg{1, 2},
 		},
 		{
 			// 0 <- 1 -> 2
@@ -234,7 +244,7 @@ func TestAllocator_assignColor(t *testing.T) {
 		a.allocatableSet = map[RealReg]struct{}{10: {}}
 		n := a.getOrAllocateNode(100)
 		n.copyFromVReg = &node{r: 10}
-		a.assignColor(n, map[RealReg]struct{}{}, nil)
+		a.assignColor(n, &a.realRegSet, nil)
 		require.Equal(t, RealReg(10), n.r)
 		_, ok := a.allocatedRegSet[n.r]
 		require.True(t, ok)
@@ -245,7 +255,8 @@ func TestAllocator_assignColor(t *testing.T) {
 		n := a.getOrAllocateNode(100)
 		n.copyFromVReg = &node{r: 10}
 		n.copyToVReg = &node{r: 20}
-		a.assignColor(n, map[RealReg]struct{}{10: {}}, nil)
+		a.realRegSet[10] = true
+		a.assignColor(n, &a.realRegSet, nil)
 		require.Equal(t, RealReg(20), n.r)
 		_, ok := a.allocatedRegSet[n.r]
 		require.True(t, ok)
@@ -257,7 +268,9 @@ func TestAllocator_assignColor(t *testing.T) {
 		n.copyFromVReg = &node{r: 10}
 		n.copyToVReg = &node{r: 20}
 		n.copyFromReal = 30
-		a.assignColor(n, map[RealReg]struct{}{10: {}, 20: {}}, nil)
+		a.realRegSet[10] = true
+		a.realRegSet[20] = true
+		a.assignColor(n, &a.realRegSet, nil)
 		require.Equal(t, RealReg(30), n.r)
 		_, ok := a.allocatedRegSet[n.r]
 		require.True(t, ok)
@@ -270,7 +283,10 @@ func TestAllocator_assignColor(t *testing.T) {
 		n.copyToVReg = &node{r: 20}
 		n.copyFromReal = 30
 		n.copyToReal = 40
-		a.assignColor(n, map[RealReg]struct{}{10: {}, 20: {}, 30: {}}, nil)
+		a.realRegSet[10] = true
+		a.realRegSet[20] = true
+		a.realRegSet[30] = true
+		a.assignColor(n, &a.realRegSet, nil)
 		require.Equal(t, RealReg(40), n.r)
 		_, ok := a.allocatedRegSet[n.r]
 		require.True(t, ok)
@@ -283,7 +299,11 @@ func TestAllocator_assignColor(t *testing.T) {
 		n.copyToVReg = &node{r: 20}
 		n.copyFromReal = 30
 		n.copyToReal = 40
-		a.assignColor(n, map[RealReg]struct{}{10: {}, 20: {}, 30: {}, 40: {}}, []RealReg{50})
+		a.realRegSet[10] = true
+		a.realRegSet[20] = true
+		a.realRegSet[30] = true
+		a.realRegSet[40] = true
+		a.assignColor(n, &a.realRegSet, []RealReg{50})
 		require.Equal(t, RealReg(50), n.r)
 		_, ok := a.allocatedRegSet[n.r]
 		require.True(t, ok)