Repository: tetratelabs/wazero
Pull request #2018 (2024-02-05 22:37:42)

wazevo(backend): PostRegAlloc on Machine interface


Patch:
diff --git a/internal/engine/wazevo/backend/compiler.go b/internal/engine/wazevo/backend/compiler.go
index 256ec3d0..ddb85893 100644
--- a/internal/engine/wazevo/backend/compiler.go
+++ b/internal/engine/wazevo/backend/compiler.go
@@ -192,8 +192,7 @@ func (c *compiler) RegAlloc() {
 
 // Finalize implements Compiler.Finalize.
 func (c *compiler) Finalize(ctx context.Context) {
-	c.mach.SetupPrologue()
-	c.mach.SetupEpilogue()
+	c.mach.PostRegAlloc()
 	c.mach.Encode(ctx)
 }
 
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue.go b/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue.go
index 913844c6..b7a86782 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue.go
@@ -4,8 +4,13 @@ import (
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
 )
 
-// SetupPrologue implements backend.Machine.
-func (m *machine) SetupPrologue() {
+// PostRegAlloc implements backend.Machine.
+func (m *machine) PostRegAlloc() {
+	m.setupPrologue()
+	m.postRegAlloc()
+}
+
+func (m *machine) setupPrologue() {
 	cur := m.ectx.RootInstr
 	prevInitInst := cur.next
 
@@ -118,8 +123,12 @@ func (m *machine) SetupPrologue() {
 	linkInstr(cur, prevInitInst)
 }
 
-// SetupEpilogue implements backend.Machine.
-func (m *machine) SetupEpilogue() {
+// postRegAlloc does multiple things while walking through the instructions:
+// 1. Inserts the epilogue code.
+// 2. Removes the redundant copy instruction.
+// 3. Inserts the dec/inc RSP instruction right before/after the call instruction.
+// 4. Lowering that is supposed to be done after regalloc.
+func (m *machine) postRegAlloc() {
 	ectx := m.ectx
 	for cur := ectx.RootInstr; cur != nil; cur = cur.next {
 		switch k := cur.kind; k {
@@ -168,7 +177,6 @@ func (m *machine) SetupEpilogue() {
 			// At this point, reg alloc is done, therefore we can safely insert dec/inc RPS instruction
 			// right before/after the call instruction. If this is done before reg alloc, the stack slot
 			// can point to the wrong location and therefore results in a wrong value.
-			// TODO: doing this in `SetupEpilogue` seems weird. Find a better home.
 			call := cur
 			next := call.next
 			if size := uint32(abi.AlignedArgResultStackSlotSize()); size > 0 {
@@ -183,7 +191,6 @@ func (m *machine) SetupEpilogue() {
 		}
 
 		// Removes the redundant copy instruction.
-		// TODO: doing this in `SetupEpilogue` seems weird. Find a better home.
 		if cur.IsCopy() && cur.op1.r.RealReg() == cur.op2.r.RealReg() {
 			prev, next := cur.prev, cur.next
 			// Remove the copy instruction.
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
index 46841e94..90c17d86 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
@@ -7,8 +7,14 @@ import (
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/wazevoapi"
 )
 
-// SetupPrologue implements backend.Machine.
-func (m *machine) SetupPrologue() {
+// PostRegAlloc implements backend.Machine.
+func (m *machine) PostRegAlloc() {
+	m.setupPrologue()
+	m.postRegAlloc()
+}
+
+// setupPrologue initializes the prologue of the function.
+func (m *machine) setupPrologue() {
 	ectx := m.executableContext
 
 	cur := ectx.RootInstr
@@ -186,8 +192,10 @@ func (m *machine) createFrameSizeSlot(cur *instruction, s int64) *instruction {
 	return cur
 }
 
-// SetupEpilogue implements backend.Machine.
-func (m *machine) SetupEpilogue() {
+// postRegAlloc does multiple things while walking through the instructions:
+// 1. Removes the redundant copy instruction.
+// 2. Inserts the epilogue.
+func (m *machine) postRegAlloc() {
 	ectx := m.executableContext
 	for cur := ectx.RootInstr; cur != nil; cur = cur.next {
 		if cur.kind == ret {
@@ -196,7 +204,6 @@ func (m *machine) SetupEpilogue() {
 		}
 
 		// Removes the redundant copy instruction.
-		// TODO: doing this in `SetupEpilogue` seems weird. Find a better home.
 		if cur.IsCopy() && cur.rn.realReg() == cur.rd.realReg() {
 			prev, next := cur.prev, cur.next
 			// Remove the copy instruction.
diff --git a/internal/engine/wazevo/backend/machine.go b/internal/engine/wazevo/backend/machine.go
index 0bbb8c91..c746cf1c 100644
--- a/internal/engine/wazevo/backend/machine.go
+++ b/internal/engine/wazevo/backend/machine.go
@@ -55,12 +55,8 @@ type (
 		// RegAlloc does the register allocation after lowering.
 		RegAlloc()
 
-		// SetupPrologue inserts the prologue after register allocations.
-		SetupPrologue()
-
-		// SetupEpilogue inserts the epilogue after register allocations.
-		// This sets up the instructions for the inverse of SetupPrologue right before
-		SetupEpilogue()
+		// PostRegAlloc does the post register allocation, e.g. setting up prologue/epilogue, redundant move elimination, etc.
+		PostRegAlloc()
 
 		// ResolveRelocations resolves the relocations after emitting machine code.
 		ResolveRelocations(refToBinaryOffset map[ssa.FuncRef]int, binary []byte, relocations []RelocationInfo)


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue_test.go b/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue_test.go
index 7ea8c6e2..d59589ac 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue_test.go
@@ -9,7 +9,7 @@ import (
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
 
-func TestMachine_SetupPrologue(t *testing.T) {
+func TestMachine_setupPrologue(t *testing.T) {
 	for _, tc := range []struct {
 		spillSlotSize int64
 		clobberedRegs []regalloc.VReg
@@ -75,7 +75,7 @@ func TestMachine_SetupPrologue(t *testing.T) {
 			root.next = udf
 			udf.prev = root
 
-			m.SetupPrologue()
+			m.setupPrologue()
 			require.Equal(t, root, m.ectx.RootInstr)
 			m.Encode(context.Background())
 			require.Equal(t, tc.exp, m.Format())
@@ -83,7 +83,7 @@ func TestMachine_SetupPrologue(t *testing.T) {
 	}
 }
 
-func TestMachine_SetupEpilogue(t *testing.T) {
+func TestMachine_postRegAlloc(t *testing.T) {
 	for _, tc := range []struct {
 		exp           string
 		abi           backend.FunctionABI
@@ -148,7 +148,7 @@ func TestMachine_SetupEpilogue(t *testing.T) {
 			ret.asRet(nil)
 			root.next = ret
 			ret.prev = root
-			m.SetupEpilogue()
+			m.postRegAlloc()
 
 			require.Equal(t, root, m.ectx.RootInstr)
 			m.Encode(context.Background())
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go
index 260cabc3..c878084a 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go
@@ -9,7 +9,7 @@ import (
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
 
-func TestMachine_SetupPrologue(t *testing.T) {
+func TestMachine_setupPrologue(t *testing.T) {
 	for _, tc := range []struct {
 		spillSlotSize int64
 		clobberedRegs []regalloc.VReg
@@ -108,7 +108,7 @@ func TestMachine_SetupPrologue(t *testing.T) {
 			root.next = udf
 			udf.prev = root
 
-			m.SetupPrologue()
+			m.setupPrologue()
 			require.Equal(t, root, m.executableContext.RootInstr)
 			m.Encode(context.Background())
 			require.Equal(t, tc.exp, m.Format())
@@ -116,7 +116,7 @@ func TestMachine_SetupPrologue(t *testing.T) {
 	}
 }
 
-func TestMachine_SetupEpilogue(t *testing.T) {
+func TestMachine_postRegAlloc(t *testing.T) {
 	for _, tc := range []struct {
 		exp           string
 		abi           backend.FunctionABI
@@ -220,7 +220,7 @@ func TestMachine_SetupEpilogue(t *testing.T) {
 			ret.asRet(nil)
 			root.next = ret
 			ret.prev = root
-			m.SetupEpilogue()
+			m.postRegAlloc()
 
 			require.Equal(t, root, m.executableContext.RootInstr)
 			m.Encode(context.Background())
diff --git a/internal/engine/wazevo/backend/machine_test.go b/internal/engine/wazevo/backend/machine_test.go
index a80d737e..d9b2b093 100644
--- a/internal/engine/wazevo/backend/machine_test.go
+++ b/internal/engine/wazevo/backend/machine_test.go
@@ -57,11 +57,8 @@ func (m mockMachine) Encode(ctx context.Context) {}
 // ResolveRelocations implements Machine.ResolveRelocations.
 func (m mockMachine) ResolveRelocations(map[ssa.FuncRef]int, []byte, []RelocationInfo) {}
 
-// SetupPrologue implements Machine.SetupPrologue.
-func (m mockMachine) SetupPrologue() {}
-
-// SetupEpilogue implements Machine.SetupEpilogue.
-func (m mockMachine) SetupEpilogue() {}
+// PostRegAlloc implements Machine.SetupPrologue.
+func (m mockMachine) PostRegAlloc() {}
 
 // Function implements Machine.Function.
 func (m mockMachine) Function() (f regalloc.Function) { return }