wazevo(regalloc): removes map usage in real reg tracking (#1801)

Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>

internal/engine/wazevo/backend/isa/arm64/abi.go

@@ -48,6 +48,12 @@ var regInfo = &regalloc.RegisterInfo{
 		v0: v0VReg, v1: v1VReg, v2: v2VReg, v3: v3VReg, v4: v4VReg, v5: v5VReg, v6: v6VReg, v7: v7VReg, v8: v8VReg, v9: v9VReg, v10: v10VReg, v11: v11VReg, v12: v12VReg, v13: v13VReg, v14: v14VReg, v15: v15VReg, v16: v16VReg, v17: v17VReg, v18: v18VReg, v19: v19VReg, v20: v20VReg, v21: v21VReg, v22: v22VReg, v23: v23VReg, v24: v24VReg, v25: v25VReg, v26: v26VReg, v27: v27VReg, v28: v28VReg, v29: v29VReg, v30: v30VReg, v31: v31VReg,
 	},
 	RealRegName: func(r regalloc.RealReg) string { return regNames[r] },
+	RealRegType: func(r regalloc.RealReg) regalloc.RegType {
+		if r < v0 {
+			return regalloc.RegTypeInt
+		}
+		return regalloc.RegTypeFloat
+	},
 }
 
 // abiImpl implements backend.FunctionABI.

internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go

