Repository: tetratelabs/wazero
Pull request #1677 (2023-08-31 08:28:36)

wazevo(regalloc): refactors assignment logic to isolate spill handling
#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr.go b/internal/engine/wazevo/backend/isa/arm64/instr.go
index 96013f97..d90c99a4 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr.go
@@ -271,59 +271,73 @@ func (i *instruction) uses(regs []regalloc.VReg) []regalloc.VReg {
 	return regs
 }
 
-func (i *instruction) assignUses(regs []regalloc.VReg) {
+func (i *instruction) assignUse(index int, reg regalloc.VReg) {
 	switch useKinds[i.kind] {
 	case useKindNone:
 	case useKindRN:
 		if rn := i.rn.reg(); rn.Valid() {
-			i.rn = i.rn.assignReg(regs[0])
+			i.rn = i.rn.assignReg(reg)
 		}
 	case useKindRNRM:
-		if rn := i.rn.reg(); rn.Valid() {
-			i.rn = i.rn.assignReg(regs[0])
-		}
-		if rm := i.rm.reg(); rm.Valid() {
-			i.rm = i.rm.assignReg(regs[1])
+		if index == 0 {
+			if rn := i.rn.reg(); rn.Valid() {
+				i.rn = i.rn.assignReg(reg)
+			}
+		} else {
+			if rm := i.rm.reg(); rm.Valid() {
+				i.rm = i.rm.assignReg(reg)
+			}
 		}
 	case useKindRNRMRA:
-		if rn := i.rn.reg(); rn.Valid() {
-			i.rn = i.rn.assignReg(regs[0])
-		}
-		if rm := i.rm.reg(); rm.Valid() {
-			i.rm = i.rm.assignReg(regs[1])
-		}
-		if ra := i.ra.reg(); ra.Valid() {
-			i.ra = i.ra.assignReg(regs[2])
+		if index == 0 {
+			if rn := i.rn.reg(); rn.Valid() {
+				i.rn = i.rn.assignReg(reg)
+			}
+		} else if index == 1 {
+			if rm := i.rm.reg(); rm.Valid() {
+				i.rm = i.rm.assignReg(reg)
+			}
+		} else {
+			if ra := i.ra.reg(); ra.Valid() {
+				i.ra = i.ra.assignReg(reg)
+			}
 		}
 	case useKindRet:
 		panic("BUG: ret instructions shouldn't be assigned")
 	case useKindAMode:
-		if amodeRN := i.amode.rn; amodeRN.Valid() {
-			i.amode.rn = regs[0]
-		}
-		if amodeRM := i.amode.rm; amodeRM.Valid() {
-			i.amode.rm = regs[1]
+		if index == 0 {
+			if amodeRN := i.amode.rn; amodeRN.Valid() {
+				i.amode.rn = reg
+			}
+		} else {
+			if amodeRM := i.amode.rm; amodeRM.Valid() {
+				i.amode.rm = reg
+			}
 		}
 	case useKindRNAMode:
-		i.rn = i.rn.assignReg(regs[0])
-		if amodeRN := i.amode.rn; amodeRN.Valid() {
-			i.amode.rn = regs[1]
-		}
-		if amodeRM := i.amode.rm; amodeRM.Valid() {
-			i.amode.rm = regs[2]
+		if index == 0 {
+			i.rn = i.rn.assignReg(reg)
+		} else if index == 1 {
+			if amodeRN := i.amode.rn; amodeRN.Valid() {
+				i.amode.rn = reg
+			}
+		} else {
+			if amodeRM := i.amode.rm; amodeRM.Valid() {
+				i.amode.rm = reg
+			}
 		}
 	case useKindCond:
 		c := cond(i.u1)
 		switch c.kind() {
 		case condKindRegisterZero:
-			i.u1 = uint64(registerAsRegZeroCond(regs[0]))
+			i.u1 = uint64(registerAsRegZeroCond(reg))
 		case condKindRegisterNotZero:
-			i.u1 = uint64(registerAsRegNotZeroCond(regs[0]))
+			i.u1 = uint64(registerAsRegNotZeroCond(reg))
 		}
 	case useKindCall:
 		panic("BUG: call instructions shouldn't be assigned")
 	case useKindCallInd:
-		i.rn = i.rn.assignReg(regs[0])
+		i.rn = i.rn.assignReg(reg)
 	default:
 		panic(fmt.Sprintf("useKind for %v not defined", i))
 	}
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go b/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
index 3940ada0..7dea1778 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
@@ -234,9 +234,9 @@ func (r *regAllocInstrImpl) IsReturn() bool {
 	return r.i.kind == ret
 }
 
-// AssignUses implements regalloc.Instr AssignUses.
-func (r *regAllocInstrImpl) AssignUses(vs []regalloc.VReg) {
-	r.i.assignUses(vs)
+// AssignUse implements regalloc.Instr AssignUse.
+func (r *regAllocInstrImpl) AssignUse(i int, v regalloc.VReg) {
+	r.i.assignUse(i, v)
 }
 
 // AssignDef implements regalloc.Instr AssignDef.
diff --git a/internal/engine/wazevo/backend/regalloc/api.go b/internal/engine/wazevo/backend/regalloc/api.go
index f6a87fd1..076f046a 100644
--- a/internal/engine/wazevo/backend/regalloc/api.go
+++ b/internal/engine/wazevo/backend/regalloc/api.go
@@ -64,9 +64,8 @@ type (
 		// Uses returns the virtual registers used by this instruction.
 		// Note: multiple returned []VReg will not be held at the same time, so it's safe to use the same slice for this.
 		Uses() []VReg
-		// AssignUses assigns the RealReg-allocated virtual registers used by this instruction.
-		// Note: input []VReg is reused, so it's not safe to hold reference to it after the end of this call.
-		AssignUses([]VReg)
+		// AssignUse assigns the RealReg-allocated virtual register used by this instruction at the given index.
+		AssignUse(index int, v VReg)
 		// AssignDef assigns a RealReg-allocated virtual register defined by this instruction.
 		// This only accepts one register because we don't allocate registers for multi-def instructions (i.e. call instruction)
 		AssignDef(VReg)
diff --git a/internal/engine/wazevo/backend/regalloc/assign.go b/internal/engine/wazevo/backend/regalloc/assign.go
index 93147a22..27b90a3b 100644
--- a/internal/engine/wazevo/backend/regalloc/assign.go
+++ b/internal/engine/wazevo/backend/regalloc/assign.go
@@ -45,11 +45,11 @@ func (a *Allocator) assignRegistersPerInstr(f Function, pc programCounter, instr
 	if indirect := instr.IsIndirectCall(); instr.IsCall() || indirect {
 		// Only take care of non-real VRegs (e.g. VReg.IsRealReg() == false) since
 		// the real VRegs are already placed in the right registers at this point.
-		a.nodes2 = a.collectActiveNonRealVRegsAt(a.nodes2,
+		a.collectActiveNonRealVRegsAt(
 			// To find the all the live registers "after" call, we need to add pcDefOffset for search.
 			pc+pcDefOffset,
 			liveNodes)
-		for _, active := range a.nodes2 {
+		for _, active := range a.nodes1 {
 			if r := active.r; a.regInfo.isCallerSaved(r) {
 				v := active.v.SetRealReg(r)
 				f.StoreRegisterBefore(v, instr)
@@ -73,10 +73,9 @@ func (a *Allocator) assignRegistersPerInstr(f Function, pc programCounter, instr
 		return
 	}
 
-	a.nodes1 = a.nodes1[:0]
-	a.vs = a.vs[:0]
+	usesSpills := a.vs[:0]
 	uses := instr.Uses()
-	for _, u := range uses {
+	for i, u := range uses {
 		if u.IsRealReg() {
 			a.vs = append(a.vs, u)
 			continue
@@ -86,42 +85,46 @@ func (a *Allocator) assignRegistersPerInstr(f Function, pc programCounter, instr
 		}
 		n := vRegIDToNode[u.ID()]
 		if !n.spill() {
-			a.vs = append(a.vs, u.SetRealReg(n.r))
+			instr.AssignUse(i, u.SetRealReg(n.r))
 		} else {
-			a.nodes1 = append(a.nodes1, n)
+			usesSpills = append(usesSpills, u)
 		}
 	}
 
-	if len(a.nodes1) == 0 { // no spill.
-		instr.AssignUses(a.vs)
-	} else {
-		panic("TODO: handle spills.")
-	}
-
 	defs := instr.Defs()
+	defSpill := VRegInvalid
 	switch len(defs) {
 	case 0:
-		return
 	case 1:
+		d := defs[0]
+		if !d.IsRealReg() {
+			if wazevoapi.RegAllocLoggingEnabled {
+				fmt.Printf("%s defines %d\n", instr, d.ID())
+			}
+
+			n := vRegIDToNode[d.ID()]
+			if !n.spill() {
+				instr.AssignDef(d.SetRealReg(n.r))
+			} else {
+				defSpill = n.v
+			}
+		}
 	default:
-		// multiple defs (== call instruction) can be special cased, and no need to assign (already real regs following the calling convention.
-		return
+		panic("BUG: multiple def instructions must be special cased")
 	}
 
-	d := defs[0]
-	if d.IsRealReg() {
-		return
-	}
-	if wazevoapi.RegAllocLoggingEnabled {
-		fmt.Printf("%s defines %d\n", instr, d.ID())
-	}
+	a.handleSpills(f, pc, instr, vRegIDToNode, liveNodes, usesSpills, defSpill)
+	a.vs = usesSpills[:0] // for reuse.
+}
 
-	n := vRegIDToNode[d.ID()]
-	if !n.spill() {
-		instr.AssignDef(d.SetRealReg(n.r))
-	} else {
-		panic("TODO: handle spills.")
+func (a *Allocator) handleSpills(
+	f Function, pc programCounter, instr Instr, vRegIDToNode []*node, liveNodes []liveNodeInBlock,
+	usesSpills []VReg, defSpill VReg,
+) {
+	if len(usesSpills) == 0 && !defSpill.Valid() {
+		return
 	}
+	panic("TODO")
 }
 
 func (a *Allocator) assignIndirectCall(f Function, instr Instr, vRegIDToNode []*node) {
@@ -165,16 +168,14 @@ func (a *Allocator) assignIndirectCall(f Function, instr Instr, vRegIDToNode []*
 	} else {
 		v = v.SetRealReg(n.r)
 	}
-	a.vs = a.vs[:0]
-	a.vs = append(a.vs, v)
-	instr.AssignUses(a.vs)
+	instr.AssignUse(0, v)
 }
 
-// collectActiveNonRealVRegsAt collects the set of active registers at the given program counter into `nodes` slice by appending
+// collectActiveNonRealVRegsAt collects the set of active registers at the given program counter into `a.nodes1` slice by appending
 // the found registers from its beginning. This excludes the VRegs backed by a real register since this is used to list the registers
 // alive but not used by a call instruction.
-func (a *Allocator) collectActiveNonRealVRegsAt(nodes []*node, pc programCounter, liveNodes []liveNodeInBlock) []*node {
-	nodes = nodes[:0]
+func (a *Allocator) collectActiveNonRealVRegsAt(pc programCounter, liveNodes []liveNodeInBlock) {
+	nodes := a.nodes1[:0]
 	for _, live := range liveNodes {
 		n := live.n
 		if n.spill() || n.v.IsRealReg() {
@@ -189,5 +190,5 @@ func (a *Allocator) collectActiveNonRealVRegsAt(nodes []*node, pc programCounter
 			nodes = append(nodes, n)
 		}
 	}
-	return nodes
+	a.nodes1 = nodes
 }
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc.go b/internal/engine/wazevo/backend/regalloc/regalloc.go
index 782e8ea2..f003082d 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc.go
@@ -64,7 +64,7 @@ type (
 		blockInfos   [] /* blockID to */ blockInfo
 		vs           []VReg
 
-		// Followings are re-used during coloring and activeRegistersAt
+		// Followings are re-used during various places e.g. coloring.
 		realRegSet map[RealReg]struct{}
 		realRegs   []RealReg
 		nodeSet    map[*node]int


Test patch:
diff --git a/internal/engine/wazevo/backend/regalloc/api_test.go b/internal/engine/wazevo/backend/regalloc/api_test.go
index 61004a0c..1e8cb6a6 100644
--- a/internal/engine/wazevo/backend/regalloc/api_test.go
+++ b/internal/engine/wazevo/backend/regalloc/api_test.go
@@ -225,9 +225,11 @@ func (m *mockInstr) IsReturn() bool { return false }
 func (m *mockBlock) Entry() bool { return m._entry }
 
 // AssignUses implements Instr.
-func (m *mockInstr) AssignUses(regs []VReg) {
-	m.uses = make([]VReg, len(regs))
-	copy(m.uses, regs)
+func (m *mockInstr) AssignUse(index int, reg VReg) {
+	if index >= len(m.uses) {
+		m.uses = append(m.uses, make([]VReg, 100)...)
+	}
+	m.uses[index] = reg
 }
 
 // AssignDef implements Instr.
diff --git a/internal/engine/wazevo/backend/regalloc/assign_test.go b/internal/engine/wazevo/backend/regalloc/assign_test.go
index db6af3c1..18154c1a 100644
--- a/internal/engine/wazevo/backend/regalloc/assign_test.go
+++ b/internal/engine/wazevo/backend/regalloc/assign_test.go
@@ -153,7 +153,8 @@ func TestAllocator_activeNonRealVRegsAt(t *testing.T) {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
 			a := NewAllocator(&RegisterInfo{})
-			ans := a.collectActiveNonRealVRegsAt(nil, tc.pc, tc.lives)
+			a.collectActiveNonRealVRegsAt(tc.pc, tc.lives)
+			ans := a.nodes1
 
 			actual := make([]VReg, len(ans))
 			for i, n := range ans {