Repository: tetratelabs/wazero
Pull request #1675 (2023-08-31 00:23:52)

wazevo(regalloc): support for spill on function pointer of indirect call
#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/regalloc/api.go b/internal/engine/wazevo/backend/regalloc/api.go
index ee39b64a..f6a87fd1 100644
--- a/internal/engine/wazevo/backend/regalloc/api.go
+++ b/internal/engine/wazevo/backend/regalloc/api.go
@@ -25,13 +25,13 @@ type (
 		ReversePostOrderBlockIteratorNext() Block
 		// ClobberedRegisters tell the clobbered registers by this function.
 		ClobberedRegisters([]VReg)
-		// StoreRegisterBefore ... TODO
+		// StoreRegisterBefore inserts store instruction(s) before the given instruction for the given virtual register.
 		StoreRegisterBefore(v VReg, instr Instr)
-		// StoreRegisterAfter ... TODO
+		// StoreRegisterAfter inserts store instruction(s) after the given instruction for the given virtual register.
 		StoreRegisterAfter(v VReg, instr Instr)
-		// ReloadRegisterBefore ... TODO
+		// ReloadRegisterBefore inserts reload instruction(s) before the given instruction for the given virtual register.
 		ReloadRegisterBefore(v VReg, instr Instr)
-		// ReloadRegisterAfter ... TODO
+		// ReloadRegisterAfter inserts reload instruction(s) after the given instruction for the given virtual register.
 		ReloadRegisterAfter(v VReg, instr Instr)
 		// Done tells the implementation that register allocation is done, and it can finalize the stack
 		Done()
@@ -77,7 +77,8 @@ type (
 		// IsCall returns true if this instruction is a call instruction. The result is used to insert
 		// caller saved register spills and restores.
 		IsCall() bool
-		// IsIndirectCall returns true if this instruction is an indirect call instruction.
+		// IsIndirectCall returns true if this instruction is an indirect call instruction which calls a function pointer.
+		//  The result is used to insert caller saved register spills and restores.
 		IsIndirectCall() bool
 		// IsReturn returns true if this instruction is a return instruction.
 		IsReturn() bool
diff --git a/internal/engine/wazevo/backend/regalloc/assign.go b/internal/engine/wazevo/backend/regalloc/assign.go
index 3e8cd5d5..93147a22 100644
--- a/internal/engine/wazevo/backend/regalloc/assign.go
+++ b/internal/engine/wazevo/backend/regalloc/assign.go
@@ -42,7 +42,7 @@ func (a *Allocator) assignRegistersPerInstr(f Function, pc programCounter, instr
 		}
 	}
 
-	if direct := instr.IsCall(); direct || instr.IsIndirectCall() {
+	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,
@@ -56,10 +56,19 @@ func (a *Allocator) assignRegistersPerInstr(f Function, pc programCounter, instr
 				f.ReloadRegisterAfter(v, instr)
 			}
 		}
-		// Direct function calls do not need assignment, while indirect one needs the assignment on the function pointer.
-		if direct {
-			return
+		if indirect {
+			// Direct function calls do not need assignment, while indirect one needs the assignment on the function pointer.
+			a.assignIndirectCall(f, instr, vRegIDToNode)
 		}
+
+		if wazevoapi.RegAllocValidationEnabled {
+			for _, def := range instr.Defs() {
+				if !def.IsRealReg() {
+					panic(fmt.Sprintf("BUG: call/indirect call instruction must define only real registers: %s", def))
+				}
+			}
+		}
+		return
 	} else if instr.IsReturn() {
 		return
 	}
@@ -115,6 +124,52 @@ func (a *Allocator) assignRegistersPerInstr(f Function, pc programCounter, instr
 	}
 }
 
