Repository: tetratelabs/wazero
Pull request #1802 (2023-10-19 07:35:08)

wazevo: add VRegTable data structure
This PR explores changing the data structure used for _last uses_ in block info.

The change takes the compilation of Python 3.11 with Wazevo from 40s to 28s.

I added a bitset to perform fast checks of virtual registers on the hot code path of `beginUpAndMarkStack`, which yielded most of the gains (changing from a map to a simple array only gave a ~17% improvement).

Note that even with this change, the hottest code path is still `beginUpAndMarkStack`, but we're starting to see a more balanced time spent in various parts:

<img width="1528" alt="image" src="https://github.com/tetratelabs/wazero/assets/865510/757af91f-9f74-4fc9-b78f-9f4a3d74278d">

Based on what I'm seeing in the current profile, we're starting to see the impact of indirect function calls through method interfaces and such (e.g., 7%+ spent in calling `Block.ID`). There is likely some room for reducing the number of those calls on the hottest code paths to amortize their cost.

My only concern is that I don't understand the lifecycle of `blockInfo` values nor how virtual register IDs are allocated, which might increase the memory footprint. We can do more investigation as part of this PR, or we might look into this in a follow-up.

Let me know what you think!

mathetake: btw can you do the same thing to liveIns, liveOuts and defs? Does that help in the performance? 
achille-roussel: That was going to be be next thing I wanted to try 👍

Patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi.go b/internal/engine/wazevo/backend/isa/arm64/abi.go
index 550b8dcb..638c8239 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi.go
@@ -13,7 +13,7 @@ import (
 const xArgRetRegMax, vArgRetRegMax = x7, v7 // x0-x7 & v0-v7.
 
 var regInfo = &regalloc.RegisterInfo{
-	AllocatableRegisters: [regalloc.RegTypeNum][]regalloc.RealReg{
+	AllocatableRegisters: [regalloc.NumRegType][]regalloc.RealReg{
 		// We don't allocate:
 		// - x18: Reserved by the macOS: https://developer.apple.com/documentation/xcode/writing-arm64-code-for-apple-platforms#Respect-the-purpose-of-specific-CPU-registers
 		// - x28: Reserved by Go runtime.
diff --git a/internal/engine/wazevo/backend/regalloc/reg.go b/internal/engine/wazevo/backend/regalloc/reg.go
index 2c5f55b2..1e8cc71c 100644
--- a/internal/engine/wazevo/backend/regalloc/reg.go
+++ b/internal/engine/wazevo/backend/regalloc/reg.go
@@ -2,6 +2,7 @@ package regalloc
 
 import (
 	"fmt"
+	"math/bits"
 
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
 )
@@ -13,6 +14,10 @@ type VReg uint64
 // VRegID is the lower 32bit of VReg, which is the pure identifier of VReg without RealReg info.
 type VRegID uint32
 
+const (
+	MaxVRegID = ^VRegID(0)
+)
+
 // RealReg returns the RealReg of this VReg.
 func (v VReg) RealReg() RealReg {
 	return RealReg(v >> 32)
@@ -58,6 +63,130 @@ func (v VReg) Valid() bool {
 	return v.ID() != vRegIDInvalid && v.RegType() != RegTypeInvalid
 }
 
+// VRegIDMinSet is used to collect the minimum ID by type in a collection of
+// virtual registers.
+//
+// We store the min values + 1 so the zero-value of the VRegIDMinSet is valid.
+type VRegIDMinSet [NumRegType]VRegID
+
+func (mins *VRegIDMinSet) Min(t RegType) VRegID {
+	return mins[t] - 1
+}
+
+func (mins *VRegIDMinSet) Observe(v VReg) {
+	if rt, id := v.RegType(), v.ID(); id < (mins[rt] - 1) {
+		mins[rt] = id + 1
+	}
+}
+
+// VRegTable is a data structure designed for fast association of program
+// counters to virtual registers.
+type VRegTable [NumRegType]VRegTypeTable
+
+func (t *VRegTable) Contains(v VReg) bool {
+	return t[v.RegType()].Contains(v.ID())
+}
+
+func (t *VRegTable) Lookup(v VReg) programCounter {
+	return t[v.RegType()].Lookup(v.ID())
+}
+
+func (t *VRegTable) Insert(v VReg, p programCounter) {
+	if v.IsRealReg() {
+		panic("BUG: cannot insert real registers in virtual register table")
+	}
+	t[v.RegType()].Insert(v.ID(), p)
+}
+
+func (t *VRegTable) Range(f func(VReg, programCounter)) {
+	for i := range t {
+		t[i].Range(func(id VRegID, p programCounter) {
+			f(VReg(id).SetRegType(RegType(i)), p)
+		})
+	}
+}
+
+func (t *VRegTable) Reset(minVRegIDs VRegIDMinSet) {
+	for i := range t {
+		t[i].Reset(minVRegIDs.Min(RegType(i)))
+	}
+}
+
+// VRegTypeTable implements a table for virtual registers of a specific type.
+//
+// The parent VRegTable uses 4 instances of this type to maintain a table for
+// each virtual register type.
+//
+// The virtual register type table uses a bitset to accelerate checking whether
+// a virtual register exists in the table. The bitset also helps identify which
+// slots in the program counter array are used, since the table is sparse and
+// some locations may not be occupied and will have the value zero even though
+// zero is a valid program counter value.
+type VRegTypeTable struct {
+	min VRegID
+	set bitset
+	pcs []programCounter
+}
+
+func (t *VRegTypeTable) Contains(id VRegID) bool {
+	return t.set.has(uint(id - t.min))
+}
+
+func (t *VRegTypeTable) Lookup(id VRegID) programCounter {
+	if id := int(id - t.min); t.set.has(uint(id)) {
+		return t.pcs[id]
+	}
+	return -1
+}
+
+func (t *VRegTypeTable) Insert(id VRegID, p programCounter) {
+	if p < 0 {
+		panic("BUG: cannot insert negative program counter in virtual register table")
+	}
+	i := int(id - t.min)
+	if len(t.pcs) <= i {
+		t.pcs = append(t.pcs, make([]programCounter, (i+1)-len(t.pcs))...)
+	}
+	t.set.set(uint(i))
+	t.pcs[i] = p
+}
+
+func (t *VRegTypeTable) Range(f func(VRegID, programCounter)) {
+	t.set.scan(func(i uint) { f(VRegID(i)+t.min, t.pcs[i]) })
+}
+
+func (t *VRegTypeTable) Reset(minVRegID VRegID) {
+	t.min = minVRegID
+	t.set = nil
+	t.pcs = nil
+}
+
+type bitset []uint64
+
+func (b bitset) scan(f func(uint)) {
+	for i, v := range b {
+		for j := uint(i * 64); v != 0; j++ {
+			n := uint(bits.TrailingZeros64(v))
+			j += n
+			v >>= (n + 1)
+			f(j)
+		}
+	}
+}
+
+func (b bitset) has(i uint) bool {
+	index, shift := i/64, i%64
+	return index < uint(len(b)) && ((b[index] & (1 << shift)) != 0)
+}
+
+func (b *bitset) set(i uint) {
+	index, shift := i/64, i%64
+	if index >= uint(len(*b)) {
+		*b = append(*b, make([]uint64, (index+1)-uint(len(*b)))...)
+	}
+	(*b)[index] |= 1 << shift
+}
+
 // RealReg represents a physical register.
 type RealReg byte
 
@@ -95,7 +224,7 @@ const (
 	RegTypeInvalid RegType = iota
 	RegTypeInt
 	RegTypeFloat
-	RegTypeNum
+	NumRegType
 )
 
 // String implements fmt.Stringer.
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc.go b/internal/engine/wazevo/backend/regalloc/regalloc.go
index 249161e9..b336ed08 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc.go
@@ -36,7 +36,7 @@ type (
 	RegisterInfo struct {
 		// 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
+		AllocatableRegisters [NumRegType][]RealReg
 		CalleeSavedRegisters [RealRegsNumMax]bool
 		CallerSavedRegisters [RealRegsNumMax]bool
 		RealRegToVReg        []VReg
@@ -78,7 +78,7 @@ type (
 		liveOuts map[VReg]struct{}
 		liveIns  map[VReg]struct{}
 		defs     map[VReg]programCounter
-		lastUses map[VReg]programCounter
+		lastUses VRegTable
 		kills    map[VReg]programCounter
 		// Pre-colored real registers can have multiple live ranges in one block.
 		realRegUses [vRegIDReservedForRealNum][]programCounter
@@ -172,6 +172,19 @@ func (a *Allocator) livenessAnalysis(f Function) {
 	for blk := f.PostOrderBlockIteratorBegin(); blk != nil; blk = f.PostOrderBlockIteratorNext() {
 		info := a.blockInfoAt(blk.ID())
 
+		// We have to do a first pass to find the lowest VRegID in the block;
+		// this is used to reduce memory utilization in the VRegTable, which
+		// can avoid allocating memory for registers zero to minVRegID-1.
+		minVRegID := VRegIDMinSet{}
+		for instr := blk.InstrIteratorBegin(); instr != nil; instr = blk.InstrIteratorNext() {
+			for _, use := range instr.Uses() {
+				if !use.IsRealReg() {
+					minVRegID.Observe(use)
+				}
+			}
+		}
+		info.lastUses.Reset(minVRegID)
+
 		var pc programCounter
 		for instr := blk.InstrIteratorBegin(); instr != nil; instr = blk.InstrIteratorNext() {
 			var srcVR, dstVR VReg
@@ -181,7 +194,7 @@ func (a *Allocator) livenessAnalysis(f Function) {
 				if use.IsRealReg() {
 					info.addRealRegUsage(use, pos)
 				} else {
-					info.lastUses[use] = pos
+					info.lastUses.Insert(use, pos)
 				}
 			}
 			for _, def := range instr.Defs() {
@@ -209,7 +222,6 @@ func (a *Allocator) livenessAnalysis(f Function) {
 			}
 			pc += pcStride
 		}
-
 		if wazevoapi.RegAllocLoggingEnabled {
 			fmt.Printf("prepared block info for block[%d]:\n%s\n\n", blk.ID(), info.Format(a.regInfo))
 		}
@@ -231,13 +243,13 @@ func (a *Allocator) livenessAnalysis(f Function) {
 	// Now that we finished gathering liveIns, liveOuts, defs, and lastUses, the only thing left is to construct kills.
 	for blk := f.PostOrderBlockIteratorBegin(); blk != nil; blk = f.PostOrderBlockIteratorNext() { // Order doesn't matter.
 		info := a.blockInfoAt(blk.ID())
-		lastUses, outs := info.lastUses, info.liveOuts
-		for use, pc := range lastUses {
+		outs := info.liveOuts
+		info.lastUses.Range(func(use VReg, pc programCounter) {
 			// Usage without live-outs is a kill.
 			if _, ok := outs[use]; !ok {
 				info.kills[use] = pc
 			}
-		}
+		})
 
 		if wazevoapi.RegAllocLoggingEnabled {
 			fmt.Printf("\nfinalized info for block[%d]:\n%s\n", blk.ID(), info.Format(a.regInfo))
@@ -251,7 +263,7 @@ func (a *Allocator) beginUpAndMarkStack(f Function, v VReg, isPhi bool, phiDefin
 			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 {
+		if !info.lastUses.Contains(v) {
 			continue
 		}
 		// TODO: we might want to avoid recursion here.
@@ -463,7 +475,7 @@ func (a *Allocator) Reset() {
 
 func (a *Allocator) allocateBlockInfo(blockID int) *blockInfo {
 	if blockID >= len(a.blockInfos) {
-		a.blockInfos = append(a.blockInfos, make([]blockInfo, blockID+1)...)
+		a.blockInfos = append(a.blockInfos, make([]blockInfo, (blockID+1)-len(a.blockInfos))...)
 	}
 	info := &a.blockInfos[blockID]
 	a.initBlockInfo(info)
@@ -543,11 +555,6 @@ func (a *Allocator) initBlockInfo(i *blockInfo) {
 	} else {
 		resetMap(a, i.defs)
 	}
-	if i.lastUses == nil {
-		i.lastUses = make(map[VReg]programCounter)
-	} else {
-		resetMap(a, i.lastUses)
-	}
 	if i.kills == nil {
 		i.kills = make(map[VReg]programCounter)
 	} else {
@@ -587,9 +594,9 @@ func (i *blockInfo) Format(ri *RegisterInfo) string {
 		buf.WriteString(fmt.Sprintf("%v@%v ", v, pos))
 	}
 	buf.WriteString("\n\tlastUses: ")
-	for v, pos := range i.lastUses {
+	i.lastUses.Range(func(v VReg, pos programCounter) {
 		buf.WriteString(fmt.Sprintf("%v@%v ", v, pos))
-	}
+	})
 	buf.WriteString("\n\tkills: ")
 	for v, pos := range i.kills {
 		buf.WriteString(fmt.Sprintf("%v@%v ", v, pos))


Test patch:
diff --git a/internal/engine/wazevo/backend/regalloc/reg_test.go b/internal/engine/wazevo/backend/regalloc/reg_test.go
index 7c48200a..84600cde 100644
--- a/internal/engine/wazevo/backend/regalloc/reg_test.go
+++ b/internal/engine/wazevo/backend/regalloc/reg_test.go
@@ -25,3 +25,33 @@ func Test_FromRealReg(t *testing.T) {
 	require.Equal(t, RealReg(5), r.RealReg())
 	require.Equal(t, VRegID(5), r.ID())
 }
+
+func TestVRegTable(t *testing.T) {
+	min := VRegIDMinSet{}
+	min.Observe(VReg(vRegIDReservedForRealNum + 2))
+	min.Observe(VReg(vRegIDReservedForRealNum + 1))
+	min.Observe(VReg(vRegIDReservedForRealNum + 0))
+
+	table := VRegTable{}
+	table.Reset(min)
+	table.Insert(VReg(vRegIDReservedForRealNum+0), 1)
+	table.Insert(VReg(vRegIDReservedForRealNum+1), 10)
+	table.Insert(VReg(vRegIDReservedForRealNum+2), 100)
+
+	vregs := map[VReg]programCounter{}
+	table.Range(func(v VReg, p programCounter) {
+		vregs[v] = p
+	})
+	require.Equal(t, 3, len(vregs))
+
+	for v, p := range vregs {
+		require.True(t, table.Contains(v))
+		require.Equal(t, p, table.Lookup(v))
+	}
+
+	table.Range(func(v VReg, p programCounter) {
+		require.Equal(t, vregs[v], p)
+		delete(vregs, v)
+	})
+	require.Equal(t, 0, len(vregs))
+}
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc_test.go b/internal/engine/wazevo/backend/regalloc/regalloc_test.go
index 0048c525..f3ab10fa 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc_test.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc_test.go
@@ -6,6 +6,18 @@ import (
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
 
+func makeVRegTable(vregs map[VReg]programCounter) (table VRegTable) {
+	min := VRegIDMinSet{}
+	for v := range vregs {
+		min.Observe(v)
+	}
+	table.Reset(min)
+	for v, p := range vregs {
+		table.Insert(v, p)
+	}
+	return table
+}
+
 func TestAllocator_livenessAnalysis(t *testing.T) {
 	const realRegID, realRegID2 = 50, 100
 	realReg, realReg2 := FromRealReg(realRegID, RegTypeInt), FromRealReg(realRegID2, RegTypeInt)
@@ -28,11 +40,12 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 			exp: map[int]*blockInfo{
 				0: {
 					defs:     map[VReg]programCounter{2: pcDefOffset + pcStride, 1: pcDefOffset},
-					lastUses: map[VReg]programCounter{1: pcStride + pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{1: pcStride + pcUseOffset}),
 					kills:    map[VReg]programCounter{1: pcStride + pcUseOffset},
 				},
 			},
 		},
+
 		{
 			name: "straight",
 			// b0 -> b1 -> b2
@@ -62,11 +75,11 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 						2:    pcDefOffset,
 						3:    pcStride*2 + pcDefOffset,
 					},
-					lastUses: map[VReg]programCounter{
+					lastUses: makeVRegTable(map[VReg]programCounter{
 						1000: pcStride + pcUseOffset,
 						1:    pcStride*2 + pcUseOffset,
 						2:    pcStride*2 + pcUseOffset,
-					},
+					}),
 					liveOuts: map[VReg]struct{}{3: {}},
 					kills: map[VReg]programCounter{
 						1000: pcStride + pcUseOffset,
@@ -77,9 +90,9 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				1: {
 					liveIns:  map[VReg]struct{}{3: {}},
 					liveOuts: map[VReg]struct{}{3: {}, 4: {}, 5: {}},
-					lastUses: map[VReg]programCounter{
+					lastUses: makeVRegTable(map[VReg]programCounter{
 						3: pcStride + pcUseOffset,
-					},
+					}),
 					defs: map[VReg]programCounter{
 						4: pcStride + pcDefOffset,
 						5: pcStride + pcDefOffset,
@@ -93,7 +106,7 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				},
 				2: {
 					liveIns:  map[VReg]struct{}{3: {}, 4: {}, 5: {}},
-					lastUses: map[VReg]programCounter{3: pcUseOffset, 4: pcUseOffset, 5: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{3: pcUseOffset, 4: pcUseOffset, 5: pcUseOffset}),
 					kills:    map[VReg]programCounter{3: pcUseOffset, 4: pcUseOffset, 5: pcUseOffset},
 				},
 			},
@@ -139,11 +152,12 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 						2:    pcStride*2 + pcDefOffset,
 					},
 					liveOuts: map[VReg]struct{}{1000: {}, 1: {}, 2: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				1: {
 					liveIns:  map[VReg]struct{}{1000: {}, 1: {}},
 					liveOuts: map[VReg]struct{}{1000: {}},
-					lastUses: map[VReg]programCounter{1: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{1: pcUseOffset}),
 					kills:    map[VReg]programCounter{1: pcUseOffset},
 					realRegDefs: [vRegIDReservedForRealNum][]programCounter{
 						realRegID:  {pcDefOffset, pcStride*4 + pcDefOffset},
@@ -157,18 +171,19 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				2: {
 					liveIns:     map[VReg]struct{}{1000: {}, 2: {}},
 					liveOuts:    map[VReg]struct{}{1000: {}},
-					lastUses:    map[VReg]programCounter{2: pcUseOffset},
+					lastUses:    makeVRegTable(map[VReg]programCounter{2: pcUseOffset}),
 					kills:       map[VReg]programCounter{2: pcUseOffset},
 					realRegUses: [vRegIDReservedForRealNum][]programCounter{realRegID2: {pcUseOffset}},
 					realRegDefs: [vRegIDReservedForRealNum][]programCounter{realRegID2: {0}},
 				},
 				3: {
 					liveIns:  map[VReg]struct{}{1000: {}},
-					lastUses: map[VReg]programCounter{1000: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{1000: pcUseOffset}),
 					kills:    map[VReg]programCounter{1000: pcUseOffset},
 				},
 			},
 		},
+
 		{
 			name: "phis",
 			//   0
@@ -203,32 +218,35 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				0: {
 					defs:     map[VReg]programCounter{1000: pcDefOffset, 2000: pcDefOffset, 3000: pcDefOffset},
 					liveOuts: map[VReg]struct{}{1000: {}, 2000: {}, 3000: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				1: {
 					liveIns:  map[VReg]struct{}{2000: {}, 3000: {}},
 					liveOuts: map[VReg]struct{}{phiVReg: {}, 3000: {}},
 					defs:     map[VReg]programCounter{phiVReg: pcDefOffset},
-					lastUses: map[VReg]programCounter{2000: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{2000: pcUseOffset}),
 					kills:    map[VReg]programCounter{2000: pcUseOffset},
 				},
 				2: {
 					liveIns:  map[VReg]struct{}{phiVReg: {}, 3000: {}},
 					liveOuts: map[VReg]struct{}{phiVReg: {}, 3000: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				3: {
 					liveIns:  map[VReg]struct{}{1000: {}, 3000: {}},
 					liveOuts: map[VReg]struct{}{phiVReg: {}, 3000: {}},
 					defs:     map[VReg]programCounter{phiVReg: pcDefOffset},
-					lastUses: map[VReg]programCounter{1000: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{1000: pcUseOffset}),
 					kills:    map[VReg]programCounter{1000: pcUseOffset},
 				},
 				4: {
 					liveIns:  map[VReg]struct{}{phiVReg: {}, 3000: {}},
-					lastUses: map[VReg]programCounter{phiVReg: pcUseOffset, 3000: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{phiVReg: pcUseOffset, 3000: pcUseOffset}),
 					kills:    map[VReg]programCounter{phiVReg: pcUseOffset, 3000: pcUseOffset},
 				},
 			},
 		},
+
 		{
 			name: "loop",
 			// 0 -> 1 -> 2
@@ -274,9 +292,9 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 						1:       pcDefOffset,
 						phiVReg: pcStride + pcDefOffset,
 					},
-					lastUses: map[VReg]programCounter{
+					lastUses: makeVRegTable(map[VReg]programCounter{
 						1: pcStride + pcUseOffset,
-					},
+					}),
 					kills: map[VReg]programCounter{
 						1: pcStride + pcUseOffset,
 					},
@@ -285,32 +303,33 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 					liveIns:  map[VReg]struct{}{phiVReg: {}},
 					liveOuts: map[VReg]struct{}{phiVReg: {}, 9999: {}},
 					defs:     map[VReg]programCounter{9999: pcDefOffset},
-					lastUses: map[VReg]programCounter{},
+					lastUses: makeVRegTable(map[VReg]programCounter{}),
 					kills:    map[VReg]programCounter{},
 				},
 				2: {
 					liveIns:  map[VReg]struct{}{phiVReg: {}, 9999: {}},
 					liveOuts: map[VReg]struct{}{100: {}},
 					defs:     map[VReg]programCounter{100: pcDefOffset},
-					lastUses: map[VReg]programCounter{phiVReg: pcUseOffset, 9999: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{phiVReg: pcUseOffset, 9999: pcUseOffset}),
 					kills:    map[VReg]programCounter{phiVReg: pcUseOffset, 9999: pcUseOffset},
 				},
 				3: {
 					liveIns:  map[VReg]struct{}{100: {}},
 					liveOuts: map[VReg]struct{}{54321: {}},
 					defs:     map[VReg]programCounter{54321: pcDefOffset},
-					lastUses: map[VReg]programCounter{100: pcStride + pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{100: pcStride + pcUseOffset}),
 					kills:    map[VReg]programCounter{100: pcStride + pcUseOffset},
 				},
 				4: {
 					liveIns:  map[VReg]struct{}{54321: {}},
 					liveOuts: map[VReg]struct{}{phiVReg: {}},
 					defs:     map[VReg]programCounter{phiVReg: pcDefOffset},
-					lastUses: map[VReg]programCounter{54321: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{54321: pcUseOffset}),
 					kills:    map[VReg]programCounter{54321: pcUseOffset},
 				},
 			},
 		},
+
 		{
 			//           -----+
 			//           v    |
@@ -336,23 +355,29 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				0: {
 					defs:     map[VReg]programCounter{99999: pcDefOffset},
 					liveOuts: map[VReg]struct{}{99999: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				1: {
 					liveIns:  map[VReg]struct{}{99999: {}},
 					liveOuts: map[VReg]struct{}{99999: {}},
-					lastUses: map[VReg]programCounter{99999: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{99999: pcUseOffset}),
 				},
 				2: {
 					liveIns:  map[VReg]struct{}{99999: {}},
 					liveOuts: map[VReg]struct{}{99999: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				3: {
 					liveIns:  map[VReg]struct{}{99999: {}},
 					liveOuts: map[VReg]struct{}{99999: {}},
+					lastUses: makeVRegTable(nil),
+				},
+				4: {
+					lastUses: makeVRegTable(nil),
 				},
-				4: {},
 			},
 		},
+
 		//      2
 		//      ^              +----+
 		//      |              v    |
@@ -405,35 +430,48 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				return newMockFunction(b0, b1, b2, b3, b4, b7, b8, b5, b6, b9)
 			},
 			exp: map[int]*blockInfo{
-				0: {},
-				1: {},
-				2: {},
+				0: {
+					lastUses: makeVRegTable(nil),
+				},
+				1: {
+					lastUses: makeVRegTable(nil),
+				},
+				2: {
+					lastUses: makeVRegTable(nil),
+				},
 				3: {
-					liveOuts: map[VReg]struct{}{100: {}},
 					defs:     map[VReg]programCounter{100: pcDefOffset},
+					liveOuts: map[VReg]struct{}{100: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				4: {
 					liveIns:  map[VReg]struct{}{100: {}},
 					liveOuts: map[VReg]struct{}{100: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				5: {
 					liveIns:  map[VReg]struct{}{100: {}},
 					liveOuts: map[VReg]struct{}{100: {}},
-					lastUses: map[VReg]programCounter{100: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{100: pcUseOffset}),
 				},
 				6: {
 					liveIns:  map[VReg]struct{}{100: {}},
 					liveOuts: map[VReg]struct{}{100: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				7: {
 					liveIns:  map[VReg]struct{}{100: {}},
 					liveOuts: map[VReg]struct{}{100: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				8: {
 					liveIns:  map[VReg]struct{}{100: {}},
 					liveOuts: map[VReg]struct{}{100: {}},
+					lastUses: makeVRegTable(nil),
+				},
+				9: {
+					lastUses: makeVRegTable(nil),
 				},
-				9: {},
 			},
 		},
 	} {
@@ -448,7 +486,7 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				initMapInInfo(exp)
 				saved := actual.intervalMng
 				actual.intervalMng = nil // Don't compare intervalManager.
-				require.Equal(t, exp, actual, "\n[exp for block[%d]]\n%s\n[actual for block[%d]]\n%s", blockID, exp, blockID, actual)
+				require.Equal(t, exp, actual, "\n[exp for block[%d]]\n%v\n[actual for block[%d]]\n%v", blockID, exp, blockID, actual)
 				actual.intervalMng = saved
 			}
 
@@ -533,9 +571,6 @@ func initMapInInfo(info *blockInfo) {
 	if info.kills == nil {
 		info.kills = make(map[VReg]programCounter)
 	}
-	if info.lastUses == nil {
-		info.lastUses = make(map[VReg]programCounter)
-	}
 }
 
 func TestNode_assignedRealReg(t *testing.T) {
diff --git a/internal/engine/wazevo/backend/regalloc/spill_handler_test.go b/internal/engine/wazevo/backend/regalloc/spill_handler_test.go
index 52800237..d8c07a92 100644
--- a/internal/engine/wazevo/backend/regalloc/spill_handler_test.go
+++ b/internal/engine/wazevo/backend/regalloc/spill_handler_test.go
@@ -29,7 +29,7 @@ func TestSpillHandler_getUnusedOrEvictReg(t *testing.T) {
 	require.Equal(t, 4, len(s.activeRegs))
 
 	regInfo := RegisterInfo{
-		AllocatableRegisters: [RegTypeNum][]RealReg{
+		AllocatableRegisters: [NumRegType][]RealReg{
 			RegTypeInt: {
 				RealReg(0xff), // unused.
 				RealReg(0), RealReg(1), RealReg(2),