@@ -227,6 +227,7 @@ func (m *machine) RegisterInfo(debug bool) *regalloc.RegisterInfo {
 		regInfoDebug.CallerSavedRegisters = regInfo.CallerSavedRegisters
 		regInfoDebug.RealRegToVReg = regInfo.RealRegToVReg
 		regInfoDebug.RealRegName = regInfo.RealRegName
+		regInfoDebug.RealRegType = regInfo.RealRegType
 		regInfoDebug.AllocatableRegisters[regalloc.RegTypeFloat] = []regalloc.RealReg{
 			v18,                            // One callee saved.
 			v7, v6, v5, v4, v3, v2, v1, v0, // Allocatable sets == Argument registers.

internal/engine/wazevo/backend/regalloc/regalloc.go

@@ -42,6 +42,7 @@ type (
 		RealRegToVReg        []VReg
 		// RealRegName returns the name of the given RealReg for debugging.
 		RealRegName func(r RealReg) string
+		RealRegType func(r RealReg) RegType
 	}
 
 	// Allocator is a register allocator.
@@ -80,8 +81,8 @@ type (
 		lastUses map[VReg]programCounter
 		kills    map[VReg]programCounter
 		// Pre-colored real registers can have multiple live ranges in one block.
-		realRegUses map[VReg][]programCounter
-		realRegDefs map[VReg][]programCounter
+		realRegUses [vRegIDReservedForRealNum][]programCounter
+		realRegDefs [vRegIDReservedForRealNum][]programCounter
 		intervalMng *intervalManager
 	}
 
@@ -185,9 +186,10 @@ func (a *Allocator) livenessAnalysis(f Function) {
 			}
 			for _, def := range instr.Defs() {
 				dstVR = def
+				defID := def.ID()
 				pos := pc + pcDefOffset
 				if def.IsRealReg() {
-					info.realRegDefs[def] = append(info.realRegDefs[def], pos)
+					info.realRegDefs[defID] = append(info.realRegDefs[defID], pos)
 				} else {
 					if _, ok := info.defs[def]; !ok {
 						// This means that this VReg is defined multiple times in a series of instructions
@@ -407,37 +409,30 @@ func (a *Allocator) buildLiveRangesForNonReals(info *blockInfo) {
 func (a *Allocator) buildLiveRangesForReals(info *blockInfo) {
 	ds, us := info.realRegDefs, info.realRegUses
 
-	// In order to do the deterministic compilation, we need to sort the registers.
-	a.vs = a.vs[:0]
-	for v := range us {
+	for i := 0; i < RealRegsNumMax; i++ {
+		r := RealReg(i)
 		// Non allocation target registers are not needed here.
-		if !a.allocatableSet[v.RealReg()] {
+		if !a.allocatableSet[r] {
 			continue
 		}
-		a.vs = append(a.vs, v)
-	}
-	sort.SliceStable(a.vs, func(i, j int) bool {
-		return a.vs[i].RealReg() < a.vs[j].RealReg()
-	})
 
-	for _, v := range a.vs {
-		uses := us[v]
-		defs, ok := ds[v]
-		if !ok || len(defs) != len(uses) {
+		uses := us[r]
+		defs := ds[r]
+		if len(defs) != len(uses) {
 			// This is likely a bug of the Instr interface implementation and/or ABI around call instructions.
 			// E.g. call or ret instructions should specify that they use all the real registers (calling convention).
 			panic(
 				fmt.Sprintf(
 					"BUG: real register (%s) is defined and used, but the number of defs and uses are different: %d (defs) != %d (uses)",
-					a.regInfo.RealRegName(v.RealReg()), len(defs), len(uses),
+					a.regInfo.RealRegName(r), len(defs), len(uses),
 				),
 			)
 		}
 
 		for i := range uses {
 			n := a.allocateNode()
-			n.r = v.RealReg()
-			n.v = v
+			n.r = r
+			n.v = FromRealReg(r, a.regInfo.RealRegType(r))
 			defined, used := defs[i], uses[i]
 			intervalNode := info.intervalMng.insert(n, defined, used)
 			n.ranges = append(n.ranges, intervalNode)
@@ -558,26 +553,22 @@ func (a *Allocator) initBlockInfo(i *blockInfo) {
 	} else {
 		resetMap(a, i.kills)
 	}
-	if i.realRegUses == nil {
-		i.realRegUses = make(map[VReg][]programCounter)
-	} else {
-		resetMap(a, i.realRegUses)
-	}
-	if i.realRegDefs == nil {
-		i.realRegDefs = make(map[VReg][]programCounter)
-	} else {
-		resetMap(a, i.realRegDefs)
+
+	for index := range i.realRegUses {
+		i.realRegUses[index] = i.realRegUses[index][:0]
+		i.realRegDefs[index] = i.realRegDefs[index][:0]
 	}
 }
 
 func (i *blockInfo) addRealRegUsage(v VReg, pc programCounter) {
-	defs := i.realRegDefs[v]
+	id := v.ID()
+	defs := i.realRegDefs[id]
 	if len(defs) == 0 {
 		// If the definition not found yet but used, this must be a function preamble,
 		// so we let's assume it is defined at the beginning.
-		i.realRegDefs[v] = append(i.realRegDefs[v], 0)
+		i.realRegDefs[id] = append(i.realRegDefs[id], 0)
 	}
-	i.realRegUses[v] = append(i.realRegUses[v], pc)
+	i.realRegUses[id] = append(i.realRegUses[id], pc)
 }
 
 // Format is for debugging.
@@ -605,11 +596,15 @@ func (i *blockInfo) Format(ri *RegisterInfo) string {
 	}
 	buf.WriteString("\n\trealRegUses: ")
 	for v, pos := range i.realRegUses {
-		buf.WriteString(fmt.Sprintf("%s@%v ", ri.RealRegName(v.RealReg()), pos))
+		if len(pos) > 0 {
+			buf.WriteString(fmt.Sprintf("%s@%v ", ri.RealRegName(RealReg(v)), pos))
+		}
 	}
 	buf.WriteString("\n\trealRegDefs: ")
 	for v, pos := range i.realRegDefs {
-		buf.WriteString(fmt.Sprintf("%s@%v ", ri.RealRegName(v.RealReg()), pos))
+		if len(pos) > 0 {
+			buf.WriteString(fmt.Sprintf("%s@%v ", ri.RealRegName(RealReg(v)), pos))
+		}
 	}
 	return buf.String()
 }

internal/engine/wazevo/backend/regalloc/regalloc_test.go

@@ -7,7 +7,8 @@ import (
 )
 
 func TestAllocator_livenessAnalysis(t *testing.T) {
-	realReg, realReg2 := FromRealReg(50, RegTypeInt), FromRealReg(100, RegTypeInt)
+	const realRegID, realRegID2 = 50, 100
+	realReg, realReg2 := FromRealReg(realRegID, RegTypeInt), FromRealReg(realRegID2, RegTypeInt)
 	const phiVReg = 12345
 	for _, tc := range []struct {
 		name  string
@@ -83,11 +84,11 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 						4: pcStride + pcDefOffset,
 						5: pcStride + pcDefOffset,
 					},
-					realRegUses: map[VReg][]programCounter{
-						realReg: {pcStride*2 + pcUseOffset},
+					realRegUses: [vRegIDReservedForRealNum][]programCounter{
+						realRegID: {pcStride*2 + pcUseOffset},
 					},
-					realRegDefs: map[VReg][]programCounter{
-						realReg: {pcDefOffset},
+					realRegDefs: [vRegIDReservedForRealNum][]programCounter{
+						realRegID: {pcDefOffset},
 					},
 				},
 				2: {
@@ -144,22 +145,22 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 					liveOuts: map[VReg]struct{}{1000: {}},
 					lastUses: map[VReg]programCounter{1: pcUseOffset},
 					kills:    map[VReg]programCounter{1: pcUseOffset},
-					realRegDefs: map[VReg][]programCounter{
-						realReg:  {pcDefOffset, pcStride*4 + pcDefOffset},
-						realReg2: {pcStride*2 + pcDefOffset},
+					realRegDefs: [vRegIDReservedForRealNum][]programCounter{
+						realRegID:  {pcDefOffset, pcStride*4 + pcDefOffset},
+						realRegID2: {pcStride*2 + pcDefOffset},
 					},
-					realRegUses: map[VReg][]programCounter{
-						realReg:  {pcStride + pcUseOffset, pcStride*5 + pcUseOffset},
-						realReg2: {pcStride*3 + pcUseOffset},
+					realRegUses: [vRegIDReservedForRealNum][]programCounter{
+						realRegID:  {pcStride + pcUseOffset, pcStride*5 + pcUseOffset},
+						realRegID2: {pcStride*3 + pcUseOffset},
 					},
 				},
 				2: {
 					liveIns:     map[VReg]struct{}{1000: {}, 2: {}},
 					liveOuts:    map[VReg]struct{}{1000: {}},
 					lastUses:    map[VReg]programCounter{2: pcUseOffset},
 					kills:       map[VReg]programCounter{2: pcUseOffset},
-					realRegUses: map[VReg][]programCounter{realReg2: {pcUseOffset}},
-					realRegDefs: map[VReg][]programCounter{realReg2: {0}},
+					realRegUses: [vRegIDReservedForRealNum][]programCounter{realRegID2: {pcUseOffset}},
+					realRegDefs: [vRegIDReservedForRealNum][]programCounter{realRegID2: {0}},
 				},
 				3: {
 					liveIns:  map[VReg]struct{}{1000: {}},
@@ -535,12 +536,6 @@ func initMapInInfo(info *blockInfo) {
 	if info.lastUses == nil {
 		info.lastUses = make(map[VReg]programCounter)
 	}
-	if info.realRegUses == nil {
-		info.realRegUses = make(map[VReg][]programCounter)
-	}
-	if info.realRegDefs == nil {
-		info.realRegDefs = make(map[VReg][]programCounter)
-	}
 }
 
 func TestNode_assignedRealReg(t *testing.T) {