+func (a *Allocator) assignIndirectCall(f Function, instr Instr, vRegIDToNode []*node) {
+	a.nodes1 = a.nodes1[:0]
+	uses := instr.Uses()
+	if wazevoapi.RegAllocValidationEnabled {
+		var nonRealRegs int
+		for _, u := range uses {
+			if !u.IsRealReg() {
+				nonRealRegs++
+			}
+		}
+		if nonRealRegs != 1 {
+			panic(fmt.Sprintf("BUG: indirect call must have only one non-real register (for function pointer): %d", nonRealRegs))
+		}
+	}
+
+	var v VReg
+	for _, u := range uses {
+		if !u.IsRealReg() {
+			v = u
+			break
+		}
+	}
+
+	if v.RegType() != RegTypeInt {
+		panic(fmt.Sprintf("BUG: function pointer for indirect call must be an integer register: %s", v))
+	}
+
+	n := vRegIDToNode[v.ID()]
+	if n.spill() {
+		// If the function pointer is spilled, we need to reload it to a register.
+		// But at this point, all the caller-saved registers are saved, we can use a callee-saved register to reload.
+		for _, r := range a.regInfo.AllocatableRegisters[RegTypeInt] {
+			if a.regInfo.isCallerSaved(r) {
+				v = v.SetRealReg(r)
+				f.ReloadRegisterBefore(v, instr)
+				break
+			}
+		}
+	} else {
+		v = v.SetRealReg(n.r)
+	}
+	a.vs = a.vs[:0]
+	a.vs = append(a.vs, v)
+	instr.AssignUses(a.vs)
+}
+
 // collectActiveNonRealVRegsAt collects the set of active registers at the given program counter into `nodes` 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.


Test patch:
diff --git a/internal/engine/wazevo/backend/regalloc/api_test.go b/internal/engine/wazevo/backend/regalloc/api_test.go
index 17175eb7..61004a0c 100644
--- a/internal/engine/wazevo/backend/regalloc/api_test.go
+++ b/internal/engine/wazevo/backend/regalloc/api_test.go
@@ -32,8 +32,8 @@ type (
 
 	// mockInstr implements Instr.
 	mockInstr struct {
-		defs, uses     []VReg
-		isCopy, isCall bool
+		defs, uses                 []VReg
+		isCopy, isCall, isIndirect bool
 	}
 )
 
@@ -102,6 +102,11 @@ func (m *mockInstr) asCall() *mockInstr {
 	return m
 }
 
