Repository: tetratelabs/wazero
Pull request #1787 (2023-10-16 06:22:15)

wazevo(regalloc): reduces map usages
#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi.go b/internal/engine/wazevo/backend/isa/arm64/abi.go
index 0b278d74..21707b78 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi.go
@@ -32,13 +32,16 @@ var regInfo = &regalloc.RegisterInfo{
 			v7, v6, v5, v4, v3, v2, v1, v0,
 		},
 	},
-	CalleeSavedRegisters: map[regalloc.RealReg]struct{}{
-		x19: {}, x20: {}, x21: {}, x22: {}, x23: {}, x24: {}, x25: {}, x26: {}, x28: {},
-		v18: {}, v19: {}, v20: {}, v21: {}, v22: {}, v23: {}, v24: {}, v25: {}, v26: {}, v27: {}, v28: {}, v29: {}, v30: {}, v31: {},
+	CalleeSavedRegisters: [regalloc.RealRegsNumMax]bool{
+		x19: true, x20: true, x21: true, x22: true, x23: true, x24: true, x25: true, x26: true, x28: true,
+		v18: true, v19: true, v20: true, v21: true, v22: true, v23: true, v24: true, v25: true, v26: true,
+		v27: true, v28: true, v29: true, v30: true, v31: true,
 	},
-	CallerSavedRegisters: map[regalloc.RealReg]struct{}{
-		x0: {}, x1: {}, x2: {}, x3: {}, x4: {}, x5: {}, x6: {}, x7: {}, x8: {}, x9: {}, x10: {}, x11: {}, x12: {}, x13: {}, x14: {}, x15: {}, x16: {}, x17: {}, x29: {}, x30: {},
-		v0: {}, v1: {}, v2: {}, v3: {}, v4: {}, v5: {}, v6: {}, v7: {}, v8: {}, v9: {}, v10: {}, v11: {}, v12: {}, v13: {}, v14: {}, v15: {}, v16: {}, v17: {},
+	CallerSavedRegisters: [regalloc.RealRegsNumMax]bool{
+		x0: true, x1: true, x2: true, x3: true, x4: true, x5: true, x6: true, x7: true, x8: true, x9: true, x10: true,
+		x11: true, x12: true, x13: true, x14: true, x15: true, x16: true, x17: true, x29: true, x30: true,
+		v0: true, v1: true, v2: true, v3: true, v4: true, v5: true, v6: true, v7: true, v8: true, v9: true, v10: true,
+		v11: true, v12: true, v13: true, v14: true, v15: true, v16: true, v17: true,
 	},
 	RealRegToVReg: []regalloc.VReg{
 		x0: x0VReg, x1: x1VReg, x2: x2VReg, x3: x3VReg, x4: x4VReg, x5: x5VReg, x6: x6VReg, x7: x7VReg, x8: x8VReg, x9: x9VReg, x10: x10VReg, x11: x11VReg, x12: x12VReg, x13: x13VReg, x14: x14VReg, x15: x15VReg, x16: x16VReg, x17: x17VReg, x18: x18VReg, x19: x19VReg, x20: x20VReg, x21: x21VReg, x22: x22VReg, x23: x23VReg, x24: x24VReg, x25: x25VReg, x26: x26VReg, x27: x27VReg, x28: x28VReg, x29: x29VReg, x30: x30VReg,
diff --git a/internal/engine/wazevo/backend/regalloc/coloring.go b/internal/engine/wazevo/backend/regalloc/coloring.go
index 67d67d8d..f5ecb033 100644
--- a/internal/engine/wazevo/backend/regalloc/coloring.go
+++ b/internal/engine/wazevo/backend/regalloc/coloring.go
@@ -228,13 +228,13 @@ func (a *Allocator) coloringFor(allocatable []RealReg) {
 	a.nodes2 = coloringStack[:0]
 }
 
-func (a *Allocator) assignColor(n *node, neighborColorsSet *[256]bool, allocatable []RealReg) {
+func (a *Allocator) assignColor(n *node, neighborColorsSet *[128]bool, allocatable []RealReg) {
 	if cfv := n.copyFromVReg; cfv != nil && cfv.r != RealRegInvalid {
 		r := cfv.r
-		if _, ok := a.allocatableSet[r]; ok {
+		if a.allocatableSet[r] {
 			if !neighborColorsSet[r] {
 				n.r = r
-				a.allocatedRegSet[r] = struct{}{}
+				a.allocatedRegSet[r] = true
 				return
 			}
 		}
@@ -242,30 +242,30 @@ func (a *Allocator) assignColor(n *node, neighborColorsSet *[256]bool, allocatab
 
 	if ctv := n.copyToVReg; ctv != nil && ctv.r != RealRegInvalid {
 		r := ctv.r
-		if _, ok := a.allocatableSet[r]; ok {
+		if a.allocatableSet[r] {
 			if !neighborColorsSet[r] {
 				n.r = r
-				a.allocatedRegSet[r] = struct{}{}
+				a.allocatedRegSet[r] = true
 				return
 			}
 		}
 	}
 
 	if r := n.copyFromReal; r != RealRegInvalid {
-		if _, ok := a.allocatableSet[r]; ok {
+		if a.allocatableSet[r] {
 			if !neighborColorsSet[r] {
 				n.r = r
-				a.allocatedRegSet[r] = struct{}{}
+				a.allocatedRegSet[r] = true
 				return
 			}
 		}
 	}
 
 	if r := n.copyToReal; r != RealRegInvalid {
-		if _, ok := a.allocatableSet[r]; ok {
+		if a.allocatableSet[r] {
 			if !neighborColorsSet[r] {
 				n.r = r
-				a.allocatedRegSet[r] = struct{}{}
+				a.allocatedRegSet[r] = true
 				return
 			}
 		}
@@ -275,7 +275,7 @@ func (a *Allocator) assignColor(n *node, neighborColorsSet *[256]bool, allocatab
 		for _, color := range allocatable {
 			if !neighborColorsSet[color] {
 				n.r = color
-				a.allocatedRegSet[color] = struct{}{}
+				a.allocatedRegSet[color] = true
 				break
 			}
 		}
diff --git a/internal/engine/wazevo/backend/regalloc/reg.go b/internal/engine/wazevo/backend/regalloc/reg.go
index 995ab188..b05edb6c 100644
--- a/internal/engine/wazevo/backend/regalloc/reg.go
+++ b/internal/engine/wazevo/backend/regalloc/reg.go
@@ -123,3 +123,5 @@ func RegTypeOf(p ssa.Type) RegType {
 		panic("invalid type")
 	}
 }
+
+const RealRegsNumMax = 128
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc.go b/internal/engine/wazevo/backend/regalloc/regalloc.go
index a134a3bc..f39de917 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc.go
@@ -20,21 +20,15 @@ import (
 // NewAllocator returns a new Allocator.
 func NewAllocator(allocatableRegs *RegisterInfo) Allocator {
 	a := Allocator{
-		regInfo:         allocatableRegs,
-		nodePool:        wazevoapi.NewPool[node](),
-		nodeSet:         make(map[*node]int),
-		allocatedRegSet: make(map[RealReg]struct{}),
-		phis:            make(map[VReg]struct{}),
-	}
-	allocatableSet := make(map[RealReg]struct{},
-		len(allocatableRegs.AllocatableRegisters[RegTypeInt])+len(allocatableRegs.AllocatableRegisters[RegTypeFloat]),
-	)
+		regInfo:  allocatableRegs,
+		nodePool: wazevoapi.NewPool[node](),
+		phis:     make(map[VReg]Block),
+	}
 	for _, regs := range allocatableRegs.AllocatableRegisters {
 		for _, r := range regs {
-			allocatableSet[r] = struct{}{}
+			a.allocatableSet[r] = true
 		}
 	}
-	a.allocatableSet = allocatableSet
 	return a
 }
 
@@ -44,8 +38,8 @@ type (
 		// AllocatableRegisters is a 2D array of allocatable RealReg, indexed by regTypeNum and regNum.
 		// The order matters: the first element is the most preferred one when allocating.
 		AllocatableRegisters [RegTypeNum][]RealReg
-		CalleeSavedRegisters map[RealReg]struct{}
-		CallerSavedRegisters map[RealReg]struct{}
+		CalleeSavedRegisters [RealRegsNumMax]bool
+		CallerSavedRegisters [RealRegsNumMax]bool
 		RealRegToVReg        []VReg
 		// RealRegName returns the name of the given RealReg for debugging.
 		RealRegName func(r RealReg) string
@@ -56,9 +50,9 @@ type (
 		// regInfo is static per ABI/ISA, and is initialized by the machine during Machine.PrepareRegisterAllocator.
 		regInfo *RegisterInfo
 		// allocatableSet is a set of allocatable RealReg derived from regInfo. Static per ABI/ISA.
-		allocatableSet map[RealReg]struct{}
+		allocatableSet [RealRegsNumMax]bool
 		// allocatedRegSet is a set of RealReg that are allocated during the allocation phase. This is reset per function.
-		allocatedRegSet          map[RealReg]struct{}
+		allocatedRegSet          [RealRegsNumMax]bool
 		allocatedCalleeSavedRegs []VReg
 		nodePool                 wazevoapi.Pool[node]
 		// vRegIDToNode maps VRegID to the node whose node.v has the VRegID.
@@ -67,12 +61,11 @@ type (
 		vs           []VReg
 		spillHandler spillHandler
 		// phis keeps track of the VRegs that are defined by phi functions.
-		phis map[VReg]struct{}
+		phis map[VReg]Block
 
 		// Followings are re-used during various places e.g. coloring.
-		realRegSet [256]bool
+		realRegSet [RealRegsNumMax]bool
 		realRegs   []RealReg
-		nodeSet    map[*node]int
 		nodes1     []*node
 		nodes2     []*node
 	}
@@ -84,10 +77,6 @@ type (
 		defs     map[VReg]programCounter
 		lastUses map[VReg]programCounter
 		kills    map[VReg]programCounter
-		// phiUses are the set of VRegs that are used in a phi function.
-		phiUses map[VReg]struct{}
-		// phiDefs are the set of VRegs that are defined in a phi function.
-		phiDefs map[VReg]struct{}
 		// Pre-colored real registers can have multiple live ranges in one block.
 		realRegUses map[VReg][]programCounter
 		realRegDefs map[VReg][]programCounter
@@ -141,9 +130,12 @@ func (a *Allocator) DoAllocation(f Function) {
 
 func (a *Allocator) determineCalleeSavedRealRegs(f Function) {
 	a.allocatedCalleeSavedRegs = a.allocatedCalleeSavedRegs[:0]
-	for r := range a.allocatedRegSet {
-		if a.regInfo.isCalleeSaved(r) {
-			a.allocatedCalleeSavedRegs = append(a.allocatedCalleeSavedRegs, a.regInfo.RealRegToVReg[r])
+	for i, allocated := range a.allocatedRegSet {
+		if allocated {
+			r := RealReg(i)
+			if a.regInfo.isCalleeSaved(r) {
+				a.allocatedCalleeSavedRegs = append(a.allocatedCalleeSavedRegs, a.regInfo.RealRegToVReg[r])
+			}
 		}
 	}
 	// In order to make the output deterministic, sort it now.
@@ -174,9 +166,8 @@ func (a *Allocator) livenessAnalysis(f Function) {
 		}
 		// If this is not the entry block, we should define phi nodes, which are not defined by instructions.
 		for _, p := range blk.BlockParams() {
-			info.phiDefs[p] = struct{}{} // Mark this block has the definition of this phi `p`.
-			info.defs[p] = 0             // Earliest definition is at the beginning of the block.
-			a.phis[p] = struct{}{}
+			info.defs[p] = 0 // Earliest definition is at the beginning of the block.
+			a.phis[p] = blk
 		}
 	}
 
@@ -209,14 +200,12 @@ func (a *Allocator) livenessAnalysis(f Function) {
 						// definition to construct live range.
 						info.defs[def] = pos
 					}
-
 					a.vs = append(a.vs, def)
 				}
 			}
 			if instr.IsCopy() {
 				if _, ok := a.phis[dstVR]; ok {
-					// This is the phi move at the end of the block, so srcVR is the phi use (input to the phi function).
-					info.phiUses[srcVR] = struct{}{}
+					info.liveOuts[srcVR] = struct{}{}
 				}
 				a.recordCopyRelation(dstVR, srcVR)
 			}
@@ -229,26 +218,15 @@ func (a *Allocator) livenessAnalysis(f Function) {
 	}
 
 	// Run the Algorithm 9.9. in the book. This will construct blockInfo.liveIns and blockInfo.liveOuts.
+	for phi, blk := range a.phis {
+		a.beginUpAndMarkStack(f, phi, true, blk)
+	}
 	for _, v := range a.vs {
 		if v.IsRealReg() {
 			// Real registers do not need to be tracked in liveOuts and liveIns because they are not allocation targets.
 			panic("BUG")
 		}
-		for blk := f.PostOrderBlockIteratorBegin(); blk != nil; blk = f.PostOrderBlockIteratorNext() {
-			if blk.Preds() == 0 && !blk.Entry() {
-				panic(fmt.Sprintf("block without predecessor must be optimized out by the compiler: %d", blk.ID()))
-			}
-			info := a.blockInfoAt(blk.ID())
-			if _, ok := info.lastUses[v]; !ok {
-				continue
-			}
-			// Phi uses are always live-outs because it is the new value coming out to the successor.
-			if _, ok := info.phiUses[v]; ok {
-				info.liveOuts[v] = struct{}{}
-			}
-			// TODO: we might want to avoid recursion here.
-			a.upAndMarkStack(blk, v, 0)
-		}
+		a.beginUpAndMarkStack(f, v, false, nil)
 	}
 
 	// Now that we finished gathering liveIns, liveOuts, defs, and lastUses, the only thing left is to construct kills.
@@ -268,16 +246,29 @@ func (a *Allocator) livenessAnalysis(f Function) {
 	}
 }
 
+func (a *Allocator) beginUpAndMarkStack(f Function, v VReg, isPhi bool, phiDefinedAt Block) {
+	for blk := f.PostOrderBlockIteratorBegin(); blk != nil; blk = f.PostOrderBlockIteratorNext() {
+		if blk.Preds() == 0 && !blk.Entry() {
+			panic(fmt.Sprintf("block without predecessor must be optimized out by the compiler: %d", blk.ID()))
+		}
+		info := a.blockInfoAt(blk.ID())
+		if _, ok := info.lastUses[v]; !ok {
+			continue
+		}
+		// TODO: we might want to avoid recursion here.
+		a.upAndMarkStack(blk, v, isPhi, phiDefinedAt, 0)
+	}
+}
+
 // upAndMarkStack is the Algorithm 9.10. in the book named Up_and_Mark_Stack(B, v).
 //
 // We recursively call this, so passing `depth` for debugging.
-func (a *Allocator) upAndMarkStack(b Block, v VReg, depth int) {
+func (a *Allocator) upAndMarkStack(b Block, v VReg, isPhi bool, phiDefinedAt Block, depth int) {
 	if wazevoapi.RegAllocLoggingEnabled {
 		fmt.Printf("%supAndMarkStack for %v at %v\n", strings.Repeat("\t", depth), v, b.ID())
 	}
 
 	info := a.blockInfoAt(b.ID())
-	_, isPhi := a.phis[v]
 	if _, ok := info.defs[v]; ok && !isPhi {
 		return // Defined in this block, so no need to go further climbing up.
 	}
@@ -293,7 +284,7 @@ func (a *Allocator) upAndMarkStack(b Block, v VReg, depth int) {
 	info.liveIns[v] = struct{}{}
 
 	// Plus if this is this block has the definition of this phi, we can stop climbing up.
-	if _, ok := info.phiDefs[v]; ok {
+	if b == phiDefinedAt {
 		return
 	}
 
@@ -309,7 +300,7 @@ func (a *Allocator) upAndMarkStack(b Block, v VReg, depth int) {
 			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)
+		a.upAndMarkStack(pred, v, isPhi, phiDefinedAt, depth+1)
 	}
 }
 
@@ -433,7 +424,7 @@ func (a *Allocator) buildLiveRangesForReals(blkID int, info *blockInfo) {
 	a.vs = a.vs[:0]
 	for v := range us {
 		// Non allocation target registers are not needed here.
-		if _, ok := a.allocatableSet[v.RealReg()]; !ok {
+		if !a.allocatableSet[v.RealReg()] {
 			continue
 		}
 		a.vs = append(a.vs, v)
@@ -474,24 +465,13 @@ func (a *Allocator) Reset() {
 	for i := range a.vRegIDToNode {
 		a.vRegIDToNode[i] = nil
 	}
-	rr := a.realRegs[:0]
-	for r := range a.allocatedRegSet {
-		rr = append(rr, r)
-	}
-	for _, r := range rr {
-		delete(a.allocatedRegSet, r)
+	for i := range a.allocatedRegSet {
+		a.allocatedRegSet[i] = false
 	}
 
-	a.nodes1 = a.nodes1[:0]
-	for n := range a.nodeSet {
-		a.nodes1 = append(a.nodes1, n)
-	}
-	for _, n := range a.nodes1 {
-		delete(a.nodeSet, n)
-	}
 	a.nodes1 = a.nodes1[:0]
 	a.nodes2 = a.nodes2[:0]
-	a.realRegs = rr[:0]
+	a.realRegs = a.realRegs[:0]
 	resetMap(a, a.phis)
 	a.vs = a.vs[:0]
 }
@@ -586,16 +566,6 @@ func (a *Allocator) initBlockInfo(i *blockInfo) {
 	} else {
 		resetMap(a, i.realRegDefs)
 	}
-	if i.phiDefs == nil {
-		i.phiDefs = make(map[VReg]struct{})
-	} else {
-		resetMap(a, i.phiDefs)
-	}
-	if i.phiUses == nil {
-		i.phiUses = make(map[VReg]struct{})
-	} else {
-		resetMap(a, i.phiUses)
-	}
 }
 
 func (i *blockInfo) addRealRegUsage(v VReg, pc programCounter) {
@@ -608,7 +578,7 @@ func (i *blockInfo) addRealRegUsage(v VReg, pc programCounter) {
 	i.realRegUses[v] = append(i.realRegUses[v], pc)
 }
 
-// String implements fmt.Stringer for debugging.
+// Format is for debugging.
 func (i *blockInfo) Format(ri *RegisterInfo) string {
 	var buf strings.Builder
 	buf.WriteString("\tliveOuts: ")
@@ -674,13 +644,11 @@ func (l *liveRange) intersects(other *liveRange) bool {
 }
 
 func (r *RegisterInfo) isCalleeSaved(reg RealReg) bool {
-	_, ok := r.CalleeSavedRegisters[reg]
-	return ok
+	return r.CalleeSavedRegisters[reg]
 }
 
 func (r *RegisterInfo) isCallerSaved(reg RealReg) bool {
-	_, ok := r.CallerSavedRegisters[reg]
-	return ok
+	return r.CallerSavedRegisters[reg]
 }
 
 // String implements fmt.Stringer for debugging.


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go b/internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go
index e2954c41..23fc2bf0 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go
@@ -13,8 +13,10 @@ import (
 
 func Test_calleeSavedRegistersSorted(t *testing.T) {
 	var exp []regalloc.VReg
-	for r := range regInfo.CalleeSavedRegisters {
-		exp = append(exp, regInfo.RealRegToVReg[r])
+	for i, r := range regInfo.CalleeSavedRegisters {
+		if r {
+			exp = append(exp, regInfo.RealRegToVReg[i])
+		}
 	}
 	sort.Slice(exp, func(i, j int) bool {
 		return exp[i].RealReg() < exp[j].RealReg()
diff --git a/internal/engine/wazevo/backend/regalloc/assign_test.go b/internal/engine/wazevo/backend/regalloc/assign_test.go
index a93d0976..543c6d67 100644
--- a/internal/engine/wazevo/backend/regalloc/assign_test.go
+++ b/internal/engine/wazevo/backend/regalloc/assign_test.go
@@ -8,7 +8,7 @@ import (
 
 func TestAllocator_assignRegistersPerInstr(t *testing.T) {
 	t.Run("call", func(t *testing.T) {
-		a := NewAllocator(&RegisterInfo{CallerSavedRegisters: map[RealReg]struct{}{1: {}, 3: {}}})
+		a := NewAllocator(&RegisterInfo{CallerSavedRegisters: [RealRegsNumMax]bool{1: true, 3: true}})
 		pc := programCounter(5)
 		liveNodes := []liveNodeInBlock{
 			{n: &node{r: 1, v: 0xa, ranges: []liveRange{{begin: 5, end: 20}}}},
@@ -26,12 +26,12 @@ func TestAllocator_assignRegistersPerInstr(t *testing.T) {
 		require.Equal(t, 2, len(f.afters))
 	})
 	t.Run("call_indirect/func_ptr not spilled", func(t *testing.T) {
-		a := NewAllocator(&RegisterInfo{CallerSavedRegisters: map[RealReg]struct{}{1: {}, 3: {}, 0xff: {}}})
+		a := NewAllocator(&RegisterInfo{CallerSavedRegisters: [RealRegsNumMax]bool{1: true, 3: true, 0xf: true}})
 		pc := programCounter(5)
 		functionPtrVRegID := 0x0
 		functionPtrVReg := VReg(functionPtrVRegID).SetRegType(RegTypeInt)
 		functionPtrLiveNode := liveNodeInBlock{
-			n: &node{r: 0xff, v: functionPtrVReg, ranges: []liveRange{{begin: 4, end: pc /* killed at this indirect call. */}}},
+			n: &node{r: 0xf, v: functionPtrVReg, ranges: []liveRange{{begin: 4, end: pc /* killed at this indirect call. */}}},
 		}
 		liveNodes := []liveNodeInBlock{
 			functionPtrLiveNode, // Function pointer, used at this PC. not save target.
@@ -48,12 +48,12 @@ func TestAllocator_assignRegistersPerInstr(t *testing.T) {
 		require.Equal(t, 2, len(f.befores))
 		require.Equal(t, 2, len(f.afters))
 		require.True(t, callInd.uses[0].IsRealReg())
-		require.Equal(t, functionPtrVReg.SetRealReg(0xff), callInd.uses[0])
+		require.Equal(t, functionPtrVReg.SetRealReg(0xf), callInd.uses[0])
 	})
 	t.Run("call_indirect/func_ptr spilled", func(t *testing.T) {
 		a := NewAllocator(&RegisterInfo{
-			CallerSavedRegisters: map[RealReg]struct{}{1: {}, 3: {}, 0xbb: {}},
-			AllocatableRegisters: [3][]RealReg{RegTypeInt: {0xff, 0xbb}},
+			CallerSavedRegisters: [RealRegsNumMax]bool{1: true, 3: true, 0xb: true},
+			AllocatableRegisters: [3][]RealReg{RegTypeInt: {0xf, 0xb}},
 		})
 		pc := programCounter(5)
 		functionPtrVRegID := 0x0
@@ -74,9 +74,9 @@ func TestAllocator_assignRegistersPerInstr(t *testing.T) {
 		require.Equal(t, 3, len(f.befores))
 		require.Equal(t, 2, len(f.afters))
 		require.Equal(t, callInd, f.befores[2].instr)
-		require.Equal(t, functionPtrVReg.SetRealReg(0xbb), f.befores[2].v)
+		require.Equal(t, functionPtrVReg.SetRealReg(0xb), f.befores[2].v)
 		require.True(t, callInd.uses[0].IsRealReg())
-		require.Equal(t, functionPtrVReg.SetRealReg(0xbb), callInd.uses[0])
+		require.Equal(t, functionPtrVReg.SetRealReg(0xb), callInd.uses[0])
 	})
 	t.Run("no spills", func(t *testing.T) {
 		r := FromRealReg(1, RegTypeInt)
@@ -94,9 +94,6 @@ func TestAllocator_assignRegistersPerInstr(t *testing.T) {
 		require.Equal(t, []VReg{r, VReg(2).SetRealReg(22), VReg(3).SetRealReg(33)}, instr.uses)
 		require.Equal(t, []VReg{VReg(4).SetRealReg(44)}, instr.defs)
 	})
-	t.Run("spills", func(t *testing.T) {
-		t.Skip("TODO")
-	})
 }
 
 func TestAllocator_activeNonRealVRegsAt(t *testing.T) {
@@ -244,7 +241,7 @@ func TestAllocator_handleSpills(t *testing.T) {
 		}
 
 		a := NewAllocator(&RegisterInfo{
-			AllocatableRegisters: [3][]RealReg{RegTypeInt: {0xb, 0xc, 0xff /* free */}},
+			AllocatableRegisters: [3][]RealReg{RegTypeInt: {0xb, 0xc, 0xf /* free */}},
 		})
 
 		f := newMockFunction(newMockBlock(0).entry())
@@ -253,11 +250,11 @@ func TestAllocator_handleSpills(t *testing.T) {
 		instr := newMockInstr().def(vr)
 		a.handleSpills(f, pc, instr, liveNodes, nil, vr)
 		require.Equal(t, 1, len(instr.defs))
-		require.Equal(t, RealReg(0xff), instr.defs[0].RealReg())
+		require.Equal(t, RealReg(0xf), instr.defs[0].RealReg())
 
 		require.Equal(t, 0, len(f.befores))
 		require.Equal(t, 1, len(f.afters))
-		requireInsertedInst(t, f, false, 0, instr, false, vr.SetRealReg(0xff))
+		requireInsertedInst(t, f, false, 0, instr, false, vr.SetRealReg(0xf))
 	})
 
 	t.Run("uses and def / not evicted / def same type", func(t *testing.T) {
@@ -269,8 +266,8 @@ func TestAllocator_handleSpills(t *testing.T) {
 
 		a := NewAllocator(&RegisterInfo{
 			AllocatableRegisters: [3][]RealReg{
-				RegTypeInt:   {0xb, 0xc, 0xaa /* free */},
-				RegTypeFloat: {0xbb /* free */},
+				RegTypeInt:   {0xb, 0xc, 0xa /* free */},
+				RegTypeFloat: {0xf /* free */},
 			},
 		})
 
@@ -281,28 +278,28 @@ func TestAllocator_handleSpills(t *testing.T) {
 		d1 := VReg(104).SetRegType(RegTypeFloat)
 		instr := newMockInstr().use(u1, u2, u3).def(d1)
 		a.handleSpills(f, pc, instr, liveNodes, []VReg{u1, u3}, d1)
-		require.Equal(t, []VReg{u1.SetRealReg(0xaa), u2, u3.SetRealReg(0xbb)}, instr.uses)
-		require.Equal(t, []VReg{d1.SetRealReg(0xbb)}, instr.defs)
+		require.Equal(t, []VReg{u1.SetRealReg(0xa), u2, u3.SetRealReg(0xf)}, instr.uses)
+		require.Equal(t, []VReg{d1.SetRealReg(0xf)}, instr.defs)
 
 		require.Equal(t, 2, len(f.befores))
-		requireInsertedInst(t, f, true, 0, instr, true, u1.SetRealReg(0xaa))
-		requireInsertedInst(t, f, true, 1, instr, true, u3.SetRealReg(0xbb))
+		requireInsertedInst(t, f, true, 0, instr, true, u1.SetRealReg(0xa))
+		requireInsertedInst(t, f, true, 1, instr, true, u3.SetRealReg(0xf))
 
 		require.Equal(t, 1, len(f.afters))
-		requireInsertedInst(t, f, false, 0, instr, false, d1.SetRealReg(0xbb))
+		requireInsertedInst(t, f, false, 0, instr, false, d1.SetRealReg(0xf))
 	})
 
 	t.Run("uses and def / not evicted / def different type", func(t *testing.T) {
 		const pc = 5
 		liveNodes := []liveNodeInBlock{
 			{n: &node{r: RealReg(0xb), v: 0xa, ranges: []liveRange{{begin: pc, end: 20}}}},
-			{n: &node{r: RealReg(0xff), v: 0xb, ranges: []liveRange{{begin: pc, end: 20}}}},
+			{n: &node{r: RealReg(0xf), v: 0xb, ranges: []liveRange{{begin: pc, end: 20}}}},
 		}
 
 		a := NewAllocator(&RegisterInfo{
 			AllocatableRegisters: [3][]RealReg{
-				RegTypeInt:   {0xb, 0xaa /* free */},
-				RegTypeFloat: {0xff},
+				RegTypeInt:   {0xb, 0xa /* free */},
+				RegTypeFloat: {0xf},
 			},
 		})
 
@@ -311,15 +308,15 @@ func TestAllocator_handleSpills(t *testing.T) {
 		d1 := VReg(104).SetRegType(RegTypeFloat)
 		instr := newMockInstr().use(u1).def(d1)
 		a.handleSpills(f, pc, instr, liveNodes, []VReg{u1}, d1)
-		require.Equal(t, []VReg{u1.SetRealReg(0xaa)}, instr.uses)
-		require.Equal(t, []VReg{d1.SetRealReg(0xff)}, instr.defs)
+		require.Equal(t, []VReg{u1.SetRealReg(0xa)}, instr.uses)
+		require.Equal(t, []VReg{d1.SetRealReg(0xf)}, instr.defs)
 
 		require.Equal(t, 2, len(f.befores))
-		requireInsertedInst(t, f, true, 0, instr, true, u1.SetRealReg(0xaa))
-		requireInsertedInst(t, f, true, 1, instr, false, liveNodes[1].n.v.SetRealReg(0xff))
+		requireInsertedInst(t, f, true, 0, instr, true, u1.SetRealReg(0xa))
+		requireInsertedInst(t, f, true, 1, instr, false, liveNodes[1].n.v.SetRealReg(0xf))
 		require.Equal(t, 2, len(f.afters))
-		requireInsertedInst(t, f, false, 0, instr, true, liveNodes[1].n.v.SetRealReg(0xff))
-		requireInsertedInst(t, f, false, 1, instr, false, d1.SetRealReg(0xff))
+		requireInsertedInst(t, f, false, 0, instr, true, liveNodes[1].n.v.SetRealReg(0xf))
+		requireInsertedInst(t, f, false, 1, instr, false, d1.SetRealReg(0xf))
 	})
 
 	t.Run("uses and def / evicted / def different type", func(t *testing.T) {
@@ -327,13 +324,13 @@ func TestAllocator_handleSpills(t *testing.T) {
 		liveNodes := []liveNodeInBlock{
 			{n: &node{r: RealReg(0xb), v: 0xa, ranges: []liveRange{{begin: pc, end: 20}}}},
 			{n: &node{r: RealReg(0xc), v: 0xa, ranges: []liveRange{{begin: pc, end: 20}}}},
-			{n: &node{r: RealReg(0xff), v: 0xb, ranges: []liveRange{{begin: pc, end: 20}}}},
+			{n: &node{r: RealReg(0xf), v: 0xb, ranges: []liveRange{{begin: pc, end: 20}}}},
 		}
 
 		a := NewAllocator(&RegisterInfo{
 			AllocatableRegisters: [3][]RealReg{
 				RegTypeInt:   {0xb, 0xc},
-				RegTypeFloat: {0xff},
+				RegTypeFloat: {0xf},
 			},
 		})
 
@@ -343,16 +340,16 @@ func TestAllocator_handleSpills(t *testing.T) {
 		instr := newMockInstr().use(u1).def(d1)
 		a.handleSpills(f, pc, instr, liveNodes, []VReg{u1}, d1)
 		require.Equal(t, []VReg{u1.SetRealReg(0xb)}, instr.uses)
-		require.Equal(t, []VReg{d1.SetRealReg(0xff)}, instr.defs)
+		require.Equal(t, []VReg{d1.SetRealReg(0xf)}, instr.defs)
 
 		require.Equal(t, 3, len(f.befores))
 		requireInsertedInst(t, f, true, 0, instr, false, liveNodes[0].n.v.SetRealReg(0xb))
 		requireInsertedInst(t, f, true, 1, instr, true, u1.SetRealReg(0xb))
-		requireInsertedInst(t, f, true, 2, instr, false, liveNodes[2].n.v.SetRealReg(0xff))
+		requireInsertedInst(t, f, true, 2, instr, false, liveNodes[2].n.v.SetRealReg(0xf))
 		require.Equal(t, 3, len(f.afters))
 		requireInsertedInst(t, f, false, 0, instr, true, liveNodes[0].n.v.SetRealReg(0xb))
-		requireInsertedInst(t, f, false, 1, instr, true, liveNodes[2].n.v.SetRealReg(0xff))
-		requireInsertedInst(t, f, false, 2, instr, false, d1.SetRealReg(0xff))
+		requireInsertedInst(t, f, false, 1, instr, true, liveNodes[2].n.v.SetRealReg(0xf))
+		requireInsertedInst(t, f, false, 2, instr, false, d1.SetRealReg(0xf))
 	})
 }
 
@@ -367,11 +364,11 @@ func TestAllocator_collectActiveNodesAt(t *testing.T) {
 	t.Run("lives", func(t *testing.T) {
 		const pc = 5
 		liveNodes := []liveNodeInBlock{
-			{n: &node{r: RealReg(0xff), v: 0xa, ranges: []liveRange{{begin: 0, end: pc - 1}}}},
+			{n: &node{r: RealReg(0xf), v: 0xa, ranges: []liveRange{{begin: 0, end: pc - 1}}}},
 			{n: &node{r: RealReg(0x1), v: 0xa, ranges: []liveRange{{begin: pc, end: 20}}}},
 			{n: &node{r: RealReg(0x2), v: 0xa, ranges: []liveRange{{begin: pc, end: 20}}}},
 			{n: &node{r: RealReg(0x4), v: 0xb, ranges: []liveRange{{begin: pc, end: 20}}}},
-			{n: &node{r: RealReg(0xff), v: 0xa, ranges: []liveRange{{begin: 1000, end: 2000000}}}},
+			{n: &node{r: RealReg(0xf), v: 0xa, ranges: []liveRange{{begin: 1000, end: 2000000}}}},
 		}
 
 		a := NewAllocator(&RegisterInfo{})
diff --git a/internal/engine/wazevo/backend/regalloc/coloring_test.go b/internal/engine/wazevo/backend/regalloc/coloring_test.go
index 6393b60c..0367c971 100644
--- a/internal/engine/wazevo/backend/regalloc/coloring_test.go
+++ b/internal/engine/wazevo/backend/regalloc/coloring_test.go
@@ -241,29 +241,32 @@ func TestAllocator_coloringFor(t *testing.T) {
 func TestAllocator_assignColor(t *testing.T) {
 	t.Run("copyFromVReg", func(t *testing.T) {
 		a := NewAllocator(&RegisterInfo{})
-		a.allocatableSet = map[RealReg]struct{}{10: {}}
+		a.allocatableSet[10] = true
 		n := a.getOrAllocateNode(100)
 		n.copyFromVReg = &node{r: 10}
 		a.assignColor(n, &a.realRegSet, nil)
 		require.Equal(t, RealReg(10), n.r)
-		_, ok := a.allocatedRegSet[n.r]
+		ok := a.allocatedRegSet[n.r]
 		require.True(t, ok)
 	})
 	t.Run("copyToVReg", func(t *testing.T) {
 		a := NewAllocator(&RegisterInfo{})
-		a.allocatableSet = map[RealReg]struct{}{10: {}, 20: {}}
+		a.allocatableSet[10] = true
+		a.allocatableSet[20] = true
 		n := a.getOrAllocateNode(100)
 		n.copyFromVReg = &node{r: 10}
 		n.copyToVReg = &node{r: 20}
 		a.realRegSet[10] = true
 		a.assignColor(n, &a.realRegSet, nil)
 		require.Equal(t, RealReg(20), n.r)
-		_, ok := a.allocatedRegSet[n.r]
+		ok := a.allocatedRegSet[n.r]
 		require.True(t, ok)
 	})
 	t.Run("copyFromReal", func(t *testing.T) {
 		a := NewAllocator(&RegisterInfo{})
-		a.allocatableSet = map[RealReg]struct{}{10: {}, 20: {}, 30: {}}
+		a.allocatableSet[10] = true
+		a.allocatableSet[20] = true
+		a.allocatableSet[30] = true
 		n := a.getOrAllocateNode(100)
 		n.copyFromVReg = &node{r: 10}
 		n.copyToVReg = &node{r: 20}
@@ -272,12 +275,15 @@ func TestAllocator_assignColor(t *testing.T) {
 		a.realRegSet[20] = true
 		a.assignColor(n, &a.realRegSet, nil)
 		require.Equal(t, RealReg(30), n.r)
-		_, ok := a.allocatedRegSet[n.r]
+		ok := a.allocatedRegSet[n.r]
 		require.True(t, ok)
 	})
 	t.Run("copyToReal", func(t *testing.T) {
 		a := NewAllocator(&RegisterInfo{})
-		a.allocatableSet = map[RealReg]struct{}{10: {}, 20: {}, 30: {}, 40: {}}
+		a.allocatableSet[10] = true
+		a.allocatableSet[20] = true
+		a.allocatableSet[30] = true
+		a.allocatableSet[40] = true
 		n := a.getOrAllocateNode(100)
 		n.copyFromVReg = &node{r: 10}
 		n.copyToVReg = &node{r: 20}
@@ -288,12 +294,16 @@ func TestAllocator_assignColor(t *testing.T) {
 		a.realRegSet[30] = true
 		a.assignColor(n, &a.realRegSet, nil)
 		require.Equal(t, RealReg(40), n.r)
-		_, ok := a.allocatedRegSet[n.r]
+		ok := a.allocatedRegSet[n.r]
 		require.True(t, ok)
 	})
 	t.Run("from allocatable sets", func(t *testing.T) {
 		a := NewAllocator(&RegisterInfo{})
-		a.allocatableSet = map[RealReg]struct{}{10: {}, 20: {}, 30: {}, 40: {}, 50: {}}
+		a.allocatableSet[10] = true
+		a.allocatableSet[20] = true
+		a.allocatableSet[30] = true
+		a.allocatableSet[40] = true
+		a.allocatableSet[50] = true
 		n := a.getOrAllocateNode(100)
 		n.copyFromVReg = &node{r: 10}
 		n.copyToVReg = &node{r: 20}
@@ -305,7 +315,7 @@ func TestAllocator_assignColor(t *testing.T) {
 		a.realRegSet[40] = true
 		a.assignColor(n, &a.realRegSet, []RealReg{50})
 		require.Equal(t, RealReg(50), n.r)
-		_, ok := a.allocatedRegSet[n.r]
+		ok := a.allocatedRegSet[n.r]
 		require.True(t, ok)
 	})
 }
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc_test.go b/internal/engine/wazevo/backend/regalloc/regalloc_test.go
index 4cefa489..aa8ff93e 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc_test.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc_test.go
@@ -611,7 +611,9 @@ func TestAllocator_buildLiveRangesForReals(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			initMapInInfo(tc.info)
 			a := NewAllocator(&RegisterInfo{})
-			a.allocatableSet = tc.allocatableRealRegs
+			for r := range tc.allocatableRealRegs {
+				a.allocatableSet[r] = true
+			}
 			a.buildLiveRangesForReals(blockID, tc.info)
 
 			actual := map[VReg][]liveRange{}
@@ -705,12 +707,6 @@ func initMapInInfo(info *blockInfo) {
 	if info.realRegDefs == nil {
 		info.realRegDefs = make(map[VReg][]programCounter)
 	}
-	if info.phiDefs == nil {
-		info.phiDefs = make(map[VReg]struct{})
-	}
-	if info.phiUses == nil {
-		info.phiUses = make(map[VReg]struct{})
-	}
 }
 
 func TestNode_assignedRealReg(t *testing.T) {
diff --git a/internal/integration_test/engine/memleak_test.go b/internal/integration_test/engine/memleak_test.go
index a807a318..d1d98abc 100644
--- a/internal/integration_test/engine/memleak_test.go
+++ b/internal/integration_test/engine/memleak_test.go
@@ -32,7 +32,7 @@ func TestMemoryLeak(t *testing.T) {
 		}
 
 		t.Run(tc.name, func(t *testing.T) {
-			duration := 15 * time.Second
+			duration := 5 * time.Second
 			t.Logf("running memory leak test for %s", duration)
 
 			ctx, cancel := context.WithTimeout(context.Background(), duration)