+func (m *mockInstr) asIndirectCall() *mockInstr {
+	m.isIndirect = true
+	return m
+}
+
 // StoreRegisterAfter implements Function.StoreRegisterAfter.
 func (m *mockFunction) StoreRegisterAfter(v VReg, instr Instr) {
 	m.storeRegisterAfter = append(m.storeRegisterAfter, storeOrReloadInfo{v, instr})
@@ -211,7 +216,7 @@ func (m *mockInstr) IsCopy() bool { return m.isCopy }
 func (m *mockInstr) IsCall() bool { return m.isCall }
 
 // IsIndirectCall implements Instr.
-func (m *mockInstr) IsIndirectCall() bool { return m.isCall }
+func (m *mockInstr) IsIndirectCall() bool { return m.isIndirect }
 
 // IsReturn implements Instr.
 func (m *mockInstr) IsReturn() bool { return false }
diff --git a/internal/engine/wazevo/backend/regalloc/assign_test.go b/internal/engine/wazevo/backend/regalloc/assign_test.go
index 494fabe0..db6af3c1 100644
--- a/internal/engine/wazevo/backend/regalloc/assign_test.go
+++ b/internal/engine/wazevo/backend/regalloc/assign_test.go
@@ -12,10 +12,10 @@ func TestAllocator_assignRegistersPerInstr(t *testing.T) {
 		pc := programCounter(5)
 		liveNodes := []liveNodeInBlock{
 			{n: &node{r: 1, v: 0xa, ranges: []liveRange{{begin: 5, end: 20}}}},
-			{n: &node{r: RealRegInvalid, v: 0xb, ranges: []liveRange{{begin: 5, end: 20}}}},           // Spill. not target.
-			{n: &node{r: 2, v: FromRealReg(1, RegTypeInt), ranges: []liveRange{{begin: 5, end: 20}}}}, // Real reg-backed VReg. not target
+			{n: &node{r: RealRegInvalid, v: 0xb, ranges: []liveRange{{begin: 5, end: 20}}}},           // Spill. not save target.
+			{n: &node{r: 2, v: FromRealReg(1, RegTypeInt), ranges: []liveRange{{begin: 5, end: 20}}}}, // Real reg-backed VReg. not save target
 			{n: &node{r: 3, v: 0xc, ranges: []liveRange{{begin: 5, end: 20}}}},
-			{n: &node{r: 4, v: 0xd, ranges: []liveRange{{begin: 5, end: 20}}}}, // real reg, but not caller saved. not target
+			{n: &node{r: 4, v: 0xd, ranges: []liveRange{{begin: 5, end: 20}}}}, // real reg, but not caller saved. not save target
 		}
 		call := newMockInstr().asCall()
 		blk := newMockBlock(0, call).entry()
@@ -25,6 +25,60 @@ func TestAllocator_assignRegistersPerInstr(t *testing.T) {
 		require.Equal(t, 2, len(f.storeRegisterBefore))
 		require.Equal(t, 2, len(f.reloadRegisterAfter))
 	})
+	t.Run("call_indirect/func_ptr not spilled", func(t *testing.T) {
+		a := NewAllocator(&RegisterInfo{CallerSavedRegisters: map[RealReg]struct{}{1: {}, 3: {}, 0xff: {}}})
+		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. */}}},
+		}
+		liveNodes := []liveNodeInBlock{
+			functionPtrLiveNode, // Function pointer, used at this PC. not save target.
+			{n: &node{r: 1, v: 0xa, ranges: []liveRange{{begin: 5, end: 20}}}},
+			{n: &node{r: 2, v: FromRealReg(1, RegTypeInt), ranges: []liveRange{{begin: 5, end: 20}}}}, // Real reg-backed VReg. not target
+			{n: &node{r: 3, v: 0xc, ranges: []liveRange{{begin: 5, end: 20}}}},
+			{n: &node{r: 4, v: 0xd, ranges: []liveRange{{begin: 5, end: 20}}}}, // real reg, but not caller saved. not save target
+		}
+		callInd := newMockInstr().asIndirectCall().use(functionPtrVReg)
+		blk := newMockBlock(0, callInd).entry()
+		f := newMockFunction(blk)
+		a.assignRegistersPerInstr(f, pc, callInd, []*node{0: functionPtrLiveNode.n}, liveNodes)
+
+		require.Equal(t, 2, len(f.storeRegisterBefore))
+		require.Equal(t, 2, len(f.reloadRegisterAfter))
+		require.True(t, callInd.uses[0].IsRealReg())
+		require.Equal(t, functionPtrVReg.SetRealReg(0xff), 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}},
+		})
+		pc := programCounter(5)
+		functionPtrVRegID := 0x0
+		functionPtrVReg := VReg(functionPtrVRegID).SetRegType(RegTypeInt)
+		liveNodes := []liveNodeInBlock{
+			{n: &node{r: 1, v: 0xa, ranges: []liveRange{{begin: 5, end: 20}}}},
+			{n: &node{r: 2, v: FromRealReg(1, RegTypeInt), ranges: []liveRange{{begin: 5, end: 20}}}}, // Real reg-backed VReg. not target
+			{n: &node{r: 3, v: 0xc, ranges: []liveRange{{begin: 5, end: 20}}}},
+			{n: &node{r: 4, v: 0xd, ranges: []liveRange{{begin: 5, end: 20}}}}, // real reg, but not caller saved. not save target
+		}
+		callInd := newMockInstr().asIndirectCall().use(functionPtrVReg)
+		blk := newMockBlock(0, callInd).entry()
+		f := newMockFunction(blk)
+		a.assignRegistersPerInstr(f, pc, callInd, []*node{
+			0: {r: RealRegInvalid},
+		}, liveNodes)
+
+		require.Equal(t, 2, len(f.storeRegisterBefore))
+		require.Equal(t, 2, len(f.reloadRegisterAfter))
+		require.Equal(t, 1, len(f.reloadRegisterBefore))
+		require.Equal(t, callInd, f.reloadRegisterBefore[0].instr)
+		require.Equal(t, functionPtrVReg.SetRealReg(0xbb), f.reloadRegisterBefore[0].v)
+		require.True(t, callInd.uses[0].IsRealReg())
+		require.Equal(t, functionPtrVReg.SetRealReg(0xbb), callInd.uses[0])
+	})
 	t.Run("no spills", func(t *testing.T) {
 		r := FromRealReg(1, RegTypeInt)
 		a := NewAllocator(&RegisterInfo{})